diff --git a/AGENTS.md b/AGENTS.md index 6b37a67d5..75b58e12f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,7 +44,9 @@ Python (editable; compiles the pybind11 extension via CMake): ```bash pip install -e . # core graph API only -pip install -e ".[cutedsl]" # + OSS CuTeDSL kernels (torch, nvidia-cutlass-dsl, cuda-python) +pip install -e ".[cutedsl]" # + OSS CuTeDSL kernels (nvidia-cutlass-dsl, cuda-python, tvm-ffi; framework-neutral) +pip install --group torch # + torch for the CuTeDSL APIs (torch, torch-c-dlpack-ext) +pip install --group jax # + jax for the CuTeDSL APIs (jax >= 0.5; XLA entry points via cutlass.jax) ``` `setup.py` honors env vars: `CUDNN_PATH`, `CUDA_PATH` / `CUDAToolkit_ROOT`, `DEBUG=1` (debug build), `CMAKE_BUILD_PARALLEL_LEVEL`, `CMAKE_GENERATOR`. @@ -61,7 +63,7 @@ cd test/python pytest # default is -m L0 (smoke level) per pytest.ini pytest -m L1 # deeper levels: L0..L4 pytest test_conv_fprop.py # one file (still filtered by -m L0 — pass -m "L0 or L1" to widen) -pytest fe_api/ # OSS kernel tests; require ".[cutedsl]" install + SM90/SM100 GPU +pytest fe_api/ # OSS kernel tests; require ".[cutedsl]" + `--group torch` (and `--group jax` for the *_jax tests) + SM90/SM100 GPU ``` Read [test/AGENTS.md](test/AGENTS.md) before touching tests — `test/python/conftest.py` has import-order and env-var requirements that are easy to break. diff --git a/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md b/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md index b80675432..5a3784cc0 100644 --- a/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md @@ -2,6 +2,12 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** in FP8 configurations: b_ptrs/sfb_ptrs as packed-uint8 (or x64 int64) pointer arrays, SFA in the physical C-contiguous atom shape, SFD outputs allocated the same way (the kernel rebuilds all SF layouts from the GEMM shapes and reads only base pointers). Packed-fp4 inputs are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `discrete_grouped_gemm_dswiglu_jax_sm100` (built on `cudnn.jax.call`; k-major weights only): all outputs (d_row/d_col, SFD tensors, amax, `dprob` as a bridge-managed zero-initialized accumulator, optional `dbias`) are XLA-managed donated buffers — no manual synchronization. Under tracing the offsets *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation. + ## Overview **Discrete Grouped GEMM + dGLU backward fusion**: A block-scaled grouped GEMM fused with a dSwiGLU/dGeGLU backward epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE workloads where each expert weight/scale lives in a separate allocation. diff --git a/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md index d89841c24..96d19ee8a 100644 --- a/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md @@ -2,6 +2,12 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** in FP8 configurations: b_ptrs/sfb_ptrs as packed-uint8 (or x64 int64) pointer arrays, SFA in the physical C-contiguous atom shape `(1, MN', K', 32, 4, 4)`, SFD outputs allocated the same way (the kernel rebuilds all SF layouts from the GEMM shapes and reads only base pointers). Column-major bias and packed-fp4 inputs are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `discrete_grouped_gemm_swiglu_jax_sm100` (built on `cudnn.jax.call`; k-major weights only): all outputs (c/d/d_col, SFD tensors, amax) are XLA-managed donated buffers — no manual synchronization. Under tracing the offsets *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation. + ## Overview **Discrete Grouped GEMM + SwiGLU fusion**: A block-scaled grouped GEMM fused with a SwiGLU/GeGLU epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE workloads where each expert weight lives in a separate allocation. diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md index 4247526ce..a84683c51 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md @@ -109,7 +109,7 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None) Two integration levels are available: -1. **`gemm_amax_jax_sm100`** (recommended for jitted programs) — an XLA custom call via [jax-tvm-ffi]. The kernel runs on XLA's compute stream (correctly ordered with surrounding ops), outputs are XLA-managed fresh arrays, and the call composes with `jax.jit`. No manual synchronization is needed. Requires the `jax` dependency group (`pip install --group jax`, which brings `jax-tvm-ffi`; Python >= 3.11). +1. **`gemm_amax_jax_sm100`** (recommended for jitted programs) — an XLA custom call built on `cudnn.jax.call` (CuTeDSL's native `cutlass.jax` bridge). The kernel runs on XLA's compute stream (correctly ordered with surrounding ops), outputs are XLA-managed fresh arrays, and the call composes with `jax.jit`. No manual synchronization is needed. Requires the `jax` dependency group (`pip install --group jax`; jax >= 0.5). ```python from cudnn import gemm_amax_jax_sm100 @@ -120,7 +120,7 @@ def quantized_matmul(a, b, sfa, sfb): return c, amax ``` -Calling it eagerly works but re-traces the `ffi_call` on every invocation; call it from inside a jitted function in hot loops. +Calling it eagerly works but re-traces the custom call on every invocation; call it from inside a jitted function in hot loops. 2. **The eager entry points below** (`gemm_amax_wrapper_sm100`, `GemmAmaxSm100`) also accept JAX arrays via DLPack. In hot loops prefer the **class API with pre-allocated output buffers** (~15 µs CPU per launch) over the wrapper — per-call `jnp` output allocation in the wrapper costs hundreds of µs of XLA dispatch. @@ -254,7 +254,7 @@ Tuple unpacking order is: `(c_tensor, amax_tensor)`. - `L == 1`; `A`/`B` k-major; `C` n-major only - `SFA`/`SFB` in the physical atom shape `(L, MN', K', 32, 4, 4)` (see "Using JAX arrays") -- Eager use only (no `jax.jit` over these entry points); synchronize before reading outputs +- The wrapper/class entry points are eager-only (use `gemm_amax_jax_sm100` under `jax.jit`); synchronize before reading outputs --- diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md index 47c1d9b63..1b96a0779 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md @@ -73,7 +73,7 @@ A (MxKxL), SFA B (NxKxL), SFB ## API Usage -The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). A `jax.jit`-compatible XLA custom-call entry point is not yet available for this kernel (its signature carries optional None-typed parameters that the jax-tvm-ffi bridge cannot supply). +The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point (`gemm_srelu_jax_sm100` / `gemm_dsrelu_jax_sm100`, built on `cudnn.jax.call`); the kernels' optional parameters are compile-time constants inside its adapter. ### High-level wrapper diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md b/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md index eed9aa9bd..f58568dc4 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md @@ -2,6 +2,14 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` (the `[in, out]` weight layout reaches the kernel through a transposed strided view, which has no row-major JAX equivalent and raises a clear error). The E8M0 scale inputs stay `uint8` as with torch. Outputs are allocated as C-contiguous `jnp` arrays. The wrapper is eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"): same contract as the wrapper with `w_out_in=True`, dispatching on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs. + +The API is compiled with `--enable-tvm-ffi`: raw framework tensors go straight to the compiled kernel (no per-call `from_dlpack` conversion), cutting per-launch CPU overhead roughly in half for torch callers as well. + ## Overview **Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize**: a persistent dense GEMM on NVIDIA Blackwell GPUs (SM100+) that projects activations, applies the Megatron MLA-YARN rotary embedding to each attention head's trailing rotary features, and MXFP8 (E4M3, block=32) quantizes the result in **both** the rowwise (D-direction) and columnwise (S-direction) layouts. Implemented with CUTLASS/CUTE. diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md index 7b0196186..2e3a56d8c 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md @@ -66,7 +66,7 @@ A (MxKxL), SFA B (NxKxL), SFB ## API Usage -The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). A `jax.jit`-compatible XLA custom-call entry point is not yet available for this kernel (its signature carries optional None-typed parameters that the jax-tvm-ffi bridge cannot supply). +The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point (`gemm_srelu_jax_sm100` / `gemm_dsrelu_jax_sm100`, built on `cudnn.jax.call`); the kernels' optional parameters are compile-time constants inside its adapter. ### High-level wrapper diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index 4131344de..46968ff81 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -68,7 +68,7 @@ Notes: The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, SF tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). -For jitted JAX programs, use **`gemm_swiglu_jax_sm100`** — an XLA custom call (via jax-tvm-ffi, `jax` dependency group) that runs on XLA's compute stream, returns fresh `(ab12, c)` arrays, and composes with `jax.jit`. It currently supports the standard (non-quantized) kernel only; use the eager wrapper for blockscaled MXFP8 inputs from JAX. `alpha` is a static (trace-time) parameter. +For jitted JAX programs, use **`gemm_swiglu_jax_sm100`** — an XLA custom call (built on `cudnn.jax.call` / CuTeDSL's native `cutlass.jax` bridge, `jax` dependency group) that runs on XLA's compute stream, returns fresh `(ab12, c)` arrays, and composes with `jax.jit`. Both the standard and the blockscaled MXFP8 quantized kernels are supported. `alpha` is a static (trace-time) parameter. ```python from cudnn import gemm_swiglu_jax_sm100 diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md index ec932bae3..9a303ec9a 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md @@ -53,6 +53,18 @@ alignment contract. The API records the pointer-array tensor on the launch stream. The caller must keep every pointed-to expert allocation alive and must not modify or free it until that stream completes. +## Using JAX arrays + +The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is imported only when torch tensors are passed, jax only when JAX arrays are passed). Because JAX arrays are always row-major, the JAX contract is narrower than torch's: + +- **Discrete weight mode only** (`b_ptrs`): dense mode's `b_tensor` uses an expert-outermost strided layout with no row-major equivalent, and `bias_tensor`'s `(n, experts)` column-major layout is likewise not expressible — both raise clear errors for JAX inputs. Each per-expert weight is a plain k-major `(n, k)` C-contiguous JAX array. +- **`b_ptrs` from JAX**: build the pointer array from `weight.unsafe_buffer_pointer()` per expert. JAX truncates int64 without x64 mode, so pass the pointers either as an int64 array (with `jax_enable_x64`) or as a **packed uint8 array** (8 little-endian bytes per pointer): `jnp.asarray(np.array(ptrs, dtype=np.int64).view(np.uint8))`. The weight arrays (and `b_ptrs`) must stay alive and un-donated until the kernel completes. +- A/offsets/alpha/prob are plain C-contiguous JAX arrays of the documented shapes; outputs are allocated as n-major C-contiguous `jnp` arrays. Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. +- The eager path launches on the **CUDA legacy default stream** (XLA does not track it): `jax.block_until_ready(...)` your inputs before calling, and synchronize the device (or the stream you passed) before reading the outputs. +- For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_jax_sm100(a_tensor, padded_offsets, alpha_tensor, b_ptrs, n, prob_tensor, ...)` (built on `cudnn.jax.call`; discrete mode, no bias): outputs are fresh XLA-managed arrays with rows at/past `padded_offsets[-1]` zero-filled, and no manual synchronization is needed. Under tracing the `padded_offsets` *values* cannot be host-validated (shapes/dtypes still are), and the per-expert weight buffers behind `b_ptrs` must stay alive and unmoved across every execution of the traced computation. + +Internal workspaces are allocated in the caller's framework allocator (torch caching allocator or XLA's pool) and written through raw pointers; they are never surfaced as arrays. + ## Wrapper API Dense mode: diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index c9c4352a2..b8a85fd2a 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -2,6 +2,12 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** on the BF16 backend in discrete weight mode (dswiglu and dgeglu), including `generate_dbias=True` and caller-provided zero-initialized `dprob`. Dense `b_tensor` and the block-scaled backend (MMA-interleaved scale-factor layouts) are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_dglu_jax_sm100` (built on `cudnn.jax.call`; discrete mode): `dprob` and (with `generate_dbias=True`) `dbias` come back as bridge-managed zero-initialized accumulator outputs — no caller-zeroed buffers, no manual synchronization. Under tracing the `padded_offsets` *values* cannot be host-validated, and the per-expert weight buffers behind `b_ptrs` must stay alive and unmoved across every execution of the traced computation. + ## Overview **Unified Grouped GEMM + dGLU fusion**: one public class and wrapper select a diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md index e78a5f0a8..ac8b0e625 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md @@ -2,6 +2,12 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** in the discrete (b_ptrs) FP8 configurations: pointer arrays as int64 (jax x64 mode) or packed uint8 (8 bytes per pointer), scale-factor tensors in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)` (the kernel rebuilds SF layouts from the GEMM shapes and reads only the base pointer), outputs allocated as C-contiguous `jnp` arrays. Dense weight mode and packed-fp4 A/B are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_dsrelu_jax_sm100` (built on `cudnn.jax.call`; discrete FP8 mode, `sf_vec_size=32`): all outputs (d/SFD tensors, `dprob`, and with `generate_dbias=True` `dbias`) are XLA-managed donated zero-initialized buffers — no manual synchronization. Under tracing the `padded_offsets` *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation. + ## Overview **Grouped GEMM + dsReLU backward fusion**: A grouped block-scaled GEMM fused with a probability-gradient backward epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. The API supports dense contiguous weights and discrete per-expert weight allocations. Groups are contiguous in the `M` dimension and described by `padded_offsets`. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md index cf7b40fa8..11dc7cbf2 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md @@ -4,6 +4,10 @@ **Legacy contiguous-only API note:** This page documents the older contiguous-only dSwiGLU API. For new integrations, prefer the unified [Grouped GEMM + dGLU](grouped_gemm_dglu.md) API, which covers dense and discrete weight layouts. +## JAX support + +JAX arrays are **not supported**: this API is dense-weight-mode only, and the expert-outermost strided B layout has no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + ## Overview **Grouped GEMM + dSwiGLU fusion**: A contiguous grouped block-scaled GEMM fused with a dSwiGLU backward epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md index 38ead6042..eacd4a3e5 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md @@ -2,6 +2,12 @@ **This is an experimental API and subject to change.** +## JAX support + +Supports **JAX arrays** on the BF16 backend in discrete weight mode (swiglu and geglu): `b_ptrs` as a packed little-endian uint8 pointer array (8 bytes per pointer; int64 accepted with jax x64 mode), outputs allocated as n-major C-contiguous `jnp` arrays. Dense `b_tensor` (expert-outermost strides), column-major `bias_tensor`, and the block-scaled backend (MMA-interleaved scale-factor layouts) are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_glu_jax_sm100` (built on `cudnn.jax.call`; discrete mode, no bias, `b_major="k"`): outputs are fresh XLA-managed arrays with rows at/past `padded_offsets[-1]` zero-filled, no manual synchronization needed. `linear_offset` is a compile-time constant (each distinct value compiles a new specialization). Under tracing the `padded_offsets` *values* cannot be host-validated, and the per-expert weight buffers behind `b_ptrs` must stay alive and unmoved across every execution of the traced computation. + ## Overview **Unified Grouped GEMM + GLU fusion**: one public class and wrapper select a diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md index 1d41af336..2166ad4e0 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md @@ -2,6 +2,10 @@ **This is an experimental API and subject to change.** +## JAX support + +JAX arrays are **not supported**: this fusion is block-scaled-only and its mandatory scale-factor inputs use an MMA-interleaved layout with no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + ## Overview **Grouped GEMM + GLU + Hadamard fusion**: A contiguous grouped block-scaled GEMM fused with a GLU epilogue, a 16-wide Hadamard transform, and per-expert `amax` reduction on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. Groups are contiguous in the `M` dimension and described by `padded_offsets`. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md index 004105c27..1e1083ce3 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md @@ -4,6 +4,10 @@ **Legacy dense-only API note:** This page documents the older dense-only grouped quant API. For new integrations, prefer the unified [Grouped GEMM + Quant (Unified)](grouped_gemm_quant_unified.md) page. +## JAX support + +JAX arrays are **not supported**: all configurations consume the SFA scale-factor tensor as an MMA-permuted strided cute tensor argument, a layout with no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + ## Overview **Grouped GEMM + Quant fusion**: A contiguous grouped block-scaled GEMM with output quantization on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md index af78b87f0..42db88288 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md @@ -2,6 +2,10 @@ **This is an experimental API and subject to change.** +## JAX support + +JAX arrays are **not supported**: both dense and discrete modes consume the SFA scale-factor tensor as an MMA-permuted strided cute tensor argument, a layout with no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + ## Overview **Grouped GEMM + sReLU fusion**: A grouped block-scaled GEMM fused with a probability-gated squared-ReLU epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. The API supports dense contiguous weights and discrete per-expert weight allocations. Groups are contiguous in the `M` dimension and described by `padded_offsets`. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md index 1ac83b772..493244140 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md @@ -4,6 +4,10 @@ **Legacy contiguous-only API note:** This page documents the older contiguous-only SwiGLU API. For new integrations, prefer the unified [Grouped GEMM + GLU](grouped_gemm_glu.md) API, which covers dense and discrete weight layouts. +## JAX support + +JAX arrays are **not supported**: this kernel consumes its scale-factor tensors as MMA-permuted strided cute tensor arguments in every configuration, a layout with no row-major (JAX) equivalent. JAX inputs raise a clear `ValueError` at the entry points. The API is otherwise type-erased and torch-lazy. + ## Overview **Grouped GEMM + SwiGLU fusion**: A contiguous grouped block-scaled GEMM fused with a SwiGLU epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index 59acdf82d..88514b2ba 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -11,6 +11,12 @@ Install the optional CuTe DSL dependencies before importing either API: pip install nvidia-cudnn-frontend[cutedsl] ``` +## JAX support + +Supports **JAX arrays** on the BF16 backend: A k-major and B n-major C-contiguous arrays, dense `(experts, m, n)` C-contiguous output or discrete output pointers (packed uint8 / int64 with jax x64 mode). The block-scaled backend's layouts are not expressible as JAX arrays and raise a clear error. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. + +For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_wgrad_jax_sm100` (built on `cudnn.jax.call`; BF16, discrete output pointers): the per-expert weight-gradient buffers behind `wgrad_ptrs` are caller-owned external memory the kernel writes through, so the entry returns a completion **token** — `jax.block_until_ready(token)` before reading them (and zero them yourself between runs unless accumulating). Under tracing the per-group offsets *values* cannot be host-validated; the external buffers must stay alive and unmoved across every execution of the traced computation. + ## Operation For expert `e`, let `begin = 0` for the first expert and diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 3661dc40e..da2101001 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -2,7 +2,11 @@ **FE-OSS APIs are experimental and subject to change.** -The GEMM CuTeDSL APIs are type-erased and torch-lazy: torch is imported only when torch tensors are passed. The dense GEMM fusions (amax, swiglu, srelu, dsrelu) additionally accept JAX arrays (see `gemm_amax.md` "Using JAX arrays" for the JAX contract), with `jax.jit`-compatible XLA custom-call entry points for amax and swiglu; the grouped / discrete-grouped / proj_rope APIs currently support torch tensors only and reject other frameworks with a clear error. +The GEMM CuTeDSL APIs are type-erased and torch-lazy: torch is imported only when torch tensors are passed. JAX arrays are additionally accepted wherever the kernel's tensor layouts are expressible as row-major arrays (each API's page has a "JAX support" section with its exact contract): + +- **Dense fusions** (amax, swiglu, srelu, dsrelu): full JAX eager support, plus `jax.jit`-compatible XLA custom-call entry points for all four (built on `cudnn.jax.call` / CuTeDSL's native `cutlass.jax` bridge; see `gemm_amax.md` "Using JAX arrays"). +- **Grouped / discrete-grouped**: JAX eager support in discrete (pointer-array) weight modes — unfused grouped GEMM, glu/dglu (BF16), dsrelu (FP8), wgrad (BF16), and discrete-grouped swiglu/dswiglu (FP8) — plus a `jax.jit`-compatible `*_jax_sm100` entry point for each of those same families (built on `cudnn.jax.call`; each API page documents its exact jit contract). Dense weight mode, column-major bias layouts, and kernels whose scale factors are MMA-permuted tensor arguments (grouped swiglu/srelu/quant/dswiglu, glu_hadamard, block-scaled glu/dglu/wgrad backends) reject JAX with clear errors. +- **proj_rope_mxfp8**: JAX eager support on both input paths with `w_out_in=True` (the transposed [in, out] weight view is torch-only), plus the `jax.jit`-compatible `gemm_proj_rope_mxfp8_jax_sm100` entry point. This folder documents the Python FE APIs implemented under `python/cudnn`. For details on currently implemented operations, see: - [GEMM + Amax](gemm_fusions/gemm_amax.md) @@ -39,6 +43,13 @@ All Frontend OSS APIs come installed with the `nvidia-cudnn-frontend` package. H pip install nvidia-cudnn-frontend[cutedsl] ``` +The `cutedsl` extra is framework-neutral (nvidia-cutlass-dsl, cuda-python, apache-tvm-ffi). Install your tensor framework separately — from a checkout, the PEP 735 dependency groups pin the right companion packages: +```bash +pip install --group torch # torch + torch-c-dlpack-ext +pip install --group jax # jax >= 0.5 (XLA entry points via cutlass.jax, shipped with nvidia-cutlass-dsl) +``` +(For the published wheel, `pip install torch torch-c-dlpack-ext` or `pip install "jax>=0.5"` directly.) + After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` ## API Usage diff --git a/pyproject.toml b/pyproject.toml index 2d6fcda2c..75a9751ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,11 +62,11 @@ cutedsl = [ # their own constraints. The engines check the version at support time and # decline when it is too old, so an older DSL costs those engines and # nothing else. See CUTEDSL_MIN_VERSION in cudnn/frost/buffers.py. + # Framework-neutral core only: the CuTeDSL APIs are type-erased, so torch (like + # jax) is opt-in via the [dependency-groups] below (`pip install --group torch`). "nvidia-cutlass-dsl[cu13]>=4.5.0", "cuda-python", - "torch", "apache-tvm-ffi>=0.1.11", - "torch-c-dlpack-ext", ] [dependency-groups] @@ -88,11 +88,11 @@ torch = [ "torch", "torch-c-dlpack-ext", ] -# jax-tvm-ffi provides the jax.jit-compatible XLA custom-call entry points -# (e.g. gemm_amax_jax_sm100) and requires Python >= 3.11. +# The jax.jit-compatible XLA custom-call entry points (cudnn.jax, e.g. +# gemm_amax_jax_sm100) build on the CuTeDSL JAX extensions (cutlass.jax, +# shipped with nvidia-cutlass-dsl), which require jax >= 0.5. jax = [ - "jax>=0.4.35", - "jax-tvm-ffi>=0.1.3", + "jax>=0.5", ] [build-system] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 9b4fee01d..ff3da8601 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -319,6 +319,8 @@ def _dlopen_cudnn(): "GemmSwigluSm100": (".gemm.cutedsl.dense.swiglu", "GemmSwigluSm100"), "gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_wrapper_sm100"), "gemm_swiglu_jax_sm100": (".gemm.cutedsl.dense.swiglu", "gemm_swiglu_jax_sm100"), + "gemm_srelu_jax_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_jax_sm100"), + "gemm_dsrelu_jax_sm100": (".gemm.cutedsl.dense.dsrelu", "gemm_dsrelu_jax_sm100"), "GemmSreluSm100": (".gemm.cutedsl.dense.srelu", "GemmSreluSm100"), "gemm_srelu_wrapper_sm100": (".gemm.cutedsl.dense.srelu", "gemm_srelu_wrapper_sm100"), "GemmDsreluSm100": (".gemm.cutedsl.dense.dsrelu", "GemmDsreluSm100"), @@ -329,11 +331,19 @@ def _dlopen_cudnn(): "GemmProjRopeMxfp8Bf16InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Bf16InSm100"), "GemmProjRopeMxfp8Mxfp8InSm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "GemmProjRopeMxfp8Mxfp8InSm100"), "gemm_proj_rope_mxfp8_wrapper_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_wrapper_sm100"), + "gemm_proj_rope_mxfp8_jax_sm100": (".gemm.cutedsl.dense.proj_rope_mxfp8", "gemm_proj_rope_mxfp8_jax_sm100"), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), "grouped_gemm": (".gemm.cutedsl.grouped", None), "GroupedGemmSm100": (".gemm.cutedsl.grouped", "GroupedGemmSm100"), "grouped_gemm_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wrapper_sm100"), + "grouped_gemm_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_jax_sm100"), + "grouped_gemm_glu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_glu_jax_sm100"), + "grouped_gemm_dglu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dglu_jax_sm100"), + "grouped_gemm_dsrelu_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_dsrelu_jax_sm100"), + "grouped_gemm_wgrad_jax_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_wgrad_jax_sm100"), + "discrete_grouped_gemm_swiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_swiglu_jax_sm100"), + "discrete_grouped_gemm_dswiglu_jax_sm100": (".gemm.cutedsl.discrete_grouped", "discrete_grouped_gemm_dswiglu_jax_sm100"), "GroupedGemmSwigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmSwigluSm100"), "grouped_gemm_swiglu_wrapper_sm100": (".gemm.cutedsl.grouped", "grouped_gemm_swiglu_wrapper_sm100"), "GroupedGemmDswigluSm100": (".gemm.cutedsl.grouped", "GroupedGemmDswigluSm100"), @@ -401,10 +411,18 @@ def __getattr__(name: str) -> Any: globals()["experimental"] = _experimental return _experimental + if name == "jax": + # `import cudnn; cudnn.jax.call` works like `import cudnn.jax`. + # Deferred so torch-only users never pay the jax import (the submodule + # itself raises a descriptive ImportError when jax >= 0.5 is missing). + _jax = importlib.import_module(".jax", __name__) + globals()["jax"] = _jax + return _jax + if name in _LAZY_OPTIONAL_IMPORTS: return _load_optional_symbol(name) - raise AttributeError(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def __dir__(): diff --git a/python/cudnn/gemm/cutedsl/_jax_ffi.py b/python/cudnn/gemm/cutedsl/_jax_ffi.py deleted file mode 100644 index 8006cc4fd..000000000 --- a/python/cudnn/gemm/cutedsl/_jax_ffi.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Shared helpers for the jax.jit-compatible entry points of the CuTeDSL GEMM APIs. - -These entry points integrate with XLA via jax-tvm-ffi: the kernel variant is compiled -with the TVM-FFI environment stream (so it runs on XLA's compute stream), all outputs -are donated pre-initialized operands (input_output_aliases + arg_spec=["args"]) matching -the kernels' destination-passing signatures, and calls compose with jax.jit. - -This module imports jax/jax_tvm_ffi at import time; it is only loaded from the -per-API jax_api modules, which are lazily exported. -""" - -from typing import Any, Callable, Tuple - -import jax_tvm_ffi - -from cudnn.api_base import TensorDesc -from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.tensor_adapter import Device - - -def c_contiguous_strides(shape: Tuple[int, ...]) -> Tuple[int, ...]: - strides, acc = [1] * len(shape), 1 - for i in range(len(shape) - 1, -1, -1): - strides[i] = acc - acc *= shape[i] - return tuple(strides) - - -def make_row_major_desc(shape: Tuple[int, ...], dtype: Any, name: str) -> TensorDesc: - """Descriptor for a C-contiguous (row-major) JAX buffer, from shape/dtype only. - - Built from aval metadata so this works for jax.jit tracers as well as concrete - arrays (tracers expose .shape/.dtype but no device or DLPack). - """ - shape = tuple(shape) - stride = c_contiguous_strides(shape) - return TensorDesc( - dtype=_convert_to_cutlass_data_type(dtype), - shape=shape, - stride=stride, - stride_order=TensorDesc._compute_stride_order(shape, stride), - device=Device("cuda", 0), - name=name, - ) - - -def get_or_register_env_stream_target( - registry: dict, - cache_key: Any, - make_gemm: Callable[[], Any], - target_prefix: str, - arg_spec: Tuple[str, ...] = ("args",), -) -> str: - """Compile the env-stream kernel variant for cache_key (once) and register it as an - XLA FFI target; return the registered target name. - - The registry keeps (target, gemm, compiled) so the compiled kernel stays alive - alongside the global registration. - """ - entry = registry.get(cache_key) - if entry is None: - gemm = make_gemm() - assert gemm.check_support() - compiled = gemm._compile_kernel(use_tvm_ffi_env_stream=True) - target = f"{target_prefix}.{len(registry)}" - jax_tvm_ffi.register_ffi_target(target, compiled, arg_spec=list(arg_spec), platform="gpu", allow_cuda_graph=True) - registry[cache_key] = (target, gemm, compiled) - entry = registry[cache_key] - return entry[0] diff --git a/python/cudnn/gemm/cutedsl/dense/amax/__init__.py b/python/cudnn/gemm/cutedsl/dense/amax/__init__.py index 5e7250a2f..8b1a1b60e 100644 --- a/python/cudnn/gemm/cutedsl/dense/amax/__init__.py +++ b/python/cudnn/gemm/cutedsl/dense/amax/__init__.py @@ -14,7 +14,7 @@ def __getattr__(name): - # Lazy: the jax entry point imports jax/jax_tvm_ffi, which must not be pulled in + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in # for torch-only users. if name == "gemm_amax_jax_sm100": from .jax_api import gemm_amax_jax_sm100 diff --git a/python/cudnn/gemm/cutedsl/dense/amax/api.py b/python/cudnn/gemm/cutedsl/dense/amax/api.py index 84eaadec9..724c3c0d2 100644 --- a/python/cudnn/gemm/cutedsl/dense/amax/api.py +++ b/python/cudnn/gemm/cutedsl/dense/amax/api.py @@ -89,23 +89,43 @@ def _check_sf_shape(self, sf_desc: TensorDesc, l: int, name: str) -> Tuple[int, for frameworks such as JAX that cannot express the permuted (strided) view. The kernel rebuilds the SF layout from the A/B shapes and consumes only the SF base - pointer, so only the element count and memory order matter. + pointer, so the memory must be exactly the C-contiguous physical allocation in both + forms: strides are validated too (a shape-matching but differently-strided tensor + would silently produce wrong results). """ shape = sf_desc.shape self._value_error_if(len(shape) != 6, f"{name} tensor must be 6-D, got shape {shape}") - if shape[0] == self.atom_m[0] and shape[1] == self.atom_m[1] and shape[3] == self.atom_k: + atom_m0, atom_m1, atom_k = self.atom_m[0], self.atom_m[1], self.atom_k + atom_elems = atom_m0 * atom_m1 * atom_k + if shape[0] == atom_m0 and shape[1] == atom_m1 and shape[3] == atom_k: mn_div_atom_m0_m1, sf_k_div_atom_k = shape[2], shape[4] - self._check_tensor_shape( + atom_shape = (atom_m0, atom_m1, mn_div_atom_m0_m1, atom_k, sf_k_div_atom_k, l) + self._check_tensor_shape(sf_desc, atom_shape, name) + _ = self._check_tensor_stride( sf_desc, - (self.atom_m[0], self.atom_m[1], mn_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l), - name, + stride=[ + canonicalize_unit_dim_strides( + atom_shape, + (atom_m1 * atom_k, atom_k, sf_k_div_atom_k * atom_elems, 1, atom_elems, mn_div_atom_m0_m1 * sf_k_div_atom_k * atom_elems), + ) + ], + name=name, + extra_error_msg=f"{name} atom view must be the (3, 4, 1, 5, 2, 0) permutation of a C-contiguous physical allocation", ) else: mn_div_atom_m0_m1, sf_k_div_atom_k = shape[1], shape[2] - self._check_tensor_shape( + physical_shape = (l, mn_div_atom_m0_m1, sf_k_div_atom_k, atom_m0, atom_m1, atom_k) + self._check_tensor_shape(sf_desc, physical_shape, name) + _ = self._check_tensor_stride( sf_desc, - (l, mn_div_atom_m0_m1, sf_k_div_atom_k, self.atom_m[0], self.atom_m[1], self.atom_k), - name, + stride=[ + canonicalize_unit_dim_strides( + physical_shape, + (mn_div_atom_m0_m1 * sf_k_div_atom_k * atom_elems, sf_k_div_atom_k * atom_elems, atom_elems, atom_m1 * atom_k, atom_k, 1), + ) + ], + name=name, + extra_error_msg=f"{name} in the physical (L, MN', K', Atom_M0, Atom_M1, Atom_K) form must be C-contiguous", ) return mn_div_atom_m0_m1, sf_k_div_atom_k @@ -309,13 +329,8 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): self._logger.debug("check_support completed successfully") return True - def _compile_kernel(self, use_tvm_ffi_env_stream: bool = False): - """Compile the kernel and return the raw TVM-FFI callable. - - With ``use_tvm_ffi_env_stream=True`` the stream argument is bound to the - TVM-FFI environment stream and dropped from the callable's signature -- - the variant used for XLA custom-call (jax.ffi) integration. - """ + def _compile_kernel(self): + """Compile the kernel and return the raw TVM-FFI callable.""" self._ensure_support_checked() gemm_amax = self._kernel( @@ -338,7 +353,7 @@ def _compile_kernel(self, use_tvm_ffi_env_stream: bool = False): sfb_cute = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) c_cute = self._make_fake_cute_tensor_from_desc(self.c_desc, assumed_align=16) 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=use_tvm_ffi_env_stream) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) return cute.compile( gemm_amax, @@ -360,7 +375,7 @@ def compile(self) -> None: self._logger.debug("Kernel already compiled; skipping recompilation") return - _compiled_kernel = self._compile_kernel(use_tvm_ffi_env_stream=False) + _compiled_kernel = self._compile_kernel() def tensor_api( a_tensor: Any, diff --git a/python/cudnn/gemm/cutedsl/dense/amax/jax_api.py b/python/cudnn/gemm/cutedsl/dense/amax/jax_api.py index 8d4284102..148b34ba1 100644 --- a/python/cudnn/gemm/cutedsl/dense/amax/jax_api.py +++ b/python/cudnn/gemm/cutedsl/dense/amax/jax_api.py @@ -3,14 +3,9 @@ """JAX-native (XLA custom call) entry point for the blockscaled GEMM + amax kernel. -Unlike the eager path (``gemm_amax_wrapper_sm100`` with JAX arrays), this entry point -integrates with XLA via jax-tvm-ffi: the kernel is compiled with the TVM-FFI environment -stream (so it runs on XLA's compute stream, correctly ordered with surrounding ops), the -output buffers are managed by XLA through donation, and the call is ``jax.jit``-compatible. -No manual synchronization or block_until_ready is needed around it. - -This module imports jax/jax_tvm_ffi at import time; it is only loaded when the -``gemm_amax_jax_sm100`` symbol is requested. +Built on :func:`cudnn.jax.call` (CuTeDSL's native JAX bridge): the kernel runs on +XLA's compute stream, outputs are XLA-managed, and the call composes with ``jax.jit`` +and CUDA graph capture. No manual synchronization is needed. """ from typing import Any, Tuple @@ -19,15 +14,26 @@ import jax.numpy as jnp import cutlass +import cutlass.cute as cute +import cutlass.utils from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.tensor_adapter import framework_dtype -from cudnn.gemm.cutedsl._jax_ffi import get_or_register_env_stream_target, make_row_major_desc as _make_desc +from cudnn.jax import call, gemm_operand_spec, row_major_desc as _make_desc, sf_atom_spec, zeros_init from .api import GemmAmaxSm100 -# cache_key -> (registered XLA target name, GemmAmaxSm100, compiled tvm-ffi callable). -# The object references keep the compiled kernel alive alongside the global registration. -_registered_targets = {} +# config_key -> (kernel instance, max_active_clusters). The instance does not vary +# with problem shapes, so the key holds only kernel-construction config — a shared +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). Shape/dtype validation is cached separately per full signature. +_kernel_cache: dict = {} +_validated_configs: set = set() + + +@cute.jit +def _amax_adapter(stream, a, b, sfa, sfb, c, amax, *, kernel, mac): + # Destination-passing kernel signature; amax arrives pre-initialized (donated input). + kernel(a, b, sfa, sfb, c, amax, mac, stream) def gemm_amax_jax_sm100( @@ -45,11 +51,7 @@ def gemm_amax_jax_sm100( Arguments are JAX arrays (or tracers): A (M, K, 1) and B (N, K, 1) k-major C-contiguous, SFA/SFB in the physical atom shape (1, MN', K', 32, 4, 4). - Returns a plain ``(c_tensor, amax_tensor)`` tuple of fresh JAX arrays (a plain - tuple, not a TupleDict, so the result is a valid JAX pytree under jit); C is n-major. - - Note: calling this eagerly re-traces the ffi_call each time -- prefer calling it - from inside a jitted function in hot loops. + Returns a plain ``(c_tensor, amax_tensor)`` tuple of fresh JAX arrays; C is n-major. """ c_dtype = _convert_to_cutlass_data_type(c_dtype) acc_dtype = _convert_to_cutlass_data_type(acc_dtype) @@ -59,7 +61,14 @@ def gemm_amax_jax_sm100( if l != 1: raise ValueError("JAX inputs must have batch dim L == 1; batch-outermost (L-major) layouts are not expressible as JAX arrays") - cache_key = ( + config_key = ( + c_dtype, + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, + ) + validation_key = ( tuple(a_tensor.shape), tuple(b_tensor.shape), tuple(sfa_tensor.shape), @@ -68,15 +77,13 @@ def gemm_amax_jax_sm100( _convert_to_cutlass_data_type(b_tensor.dtype), _convert_to_cutlass_data_type(sfa_tensor.dtype), _convert_to_cutlass_data_type(sfb_tensor.dtype), - c_dtype, - acc_dtype, - mma_tiler_mn, - cluster_shape_mn, - sf_vec_size, + config_key, ) - def make_gemm(): - return GemmAmaxSm100( + if validation_key not in _validated_configs: + # Validation reuses the class API's check_support on metadata-only descriptors + # (works for jax.jit tracers, which expose only .shape/.dtype). + gemm = GemmAmaxSm100( sample_a=_make_desc(tuple(a_tensor.shape), a_tensor.dtype, "sample_a"), sample_b=_make_desc(tuple(b_tensor.shape), b_tensor.dtype, "sample_b"), sample_sfa=_make_desc(tuple(sfa_tensor.shape), sfa_tensor.dtype, "sample_sfa"), @@ -88,25 +95,35 @@ def make_gemm(): cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, ) - - # arg_spec=["args"]: only the operands are passed to the kernel; the result - # buffers arrive as the two donated trailing operands (input_output_aliases below), - # matching the kernel's destination-passing signature (a, b, sfa, sfb, c, amax). - target = get_or_register_env_stream_target(_registered_targets, cache_key, make_gemm, "cudnn.gemm_amax_sm100") - - c_jax_dtype = framework_dtype(c_dtype, "jax") - c_buf = jnp.zeros((m, n, l), dtype=c_jax_dtype) - # Zero-init is a valid amax identity: the kernel accumulates max(|c|) >= 0 via a - # signed-integer atomic max of non-negative float bit patterns. - amax_buf = jnp.zeros((1, 1, 1), dtype=jnp.float32) - - c_tensor, amax_tensor = jax.ffi.ffi_call( - target, - ( - jax.ShapeDtypeStruct((m, n, l), c_jax_dtype), + assert gemm.check_support() + if config_key not in _kernel_cache: + kernel = gemm._kernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin + _kernel_cache[config_key] = (kernel, mac) + _validated_configs.add(validation_key) + kernel, mac = _kernel_cache[config_key] + + operand = gemm_operand_spec() + sf = sf_atom_spec() + c_tensor, amax_tensor = call( + _amax_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n, l), framework_dtype(c_dtype, "jax")), jax.ShapeDtypeStruct((1, 1, 1), jnp.float32), ), - input_output_aliases={4: 0, 5: 1}, - )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_buf, amax_buf) + input_spec=(operand, operand, sf, sf), + output_spec=(operand, None), + # Both outputs are donated pre-initialized inputs: amax needs the zero identity + # (signed-int atomic max of non-negative values), and routing c through the + # donated-input path gives it the explicit (1, 0, 2) layout spec -- the bridge's + # leading-dim inference rejects trailing-unit-dim buffers on pure results. + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor) return c_tensor, amax_tensor diff --git a/python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py b/python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py index dab1e63a6..648ed62f3 100644 --- a/python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py +++ b/python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py @@ -10,3 +10,16 @@ "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100", ] + + +__all__.append("gemm_dsrelu_jax_sm100") + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "gemm_dsrelu_jax_sm100": + from ..srelu.jax_api import gemm_dsrelu_jax_sm100 + + return gemm_dsrelu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__init__.py b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__init__.py index fd2c8ff70..2e0fb0c94 100644 --- a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__init__.py +++ b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__init__.py @@ -13,4 +13,15 @@ "GemmProjRopeMxfp8Mxfp8InSm100", "gemm_proj_rope_mxfp8_wrapper_sm100", "gemm_proj_rope_mxfp8_reference", + "gemm_proj_rope_mxfp8_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "gemm_proj_rope_mxfp8_jax_sm100": + from .jax_api import gemm_proj_rope_mxfp8_jax_sm100 + + return gemm_proj_rope_mxfp8_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py index dd8eec740..b210477e0 100644 --- a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py +++ b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py @@ -31,11 +31,28 @@ import logging from typing import Optional +import cutlass import cutlass.utils import cutlass.cute as cute -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import from_dlpack, make_fake_stream from cudnn.api_base import APIBase, TensorDesc, TupleDict +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_device, + get_shape, + is_torch_tensor, +) + + +def _maybe_detach(tensor): + """Detach torch autograd-tracked tensors before DLPack export; no-op for other frameworks.""" + return tensor.detach() if is_torch_tensor(tensor) else tensor # ====================================================================================== @@ -56,24 +73,27 @@ def __init__( sample_out_scales_col: torch.Tensor, w_out_in: bool = False, ): - from cudnn.tensor_adapter import is_torch_tensor - - if sample_x is not None and not is_torch_tensor(sample_x): - raise ValueError("GemmProjRopeMxfp8Bf16InSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") super().__init__() self._warn_experimental_api() self._logger.debug("Entering __init__") - self.x_desc = self._make_tensor_desc(sample_x, name="sample_x") - self.w_desc = self._make_tensor_desc(sample_w, name="sample_w") - self.cos_desc = self._make_tensor_desc(sample_cos, name="sample_cos") - self.sin_desc = self._make_tensor_desc(sample_sin, name="sample_sin") - self.out_fp8_row_desc = self._make_tensor_desc(sample_out_fp8_row, name="sample_out_fp8_row") - self.out_scales_row_desc = self._make_tensor_desc(sample_out_scales_row, name="sample_out_scales_row") - self.out_fp8_col_desc = self._make_tensor_desc(sample_out_fp8_col, name="sample_out_fp8_col") - self.out_scales_col_desc = self._make_tensor_desc(sample_out_scales_col, name="sample_out_scales_col") + self._framework = detect_framework(sample_x) + self.x_desc = self._make_tensor_desc(sample_x, name="sample_x", canonical=True) + self.w_desc = self._make_tensor_desc(sample_w, name="sample_w", canonical=True) + self.cos_desc = self._make_tensor_desc(sample_cos, name="sample_cos", canonical=True) + self.sin_desc = self._make_tensor_desc(sample_sin, name="sample_sin", canonical=True) + self.out_fp8_row_desc = self._make_tensor_desc(sample_out_fp8_row, name="sample_out_fp8_row", canonical=True) + self.out_scales_row_desc = self._make_tensor_desc(sample_out_scales_row, name="sample_out_scales_row", canonical=True) + self.out_fp8_col_desc = self._make_tensor_desc(sample_out_fp8_col, name="sample_out_fp8_col", canonical=True) + self.out_scales_col_desc = self._make_tensor_desc(sample_out_scales_col, name="sample_out_scales_col", canonical=True) self.w_out_in = bool(w_out_in) + if self._framework == "jax" and not self.w_out_in: + raise ValueError( + "w_out_in=False is not expressible as JAX arrays for GemmProjRopeMxfp8Bf16InSm100 " + "(the [in, out] weight is presented to the kernel through a transposed strided view); " + "pass the weight as [out, in] with w_out_in=True" + ) self.tokens = int(sample_x.shape[0]) proj_dim = int(sample_w.shape[0] if self.w_out_in else sample_w.shape[1]) self.num_heads = proj_dim // HEAD_DIM @@ -91,18 +111,16 @@ def __init__( self._logger.debug(f"__init__ completed: x {self.x_desc.shape}, w {self.w_desc.shape}, w_out_in {self.w_out_in}") def check_support(self) -> bool: - import torch - self._logger.debug("Entering check_support") - 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.cos_desc, dtype=torch.bfloat16, name="cos") - self._check_dtype(self.sin_desc, dtype=torch.bfloat16, name="sin") - self._check_dtype(self.out_fp8_row_desc, dtype=torch.float8_e4m3fn, name="out_fp8_row") - self._check_dtype(self.out_fp8_col_desc, dtype=torch.float8_e4m3fn, name="out_fp8_col") - self._check_dtype(self.out_scales_row_desc, dtype=torch.uint8, name="out_scales_row") - self._check_dtype(self.out_scales_col_desc, dtype=torch.uint8, name="out_scales_col") + self._check_dtype(self.x_desc, dtype=cutlass.BFloat16, name="x") + self._check_dtype(self.w_desc, dtype=cutlass.BFloat16, name="w") + self._check_dtype(self.cos_desc, dtype=cutlass.BFloat16, name="cos") + self._check_dtype(self.sin_desc, dtype=cutlass.BFloat16, name="sin") + self._check_dtype(self.out_fp8_row_desc, dtype=cutlass.Float8E4M3FN, name="out_fp8_row") + self._check_dtype(self.out_fp8_col_desc, dtype=cutlass.Float8E4M3FN, name="out_fp8_col") + self._check_dtype(self.out_scales_row_desc, dtype=cutlass.Uint8, name="out_scales_row") + self._check_dtype(self.out_scales_col_desc, dtype=cutlass.Uint8, name="out_scales_col") self._value_error_if( self.tokens % TILE_M != 0, @@ -171,14 +189,15 @@ def check_support(self) -> bool: return True def _to_cute_tensors(self, x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col): - """``w`` may be [in, out] (default) or TE-native [out, in] (``w_out_in``); both present as [out, in].""" - mA = from_dlpack(x.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + """Compile-time sample conversion. ``w`` may be [in, out] (default) or TE-native + [out, in] (``w_out_in``); both present as [out, in].""" + mA = from_dlpack(_maybe_detach(x), assumed_align=16).mark_layout_dynamic(leading_dim=1) if self.w_out_in: - mB = from_dlpack(w.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mB = from_dlpack(_maybe_detach(w), assumed_align=16).mark_layout_dynamic(leading_dim=1) else: - mB = from_dlpack(w.detach().transpose(0, 1), assumed_align=16).mark_layout_dynamic(leading_dim=0) - mCos = from_dlpack(cos.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) - mSin = from_dlpack(sin.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mB = from_dlpack(_maybe_detach(w).transpose(0, 1), assumed_align=16).mark_layout_dynamic(leading_dim=0) + mCos = from_dlpack(_maybe_detach(cos), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mSin = from_dlpack(_maybe_detach(sin), assumed_align=16).mark_layout_dynamic(leading_dim=1) mQrow = from_dlpack(out_fp8_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) mSrow = from_dlpack(out_scales_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) mQcol = from_dlpack(out_fp8_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) @@ -186,8 +205,6 @@ def _to_cute_tensors(self, x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_ return mA, mB, mCos, mSin, mQrow, mSrow, mQcol, mScol def compile(self) -> None: - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -196,7 +213,7 @@ def compile(self) -> None: grid_m = self.tokens // TILE_M max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters(1) swizzle_size = 8 - compile_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) self._compiled_kernel = cute.compile( _bf16in_host, *cute_tensors, @@ -204,7 +221,8 @@ def compile(self) -> None: self.num_heads, max_active_clusters, swizzle_size, - compile_stream, + fake_stream, + options="--enable-tvm-ffi", ) self._samples = None self._logger.debug("Kernel compiled successfully") @@ -221,13 +239,31 @@ def execute( out_scales_col, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(x)) self._runtime_error_if( self._compiled_kernel is None, "GemmProjRopeMxfp8Bf16InSm100 kernel not compiled; call compile() first", ) - cute_tensors = self._to_cute_tensors(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col) - self._compiled_kernel(*cute_tensors, current_stream) + # TVM-FFI entry point: raw DLPack-capable tensors go straight to the compiled + # kernel (no per-call from_dlpack). torch inputs are detached views so autograd- + # tracked tensors stay exportable; w in [in, out] layout is a transposed view. + w_arg = _maybe_detach(w) + if not self.w_out_in: + w_arg = w_arg.transpose(0, 1) + self._compiled_kernel( + _maybe_detach(x), + w_arg, + _maybe_detach(cos), + _maybe_detach(sin), + out_fp8_row, + out_scales_row, + out_fp8_col, + out_scales_col, + current_stream, + ) # ====================================================================================== @@ -254,24 +290,21 @@ def __init__( sample_out_fp8_col: torch.Tensor, sample_out_scales_col: torch.Tensor, ): - from cudnn.tensor_adapter import is_torch_tensor - - if sample_x_code is not None and not is_torch_tensor(sample_x_code): - raise ValueError("GemmProjRopeMxfp8Mxfp8InSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") super().__init__() self._warn_experimental_api() self._logger.debug("Entering __init__") - self.x_code_desc = self._make_tensor_desc(sample_x_code, name="sample_x_code") - self.x_scale_desc = self._make_tensor_desc(sample_x_scale, name="sample_x_scale") - self.w_code_desc = self._make_tensor_desc(sample_w_code, name="sample_w_code") - self.w_scale_desc = self._make_tensor_desc(sample_w_scale, name="sample_w_scale") - self.cos_desc = self._make_tensor_desc(sample_cos, name="sample_cos") - self.sin_desc = self._make_tensor_desc(sample_sin, name="sample_sin") - self.out_fp8_row_desc = self._make_tensor_desc(sample_out_fp8_row, name="sample_out_fp8_row") - self.out_scales_row_desc = self._make_tensor_desc(sample_out_scales_row, name="sample_out_scales_row") - self.out_fp8_col_desc = self._make_tensor_desc(sample_out_fp8_col, name="sample_out_fp8_col") - self.out_scales_col_desc = self._make_tensor_desc(sample_out_scales_col, name="sample_out_scales_col") + self._framework = detect_framework(sample_x_code) + self.x_code_desc = self._make_tensor_desc(sample_x_code, name="sample_x_code", canonical=True) + self.x_scale_desc = self._make_tensor_desc(sample_x_scale, name="sample_x_scale", canonical=True) + self.w_code_desc = self._make_tensor_desc(sample_w_code, name="sample_w_code", canonical=True) + self.w_scale_desc = self._make_tensor_desc(sample_w_scale, name="sample_w_scale", canonical=True) + self.cos_desc = self._make_tensor_desc(sample_cos, name="sample_cos", canonical=True) + self.sin_desc = self._make_tensor_desc(sample_sin, name="sample_sin", canonical=True) + self.out_fp8_row_desc = self._make_tensor_desc(sample_out_fp8_row, name="sample_out_fp8_row", canonical=True) + self.out_scales_row_desc = self._make_tensor_desc(sample_out_scales_row, name="sample_out_scales_row", canonical=True) + self.out_fp8_col_desc = self._make_tensor_desc(sample_out_fp8_col, name="sample_out_fp8_col", canonical=True) + self.out_scales_col_desc = self._make_tensor_desc(sample_out_scales_col, name="sample_out_scales_col", canonical=True) self.tokens = int(sample_x_code.shape[0]) proj_dim = int(sample_w_code.shape[0]) # weight is [N, K] @@ -293,20 +326,18 @@ def __init__( self._logger.debug(f"__init__ completed: x_code {self.x_code_desc.shape}, w_code {self.w_code_desc.shape}") def check_support(self) -> bool: - import torch - self._logger.debug("Entering check_support") - self._check_dtype(self.x_code_desc, dtype=torch.float8_e4m3fn, name="x_code") - self._check_dtype(self.w_code_desc, dtype=torch.float8_e4m3fn, name="w_code") - self._check_dtype(self.x_scale_desc, dtype=torch.uint8, name="x_scale") - self._check_dtype(self.w_scale_desc, dtype=torch.uint8, name="w_scale") - self._check_dtype(self.cos_desc, dtype=torch.bfloat16, name="cos") - self._check_dtype(self.sin_desc, dtype=torch.bfloat16, name="sin") - self._check_dtype(self.out_fp8_row_desc, dtype=torch.float8_e4m3fn, name="out_fp8_row") - self._check_dtype(self.out_fp8_col_desc, dtype=torch.float8_e4m3fn, name="out_fp8_col") - self._check_dtype(self.out_scales_row_desc, dtype=torch.uint8, name="out_scales_row") - self._check_dtype(self.out_scales_col_desc, dtype=torch.uint8, name="out_scales_col") + self._check_dtype(self.x_code_desc, dtype=cutlass.Float8E4M3FN, name="x_code") + self._check_dtype(self.w_code_desc, dtype=cutlass.Float8E4M3FN, name="w_code") + self._check_dtype(self.x_scale_desc, dtype=cutlass.Uint8, name="x_scale") + self._check_dtype(self.w_scale_desc, dtype=cutlass.Uint8, name="w_scale") + self._check_dtype(self.cos_desc, dtype=cutlass.BFloat16, name="cos") + self._check_dtype(self.sin_desc, dtype=cutlass.BFloat16, name="sin") + self._check_dtype(self.out_fp8_row_desc, dtype=cutlass.Float8E4M3FN, name="out_fp8_row") + self._check_dtype(self.out_fp8_col_desc, dtype=cutlass.Float8E4M3FN, name="out_fp8_col") + self._check_dtype(self.out_scales_row_desc, dtype=cutlass.Uint8, name="out_scales_row") + self._check_dtype(self.out_scales_col_desc, dtype=cutlass.Uint8, name="out_scales_col") self._value_error_if( self.tokens % TILE_M != 0, @@ -401,12 +432,13 @@ def _grid_params(self): return grid_m, t2r_x8, swizzle_size def _to_cute_tensors(self, x_code, x_scale, w_code, w_scale, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col): - mA = from_dlpack(x_code.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + """Compile-time sample conversion.""" + mA = from_dlpack(_maybe_detach(x_code), assumed_align=16).mark_layout_dynamic(leading_dim=1) mSFA = _mxfp8_as_e8m0(x_scale) - mB = from_dlpack(w_code.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mB = from_dlpack(_maybe_detach(w_code), assumed_align=16).mark_layout_dynamic(leading_dim=1) mSFB = _mxfp8_as_e8m0(w_scale) - mCos = from_dlpack(cos.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) - mSin = from_dlpack(sin.detach(), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mCos = from_dlpack(_maybe_detach(cos), assumed_align=16).mark_layout_dynamic(leading_dim=1) + mSin = from_dlpack(_maybe_detach(sin), assumed_align=16).mark_layout_dynamic(leading_dim=1) mQrow = from_dlpack(out_fp8_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) mSrow = from_dlpack(out_scales_row, assumed_align=16).mark_layout_dynamic(leading_dim=2) mQcol = from_dlpack(out_fp8_col, assumed_align=16).mark_layout_dynamic(leading_dim=2) @@ -414,8 +446,6 @@ def _to_cute_tensors(self, x_code, x_scale, w_code, w_scale, cos, sin, out_fp8_r return mA, mSFA, mB, mSFB, mCos, mSin, mQrow, mSrow, mQcol, mScol def compile(self) -> None: - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -424,7 +454,7 @@ def compile(self) -> None: grid_m, t2r_x8, swizzle_size = self._grid_params() max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters(1) k_scale_words = self.k_dim // 128 # compact-scale uint32 words per row = K // 128 (deduced, not hardcoded) - compile_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) self._compiled_kernel = cute.compile( _mxfp8in_host, *cute_tensors, @@ -434,7 +464,8 @@ def compile(self) -> None: swizzle_size, t2r_x8, k_scale_words, - compile_stream, + fake_stream, + options="--enable-tvm-ffi", ) self._samples = None self._logger.debug("Kernel compiled successfully") @@ -453,24 +484,30 @@ def execute( out_scales_col, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(x_code)) self._runtime_error_if( self._compiled_kernel is None, "GemmProjRopeMxfp8Mxfp8InSm100 kernel not compiled; call compile() first", ) - cute_tensors = self._to_cute_tensors( - x_code, - x_scale, - w_code, - w_scale, - cos, - sin, + # TVM-FFI entry point: raw DLPack-capable tensors go straight to the compiled + # kernel. The uint8 scale tensors keep the per-call e8m0 element-type reinterpret + # (the ABI validates the compiled e8m0 dtype). + self._compiled_kernel( + _maybe_detach(x_code), + _mxfp8_as_e8m0(x_scale), + _maybe_detach(w_code), + _mxfp8_as_e8m0(w_scale), + _maybe_detach(cos), + _maybe_detach(sin), out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, + current_stream, ) - self._compiled_kernel(*cute_tensors, current_stream) # ====================================================================================== @@ -496,15 +533,12 @@ def _check_contiguous(api, **named_descs): def _check_sm100(api): - import torch - - api._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) + api._runtime_error_if(not cuda_is_available(), "CUDA is not available") + major, minor = get_compute_capability() compute_capability = major * 10 + minor api._runtime_error_if( compute_capability < 100, - f"GemmProjRopeMxfp8 requires SM100+ compute capability, but found SM{compute_capability} on device {device}", + f"GemmProjRopeMxfp8 requires SM100+ compute capability, but found SM{compute_capability} on the current device", ) @@ -541,27 +575,50 @@ def gemm_proj_rope_mxfp8_wrapper_sm100( Returns: ``TupleDict(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)``. """ - from cudnn.tensor_adapter import is_torch_tensor - - if x is not None and not is_torch_tensor(x): - raise ValueError("gemm_proj_rope_mxfp8_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch + framework = detect_framework(x) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for gemm_proj_rope_mxfp8_wrapper_sm100; pass torch tensors or JAX arrays") - assert x.dtype == w.dtype, f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}" + x_cutlass_dtype = _convert_to_cutlass_data_type(x.dtype) + assert x_cutlass_dtype == _convert_to_cutlass_data_type( + w.dtype + ), f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}" - tokens = x.shape[0] + tokens = get_shape(x)[0] device = x.device - proj_dim = w.shape[0] if w_out_in else w.shape[1] + proj_dim = get_shape(w)[0] if w_out_in else get_shape(w)[1] num_heads = proj_dim // HEAD_DIM - out_fp8_row = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) - out_scales_row = torch.empty(tokens, num_heads, HEAD_DIM // BLOCK, dtype=torch.uint8, device=device) - out_fp8_col = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) - out_scales_col = torch.empty(tokens // BLOCK, num_heads, HEAD_DIM, dtype=torch.uint8, device=device) + if framework == "torch": + import torch - if x.dtype == torch.bfloat16: + out_fp8_row = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + out_scales_row = torch.empty(tokens, num_heads, HEAD_DIM // BLOCK, dtype=torch.uint8, device=device) + out_fp8_col = torch.empty(tokens, num_heads, HEAD_DIM, dtype=torch.float8_e4m3fn, device=device) + out_scales_col = torch.empty(tokens // BLOCK, num_heads, HEAD_DIM, dtype=torch.uint8, device=device) + else: + import jax + import jax.numpy as jnp + + if not w_out_in: + raise ValueError( + "w_out_in=False is not expressible as JAX arrays for gemm_proj_rope_mxfp8_wrapper_sm100 " + "(the [in, out] weight is presented to the kernel through a transposed strided view); " + "pass the weight as [out, in] with w_out_in=True" + ) + fp8 = framework_dtype(cutlass.Float8E4M3FN, "jax") + u8 = framework_dtype(cutlass.Uint8, "jax") + out_fp8_row = jnp.empty((tokens, num_heads, HEAD_DIM), dtype=fp8, device=device) + out_scales_row = jnp.empty((tokens, num_heads, HEAD_DIM // BLOCK), dtype=u8, device=device) + out_fp8_col = jnp.empty((tokens, num_heads, HEAD_DIM), dtype=fp8, device=device) + out_scales_col = jnp.empty((tokens // BLOCK, num_heads, HEAD_DIM), dtype=u8, device=device) + # The kernel writes into these buffers on the launch stream; make sure XLA has + # finished materializing them before the kernel runs. + jax.block_until_ready((out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)) + + if x_cutlass_dtype is cutlass.BFloat16: assert x_scale is None and w_scale is None, "bf16 inputs must not be given MXFP8 scales (x_scale/w_scale); those are for the float8_e4m3fn path" - key = (tuple(x.shape), tuple(w.shape), bool(w_out_in), device) + key = (get_shape(x), get_shape(w), bool(w_out_in), get_device(x)) obj = _bf16in_obj_cache.get(key) if obj is None: obj = GemmProjRopeMxfp8Bf16InSm100( @@ -580,14 +637,14 @@ def gemm_proj_rope_mxfp8_wrapper_sm100( _bf16in_obj_cache[key] = obj obj.execute(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col, current_stream=stream) - elif x.dtype == torch.float8_e4m3fn: + elif x_cutlass_dtype is cutlass.Float8E4M3FN: assert x_scale is not None and w_scale is not None, "MXFP8 (float8_e4m3fn) inputs require x_scale and w_scale (E8M0 rowwise block scales)" # the mxfp8in kernel expects the weight as [out, in]; transpose code + scale for [in, out]. if w_out_in: wc, ws = w, w_scale else: wc, ws = w.T.contiguous(), w_scale.T.contiguous() - key = (tuple(x.shape), tuple(wc.shape), device) + key = (get_shape(x), get_shape(wc), get_device(x)) obj = _mxfp8in_obj_cache.get(key) if obj is None: obj = GemmProjRopeMxfp8Mxfp8InSm100( diff --git a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py index 00f1cdd60..cb5d2c1d9 100644 --- a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py +++ b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py @@ -650,6 +650,8 @@ def gemm_proj_rope_mxfp8_host( def _as_e8m0(t): - ct = from_dlpack(t.detach(), assumed_align=16) + from cudnn.tensor_adapter import is_torch_tensor + + ct = from_dlpack(t.detach() if is_torch_tensor(t) else t, assumed_align=16, enable_tvm_ffi=True) ct.element_type = cutlass.Float8E8M0FNU return ct.mark_layout_dynamic(leading_dim=1) diff --git a/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py new file mode 100644 index 000000000..153fed686 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the fused projection GEMM + RoPE + +dual-direction MXFP8 quantize, built on :func:`cudnn.jax.call`.""" + +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np + +import cutlass +import cutlass.cute as cute +import cutlass.utils + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.jax import call, row_major_desc as _make_desc +from .api import GemmProjRopeMxfp8Bf16InSm100, GemmProjRopeMxfp8Mxfp8InSm100 +from .gemm_proj_rope_mxfp8_bf16in import HEAD_DIM, BLOCK, TILE_M, gemm_proj_rope_mxfp8_host as _bf16in_host +from .gemm_proj_rope_mxfp8_mxfp8in import gemm_proj_rope_mxfp8_host as _mxfp8in_host + +_bf16in_grid_cache: dict = {} +_mxfp8in_grid_cache: dict = {} + + +@cute.jit +def _proj_rope_bf16in_adapter(stream, x, w, cos, sin, qrow, srow, qcol, scol, *, grid_m, num_heads, mac, swizzle_size): + _bf16in_host(x, w, cos, sin, qrow, srow, qcol, scol, grid_m, num_heads, mac, swizzle_size, stream) + + +@cute.jit +def _proj_rope_mxfp8in_adapter( + stream, x, x_scale, w, w_scale, cos, sin, qrow, srow, qcol, scol, *, grid_m, num_heads, mac, swizzle_size, t2r_x8, k_scale_words +): + _mxfp8in_host(x, x_scale, w, w_scale, cos, sin, qrow, srow, qcol, scol, grid_m, num_heads, mac, swizzle_size, t2r_x8, k_scale_words, stream) + + +def _as_e8m0_array(scale: Any) -> Any: + """Present a uint8 E8M0-bit-pattern array as float8_e8m0fnu (free bitcast, jit-safe).""" + if _convert_to_cutlass_data_type(scale.dtype) is cutlass.Uint8: + import ml_dtypes + + return scale.view(ml_dtypes.float8_e8m0fnu) + return scale + + +def gemm_proj_rope_mxfp8_jax_sm100( + x: Any, + w: Any, + cos: Any, + sin: Any, + x_scale: Optional[Any] = None, + w_scale: Optional[Any] = None, +) -> Tuple[Any, Any, Any, Any]: + """Projection GEMM + RoPE + dual-direction MXFP8 quantize as an XLA custom call. + + Same contract as the eager wrapper with ``w_out_in=True`` (the only JAX-expressible + weight layout): ``x [tokens, K]`` / ``w [out, in]`` bfloat16 for the BF16 GEMM, or + ``float8_e4m3fn`` codes plus E8M0 block scales (``uint8`` bit patterns or + ``float8_e8m0fnu``) for the MXFP8 GEMM. Returns + ``(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)``. + """ + x_cutlass_dtype = _convert_to_cutlass_data_type(x.dtype) + tokens = x.shape[0] + proj_dim = w.shape[0] + num_heads = proj_dim // HEAD_DIM + + out_types = ( + jax.ShapeDtypeStruct((tokens, num_heads, HEAD_DIM), np.dtype("float8_e4m3fn")), + jax.ShapeDtypeStruct((tokens, num_heads, HEAD_DIM // BLOCK), np.uint8), + jax.ShapeDtypeStruct((tokens, num_heads, HEAD_DIM), np.dtype("float8_e4m3fn")), + jax.ShapeDtypeStruct((tokens // BLOCK, num_heads, HEAD_DIM), np.uint8), + ) + + if x_cutlass_dtype is cutlass.BFloat16: + if x_scale is not None or w_scale is not None: + raise ValueError("bf16 inputs must not be given MXFP8 scales (x_scale/w_scale); those are for the float8_e4m3fn path") + cache_key = (tuple(x.shape), tuple(w.shape)) + entry = _bf16in_grid_cache.get(cache_key) + if entry is None: + obj = GemmProjRopeMxfp8Bf16InSm100( + sample_x=_make_desc(tuple(x.shape), x.dtype, "sample_x"), + sample_w=_make_desc(tuple(w.shape), w.dtype, "sample_w"), + sample_cos=_make_desc(tuple(cos.shape), cos.dtype, "sample_cos"), + sample_sin=_make_desc(tuple(sin.shape), sin.dtype, "sample_sin"), + sample_out_fp8_row=_make_desc(out_types[0].shape, cutlass.Float8E4M3FN, "sample_out_fp8_row"), + sample_out_scales_row=_make_desc(out_types[1].shape, cutlass.Uint8, "sample_out_scales_row"), + sample_out_fp8_col=_make_desc(out_types[2].shape, cutlass.Float8E4M3FN, "sample_out_fp8_col"), + sample_out_scales_col=_make_desc(out_types[3].shape, cutlass.Uint8, "sample_out_scales_col"), + w_out_in=True, + ) + assert obj.check_support() + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(1) + entry = (tokens // TILE_M, num_heads, mac, 8) + _bf16in_grid_cache[cache_key] = entry + grid_m, heads, mac, swizzle = entry + + return call( + _proj_rope_bf16in_adapter, + output_shape_dtype=out_types, + grid_m=grid_m, + num_heads=heads, + mac=mac, + swizzle_size=swizzle, + )(x, w, cos, sin) + + if x_cutlass_dtype is cutlass.Float8E4M3FN: + if x_scale is None or w_scale is None: + raise ValueError("MXFP8 (float8_e4m3fn) inputs require x_scale and w_scale (E8M0 rowwise block scales)") + x_scale = _as_e8m0_array(x_scale) + w_scale = _as_e8m0_array(w_scale) + cache_key = (tuple(x.shape), tuple(w.shape)) + entry = _mxfp8in_grid_cache.get(cache_key) + if entry is None: + obj = GemmProjRopeMxfp8Mxfp8InSm100( + sample_x_code=_make_desc(tuple(x.shape), x.dtype, "sample_x_code"), + sample_x_scale=_make_desc(tuple(x_scale.shape), cutlass.Uint8, "sample_x_scale"), + sample_w_code=_make_desc(tuple(w.shape), w.dtype, "sample_w_code"), + sample_w_scale=_make_desc(tuple(w_scale.shape), cutlass.Uint8, "sample_w_scale"), + sample_cos=_make_desc(tuple(cos.shape), cos.dtype, "sample_cos"), + sample_sin=_make_desc(tuple(sin.shape), sin.dtype, "sample_sin"), + sample_out_fp8_row=_make_desc(out_types[0].shape, cutlass.Float8E4M3FN, "sample_out_fp8_row"), + sample_out_scales_row=_make_desc(out_types[1].shape, cutlass.Uint8, "sample_out_scales_row"), + sample_out_fp8_col=_make_desc(out_types[2].shape, cutlass.Float8E4M3FN, "sample_out_fp8_col"), + sample_out_scales_col=_make_desc(out_types[3].shape, cutlass.Uint8, "sample_out_scales_col"), + ) + assert obj.check_support() + grid_m, t2r_x8, swizzle = obj._grid_params() + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(1) + entry = (grid_m, num_heads, mac, swizzle, t2r_x8, int(x.shape[1]) // 128) + _mxfp8in_grid_cache[cache_key] = entry + grid_m, heads, mac, swizzle, t2r_x8, k_scale_words = entry + + return call( + _proj_rope_mxfp8in_adapter, + output_shape_dtype=out_types, + grid_m=grid_m, + num_heads=heads, + mac=mac, + swizzle_size=swizzle, + t2r_x8=t2r_x8, + k_scale_words=k_scale_words, + )(x, x_scale, w, w_scale, cos, sin) + + raise ValueError(f"unsupported input dtype {x.dtype}; expected bfloat16 (BF16 GEMM) or float8_e4m3fn (MXFP8 GEMM)") diff --git a/python/cudnn/gemm/cutedsl/dense/srelu/__init__.py b/python/cudnn/gemm/cutedsl/dense/srelu/__init__.py index 9ca3e844c..4acc96f4d 100644 --- a/python/cudnn/gemm/cutedsl/dense/srelu/__init__.py +++ b/python/cudnn/gemm/cutedsl/dense/srelu/__init__.py @@ -10,3 +10,16 @@ "GemmSreluSm100", "gemm_srelu_wrapper_sm100", ] + + +__all__.append("gemm_srelu_jax_sm100") + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "gemm_srelu_jax_sm100": + from .jax_api import gemm_srelu_jax_sm100 + + return gemm_srelu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py b/python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py new file mode 100644 index 000000000..3781067e5 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py @@ -0,0 +1,291 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry points for the GEMM + sReLU forward and +backward (dsReLU) kernels, built on :func:`cudnn.jax.call`. + +The kernels' optional sfd/amax/norm_const parameters are compile-time ``None`` +constants inside the ``@cute.jit`` adapters for the JAX-reachable (non-fp8-D) +configurations. +""" + +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import TensorSpec, call, gemm_operand_spec, row_major_desc as _make_desc, sf_atom_spec, zeros_init +from .api import GemmSreluSm100 +from ..dsrelu.api import GemmDsreluSm100 + +# config_key -> (kernel instance, max_active_clusters). The instances do not vary +# with problem shapes, so the keys hold only kernel-construction config — a shared +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). Shape/dtype validation is cached separately per full signature. +_srelu_kernel_cache: dict = {} +_srelu_validated_configs: set = set() +_dsrelu_kernel_cache: dict = {} +_dsrelu_validated_configs: set = set() + +# Constexpr epilogues, identical to the class APIs' compile() +_SRELU_EPILOGUE = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 # noqa: E731 +_DSRELU_EPILOGUE = lambda x, y: cute.where(x > 0, x, cute.full_like(x, 0)) * 2 * y # noqa: E731 + + +def _prob_spec() -> TensorSpec: + # (m, 1, 1) with m innermost: explicit ranks because trailing unit dims make + # leading-dim inference ambiguous + return TensorSpec(layout=(0, 1, 2)) + + +@cute.jit +def _srelu_adapter(stream, a, b, sfa, sfb, prob, c, d, *, kernel, mac, alpha): + # sfd/amax/norm_const are compile-time Nones (fp8-D configs are not reachable from JAX) + kernel( + a_tensor=a, + b_tensor=b, + sfa_tensor=sfa, + sfb_tensor=sfb, + c_tensor=c, + d_tensor=d, + prob_tensor=prob, + amax_tensor=None, + sfd_tensor=None, + norm_const_tensor=None, + alpha=alpha, + max_active_clusters=mac, + stream=stream, + epilogue_op=_SRELU_EPILOGUE, + ) + + +@cute.jit +def _dsrelu_adapter(stream, a, b, sfa, sfb, c, prob, d, dprob, *, kernel, mac, alpha): + kernel( + a_tensor=a, + b_tensor=b, + sfa_tensor=sfa, + sfb_tensor=sfb, + c_tensor=c, + d_tensor=d, + prob_tensor=prob, + dprob_tensor=dprob, + amax_tensor=None, + sfd_tensor=None, + norm_const_tensor=None, + alpha=alpha, + max_active_clusters=mac, + stream=stream, + epilogue_op=_DSRELU_EPILOGUE, + ) + + +def gemm_srelu_jax_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + c_dtype: Any = cutlass.BFloat16, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + 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, +) -> Tuple[Any, Any]: + """Blockscaled GEMM + sReLU as an XLA custom call; usable eagerly or under jax.jit. + + A (M, K, 1) / B (N, K, 1) k-major C-contiguous fp8 arrays, SFA/SFB in the physical + atom shape (1, MN', K', 32, 4, 4), prob (M, 1, 1) float32. Returns a plain + ``(c_tensor, d_tensor)`` tuple of fresh n-major JAX arrays. fp8 ``d_dtype`` + configurations (which produce sfd/norm_const outputs) are not reachable from JAX. + """ + c_dtype = _convert_to_cutlass_data_type(c_dtype) + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + m, _, l = a_tensor.shape + n, _, _ = b_tensor.shape + if l != 1: + raise ValueError("JAX inputs must have batch dim L == 1; batch-outermost (L-major) layouts are not expressible as JAX arrays") + if d_dtype in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): + raise ValueError("fp8 d_dtype requires sfd/norm_const outputs, which are not reachable from JAX; use the eager wrapper with torch tensors") + + config_key = ( + c_dtype, + d_dtype, + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, + vector_f32, + ) + validation_key = ( + tuple(a_tensor.shape), + tuple(b_tensor.shape), + tuple(sfa_tensor.shape), + tuple(sfb_tensor.shape), + _convert_to_cutlass_data_type(a_tensor.dtype), + _convert_to_cutlass_data_type(b_tensor.dtype), + _convert_to_cutlass_data_type(sfa_tensor.dtype), + _convert_to_cutlass_data_type(sfb_tensor.dtype), + config_key, + ) + if validation_key not in _srelu_validated_configs: + gemm = GemmSreluSm100( + sample_a=_make_desc(tuple(a_tensor.shape), a_tensor.dtype, "sample_a"), + sample_b=_make_desc(tuple(b_tensor.shape), b_tensor.dtype, "sample_b"), + sample_c=_make_desc((m, n, l), c_dtype, "sample_c"), + sample_d=_make_desc((m, n, l), d_dtype, "sample_d"), + sample_sfa=_make_desc(tuple(sfa_tensor.shape), sfa_tensor.dtype, "sample_sfa"), + sample_sfb=_make_desc(tuple(sfb_tensor.shape), sfb_tensor.dtype, "sample_sfb"), + sample_prob=_make_desc((m, 1, 1), cutlass.Float32, "sample_prob"), + alpha=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, + ) + assert gemm.check_support() + if config_key not in _srelu_kernel_cache: + kernel = gemm._kernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=gemm.cluster_shape_mn, + vector_f32=vector_f32, + ) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(gemm.cluster_shape_mn[0] * gemm.cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin + _srelu_kernel_cache[config_key] = (kernel, mac) + _srelu_validated_configs.add(validation_key) + kernel, mac = _srelu_kernel_cache[config_key] + + operand = gemm_operand_spec() + sf = sf_atom_spec() + c_tensor, d_tensor = call( + _srelu_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n, l), framework_dtype(c_dtype, "jax")), + jax.ShapeDtypeStruct((m, n, l), framework_dtype(d_dtype, "jax")), + ), + input_spec=(operand, operand, sf, sf, _prob_spec()), + output_spec=(operand, operand), + # Donated pre-initialized outputs: the bridge's leading-dim inference rejects + # trailing-unit-dim buffers on pure results. + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + alpha=float(alpha), + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor) + + return c_tensor, d_tensor + + +def gemm_dsrelu_jax_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + 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, +) -> Tuple[Any, Any]: + """Blockscaled GEMM + dsReLU (backward) as an XLA custom call. + + Returns a plain ``(d_tensor, dprob_tensor)`` tuple; dprob is a zero-initialized + atomic-add accumulator managed by the bridge. + """ + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + m, _, l = a_tensor.shape + n, _, _ = b_tensor.shape + if l != 1: + raise ValueError("JAX inputs must have batch dim L == 1; batch-outermost (L-major) layouts are not expressible as JAX arrays") + if d_dtype in (cutlass.Float8E4M3FN, cutlass.Float8E5M2): + raise ValueError("fp8 d_dtype requires sfd/norm_const outputs, which are not reachable from JAX; use the eager wrapper with torch tensors") + + config_key = ( + d_dtype, + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, + vector_f32, + ) + validation_key = ( + tuple(a_tensor.shape), + tuple(b_tensor.shape), + tuple(c_tensor.shape), + tuple(sfa_tensor.shape), + tuple(sfb_tensor.shape), + _convert_to_cutlass_data_type(a_tensor.dtype), + _convert_to_cutlass_data_type(b_tensor.dtype), + _convert_to_cutlass_data_type(c_tensor.dtype), + _convert_to_cutlass_data_type(sfa_tensor.dtype), + _convert_to_cutlass_data_type(sfb_tensor.dtype), + config_key, + ) + if validation_key not in _dsrelu_validated_configs: + gemm = GemmDsreluSm100( + sample_a=_make_desc(tuple(a_tensor.shape), a_tensor.dtype, "sample_a"), + sample_b=_make_desc(tuple(b_tensor.shape), b_tensor.dtype, "sample_b"), + sample_c=_make_desc(tuple(c_tensor.shape), c_tensor.dtype, "sample_c"), + sample_d=_make_desc((m, n, l), d_dtype, "sample_d"), + sample_dprob=_make_desc((m, 1, l), cutlass.Float32, "sample_dprob"), + sample_sfa=_make_desc(tuple(sfa_tensor.shape), sfa_tensor.dtype, "sample_sfa"), + sample_sfb=_make_desc(tuple(sfb_tensor.shape), sfb_tensor.dtype, "sample_sfb"), + sample_prob=_make_desc((m, 1, 1), cutlass.Float32, "sample_prob"), + alpha=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, + ) + assert gemm.check_support() + if config_key not in _dsrelu_kernel_cache: + kernel = gemm._kernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=gemm.cluster_shape_mn, + vector_f32=vector_f32, + ) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(gemm.cluster_shape_mn[0] * gemm.cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin + _dsrelu_kernel_cache[config_key] = (kernel, mac) + _dsrelu_validated_configs.add(validation_key) + kernel, mac = _dsrelu_kernel_cache[config_key] + + operand = gemm_operand_spec() + sf = sf_atom_spec() + d_tensor, dprob_tensor = call( + _dsrelu_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n, l), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, 1, l), jnp.float32), + ), + input_spec=(operand, operand, sf, sf, operand, _prob_spec()), + output_spec=(operand, _prob_spec()), + # d is donated for the trailing-unit-dim layout spec; dprob is a genuine + # zero-initialized atomic-add accumulator. + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + alpha=float(alpha), + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, prob_tensor) + + return d_tensor, dprob_tensor diff --git a/python/cudnn/gemm/cutedsl/dense/swiglu/__init__.py b/python/cudnn/gemm/cutedsl/dense/swiglu/__init__.py index 61847af94..1e8168b18 100644 --- a/python/cudnn/gemm/cutedsl/dense/swiglu/__init__.py +++ b/python/cudnn/gemm/cutedsl/dense/swiglu/__init__.py @@ -14,7 +14,7 @@ def __getattr__(name): - # Lazy: the jax entry point imports jax/jax_tvm_ffi, which must not be pulled in + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in # for torch-only users. if name == "gemm_swiglu_jax_sm100": from .jax_api import gemm_swiglu_jax_sm100 diff --git a/python/cudnn/gemm/cutedsl/dense/swiglu/api.py b/python/cudnn/gemm/cutedsl/dense/swiglu/api.py index c87c805eb..5ae3924cc 100644 --- a/python/cudnn/gemm/cutedsl/dense/swiglu/api.py +++ b/python/cudnn/gemm/cutedsl/dense/swiglu/api.py @@ -381,13 +381,8 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): self._logger.debug("check_support completed successfully") return True - def _compile_kernel(self, use_tvm_ffi_env_stream: bool = False): - """Compile the kernel and return the raw TVM-FFI callable. - - With ``use_tvm_ffi_env_stream=True`` the stream argument is bound to the - TVM-FFI environment stream and dropped from the callable's signature -- - the variant used for XLA custom-call (jax.ffi) integration. - """ + def _compile_kernel(self): + """Compile the kernel and return the raw TVM-FFI callable.""" self._ensure_support_checked() if self._kernel is PersistentDenseGemmKernel: @@ -416,7 +411,7 @@ def _compile_kernel(self, use_tvm_ffi_env_stream: bool = False): "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", ) - fake_stream = make_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) if self._kernel is PersistentDenseGemmKernel: self._logger.debug("Compiling gemm_swiglu") @@ -456,7 +451,7 @@ def compile(self) -> None: self._logger.debug("Kernel already compiled; skipping recompilation") return - _compiled_kernel = self._compile_kernel(use_tvm_ffi_env_stream=False) + _compiled_kernel = self._compile_kernel() if self._kernel is PersistentDenseGemmKernel: diff --git a/python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py b/python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py index 5fcd2b409..0bf253c59 100644 --- a/python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py +++ b/python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py @@ -3,9 +3,10 @@ """JAX-native (XLA custom call) entry point for the GEMM + SwiGLU kernels. -See dense/amax/jax_api.py for the integration pattern: the kernel variant is compiled -with the TVM-FFI environment stream (runs on XLA's compute stream), all outputs are -donated pre-initialized operands, and the call composes with jax.jit. +Built on :func:`cudnn.jax.call` (CuTeDSL's native JAX bridge). Both the standard and +the blockscaled (MXFP8) quantized kernels are supported: the quantized kernel's +optional amax/sfc/norm_const parameters are compile-time ``None`` constants inside the +``@cute.jit`` adapter for the JAX-reachable configurations. """ from typing import Any, Optional, Tuple @@ -14,13 +15,32 @@ import jax.numpy as jnp import cutlass +import cutlass.cute as cute +import cutlass.utils from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.tensor_adapter import framework_dtype -from cudnn.gemm.cutedsl._jax_ffi import get_or_register_env_stream_target, make_row_major_desc as _make_desc +from cudnn.jax import call, gemm_operand_spec, row_major_desc as _make_desc, sf_atom_spec, zeros_init from .api import GemmSwigluSm100 -_registered_targets = {} +# config_key -> (kernel instance, max_active_clusters). The instance does not vary +# with problem shapes, so the key holds only kernel-construction config — a shared +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). Shape/dtype validation is cached separately per full signature. +_kernel_cache: dict = {} +_validated_configs: set = set() + + +@cute.jit +def _swiglu_adapter(stream, a, b, ab12, c, *, kernel, mac, alpha): + kernel(a, b, ab12, c, alpha, mac, stream) + + +@cute.jit +def _swiglu_quant_adapter(stream, a, b, sfa, sfb, c, ab12, *, kernel, mac, alpha): + # amax/sfc/norm_const are compile-time Nones: fp8-C and fp4-A/B configurations + # (which would produce those outputs) are not reachable from JAX. + kernel(a, b, sfa, sfb, c, ab12, None, None, None, alpha, mac, stream) def gemm_swiglu_jax_sm100( @@ -41,13 +61,10 @@ def gemm_swiglu_jax_sm100( ) -> Tuple[Any, Any]: """GEMM + SwiGLU as an XLA custom call; usable eagerly or under jax.jit. - A (M, K, 1) and B (N, K, 1) are k-major C-contiguous JAX arrays (or tracers). - Returns a plain ``(ab12_tensor, c_tensor)`` tuple of fresh n-major JAX arrays. - ``alpha`` is a static (trace-time) parameter. - - Supports the non-quantized kernel only: the quantized kernel's compiled signature - carries None-typed parameters the XLA FFI bridge cannot supply -- use - ``gemm_swiglu_wrapper_sm100`` (eager) for blockscaled MXFP8 inputs from JAX. + A (M, K, 1) and B (N, K, 1) are k-major C-contiguous JAX arrays (or tracers); + optional SFA/SFB (blockscaled MXFP8 path) in the physical atom shape + (1, MN', K', 32, 4, 4). Returns a plain ``(ab12_tensor, c_tensor)`` tuple of + fresh n-major JAX arrays. ``alpha`` is a static (trace-time) parameter. """ ab12_dtype = _convert_to_cutlass_data_type(ab12_dtype) c_dtype = _convert_to_cutlass_data_type(c_dtype) @@ -58,29 +75,35 @@ def gemm_swiglu_jax_sm100( if l != 1: raise ValueError("JAX inputs must have batch dim L == 1; batch-outermost (L-major) layouts are not expressible as JAX arrays") - if sfa_tensor is not None or sfb_tensor is not None: - # The quantized kernel's compiled signature carries explicit None-typed - # parameters (amax/sfc/norm_const) that the XLA FFI bridge cannot supply. - raise NotImplementedError( - "gemm_swiglu_jax_sm100 currently supports the non-quantized kernel only; " - "use gemm_swiglu_wrapper_sm100 (eager) for blockscaled MXFP8 inputs from JAX" - ) + is_quantized = sfa_tensor is not None and sfb_tensor is not None + if (sfa_tensor is None) != (sfb_tensor is None): + raise ValueError("Provide both sfa_tensor and sfb_tensor for the quantized kernel, or neither") - cache_key = ( - tuple(a_tensor.shape), - tuple(b_tensor.shape), - _convert_to_cutlass_data_type(a_tensor.dtype), - _convert_to_cutlass_data_type(b_tensor.dtype), - alpha, + config_key = ( + is_quantized, ab12_dtype, c_dtype, acc_dtype, mma_tiler_mn, cluster_shape_mn, + sf_vec_size, + vector_f32, + ab12_stages, + ) + validation_key = ( + tuple(a_tensor.shape), + tuple(b_tensor.shape), + _convert_to_cutlass_data_type(a_tensor.dtype), + _convert_to_cutlass_data_type(b_tensor.dtype), + tuple(sfa_tensor.shape) if is_quantized else None, + tuple(sfb_tensor.shape) if is_quantized else None, + _convert_to_cutlass_data_type(sfa_tensor.dtype) if is_quantized else None, + _convert_to_cutlass_data_type(sfb_tensor.dtype) if is_quantized else None, + config_key, ) - def make_gemm(): - return GemmSwigluSm100( + if validation_key not in _validated_configs: + gemm = GemmSwigluSm100( sample_a=_make_desc(tuple(a_tensor.shape), a_tensor.dtype, "sample_a"), sample_b=_make_desc(tuple(b_tensor.shape), b_tensor.dtype, "sample_b"), sample_ab12=_make_desc((m, n, l), ab12_dtype, "sample_ab12"), @@ -89,32 +112,64 @@ def make_gemm(): acc_dtype=acc_dtype, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + sample_sfa=_make_desc(tuple(sfa_tensor.shape), sfa_tensor.dtype, "sample_sfa") if is_quantized else None, + sample_sfb=_make_desc(tuple(sfb_tensor.shape), sfb_tensor.dtype, "sample_sfb") if is_quantized else None, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + ab12_stages=ab12_stages, ) - - # arg_spec: operands first (the donated output buffers land in the kernel's - # destination-passing slots), then alpha as an XLA call attribute. - target = get_or_register_env_stream_target( - _registered_targets, - cache_key, - make_gemm, - "cudnn.gemm_swiglu_sm100", - arg_spec=("args", "attrs.alpha"), - ) - - ab12_jax_dtype = framework_dtype(ab12_dtype, "jax") - c_jax_dtype = framework_dtype(c_dtype, "jax") - ab12_buf = jnp.zeros((m, n, l), dtype=ab12_jax_dtype) - c_buf = jnp.zeros((m, n // 2, l), dtype=c_jax_dtype) + assert gemm.check_support() + if config_key not in _kernel_cache: + if is_quantized: + kernel = gemm._kernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=gemm.cluster_shape_mn, + vector_f32=vector_f32, + ab12_stages=ab12_stages, + ) + else: + kernel = gemm._kernel( + acc_dtype=acc_dtype, + use_2cta_instrs=(mma_tiler_mn[0] == 256), + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=gemm.cluster_shape_mn, + ) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(gemm.cluster_shape_mn[0] * gemm.cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin + _kernel_cache[config_key] = (kernel, mac) + _validated_configs.add(validation_key) + kernel, mac = _kernel_cache[config_key] + + operand = gemm_operand_spec() out_types = ( - jax.ShapeDtypeStruct((m, n, l), ab12_jax_dtype), - jax.ShapeDtypeStruct((m, n // 2, l), c_jax_dtype), + jax.ShapeDtypeStruct((m, n, l), framework_dtype(ab12_dtype, "jax")), # ab12 + jax.ShapeDtypeStruct((m, n // 2, l), framework_dtype(c_dtype, "jax")), # c ) - - # Kernel signature: (a, b, ab12, c, alpha[, stream from the TVM-FFI environment]). - ab12_tensor, c_tensor = jax.ffi.ffi_call( - target, - out_types, - input_output_aliases={2: 0, 3: 1}, - )(a_tensor, b_tensor, ab12_buf, c_buf, alpha=float(alpha)) + # Outputs are donated pre-initialized inputs: the bridge's leading-dim inference + # rejects trailing-unit-dim buffers on pure results, and the donated-input path + # carries the explicit (1, 0, 2) layout spec. + if not is_quantized: + ab12_tensor, c_tensor = call( + _swiglu_adapter, + output_shape_dtype=out_types, + input_spec=(operand, operand), + output_spec=(operand, operand), + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + alpha=float(alpha), + )(a_tensor, b_tensor) + else: + sf = sf_atom_spec() + c_tensor, ab12_tensor = call( + _swiglu_quant_adapter, + output_shape_dtype=(out_types[1], out_types[0]), + input_spec=(operand, operand, sf, sf), + output_spec=(operand, operand), + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + alpha=float(alpha), + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor) return ab12_tensor, c_tensor diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/__init__.py b/python/cudnn/gemm/cutedsl/discrete_grouped/__init__.py index 1706b27c4..94003f24d 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/__init__.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/__init__.py @@ -23,4 +23,22 @@ "discrete_grouped_gemm_swiglu_wrapper_sm100", "DiscreteGroupedGemmDswigluSm100", "discrete_grouped_gemm_dswiglu_wrapper_sm100", + "discrete_grouped_gemm_swiglu_jax_sm100", + "discrete_grouped_gemm_dswiglu_jax_sm100", ] + +# Lazy: the jax entry points import jax/cutlass.jax, which must not be pulled in +# for torch-only users. +_JAX_LAZY_EXPORTS = { + "discrete_grouped_gemm_swiglu_jax_sm100": ".swiglu", + "discrete_grouped_gemm_dswiglu_jax_sm100": ".dswiglu", +} + + +def __getattr__(name): + module_name = _JAX_LAZY_EXPORTS.get(name) + if module_name is not None: + import importlib + + return getattr(importlib.import_module(module_name, __name__), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py b/python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py index 4c5fc9507..766152325 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.py @@ -44,6 +44,12 @@ def _require_pointer_tensor(ptrs: torch.Tensor, name: str, expected_len: int | None = None) -> None: + """Validate a torch int64 device pointer-array tensor (torch-only contract). + + Type-erased APIs that also accept JAX arrays should use the framework-neutral + ``_validate_pointer_tensor`` / ``_pointer_values`` helpers from + ``cudnn.gemm.cutedsl.grouped.unfused._bf16_api`` instead. + """ import torch if ptrs.dtype != torch.int64: diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.py b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.py index 3b52720f3..50ef81d15 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "DiscreteGroupedGemmDswigluSm100", "discrete_grouped_gemm_dswiglu_wrapper_sm100", + "discrete_grouped_gemm_dswiglu_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "discrete_grouped_gemm_dswiglu_jax_sm100": + from .jax_api import discrete_grouped_gemm_dswiglu_jax_sm100 + + return discrete_grouped_gemm_dswiglu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py index c9b7433f6..bcd8e439d 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py @@ -24,7 +24,20 @@ from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, + get_shape, + get_strides, + is_torch_tensor, +) class DiscreteGroupedGemmDswigluSm100(APIBase): @@ -106,40 +119,38 @@ def __init__( :param epilogue_op: Optional epilogue operation ("relu", "srelu", or None) :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("DiscreteGroupedGemmDswigluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() self._logger.debug("Entering __init__") + self._framework = detect_framework(sample_a) self._value_error_if(num_experts == 0, "num_experts must be > 0") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") - self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") - - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + # Set before building descriptors so uint8 sample tensors (the packed-fp4 + # container dtype, e.g. from JAX which has no fp4 dtype) get fp4 logical shapes. + self._interpret_uint8_as_fp4x2 = True + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row", canonical=True) + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob", canonical=True) + self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias", canonical=True) + + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) @@ -150,10 +161,10 @@ def __init__( f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) must equal expert_cnt ({self.expert_cnt})", ) - self.b_dtype = b_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) self.b_shape = b_shape - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -183,13 +194,42 @@ def __init__( self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") self._workspace = None + self._compile_b_ptrs = None + self._compile_sfb_ptrs = None + self._live_ptrs = None self._logger.debug("__init__ completed") + def _check_sf_shape(self, desc, mn_div_128: int, rest: int, name: str) -> bool: + """Validate a 6-D scale-factor descriptor; returns True for the physical form. + + Accepts the torch-style permuted atom view (32, 4, MN', 4, K', 1) or the + physical C-contiguous allocation (1, MN', K', 32, 4, 4) -- byte-identical + memory, and the kernel consumes only the SF base pointer. + """ + if desc is None: + return False + is_physical = desc.ndim == 6 and desc.shape[0] == 1 and desc.shape[3] == 32 + self._check_tensor_shape( + desc, + (1, mn_div_128, rest, 32, 4, 4) if is_physical else (32, 4, mn_div_128, 4, rest, 1), + name, + ) + # The kernel consumes only the SF base pointer and rebuilds the layout from the + # GEMM shapes, so both forms must be exactly the C-contiguous physical allocation + # in memory: validate strides too (a shape-matching but differently-strided tensor + # would silently produce wrong results). + if is_physical: + expected = canonicalize_unit_dim_strides((1, mn_div_128, rest, 32, 4, 4), (mn_div_128 * rest * 512, rest * 512, 512, 16, 4, 1)) + extra = f"{name} in the physical (1, MN', K', 32, 4, 4) form must be C-contiguous" + else: + expected = canonicalize_unit_dim_strides((32, 4, mn_div_128, 4, rest, 1), (16, 4, rest * 512, 1, 512, mn_div_128 * rest * 512)) + extra = f"{name} atom view must be the (3, 4, 1, 5, 2, 0) permutation of a C-contiguous (1, MN', K', 32, 4, 4) allocation" + _ = self._check_tensor_stride(desc, stride=[expected], name=name, extra_error_msg=extra) + return is_physical + def check_support(self) -> bool: """Check if the kernel configuration is supported.""" - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -215,11 +255,17 @@ def check_support(self) -> bool: self._check_tensor_shape(self.d_col_desc, (tensor_m, n_out, 1), "D_col") rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") + # SF tensors are accepted in either of two byte-identical forms: the torch-style + # permuted atom view (32, 4, MN', 4, K', 1), or the physical C-contiguous + # allocation (1, MN', K', 32, 4, 4) for frameworks such as JAX that cannot + # express the permuted (strided) view. This is safe because the kernel rebuilds + # every SF layout from the A/D shapes (tile_atom_to_shape_SF) and consumes only + # the SF base pointers. + self._sfa_is_physical = self._check_sf_shape(self.sfa_desc, ceil_div(tensor_m, 128), rest_k, "SFA") rest_n_out = ceil_div(ceil_div(n_out, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_row_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_n_out, 1), "SFD_row") + self._sfd_row_is_physical = self._check_sf_shape(self.sfd_row_desc, ceil_div(tensor_m, 128), rest_n_out, "SFD_row") rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n_out, 128), 4, rest_m, 1), "SFD_col") + self._sfd_col_is_physical = self._check_sf_shape(self.sfd_col_desc, ceil_div(n_out, 128), rest_m, "SFD_col") self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") self._check_tensor_shape(self.beta_desc, (self.expert_cnt,), "beta") @@ -255,10 +301,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -270,7 +316,7 @@ def check_support(self) -> bool: self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD", ) self._check_dtype( @@ -291,7 +337,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype in [cutlass.Float8E4M3FN] and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} not supported", ) self._value_error_if( @@ -299,42 +345,42 @@ def check_support(self) -> bool: f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} not supported", ) - self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator") + self._check_dtype(self.acc_dtype, dtype=cutlass.Float32, name="Accumulator") self._check_dtype( self.prob_desc, - dtype=torch.float32, + dtype=cutlass.Float32, name="prob", extra_error_msg="prob must be float32", ) self._check_dtype( self.dprob_desc, - dtype=torch.float32, + dtype=cutlass.Float32, name="dprob", extra_error_msg="dprob must be float32", ) self._check_dtype( self.dbias_desc, - dtype=torch.bfloat16, + dtype=cutlass.BFloat16, name="Dbias", extra_error_msg="dbias must be bfloat16", ) - self.c_dtype = self._check_dtype(self.c_desc, dtype=[torch.float32, torch.float16, torch.bfloat16], name="C") + self.c_dtype = self._check_dtype(self.c_desc, dtype=[cutlass.Float32, cutlass.Float16, cutlass.BFloat16], name="C") if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_row_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D", ) else: self.d_dtype = self._check_dtype( self.d_row_desc, dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="D", ) @@ -345,7 +391,7 @@ def check_support(self) -> bool: extra_error_msg="D_col must have the same dtype as D", ) - kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype == torch.float8_e8m0fnu and self._is_fp8(self.d_dtype) + kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype is cutlass.Float8E8M0FNU and self._is_fp8(self.d_dtype) self._value_error_if( kernel_generate_sfd and not self._user_requested_sfd, "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=torch.float8_e8m0fnu", @@ -437,10 +483,9 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): f"expert_cnt must be <= 1024, got {self.expert_cnt}", ) - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() if major * 10 + minor < 100: raise RuntimeError(f"DiscreteGroupedGemmDswiglu requires SM100+, found SM{major * 10 + minor}") @@ -450,8 +495,6 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): def compile(self) -> None: """Compile the backward kernel from tensor descriptors captured in __init__.""" - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -494,7 +537,9 @@ def compile(self) -> None: fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) workspace_bytes = gemm_dglu.get_workspace_bytes() - self._workspace = torch.empty(workspace_bytes, dtype=torch.uint8, device="cuda") + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) align = 32 if ab_cutlass_dtype.width == 4 else 16 @@ -523,37 +568,72 @@ def compile(self) -> None: ) tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[5] = stride_tensor_m_128 - sfa_tensor = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - assumed_align=16, - ) - sfd_row_tensor = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + if self._sfa_is_physical: + # SFA in the physical C-contiguous atom shape (1, M', K', 32, 4, 4); + # the kernel rebuilds the SF layout from A's shape and consumes only + # the SFA base pointer, so only the calling convention differs. + sfa_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfa_desc.dtype, + shape=self.sfa_desc.shape, + stride_order=self.sfa_desc.stride_order, assumed_align=16, + dynamic_mode=1, + divisibility=1, ) - sfd_col_tensor = None - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + else: + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[5] = stride_tensor_m_128 + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), assumed_align=16, ) + sfd_row_tensor = None + if self.sfd_row_desc is not None: + if self._sfd_row_is_physical: + # Physical C-contiguous atom shape (1, M', N2', 32, 4, 4); pointer-only. + sfd_row_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfd_row_desc.dtype, + shape=self.sfd_row_desc.shape, + stride_order=self.sfd_row_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=1, + ) + else: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + assumed_align=16, + ) + sfd_col_tensor = None + if self.sfd_col_desc is not None: + if self._sfd_col_is_physical: + # Physical C-contiguous atom shape (1, N2', M'', 32, 4, 4); pointer-only. + sfd_col_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfd_col_desc.dtype, + shape=self.sfd_col_desc.shape, + stride_order=self.sfd_col_desc.stride_order, + assumed_align=16, + dynamic_mode=2, + divisibility=1, + ) + else: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + assumed_align=16, + ) amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) @@ -573,12 +653,18 @@ def compile(self) -> None: ) dbias_tensor = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) - # Use internal device-resident int64 arrays to provide valid pointer-like - # compile-time placeholders for b_ptrs/sfb_ptrs (required by kernel __call__). - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + # Use internal device-resident buffers to provide valid pointer-like compile-time + # placeholders for b_ptrs/sfb_ptrs (required by kernel __call__): real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, retyped + # to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + self._compile_sfb_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + b_ptrs_placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + b_ptrs_placeholder.element_type = cutlass.Int64 + b_ptrs_cute = b_ptrs_placeholder.iterator + sfb_ptrs_placeholder = from_dlpack(self._compile_sfb_ptrs, assumed_align=8) + sfb_ptrs_placeholder.element_type = cutlass.Int64 + sfb_ptrs_cute = sfb_ptrs_placeholder.iterator workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator # linear_offset, geglu_alpha, glu_clamp_max, and glu_clamp_min are runtime @@ -657,8 +743,8 @@ def tensor_api( glu_clamp_min: float = -7.0, ): norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + b_ptrs_addr = int(get_data_ptr(b_ptrs_device)) + sfb_ptrs_addr = int(get_data_ptr(sfb_ptrs_device)) _compiled_kernel( a_tensor, @@ -744,12 +830,19 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) - if a_tensor.shape[0] == 0: + if get_shape(a_tensor)[0] == 0: self._logger.debug("execute: valid_m is zero, skipping") return self._runtime_error_if(self._compiled_kernel is None, "Kernel not compiled; call compile() first") + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # arrays referenced until the next execute so their buffers outlive the launch. + self._live_ptrs = (b_ptrs, sfb_ptrs) # Resolve linear_offset default: None -> activation-derived legacy value # (1.0 for dgeglu, 0.0 for dswiglu) for backwards compatibility with @@ -862,107 +955,155 @@ def discrete_grouped_gemm_dswiglu_wrapper_sm100( TupleDict with keys: d_row_tensor, d_col_tensor, dprob_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor """ - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("discrete_grouped_gemm_dswiglu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if d_dtype is None: - d_dtype = torch.bfloat16 + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for discrete_grouped_gemm_dswiglu_wrapper_sm100; pass torch tensors or JAX arrays") + + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) + ab_dtype = _convert_to_cutlass_data_type(a_tensor.dtype) + + if framework == "jax" and ab_dtype in (cutlass.Uint8, cutlass.Float4E2M1FN): + raise ValueError( + "packed-fp4 inputs are not expressible as JAX arrays for this API " + "(JAX has no packed fp4 dtype and the compiled kernel entry point requires float4_e2m1fn_x2 tensors); " + "use torch tensors for FP4, or FP8 inputs for JAX" + ) # Resolve linear_offset default: None means "use the activation-derived legacy # default" (1.0 for dgeglu, 0.0 for dswiglu) for backwards compatibility. if linear_offset is None: linear_offset = 1.0 if act_func == "dgeglu" else 0.0 - valid_m, k_physical, _ = a_tensor.shape - _require_pointer_tensor(b_ptrs, "b_ptrs") - num_experts = b_ptrs.shape[0] - _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) - k_logical = k_physical * 2 if b_dtype in (torch.float4_e2m1fn_x2, torch.uint8) else k_physical + valid_m, k_physical, _ = get_shape(a_tensor) + num_experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + _validate_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) + k_logical = k_physical * 2 if b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) else k_physical b_shape = (n, k_logical) if cd_major != "n": raise ValueError(f"cd_major must be 'n', got {cd_major}") n_out = 2 * n - d_row_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) - d_col_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) + if framework == "torch": + import torch + + d_row_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) + d_col_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) + else: + import jax + import jax.numpy as jnp + + # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. + # The kernel writes into these buffers on the launch stream; materialize them first. + d_row_tensor = jnp.empty((valid_m, n_out, 1), dtype=framework_dtype(d_dtype, "jax"), device=a_tensor.device) + d_col_tensor = jnp.empty((valid_m, n_out, 1), dtype=framework_dtype(d_dtype, "jax"), device=a_tensor.device) + jax.block_until_ready((d_row_tensor, d_col_tensor)) sfd_row_tensor = None sfd_col_tensor = None amax_tensor = None dbias_tensor = None - if a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [torch.float8_e8m0fnu, torch.float8_e4m3fn]: + if ab_dtype in [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] and _convert_to_cutlass_data_type( + sfa_tensor.dtype + ) in [cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN]: sf_dtype = sfa_tensor.dtype mma_permute_order = (3, 4, 1, 5, 2, 0) sf_k_row = ceil_div(n_out, sf_vec_size) mma_shape_row = (1, ceil_div(valid_m, 128), ceil_div(sf_k_row, 4), 32, 4, 4) - sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - sf_k_col = ceil_div(valid_m, sf_vec_size) mma_shape_col = (1, ceil_div(n_out, 128), ceil_div(sf_k_col, 4), 32, 4, 4) - sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - - if d_dtype in [torch.bfloat16, torch.float16]: - amax_tensor = torch.full( - (num_experts, 2, 1), - float("-inf"), - dtype=torch.float32, - device=a_tensor.device, - ) + + if framework == "torch": + import torch + + sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + else: + import jax + import jax.numpy as jnp + + # Physical C-contiguous atom-shape allocations (the kernel rebuilds the SF + # layouts from the output shape and consumes only the base pointers). + sfd_row_tensor = jnp.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device) + sfd_col_tensor = jnp.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device) + jax.block_until_ready((sfd_row_tensor, sfd_col_tensor)) + + if d_dtype in [cutlass.BFloat16, cutlass.Float16]: + if framework == "torch": + import torch + + amax_tensor = torch.full( + (num_experts, 2, 1), + float("-inf"), + dtype=torch.float32, + device=a_tensor.device, + ) + else: + import jax + import jax.numpy as jnp + + amax_tensor = jax.block_until_ready(jnp.full((num_experts, 2, 1), float("-inf"), dtype=jnp.float32, device=a_tensor.device)) if generate_dbias: - dbias_tensor = torch.zeros((num_experts, n_out, 1), dtype=torch.bfloat16, device=a_tensor.device) + if framework == "torch": + import torch + + dbias_tensor = torch.zeros((num_experts, n_out, 1), dtype=torch.bfloat16, device=a_tensor.device) + else: + import jax + import jax.numpy as jnp + + dbias_tensor = jax.block_until_ready(jnp.zeros((num_experts, n_out, 1), dtype=framework_dtype(cutlass.BFloat16, "jax"), device=a_tensor.device)) def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(i for i, s in sorted(enumerate(tensor.stride()), key=lambda x: x[1])) + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + return tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: (x[1], tensor_shape[x[0]]))) def tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype + tensor_shape = get_shape(tensor) + return tensor_shape, canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)), _convert_to_cutlass_data_type(tensor.dtype) def dynamic_m_tensor_signature( tensor: Optional[torch.Tensor], static_shape_suffix: Optional[Tuple[int, ...]], dynamic_stride_dims: Tuple[int, ...] = () ) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(tensor.stride())) - return static_shape_suffix, stride_signature, tensor.dtype + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(tensor_stride)) + return static_shape_suffix, stride_signature, _convert_to_cutlass_data_type(tensor.dtype) cache_key = ( - a_tensor.shape[1:], + get_shape(a_tensor)[1:], stride_order(a_tensor), - a_tensor.dtype, + ab_dtype, b_shape, b_dtype, - c_tensor.shape[1:], + get_shape(c_tensor)[1:], stride_order(c_tensor), - c_tensor.dtype, - *dynamic_m_tensor_signature(sfa_tensor, (sfa_tensor.shape[4], 1) if sfa_tensor is not None else None, dynamic_stride_dims=(5,)), + _convert_to_cutlass_data_type(c_tensor.dtype), + *dynamic_m_tensor_signature( + sfa_tensor, + (get_shape(sfa_tensor)[4], 1) if sfa_tensor is not None else None, + dynamic_stride_dims=(0, 1, 5), + ), *tensor_signature(alpha_tensor), *tensor_signature(beta_tensor), - *dynamic_m_tensor_signature(prob_tensor, (1, 1)), - *dynamic_m_tensor_signature(dprob_tensor, (1, 1)), + *dynamic_m_tensor_signature(prob_tensor, (1, 1), dynamic_stride_dims=(1, 2)), + *dynamic_m_tensor_signature(dprob_tensor, (1, 1), dynamic_stride_dims=(1, 2)), *tensor_signature(dbias_tensor), *tensor_signature(norm_const_tensor), - tuple(b_ptrs.shape), - tuple(b_ptrs.stride()), - b_ptrs.dtype, - tuple(sfb_ptrs.shape), - tuple(sfb_ptrs.stride()), - sfb_ptrs.dtype, - tuple(padded_offsets.shape), - tuple(padded_offsets.stride()), - padded_offsets.dtype, + *tensor_signature(b_ptrs), + *tensor_signature(sfb_ptrs), + *tensor_signature(padded_offsets), acc_dtype, d_dtype, cd_major, diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py new file mode 100644 index 000000000..78564668b --- /dev/null +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py @@ -0,0 +1,399 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the discrete-weight block-scaled +grouped GEMM dGLU backward (dSwiGLU/dGeGLU), built on :func:`cudnn.jax.call`. + +FP8 inputs only (the packed-fp4 uint8 container dtype is not expressible as JAX +arrays). This discrete kernel is JAX-expressible where the contiguous block-scaled +dGLU kernel is not: SFB arrives as per-expert base pointers and SFA/SFD travel in +the physical C-contiguous atom shape ``(1, MN', K', 32, 4, 4)`` — the kernel +rebuilds every scale-factor layout from the A/D shapes and consumes only the SF +base pointers. The per-expert B/SFB pointers travel as regular device arrays whose +*values* are raw addresses — the referenced weight/scale buffers are not visible +to XLA, so the caller must keep them alive (and unmoved) across every execution of +the traced computation. ``dprob`` and ``dbias`` are kernel-accumulated outputs, so +unlike the eager wrapper they are not caller-provided buffers here: both are +donated zero-initialized outputs of the custom call. ``padded_offsets`` values +cannot be host-validated under tracing; malformed offsets are the caller's +responsibility here (the eager wrapper validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.api_base import ceil_div +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import TensorSpec, call, gemm_operand_spec, neg_inf_init, zeros_init +from ...grouped.unfused.jax_api import _pointer_count, _prob_spec +from ..swiglu.jax_api import _as_e8m0_array, _sf_atom_spec, _sf_byte_zeros_init, _JAX_FP4_ERROR +from .discrete_B_blockscaled_grouped_gemm_dglu_dbias import BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_fp8_dtypes = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_c_dtypes = (cutlass.Float32, cutlass.Float16, cutlass.BFloat16) +_d_dtypes = (cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_amax_dtypes = (cutlass.BFloat16, cutlass.Float16) + + +def _epilogue_identity(x): + return x + + +def _epilogue_relu(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + + +def _epilogue_srelu(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + +_EPILOGUE_OPS = { + None: _epilogue_identity, + "none": _epilogue_identity, + "identity": _epilogue_identity, + "relu": _epilogue_relu, + "srelu": _epilogue_srelu, +} + + +@cute.jit +def _discrete_dswiglu_adapter( + stream, + a, + b_ptrs, + sfb_ptrs, + c, + sfa, + padded_offsets, + alpha, + beta, + prob, + norm_const, + d_row, + d_col, + dprob, + sfd_row, + sfd_col, + amax, + dbias, + workspace, + *, + kernel, + n, + k, + mac, + has_amax, + has_dbias, + epilogue, + linear_offset, + geglu_alpha, + glu_clamp_max, + glu_clamp_min, +): + # Discrete-mode b/sfb are raw pointers to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffers recast for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + sfb_arg = cute.recast_ptr(sfb_ptrs.iterator, dtype=cutlass.Int64) + # amax exists only for 16-bit d and dbias only with generate_dbias (eager + # parity: the other configurations compile with compile-time Nones); the + # dummy placeholder buffers are never passed to the kernel. + amax_arg = None + if cutlass.const_expr(has_amax): + amax_arg = amax + dbias_arg = None + if cutlass.const_expr(has_dbias): + dbias_arg = dbias + kernel( + a=a, + b_ptrs=b_arg, + sfb_ptrs=sfb_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d_row, + d_col=d_col, + sfa=sfa, + sfd_row_tensor=sfd_row, + sfd_col_tensor=sfd_col, + amax_tensor=amax_arg, + norm_const_tensor=norm_const, + padded_offsets=padded_offsets, + alpha=alpha, + beta=beta, + prob=prob, + dprob=dprob, + linear_offset=cutlass.Float32(linear_offset), + dbias_tensor=dbias_arg, + max_active_clusters=mac, + stream=stream, + epilogue_op=epilogue, + geglu_alpha=cutlass.Float32(geglu_alpha), + glu_clamp_max=cutlass.Float32(glu_clamp_max), + glu_clamp_min=cutlass.Float32(glu_clamp_min), + ) + + +def discrete_grouped_gemm_dswiglu_jax_sm100( + a_tensor: Any, + b_ptrs: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_ptrs: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any, + n: int, + generate_dbias: bool = False, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + 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, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + epilogue_op: Optional[str] = None, + 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, +) -> dict: + """Discrete-weight block-scaled grouped GEMM dGLU backward as an XLA custom call. + + Same contract as the eager wrapper's FP8 JAX mode: A ``(m, k, 1)`` k-major + C-contiguous fp8 (e4m3/e5m2) gradient input, C ``(m, 2n, 1)`` n-major forward + activations (fp32/fp16/bf16), SFA in the physical C-contiguous E8M0 atom shape + ``(1, m/128, ceil(ceil(k/sf_vec_size)/4), 32, 4, 4)`` (``float8_e8m0fnu``, or + ``uint8`` bit patterns), ``padded_offsets (experts,)`` int32 cumulative + 256-aligned row offsets, ``alpha``/``beta (experts,)`` float32, ``prob + (m, 1, 1)`` float32, ``norm_const (1,)`` float32, and ``b_ptrs``/``sfb_ptrs`` + holding per-expert ``(n, k)`` k-major weight / SFB atom base addresses (packed + little-endian uint8, 8 bytes per pointer — or int64 with x64 mode). ``n`` is + the per-expert weight N (half the activation width); ``d_row``/``d_col`` come + back ``(m, 2n, 1)``. The scalar GLU knobs (``linear_offset`` — defaulting per + ``act_func`` — ``geglu_alpha``, ``glu_clamp_max``, ``glu_clamp_min``) and + ``epilogue_op`` are compile-time constants of the traced call. Rows at/past + ``padded_offsets[-1]`` come back zero-filled (the outputs are donated + zero-initialized buffers, matching the eager contract of a caller-zeroed + ``dprob``); ``dprob`` accumulates through floating-point atomics, so its + values are not bitwise-deterministic across runs. + + Returns a dict with the eager wrapper's keys: ``d_row_tensor``/``d_col_tensor + (m, 2n, 1)``, ``dprob_tensor ((m, 1, 1) float32)``, ``dbias_tensor + ((experts, 2n, 1) bfloat16, None unless generate_dbias)``, ``amax_tensor + ((experts, 2, 1) float32, -inf-initialized, None unless d_dtype is + bf16/fp16)``, and ``sfd_row_tensor``/``sfd_col_tensor`` (physical E8M0 atom + shape; written only for fp8 ``d_dtype``, zero bytes otherwise). + + Rejected for JAX (as in the eager wrapper): packed-fp4 inputs/outputs and the + contiguous kernel's non-k-major weight layouts. + """ + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + ab_dtype = _convert_to_cutlass_data_type(a_tensor.dtype) + if ab_dtype in (cutlass.Uint8, cutlass.Float4E2M1FN): + raise ValueError(_JAX_FP4_ERROR) + if ab_dtype not in _fp8_dtypes: + raise ValueError(f"a_tensor must be float8_e4m3fn or float8_e5m2, got {a_tensor.dtype}") + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + if n is None or n <= 0: + raise ValueError(f"n must be > 0, got {n}") + n_out = 2 * n + if tuple(c_tensor.shape) != (m, n_out, 1): + raise ValueError(f"c_tensor must have shape ({m}, {n_out}, 1), got {tuple(c_tensor.shape)}") + c_dtype = _convert_to_cutlass_data_type(c_tensor.dtype) + if c_dtype not in _c_dtypes: + raise ValueError(f"c_tensor must be FP32, FP16, or BF16, got {c_dtype}") + if d_dtype not in _d_dtypes: + raise ValueError(f"d_dtype must be FP16, BF16, or FP8 for JAX (packed fp4 has no JAX dtype), got {d_dtype}") + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + if sf_vec_size != 32: + raise ValueError(f"fp8 inputs require sf_vec_size 32, got {sf_vec_size}") + if act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {act_func}") + if epilogue_op not in _EPILOGUE_OPS: + raise ValueError(f"Invalid epilogue operation: {epilogue_op}. Valid: None, 'relu', 'srelu'") + if linear_offset is None: + linear_offset = 1.0 if act_func == "dgeglu" else 0.0 + + sfa_tensor = _as_e8m0_array(sfa_tensor) + if _convert_to_cutlass_data_type(sfa_tensor.dtype) is not cutlass.Float8E8M0FNU: + raise ValueError(f"sfa_tensor must be float8_e8m0fnu (or uint8 bit patterns) for fp8 inputs, got {sfa_tensor.dtype}") + rest_k = ceil_div(ceil_div(k, sf_vec_size), 4) + expected_sfa = (1, ceil_div(m, 128), rest_k, 32, 4, 4) + if tuple(sfa_tensor.shape) != expected_sfa: + raise ValueError(f"sfa_tensor must use the physical C-contiguous atom shape {expected_sfa}, got {tuple(sfa_tensor.shape)}") + + expert_cnt = _pointer_count(b_ptrs) + sfb_cnt = _pointer_count(sfb_ptrs, "sfb_ptrs") + if sfb_cnt != expert_cnt: + raise ValueError(f"sfb_ptrs length mismatch: expected {expert_cnt} pointers, got {sfb_cnt}") + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(beta_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(beta_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"beta_tensor must be ({expert_cnt},) float32, got {tuple(beta_tensor.shape)} {beta_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + if tuple(norm_const_tensor.shape) != (1,) or _convert_to_cutlass_data_type(norm_const_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"norm_const_tensor must be (1,) float32, got {tuple(norm_const_tensor.shape)} {norm_const_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + if mma_tiler_mn[1] != 256: + raise ValueError(f"MMA tiler N must be 256, got {mma_tiler_mn[1]}") + + if not BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel.can_implement( + ab_dtype, + cutlass.Float8E8M0FNU, + sf_vec_size, + acc_dtype, + d_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + "k", + "n", + BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel.FIX_PAD_SIZE, + act_func, + ): + raise ValueError("Unsupported discrete grouped GEMM dSwiGLU tile, cluster, alignment, or layout configuration") + + # SFD is generated only for the fp8-in/fp8-out E8M0 configuration; 16-bit d + # accumulates per-expert amax instead (both mirror the eager wrapper). + has_amax = d_dtype in _amax_dtypes + if not (d_dtype in _fp8_dtypes) and discrete_col_sfd: + discrete_col_sfd = False # eager parity: ignored when SFD is not generated + + cache_key = ( + expert_cnt, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + sf_vec_size, + vector_f32, + discrete_col_sfd, + act_func, + use_dynamic_sched, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel( + sf_vec_size=sf_vec_size, + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + sf_jax_dtype = framework_dtype(cutlass.Float8E8M0FNU, "jax") + operand = gemm_operand_spec() + prob_spec = _prob_spec() + sf_spec = _sf_atom_spec() + # amax/dbias keep the adapter arity fixed across configurations: when disabled, + # a minimal placeholder buffer is donated (amax keeps its tiny real shape) and + # the kernel receives a compile-time None instead. + dbias_shape = (expert_cnt, n_out, 1) if generate_dbias else (1, 1, 1) + outputs = call( + _discrete_dswiglu_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n_out, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, n_out, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, 1, 1), jnp.float32), # dprob (kernel-accumulated) + jax.ShapeDtypeStruct((1, ceil_div(m, 128), ceil_div(ceil_div(n_out, sf_vec_size), 4), 32, 4, 4), sf_jax_dtype), + jax.ShapeDtypeStruct((1, ceil_div(n_out, 128), ceil_div(ceil_div(m, sf_vec_size), 4), 32, 4, 4), sf_jax_dtype), + jax.ShapeDtypeStruct((expert_cnt, 2, 1), jnp.float32), + jax.ShapeDtypeStruct(dbias_shape, framework_dtype(cutlass.BFloat16, "jax")), + jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8), + ), + input_spec=(operand, None, None, operand, sf_spec, None, None, None, prob_spec, None), + output_spec=(operand, operand, prob_spec, sf_spec, sf_spec, operand, operand, None), + # All outputs donated: d_row/d_col/sfd for the unit-dim layout specs (and + # defined bytes past the last offset); dprob/dbias because the kernel + # accumulates into them (atomic add) and expects zeroed buffers; amax + # because the kernel accumulates into it (atomic max over a + # -inf-initialized buffer, eager parity); the workspace because the helper + # kernel writes the per-expert TMA descriptors into it (XLA inputs are + # immutable). + initialized_outputs={ + 0: zeros_init, + 1: zeros_init, + 2: zeros_init, + 3: _sf_byte_zeros_init, + 4: _sf_byte_zeros_init, + 5: neg_inf_init if has_amax else zeros_init, + 6: zeros_init, + 7: zeros_init, + }, + kernel=kernel, + n=int(n), + k=int(k), + mac=mac, + has_amax=has_amax, + has_dbias=bool(generate_dbias), + epilogue=_EPILOGUE_OPS[epilogue_op], + linear_offset=float(linear_offset), + geglu_alpha=float(geglu_alpha), + glu_clamp_max=float(glu_clamp_max), + glu_clamp_min=float(glu_clamp_min), + )(a_tensor, b_ptrs, sfb_ptrs, c_tensor, sfa_tensor, padded_offsets, alpha_tensor, beta_tensor, prob_tensor, norm_const_tensor) + + d_row_out, d_col_out, dprob_out, sfd_row_out, sfd_col_out, amax_out, dbias_out, _workspace = outputs + return { + "d_row_tensor": d_row_out, + "d_col_tensor": d_col_out, + "dprob_tensor": dprob_out, + "dbias_tensor": dbias_out if generate_dbias else None, + "amax_tensor": amax_out if has_amax else None, + "sfd_row_tensor": sfd_row_out, + "sfd_col_tensor": sfd_col_out, + } diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.py b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.py index b9e1c2d35..bd1f7a4ad 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "DiscreteGroupedGemmSwigluSm100", "discrete_grouped_gemm_swiglu_wrapper_sm100", + "discrete_grouped_gemm_swiglu_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "discrete_grouped_gemm_swiglu_jax_sm100": + from .jax_api import discrete_grouped_gemm_swiglu_jax_sm100 + + return discrete_grouped_gemm_swiglu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py index 5b85f781c..779e947d8 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py @@ -29,7 +29,20 @@ from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, + get_shape, + get_strides, + is_torch_tensor, +) class DiscreteGroupedGemmSwigluSm100(APIBase): @@ -114,41 +127,39 @@ def __init__( :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("DiscreteGroupedGemmSwigluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() self._logger.debug("Entering __init__") + self._framework = detect_framework(sample_a) self._value_error_if(num_experts == 0, "num_experts must be > 0") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + # Set before building descriptors so uint8 sample tensors (the packed-fp4 + # container dtype, e.g. from JAX which has no fp4 dtype) get fp4 logical shapes. + self._interpret_uint8_as_fp4x2 = True + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias", canonical=True) + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) self.expert_cnt = num_experts self._value_error_if( @@ -156,10 +167,10 @@ def __init__( f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) must equal expert_cnt ({self.expert_cnt})", ) - self.b_dtype = b_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) self.b_shape = b_shape - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -182,16 +193,45 @@ def __init__( self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") self._workspace = None + self._compile_b_ptrs = None + self._compile_sfb_ptrs = None + self._live_ptrs = None self._logger.debug("__init__ completed") + def _check_sf_shape(self, desc, mn_div_128: int, rest: int, name: str) -> bool: + """Validate a 6-D scale-factor descriptor; returns True for the physical form. + + Accepts the torch-style permuted atom view (32, 4, MN', 4, K', 1) or the + physical C-contiguous allocation (1, MN', K', 32, 4, 4) -- byte-identical + memory, and the kernel consumes only the SF base pointer. + """ + if desc is None: + return False + is_physical = desc.ndim == 6 and desc.shape[0] == 1 and desc.shape[3] == 32 + self._check_tensor_shape( + desc, + (1, mn_div_128, rest, 32, 4, 4) if is_physical else (32, 4, mn_div_128, 4, rest, 1), + name, + ) + # The kernel consumes only the SF base pointer and rebuilds the layout from the + # GEMM shapes, so both forms must be exactly the C-contiguous physical allocation + # in memory: validate strides too (a shape-matching but differently-strided tensor + # would silently produce wrong results). + if is_physical: + expected = canonicalize_unit_dim_strides((1, mn_div_128, rest, 32, 4, 4), (mn_div_128 * rest * 512, rest * 512, 512, 16, 4, 1)) + extra = f"{name} in the physical (1, MN', K', 32, 4, 4) form must be C-contiguous" + else: + expected = canonicalize_unit_dim_strides((32, 4, mn_div_128, 4, rest, 1), (16, 4, rest * 512, 1, 512, mn_div_128 * rest * 512)) + extra = f"{name} atom view must be the (3, 4, 1, 5, 2, 0) permutation of a C-contiguous (1, MN', K', 32, 4, 4) allocation" + _ = self._check_tensor_stride(desc, stride=[expected], name=name, extra_error_msg=extra) + return is_physical + def check_support(self) -> bool: """Check if the kernel configuration is supported. :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -223,15 +263,17 @@ def check_support(self) -> bool: self._check_tensor_shape(self.bias_desc, (n, self.expert_cnt), "bias") rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") + # SF tensors are accepted in either of two byte-identical forms: the torch-style + # permuted atom view (32, 4, MN', 4, K', 1), or the physical C-contiguous + # allocation (1, MN', K', 32, 4, 4) for frameworks such as JAX that cannot + # express the permuted (strided) view. This is safe because the kernel rebuilds + # every SF layout from the A/D shapes (tile_atom_to_shape_SF) and consumes only + # the SF base pointers. + self._sfa_is_physical = self._check_sf_shape(self.sfa_desc, ceil_div(tensor_m, 128), rest_k, "SFA") rest_n2 = ceil_div(ceil_div(n // 2, self.sf_vec_size), 4) - self._check_tensor_shape( - self.sfd_row_desc, - (32, 4, ceil_div(tensor_m, 128), 4, rest_n2, 1), - "SFD_row", - ) + self._sfd_row_is_physical = self._check_sf_shape(self.sfd_row_desc, ceil_div(tensor_m, 128), rest_n2, "SFD_row") rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n // 2, 128), 4, rest_m, 1), "SFD_col") + self._sfd_col_is_physical = self._check_sf_shape(self.sfd_col_desc, ceil_div(n // 2, 128), rest_m, "SFD_col") self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") @@ -268,10 +310,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -283,7 +325,7 @@ def check_support(self) -> bool: self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD", ) self._check_dtype( @@ -300,7 +342,7 @@ def check_support(self) -> bool: ) self._check_dtype( self.bias_desc, - dtype=[torch.bfloat16, torch.float16], + dtype=[cutlass.BFloat16, cutlass.Float16], name="bias", extra_error_msg="bias must be fp16 or bfloat16", ) @@ -310,7 +352,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype in [cutlass.Float8E4M3FN] and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) self._value_error_if( @@ -320,19 +362,19 @@ def check_support(self) -> bool: self._check_dtype( self.acc_dtype, - dtype=torch.float32, + dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32", ) self.c_dtype = self._check_dtype( self.c_desc, dtype=[ - torch.float32, - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="C", ) @@ -340,7 +382,7 @@ def check_support(self) -> bool: if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D", extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -348,11 +390,11 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_desc, dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="D", ) @@ -364,7 +406,7 @@ def check_support(self) -> bool: ) self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, + self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype is cutlass.Float32, "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported", ) @@ -453,7 +495,7 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): "Invalid configuration: fp8 ab_dtype with mma_tiler_mn[1] == 128 and fp8 d_dtype is not supported", ) self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [torch.float16, torch.bfloat16]), + self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [cutlass.Float16, cutlass.BFloat16]), f"Invalid configuration: for fp4 ab_dtype, c_dtype must be float16 or bfloat16, got {self.c_dtype}", ) self._not_implemented_error_if( @@ -461,10 +503,9 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): "Discrete bias fusion currently requires mma_tiler_mn[1] == 256", ) - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: raise RuntimeError(f"DiscreteGroupedGemmSwiglu requires SM100+, but found SM{compute_capability}") @@ -475,8 +516,6 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): def compile(self) -> None: """Compile the kernel from tensor descriptors captured in __init__.""" - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -524,7 +563,9 @@ def compile(self) -> None: fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) workspace_bytes = gemm_glu.get_workspace_bytes() - self._workspace = torch.empty(workspace_bytes, dtype=torch.uint8, device="cuda") + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) align = 32 if ab_cutlass_dtype.width == 4 else 16 @@ -553,37 +594,72 @@ def compile(self) -> None: ) tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[5] = stride_tensor_m_128 - sfa_tensor = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - assumed_align=16, - ) - sfd_row_tensor = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + if self._sfa_is_physical: + # SFA in the physical C-contiguous atom shape (1, M', K', 32, 4, 4); + # the kernel rebuilds the SF layout from A's shape and consumes only + # the SFA base pointer, so only the calling convention differs. + sfa_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfa_desc.dtype, + shape=self.sfa_desc.shape, + stride_order=self.sfa_desc.stride_order, assumed_align=16, + dynamic_mode=1, + divisibility=1, ) - sfd_col_tensor = None - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + else: + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[5] = stride_tensor_m_128 + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), assumed_align=16, ) + sfd_row_tensor = None + if self.sfd_row_desc is not None: + if self._sfd_row_is_physical: + # Physical C-contiguous atom shape (1, M', N2', 32, 4, 4); pointer-only. + sfd_row_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfd_row_desc.dtype, + shape=self.sfd_row_desc.shape, + stride_order=self.sfd_row_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=1, + ) + else: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + assumed_align=16, + ) + sfd_col_tensor = None + if self.sfd_col_desc is not None: + if self._sfd_col_is_physical: + # Physical C-contiguous atom shape (1, N2', M'', 32, 4, 4); pointer-only. + sfd_col_tensor = self._make_fake_cute_compact_tensor( + dtype=self.sfd_col_desc.dtype, + shape=self.sfd_col_desc.shape, + stride_order=self.sfd_col_desc.stride_order, + assumed_align=16, + dynamic_mode=2, + divisibility=1, + ) + else: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + assumed_align=16, + ) amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) @@ -598,12 +674,18 @@ def compile(self) -> None: ) bias_tensor = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) - # Use internal device-resident int64 arrays to provide valid pointer-like - # compile-time placeholders for b_ptrs/sfb_ptrs (required by kernel __call__). - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + # Use internal device-resident buffers to provide valid pointer-like compile-time + # placeholders for b_ptrs/sfb_ptrs (required by kernel __call__): real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, retyped + # to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + self._compile_sfb_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + b_ptrs_placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + b_ptrs_placeholder.element_type = cutlass.Int64 + b_ptrs_cute = b_ptrs_placeholder.iterator + sfb_ptrs_placeholder = from_dlpack(self._compile_sfb_ptrs, assumed_align=8) + sfb_ptrs_placeholder.element_type = cutlass.Int64 + sfb_ptrs_cute = sfb_ptrs_placeholder.iterator workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator @@ -680,8 +762,8 @@ def tensor_api( glu_clamp_min: float = -7.0, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + b_ptrs_addr = int(get_data_ptr(b_ptrs_device)) + sfb_ptrs_addr = int(get_data_ptr(sfb_ptrs_device)) _compiled_kernel( a_tensor, @@ -760,15 +842,22 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) - if a_tensor.shape[0] == 0: + if get_shape(a_tensor)[0] == 0: self._logger.debug("execute: valid_m is zero, skipping kernel execution") return self._runtime_error_if( self._compiled_kernel is None, "Kernel not compiled; call compile() first", ) + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # arrays referenced until the next execute so their buffers outlive the launch. + self._live_ptrs = (b_ptrs, sfb_ptrs) # Resolve linear_offset default: None -> activation-derived legacy value # (1.0 for geglu, 0.0 for swiglu) for backwards compatibility with callers @@ -897,118 +986,162 @@ def discrete_grouped_gemm_swiglu_wrapper_sm100( TupleDict with keys: c_tensor, d_tensor, d_col_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor """ - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("discrete_grouped_gemm_swiglu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if c_dtype is None: - c_dtype = torch.bfloat16 - if d_dtype is None: - d_dtype = torch.bfloat16 + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for discrete_grouped_gemm_swiglu_wrapper_sm100; pass torch tensors or JAX arrays") + + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + c_dtype = _convert_to_cutlass_data_type(c_dtype) if c_dtype is not None else cutlass.BFloat16 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) + ab_dtype = _convert_to_cutlass_data_type(a_tensor.dtype) + + if framework == "jax": + if bias_tensor is not None: + raise ValueError( + "bias_tensor is not expressible as JAX arrays (its (n, experts) column-major layout has no row-major equivalent); " "omit bias for JAX inputs" + ) + if ab_dtype in (cutlass.Uint8, cutlass.Float4E2M1FN): + raise ValueError( + "packed-fp4 inputs are not expressible as JAX arrays for this API " + "(JAX has no packed fp4 dtype and the compiled kernel entry point requires float4_e2m1fn_x2 tensors); " + "use torch tensors for FP4, or FP8 inputs for JAX" + ) # Resolve linear_offset default: None means "use the activation-derived legacy # default" (1.0 for geglu, 0.0 for swiglu) for backwards compatibility. if linear_offset is None: linear_offset = 1.0 if act_func == "geglu" else 0.0 - valid_m, k_physical, _ = a_tensor.shape - _require_pointer_tensor(b_ptrs, "b_ptrs") - num_experts = b_ptrs.shape[0] - _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) + valid_m, k_physical, _ = get_shape(a_tensor) + num_experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + _validate_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) n_out = n // 2 - k_logical = k_physical * 2 if b_dtype in (torch.float4_e2m1fn_x2, torch.uint8) else k_physical + k_logical = k_physical * 2 if b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) else k_physical b_shape = (n, k_logical) - if bias_tensor is not None and tuple(bias_tensor.shape) != (n, num_experts): - raise ValueError(f"bias_tensor must have shape {(n, num_experts)}, got {tuple(bias_tensor.shape)}") + if bias_tensor is not None and get_shape(bias_tensor) != (n, num_experts): + raise ValueError(f"bias_tensor must have shape {(n, num_experts)}, got {get_shape(bias_tensor)}") _logger.debug("discrete_grouped_gemm_swiglu_wrapper_sm100: Creating output tensors") - if cd_major == "n": - c_tensor = torch.empty_strided((valid_m, n, 1), (n, 1, valid_m * n), dtype=c_dtype, device=a_tensor.device) + if cd_major != "n": + raise ValueError(f"cd_major must be 'n', got {cd_major}") + + if framework == "torch": + import torch + + c_tensor = torch.empty_strided((valid_m, n, 1), (n, 1, valid_m * n), dtype=framework_dtype(c_dtype, "torch"), device=a_tensor.device) d_tensor = torch.empty_strided( (valid_m, n_out, 1), (n_out, 1, valid_m * n_out), - dtype=d_dtype, + dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device, ) d_col_tensor = torch.empty_strided( (valid_m, n_out, 1), (n_out, 1, valid_m * n_out), - dtype=d_dtype, + dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device, ) else: - raise ValueError(f"cd_major must be 'n', got {cd_major}") + import jax + import jax.numpy as jnp + + # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. + # The kernel writes into these buffers on the launch stream; materialize them first. + c_tensor = jnp.empty((valid_m, n, 1), dtype=framework_dtype(c_dtype, "jax"), device=a_tensor.device) + d_tensor = jnp.empty((valid_m, n_out, 1), dtype=framework_dtype(d_dtype, "jax"), device=a_tensor.device) + d_col_tensor = jnp.empty((valid_m, n_out, 1), dtype=framework_dtype(d_dtype, "jax"), device=a_tensor.device) + jax.block_until_ready((c_tensor, d_tensor, d_col_tensor)) sfd_row_tensor = None sfd_col_tensor = None amax_tensor = None - if a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [torch.float8_e8m0fnu, torch.float8_e4m3fn]: + if ab_dtype in [ + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ] and _convert_to_cutlass_data_type( + sfa_tensor.dtype + ) in [cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN]: _logger.debug("discrete_grouped_gemm_swiglu_wrapper_sm100: Detected fp8 config, constructing sfd tensors") - sf_dtype = sfa_tensor.dtype mma_permute_order = (3, 4, 1, 5, 2, 0) sf_k_row = ceil_div(n_out, sf_vec_size) mma_shape_row = (1, ceil_div(valid_m, 128), ceil_div(sf_k_row, 4), 32, 4, 4) - sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - sf_k_col = ceil_div(valid_m, sf_vec_size) mma_shape_col = (1, ceil_div(n_out, 128), ceil_div(sf_k_col, 4), 32, 4, 4) - sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - if d_dtype in [torch.bfloat16, torch.float16]: + if framework == "torch": + import torch + + sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + else: + import jax + import jax.numpy as jnp + + # Physical C-contiguous atom-shape allocations (the kernel rebuilds the SF + # layouts from the output shape and consumes only the base pointers). + sfd_row_tensor = jnp.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device) + sfd_col_tensor = jnp.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device) + jax.block_until_ready((sfd_row_tensor, sfd_col_tensor)) + + if d_dtype in [cutlass.BFloat16, cutlass.Float16]: _logger.debug("discrete_grouped_gemm_swiglu_wrapper_sm100: Constructing amax_tensor") - amax_tensor = torch.full((num_experts, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) + if framework == "torch": + import torch + + amax_tensor = torch.full((num_experts, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) + else: + import jax + import jax.numpy as jnp + + amax_tensor = jax.block_until_ready(jnp.full((num_experts, 1), float("-inf"), dtype=jnp.float32, device=a_tensor.device)) def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(i for i, s in sorted(enumerate(tensor.stride()), key=lambda x: x[1])) + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + return tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: (x[1], tensor_shape[x[0]]))) def tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype + tensor_shape = get_shape(tensor) + return tensor_shape, canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)), _convert_to_cutlass_data_type(tensor.dtype) def dynamic_m_tensor_signature( tensor: Optional[torch.Tensor], static_shape_suffix: Optional[Tuple[int, ...]], dynamic_stride_dims: Tuple[int, ...] = () ) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(tensor.stride())) - return static_shape_suffix, stride_signature, tensor.dtype + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(tensor_stride)) + return static_shape_suffix, stride_signature, _convert_to_cutlass_data_type(tensor.dtype) cache_key = ( - a_tensor.shape[1:], + get_shape(a_tensor)[1:], stride_order(a_tensor), - a_tensor.dtype, + ab_dtype, b_shape, b_dtype, - c_tensor.shape[1:], + get_shape(c_tensor)[1:], stride_order(c_tensor), - c_tensor.dtype, - *dynamic_m_tensor_signature(sfa_tensor, (sfa_tensor.shape[4], 1) if sfa_tensor is not None else None, dynamic_stride_dims=(5,)), + _convert_to_cutlass_data_type(c_tensor.dtype), + *dynamic_m_tensor_signature( + sfa_tensor, + (get_shape(sfa_tensor)[4], 1) if sfa_tensor is not None else None, + dynamic_stride_dims=(0, 1, 5), + ), *tensor_signature(alpha_tensor), *tensor_signature(bias_tensor), *tensor_signature(norm_const_tensor), - *dynamic_m_tensor_signature(prob_tensor, (1, 1)), - tuple(b_ptrs.shape), - tuple(b_ptrs.stride()), - b_ptrs.dtype, - tuple(sfb_ptrs.shape), - tuple(sfb_ptrs.stride()), - sfb_ptrs.dtype, - tuple(padded_offsets.shape), - tuple(padded_offsets.stride()), - padded_offsets.dtype, + *dynamic_m_tensor_signature(prob_tensor, (1, 1), dynamic_stride_dims=(1, 2)), + *tensor_signature(b_ptrs), + *tensor_signature(sfb_ptrs), + *tensor_signature(padded_offsets), acc_dtype, c_dtype, d_dtype, diff --git a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py new file mode 100644 index 000000000..b4b65ae52 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the discrete-weight block-scaled +grouped GEMM GLU forward (SwiGLU/GeGLU), built on :func:`cudnn.jax.call`. + +FP8 inputs only (the packed-fp4 uint8 container dtype is not expressible as JAX +arrays) and no bias (its (n, experts) column-major layout is likewise +inexpressible; ``bias`` is a compile-time ``None`` inside the adapter). This +discrete kernel is JAX-expressible where the contiguous block-scaled GLU kernel is +not: SFB arrives as per-expert base pointers and SFA/SFD travel in the physical +C-contiguous atom shape ``(1, MN', K', 32, 4, 4)`` — the kernel rebuilds every +scale-factor layout from the A/D shapes and consumes only the SF base pointers. +The per-expert B/SFB pointers travel as regular device arrays whose *values* are +raw addresses — the referenced weight/scale buffers are not visible to XLA, so the +caller must keep them alive (and unmoved) across every execution of the traced +computation. ``padded_offsets`` values cannot be host-validated under tracing; +malformed offsets are the caller's responsibility here (the eager wrapper +validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.api_base import ceil_div +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import TensorSpec, call, gemm_operand_spec, neg_inf_init, zeros_init +from ...grouped.unfused.jax_api import _pointer_count, _prob_spec +from .discrete_B_blockscaled_grouped_gemm_glu_bias import BlockScaledDiscreteWeightGroupedGemmBiasKernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_fp8_dtypes = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_c_dtypes = (cutlass.Float32, cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_d_dtypes = (cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_amax_dtypes = (cutlass.BFloat16, cutlass.Float16) + +_JAX_FP4_ERROR = ( + "packed-fp4 inputs are not expressible as JAX arrays for this API " + "(JAX has no packed fp4 dtype and the compiled kernel entry point requires float4_e2m1fn_x2 tensors); " + "use torch tensors for FP4, or FP8 inputs for JAX" +) + + +def _sf_atom_spec() -> TensorSpec: + # Physical C-contiguous (1, MN', K', 32, 4, 4) scale-factor atom allocation: + # explicit row-major ranks because the unit dim makes leading-dim inference + # ambiguous. The kernel rebuilds the SF layout from the GEMM shapes and consumes + # only the base pointer, so no permuted view is needed. + return TensorSpec(layout=(5, 4, 3, 2, 1, 0)) + + +def _sf_byte_zeros_init(shape_dtype: jax.ShapeDtypeStruct) -> jax.Array: + # float8_e8m0fnu has no representable zero; zero the raw bytes instead + # (byte 0x00 decodes to 2^-127) via a free uint8 bitcast. + return jnp.zeros(shape_dtype.shape, jnp.uint8).view(shape_dtype.dtype) + + +def _as_e8m0_array(scale: Any) -> Any: + """Present a uint8 E8M0-bit-pattern array as float8_e8m0fnu (free bitcast, jit-safe).""" + if _convert_to_cutlass_data_type(scale.dtype) is cutlass.Uint8: + import ml_dtypes + + return scale.view(ml_dtypes.float8_e8m0fnu) + return scale + + +@cute.jit +def _discrete_swiglu_adapter( + stream, + a, + b_ptrs, + sfb_ptrs, + sfa, + padded_offsets, + alpha, + prob, + norm_const, + c, + d, + d_col, + sfd_row, + sfd_col, + amax, + workspace, + *, + kernel, + n, + k, + mac, + has_amax, + linear_offset, + geglu_alpha, + glu_clamp_max, + glu_clamp_min, +): + # Discrete-mode b/sfb are raw pointers to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffers recast for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + sfb_arg = cute.recast_ptr(sfb_ptrs.iterator, dtype=cutlass.Int64) + # amax exists only for 16-bit d (eager parity: fp8 d compiles with amax=None); + # the dummy fp8-mode buffer is never passed to the kernel. + amax_arg = None + if cutlass.const_expr(has_amax): + amax_arg = amax + kernel( + a=a, + b_ptrs=b_arg, + sfb_ptrs=sfb_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d, + d_col=d_col, + sfa=sfa, + sfd_row_tensor=sfd_row, + sfd_col_tensor=sfd_col, + amax_tensor=amax_arg, + norm_const_tensor=norm_const, + padded_offsets=padded_offsets, + alpha=alpha, + prob=prob, + bias=None, + max_active_clusters=mac, + stream=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 discrete_grouped_gemm_swiglu_jax_sm100( + a_tensor: Any, + b_ptrs: Any, + sfa_tensor: Any, + sfb_ptrs: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any, + n: int, + bias_tensor: Optional[Any] = None, + c_dtype: Any = cutlass.BFloat16, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + 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, + 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, +) -> dict: + """Discrete-weight block-scaled grouped GEMM GLU forward as an XLA custom call. + + Same contract as the eager wrapper's FP8 JAX mode: A ``(m, k, 1)`` k-major + C-contiguous fp8 (e4m3/e5m2), SFA in the physical C-contiguous E8M0 atom shape + ``(1, m/128, ceil(ceil(k/sf_vec_size)/4), 32, 4, 4)`` (``float8_e8m0fnu``, or + ``uint8`` bit patterns), ``padded_offsets (experts,)`` int32 cumulative + 256-aligned row offsets, ``alpha (experts,)`` float32, ``prob (m, 1, 1)`` + float32, ``norm_const (1,)`` float32, and ``b_ptrs``/``sfb_ptrs`` holding + per-expert ``(n, k)`` k-major weight / SFB atom base addresses (packed + little-endian uint8, 8 bytes per pointer — or int64 with x64 mode). ``n`` is + the full weight N before the GLU split; ``d``/``d_col`` come back + ``(m, n // 2, 1)``. The scalar GLU knobs (``linear_offset`` — defaulting per + ``act_func`` — ``geglu_alpha``, ``glu_clamp_max``, ``glu_clamp_min``) are + compile-time constants of the traced call. Rows at/past ``padded_offsets[-1]`` + come back zero-filled (the outputs are donated zero-initialized buffers). + + Returns a dict with the eager wrapper's keys: ``c_tensor (m, n, 1)``, + ``d_tensor``/``d_col_tensor (m, n//2, 1)``, ``sfd_row_tensor``/``sfd_col_tensor`` + (physical E8M0 atom shape), and ``amax_tensor ((experts, 1) float32, + -inf-initialized)`` — ``None`` unless ``d_dtype`` is bf16/fp16. + + Rejected for JAX (as in the eager wrapper): packed-fp4 inputs/outputs, bias, + and the contiguous kernel's non-k-major weight layouts. + """ + c_dtype = _convert_to_cutlass_data_type(c_dtype) + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + if bias_tensor is not None: + raise ValueError( + "bias_tensor is not expressible as JAX arrays (its (n, experts) column-major layout has no row-major equivalent); omit bias for JAX inputs" + ) + ab_dtype = _convert_to_cutlass_data_type(a_tensor.dtype) + if ab_dtype in (cutlass.Uint8, cutlass.Float4E2M1FN): + raise ValueError(_JAX_FP4_ERROR) + if ab_dtype not in _fp8_dtypes: + raise ValueError(f"a_tensor must be float8_e4m3fn or float8_e5m2, got {a_tensor.dtype}") + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + if n is None or n <= 0 or n % 2 != 0: + raise ValueError(f"n must be positive and even (gate+up combined width), got {n}") + n_out = n // 2 + + if c_dtype not in _c_dtypes: + raise ValueError(f"c_dtype must be FP32, FP16, BF16, or FP8 for JAX (packed fp4 has no JAX dtype), got {c_dtype}") + if d_dtype not in _d_dtypes: + raise ValueError(f"d_dtype must be FP16, BF16, or FP8 for JAX (packed fp4 has no JAX dtype), got {d_dtype}") + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + if sf_vec_size != 32: + raise ValueError(f"fp8 inputs require sf_vec_size 32, got {sf_vec_size}") + if act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {act_func}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "geglu" else 0.0 + + sfa_tensor = _as_e8m0_array(sfa_tensor) + if _convert_to_cutlass_data_type(sfa_tensor.dtype) is not cutlass.Float8E8M0FNU: + raise ValueError(f"sfa_tensor must be float8_e8m0fnu (or uint8 bit patterns) for fp8 inputs, got {sfa_tensor.dtype}") + rest_k = ceil_div(ceil_div(k, sf_vec_size), 4) + expected_sfa = (1, ceil_div(m, 128), rest_k, 32, 4, 4) + if tuple(sfa_tensor.shape) != expected_sfa: + raise ValueError(f"sfa_tensor must use the physical C-contiguous atom shape {expected_sfa}, got {tuple(sfa_tensor.shape)}") + + expert_cnt = _pointer_count(b_ptrs) + sfb_cnt = _pointer_count(sfb_ptrs, "sfb_ptrs") + if sfb_cnt != expert_cnt: + raise ValueError(f"sfb_ptrs length mismatch: expected {expert_cnt} pointers, got {sfb_cnt}") + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + if tuple(norm_const_tensor.shape) != (1,) or _convert_to_cutlass_data_type(norm_const_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"norm_const_tensor must be (1,) float32, got {tuple(norm_const_tensor.shape)} {norm_const_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + if mma_tiler_mn[1] != 256: + raise ValueError(f"MMA tiler N must be 256, got {mma_tiler_mn[1]}") + + if not BlockScaledDiscreteWeightGroupedGemmBiasKernel.can_implement( + ab_dtype, + cutlass.Float8E8M0FNU, + sf_vec_size, + acc_dtype, + d_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + "k", + "n", + BlockScaledDiscreteWeightGroupedGemmBiasKernel.FIX_PAD_SIZE, + ): + raise ValueError("Unsupported discrete grouped GEMM SwiGLU tile, cluster, alignment, or layout configuration") + + has_amax = d_dtype in _amax_dtypes + + cache_key = ( + expert_cnt, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + sf_vec_size, + vector_f32, + discrete_col_sfd, + act_func, + use_dynamic_sched, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = BlockScaledDiscreteWeightGroupedGemmBiasKernel( + sf_vec_size=sf_vec_size, + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=True, # the fp8 JAX path always generates SFD (E8M0 scale factors) + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + act_func=act_func, + enable_bias=False, + use_dynamic_sched=use_dynamic_sched, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + sf_jax_dtype = framework_dtype(cutlass.Float8E8M0FNU, "jax") + operand = gemm_operand_spec() + sf_spec = _sf_atom_spec() + c_out, d_out, d_col_out, sfd_row_out, sfd_col_out, amax_out, _workspace = call( + _discrete_swiglu_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n, 1), framework_dtype(c_dtype, "jax")), + jax.ShapeDtypeStruct((m, n_out, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, n_out, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((1, ceil_div(m, 128), ceil_div(ceil_div(n_out, sf_vec_size), 4), 32, 4, 4), sf_jax_dtype), + jax.ShapeDtypeStruct((1, ceil_div(n_out, 128), ceil_div(ceil_div(m, sf_vec_size), 4), 32, 4, 4), sf_jax_dtype), + jax.ShapeDtypeStruct((expert_cnt, 1), jnp.float32), + jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8), + ), + input_spec=(operand, None, None, sf_spec, None, None, _prob_spec(), None), + output_spec=(operand, operand, operand, sf_spec, sf_spec, TensorSpec(layout=(1, 0)), None), + # All outputs donated: c/d/d_col/sfd for the unit-dim layout specs (and + # defined bytes past the last offset); amax because the kernel accumulates + # into it (atomic max over a -inf-initialized buffer, eager parity); the + # workspace because the helper kernel writes the per-expert TMA descriptors + # into it (XLA inputs are immutable). + initialized_outputs={ + 0: zeros_init, + 1: zeros_init, + 2: zeros_init, + 3: _sf_byte_zeros_init, + 4: _sf_byte_zeros_init, + 5: neg_inf_init if has_amax else zeros_init, + 6: zeros_init, + }, + kernel=kernel, + n=int(n), + k=int(k), + mac=mac, + has_amax=has_amax, + linear_offset=float(linear_offset), + geglu_alpha=float(geglu_alpha), + glu_clamp_max=float(glu_clamp_max), + glu_clamp_min=float(glu_clamp_min), + )(a_tensor, b_ptrs, sfb_ptrs, sfa_tensor, padded_offsets, alpha_tensor, prob_tensor, norm_const_tensor) + + return { + "c_tensor": c_out, + "d_tensor": d_out, + "d_col_tensor": d_col_out, + "amax_tensor": amax_out if has_amax else None, + "sfd_row_tensor": sfd_row_out, + "sfd_col_tensor": sfd_col_out, + } diff --git a/python/cudnn/gemm/cutedsl/grouped/__init__.py b/python/cudnn/gemm/cutedsl/grouped/__init__.py index d37a73a0d..8cfba7f10 100644 --- a/python/cudnn/gemm/cutedsl/grouped/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/__init__.py @@ -72,4 +72,28 @@ "grouped_gemm_wgrad_wrapper_sm100", "GroupedGemmSm100", "grouped_gemm_wrapper_sm100", + "grouped_gemm_jax_sm100", + "grouped_gemm_glu_jax_sm100", + "grouped_gemm_dglu_jax_sm100", + "grouped_gemm_dsrelu_jax_sm100", + "grouped_gemm_wgrad_jax_sm100", ] + +# Lazy: the jax entry points import jax/cutlass.jax, which must not be pulled in +# for torch-only users. +_JAX_LAZY_EXPORTS = { + "grouped_gemm_jax_sm100": ".unfused", + "grouped_gemm_glu_jax_sm100": ".glu", + "grouped_gemm_dglu_jax_sm100": ".dglu", + "grouped_gemm_dsrelu_jax_sm100": ".dsrelu", + "grouped_gemm_wgrad_jax_sm100": ".wgrad", +} + + +def __getattr__(name): + module_name = _JAX_LAZY_EXPORTS.get(name) + if module_name is not None: + import importlib + + return getattr(importlib.import_module(module_name, __name__), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/backend_utils.py b/python/cudnn/gemm/cutedsl/grouped/backend_utils.py index dc12ad0df..93c4edaf6 100644 --- a/python/cudnn/gemm/cutedsl/grouped/backend_utils.py +++ b/python/cudnn/gemm/cutedsl/grouped/backend_utils.py @@ -17,7 +17,11 @@ class GroupedGemmBackend(str, Enum): @contextmanager def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]: - """Run PyTorch work on the CUDA stream used for the kernel launch.""" + """Run PyTorch work on the CUDA stream used for the kernel launch. + + torch-only: callers must guard this context so non-torch (e.g. JAX) code paths + never enter it -- it imports torch and interprets ``device`` as a torch device. + """ import torch if current_stream is None: @@ -44,9 +48,15 @@ def select_grouped_gemm_backend( scale_controls, block_scaled_dtype_pairs, ): - import torch + # Compare in canonical (cutlass) dtype space so torch/jax/numpy/str dtypes all + # resolve; dtypes with no cutlass mapping fall through to the unsupported-pair error. + import cutlass + + from cudnn.datatypes import _convert_to_cutlass_data_type_or_none - bf16_operands = (a_dtype == torch.bfloat16, b_dtype == torch.bfloat16) + a_dtype_canonical = _convert_to_cutlass_data_type_or_none(a_dtype) + b_dtype_canonical = _convert_to_cutlass_data_type_or_none(b_dtype) + bf16_operands = (a_dtype_canonical is cutlass.BFloat16, b_dtype_canonical is cutlass.BFloat16) if any(bf16_operands): if not all(bf16_operands): raise ValueError(f"{operation}: mixed dtype families: a_dtype={a_dtype}, " f"b_dtype={b_dtype}") @@ -54,7 +64,10 @@ def select_grouped_gemm_backend( if forbidden: raise ValueError(f"{operation}: BF16 forbids scale control {forbidden[0]}") return GroupedGemmBackend.BF16 - if (a_dtype, b_dtype) in block_scaled_dtype_pairs: + canonical_pairs = { + (_convert_to_cutlass_data_type_or_none(pair_a), _convert_to_cutlass_data_type_or_none(pair_b)) for pair_a, pair_b in block_scaled_dtype_pairs + } + if (a_dtype_canonical, b_dtype_canonical) in canonical_pairs: return GroupedGemmBackend.BLOCK_SCALED raise ValueError(f"{operation}: unsupported dtype pair a_dtype={a_dtype}, " f"b_dtype={b_dtype}") diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/__init__.py b/python/cudnn/gemm/cutedsl/grouped/dglu/__init__.py index a21a85e06..6a4a533c1 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100", + "grouped_gemm_dglu_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "grouped_gemm_dglu_jax_sm100": + from .jax_api import grouped_gemm_dglu_jax_sm100 + + return grouped_gemm_dglu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py index 37ea0e3fc..1b244dfdb 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py @@ -17,21 +17,27 @@ from cudnn.api_base import APIBase, TensorDesc from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _pointer_values, _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + get_compute_capability, + get_data_ptr, + get_device, + get_version, + is_torch_tensor, + to_host_list, +) from ..moe_utils import MoEWeightMode from .moe_grouped_gemm_dglu_dbias import MoEGroupedGemmDgluDbiasBf16Kernel -_OUTPUT_DTYPES = None - def _output_dtypes(): - global _OUTPUT_DTYPES - if _OUTPUT_DTYPES is None: - import torch - - _OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] - return _OUTPUT_DTYPES + return [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] class GroupedGemmDgluBf16API(APIBase): @@ -61,12 +67,11 @@ def __init__( b_major: str = "k", use_dynamic_sched: bool = False, ) -> None: - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() + self._framework = detect_framework(sample_a) if sample_b is not None and num_experts is None: self.weight_mode = MoEWeightMode.DENSE @@ -77,22 +82,22 @@ def __init__( else: raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") - self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob", canonical=True) + self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias", canonical=True) self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) self._sample_offsets_ref = weakref.ref(sample_padded_offsets) - self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_offsets_version = get_version(sample_padded_offsets) self._sample_data_ptrs = { - name: tensor.data_ptr() + name: get_data_ptr(tensor) for name, tensor in ( ("sample_a", sample_a), ("sample_b", sample_b), @@ -110,8 +115,8 @@ def __init__( self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) self.b_shape = tuple(b_shape) if b_shape is not None else None - self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype - self.acc_dtype = acc_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = tuple(mma_tiler_mn) self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) @@ -123,7 +128,8 @@ def __init__( self._has_dbias = self.dbias_desc is not None self._kernel = MoEGroupedGemmDgluDbiasBf16Kernel self._workspace: Optional[torch.Tensor] = None - self._compile_b_ptrs: Optional[torch.Tensor] = None + self._live_b_ptrs = None + self._compile_b_ptrs = None self._validated_offsets: dict[int, tuple] = {} self._validated_pointer_values: dict[int, tuple] = {} self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) @@ -145,12 +151,12 @@ def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: @staticmethod def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(int(value) for value in tensor.detach().cpu().tolist()) + return tuple(int(value) for value in to_host_list(tensor)) @staticmethod def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: cached = cache.get(id(tensor)) - return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + return bool(cached and cached[0]() is tensor and cached[1] == get_version(tensor) and cached[2] == extra) @staticmethod def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: @@ -159,7 +165,7 @@ def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) - def discard(_reference, *, cache=cache, key=key): cache.pop(key, None) - cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + cache[key] = (weakref.ref(tensor, discard), get_version(tensor), extra) @staticmethod def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None: @@ -186,23 +192,27 @@ def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> Non def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): return - values = self._copy_values_to_host(b_ptrs) + values = _pointer_values(b_ptrs) if any(value == 0 or value % 16 != 0 for value in values): raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) @staticmethod def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: - if tensor.data_ptr() % 16 != 0: + if get_data_ptr(tensor) % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") @staticmethod def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: - if tensor.data_ptr() % 8 != 0: + if get_data_ptr(tensor) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") - @staticmethod - def _record_pointer_stream(b_ptrs: torch.Tensor, stream: cuda.CUstream) -> None: + def _record_pointer_stream(self, b_ptrs: torch.Tensor, stream: cuda.CUstream) -> None: + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # array referenced until the next execute so its buffer outlives the launch. + self._live_b_ptrs = b_ptrs + return import torch handle = int(stream) @@ -217,8 +227,6 @@ def _record_pointer_stream(b_ptrs: torch.Tensor, stream: cuda.CUstream) -> None: b_ptrs.record_stream(launch_stream) def check_support(self) -> bool: - import torch - if self.a_desc.ndim != 3: raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") tensor_m, k, one = self.a_desc.shape @@ -260,20 +268,20 @@ def check_support(self) -> bool: self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") self._expect_stride(self.alpha_desc, (1,), "sample_alpha") self._expect_stride(self.beta_desc, (1,), "sample_beta") - self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") - self._expect_stride(self.dprob_desc, (1, 1, 1), "sample_dprob") + self._expect_stride(self.prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_prob") + self._expect_stride(self.dprob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_dprob") - self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + self._check_dtype(self.a_desc, cutlass.BFloat16, "sample_a") if self.b_desc is not None: - self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") - self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.b_desc, cutlass.BFloat16, "sample_b") + self._check_dtype(self.b_dtype, cutlass.BFloat16, "b_dtype") self._check_dtype(self.c_desc, _output_dtypes(), "sample_c") self._check_dtype(self.d_row_desc, _output_dtypes(), "sample_d_row") - self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") - self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") - self._check_dtype(self.beta_desc, torch.float32, "sample_beta") - self._check_dtype(self.prob_desc, torch.float32, "sample_prob") - self._check_dtype(self.dprob_desc, torch.float32, "sample_dprob") + self._check_dtype(self.padded_offsets_desc, cutlass.Int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, cutlass.Float32, "sample_alpha") + self._check_dtype(self.beta_desc, cutlass.Float32, "sample_beta") + self._check_dtype(self.prob_desc, cutlass.Float32, "sample_prob") + self._check_dtype(self.dprob_desc, cutlass.Float32, "sample_dprob") device = self.a_desc.device for desc, name in ( @@ -291,15 +299,15 @@ def check_support(self) -> bool: if self.dbias_desc is not None: self._expect_shape(self.dbias_desc, (self.expert_cnt, two_n, 1), "sample_dbias") - self._expect_stride(self.dbias_desc, (two_n, 1, 1), "sample_dbias") - self._check_dtype(self.dbias_desc, torch.bfloat16, "sample_dbias") + self._expect_stride(self.dbias_desc, canonicalize_unit_dim_strides((self.expert_cnt, two_n, 1), (two_n, 1, 1)), "sample_dbias") + self._check_dtype(self.dbias_desc, cutlass.BFloat16, "sample_dbias") self._expect_device(self.dbias_desc, device, "sample_dbias") for name, pointer in self._sample_data_ptrs.items(): if pointer % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") - if self.acc_dtype != torch.float32: + if self.acc_dtype is not cutlass.Float32: raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") if self.m_aligned != 256: raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") @@ -314,13 +322,13 @@ def check_support(self) -> bool: self._validate_offset_sequence(self._sample_offset_values, expert_cnt=self.expert_cnt, tensor_m=tensor_m) sample_offsets = self._sample_offsets_ref() - if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + if sample_offsets is not None and get_version(sample_offsets) == self._sample_offsets_version: self._remember_validation(self._validated_offsets, sample_offsets, (self.expert_cnt, tensor_m)) elif sample_offsets is not None: self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) if not self._kernel.can_implement( - _convert_to_cutlass_data_type(torch.bfloat16), + cutlass.BFloat16, _convert_to_cutlass_data_type(self.c_desc.dtype), _convert_to_cutlass_data_type(self.d_row_desc.dtype), _convert_to_cutlass_data_type(self.acc_dtype), @@ -339,18 +347,16 @@ def check_support(self) -> bool: ): raise ValueError("Unsupported BF16 grouped GEMM dGLU configuration") - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(self.a_desc.device) + major, minor = get_compute_capability() capability = major * 10 + minor if capability < 100: - raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability} on {self.a_desc.device}") + raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability}") self._is_supported = True return True def compile(self) -> None: - import torch - self._ensure_support_checked() if self._compiled_kernel is not None: return @@ -372,8 +378,10 @@ def compile(self) -> None: raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") workspace_bytes = kernel.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) - if self._workspace.data_ptr() % 128 != 0: + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) + if get_data_ptr(self._workspace) % 128 != 0: raise RuntimeError("workspace allocation must be 128-byte aligned") workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -410,9 +418,14 @@ def compile(self) -> None: b_stride = cutlass.Int64(0) b_major_mode = OperandMajorMode.K else: - self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + # Compile-time placeholder for the pointer-array argument: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) self._validate_pointer_array_alignment(self._compile_b_ptrs) - b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + placeholder.element_type = cutlass.Int64 + b_fake = placeholder.iterator n, k = self.b_shape[:2] n_value = cutlass.Int32(n) k_value = cutlass.Int32(k) @@ -458,7 +471,7 @@ def tensor_api( stream, linear_offset, ) -> None: - b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(get_data_ptr(b_ptrs)) raw_compiled( a_tensor, b_arg, @@ -481,7 +494,7 @@ def tensor_api( self._compiled_kernel = tensor_api def _validate_live_tensor(self, tensor: torch.Tensor, sample: TensorDesc, name: str, *, dynamic_m: bool = False) -> TensorDesc: - desc = self._make_tensor_desc(tensor, name=name) + desc = self._make_tensor_desc(tensor, name=name, canonical=True) if desc.dtype != sample.dtype: raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") if desc.device != sample.device: @@ -511,7 +524,10 @@ def execute( linear_offset: float = 0.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if self._compiled_kernel is None: raise RuntimeError("Kernel not compiled; call compile() first") @@ -535,8 +551,8 @@ def execute( self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") self._expect_stride(c_desc, (two_n, 1, tensor_m * two_n), "c_tensor") self._expect_stride(d_desc, (two_n, 1, tensor_m * two_n), "d_row_tensor") - self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") - self._expect_stride(dprob_desc, (1, 1, 1), "dprob_tensor") + self._expect_stride(prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "prob_tensor") + self._expect_stride(dprob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "dprob_tensor") self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) for tensor, name in ( @@ -567,10 +583,10 @@ def execute( else: if b_tensor is not None or b_ptrs is None: raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") - _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) - if b_ptrs.device != self.a_desc.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor ({self.a_desc.device}), got {b_ptrs.device}") - if b_ptrs.data_ptr() % 8 != 0: + _validate_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if get_device(b_ptrs) != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor ({self.a_desc.device}), got {get_device(b_ptrs)}") + if get_data_ptr(b_ptrs) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") self._validate_pointer_values_once(b_ptrs) self._record_pointer_stream(b_ptrs, current_stream) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py index 025f71ed2..8b02e181c 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py @@ -184,6 +184,13 @@ def __init__( :param glu_clamp_min: Compile-time dGeGLU lower clamp. Ignored when ``act_func == "dswiglu"``. """ + from cudnn.tensor_adapter import detect_framework + + if sample_a is not None and detect_framework(sample_a) != "torch": + raise ValueError( + "GroupedGemmDgluBlockScaledAPI supports torch tensors only: the block-scaled " + "scale-factor tensors use an MMA-interleaved layout that is not expressible as JAX arrays" + ) import torch if acc_dtype is None: diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py index 081dc51b1..e5eebd985 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py @@ -33,26 +33,44 @@ import os from typing import Any, Tuple, Optional, overload +import cutlass + from cudnn.api_base import APIBase, TupleDict, ceil_div, get_device_type +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + cuda_is_available, + detect_framework, + framework_dtype, + get_compute_capability, + get_device, + get_shape, + get_strides, +) -_BLOCK_SCALED_DTYPE_PAIRS = None +_JAX_DENSE_B_ERROR = ( + "Dense weight mode (b_tensor) is not expressible as JAX arrays " + "(the expert-outermost strided B layout has no row-major equivalent); " + "use discrete mode (b_ptrs) with per-expert weight pointers" +) +_JAX_BLOCK_SCALED_ERROR = ( + "The block-scaled grouped GEMM dGLU backend is not expressible as JAX arrays " + "(its scale-factor tensors use an MMA-interleaved layout with no row-major equivalent); " + "only the BF16 backend supports JAX inputs" +) def _block_scaled_dtype_pairs(): - global _BLOCK_SCALED_DTYPE_PAIRS - if _BLOCK_SCALED_DTYPE_PAIRS is None: - import torch - - _BLOCK_SCALED_DTYPE_PAIRS = { - (dtype, dtype) - for dtype in ( - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ) - } - return _BLOCK_SCALED_DTYPE_PAIRS + # Canonical (cutlass) dtype vocabulary; select_grouped_gemm_backend canonicalizes + # the caller's dtypes so torch/jax/numpy/str dtypes all compare against these. + return { + (dtype, dtype) + for dtype in ( + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + ) + } from ._bf16_api import GroupedGemmDgluBf16API @@ -188,14 +206,11 @@ def __init__( self._pending_init_kwargs = dict(locals()) self._pending_init_kwargs.pop("self") self._pending_init_kwargs.pop("__class__", None) - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmDgluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if sample_a is not None and framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmDgluSm100; pass torch tensors or JAX arrays") if acc_dtype is None: - import torch - - self._pending_init_kwargs["acc_dtype"] = torch.float32 + self._pending_init_kwargs["acc_dtype"] = cutlass.Float32 self._implementation = None def check_support(self) -> bool: @@ -256,7 +271,14 @@ def check_support(self) -> bool: use_dynamic_sched=kwargs["use_dynamic_sched"], ) else: - self._implementation = GroupedGemmDgluBlockScaledAPI(**kwargs) + if detect_framework(kwargs["sample_a"]) == "jax": + raise ValueError(_JAX_BLOCK_SCALED_ERROR) + block_kwargs = dict(kwargs) + # The block-scaled implementation is torch-native: hand it torch dtypes. + block_kwargs["acc_dtype"] = framework_dtype(block_kwargs["acc_dtype"], "torch") + if block_kwargs.get("b_dtype") is not None: + block_kwargs["b_dtype"] = framework_dtype(block_kwargs["b_dtype"], "torch") + self._implementation = GroupedGemmDgluBlockScaledAPI(**block_kwargs) self._kernel = self._implementation._kernel self.weight_mode = self._implementation.weight_mode supported = self._implementation.check_support() @@ -492,12 +514,14 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: b_ptrs = call.b_ptrs sfb_ptrs = call.sfb_ptrs n = call.n - b_dtype = call.b_dtype b_major = call.b_major norm_const_tensor = call.norm_const_tensor - acc_dtype = call.acc_dtype - d_dtype = call.d_dtype cd_major = call.cd_major + # The block-scaled path is torch-native (torch-only allocations and kernels); + # the normalized call carries canonical (cutlass) dtypes, so map them back. + acc_dtype = framework_dtype(call.acc_dtype, "torch") + d_dtype = framework_dtype(call.d_dtype, "torch") + b_dtype = framework_dtype(call.b_dtype, "torch") if call.b_dtype is not None else None mma_tiler_mn = call.mma_tiler_mn cluster_shape_mn = call.cluster_shape_mn sf_vec_size = call.sf_vec_size @@ -864,16 +888,14 @@ def dynamic_m_tensor_signature( def _normalize_dglu_call( call: DgluCall, ) -> tuple[DgluCall, GroupedGemmBackend]: - import torch - - from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor + from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor - if call.acc_dtype is None or call.d_dtype is None: - call = replace( - call, - acc_dtype=call.acc_dtype if call.acc_dtype is not None else torch.float32, - d_dtype=call.d_dtype if call.d_dtype is not None else torch.bfloat16, - ) + call = replace( + call, + acc_dtype=_convert_to_cutlass_data_type(call.acc_dtype) if call.acc_dtype is not None else cutlass.Float32, + d_dtype=_convert_to_cutlass_data_type(call.d_dtype) if call.d_dtype is not None else cutlass.BFloat16, + b_dtype=_convert_to_cutlass_data_type(call.b_dtype) if call.b_dtype is not None else None, + ) is_dense = call.b_tensor is not None is_discrete = call.b_ptrs is not None @@ -881,14 +903,16 @@ def _normalize_dglu_call( raise ValueError("Provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs), not both") if not is_dense and not is_discrete: raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") - if call.a_tensor.ndim != 3 or call.a_tensor.shape[2] != 1: - raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(call.a_tensor.shape)}") + a_shape = get_shape(call.a_tensor) + if len(a_shape) != 3 or a_shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {a_shape}") - valid_m, k, _ = call.a_tensor.shape + valid_m, k, _ = a_shape if is_dense: - if call.b_tensor.ndim != 3: - raise ValueError(f"b_tensor must have shape (n, k, experts), got {tuple(call.b_tensor.shape)}") - n_weight, b_k, num_experts = call.b_tensor.shape + b_full_shape = get_shape(call.b_tensor) + if len(b_full_shape) != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got {b_full_shape}") + n_weight, b_k, num_experts = b_full_shape if b_k != k: raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") defining_b_dtype = call.b_tensor.dtype @@ -897,8 +921,7 @@ def _normalize_dglu_call( if call.n is not None or call.b_dtype is not None: raise ValueError("Dense mode forbids n and b_dtype") else: - _require_pointer_tensor(call.b_ptrs, "b_ptrs") - num_experts = call.b_ptrs.numel() + num_experts = _validate_pointer_tensor(call.b_ptrs, "b_ptrs") if call.n is None or call.b_dtype is None: raise ValueError("n and b_dtype are required for discrete mode") n_weight = call.n @@ -958,7 +981,7 @@ def _normalize_dglu_call( raise ValueError(f"cd_major must be 'n', got {call.cd_major}") if call.act_func not in ("dswiglu", "dgeglu"): raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {call.act_func}") - if call.d_dtype not in (torch.bfloat16, torch.float16, torch.float32): + if call.d_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {call.d_dtype}") if call.m_aligned != 256: raise ValueError(f"m_aligned must be 256, got {call.m_aligned}") @@ -967,30 +990,29 @@ def _normalize_dglu_call( if n_weight <= 0 or n_weight % 32 != 0: raise ValueError(f"N must be positive and divisible by 32, got {n_weight}") two_n = 2 * n_weight - if tuple(call.c_tensor.shape) != (valid_m, two_n, 1): - raise ValueError(f"c_tensor must have shape {(valid_m, two_n, 1)}, got {tuple(call.c_tensor.shape)}") - if tuple(call.prob_tensor.shape) != (valid_m, 1, 1): - raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got {tuple(call.prob_tensor.shape)}") - if tuple(call.dprob_tensor.shape) != (valid_m, 1, 1): - raise ValueError(f"dprob_tensor must have shape {(valid_m, 1, 1)}, got {tuple(call.dprob_tensor.shape)}") - if call.dprob_tensor.dtype != torch.float32: + if get_shape(call.c_tensor) != (valid_m, two_n, 1): + raise ValueError(f"c_tensor must have shape {(valid_m, two_n, 1)}, got {get_shape(call.c_tensor)}") + if get_shape(call.prob_tensor) != (valid_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got {get_shape(call.prob_tensor)}") + if get_shape(call.dprob_tensor) != (valid_m, 1, 1): + raise ValueError(f"dprob_tensor must have shape {(valid_m, 1, 1)}, got {get_shape(call.dprob_tensor)}") + if _convert_to_cutlass_data_type(call.dprob_tensor.dtype) is not cutlass.Float32: raise ValueError(f"dprob_tensor must have dtype torch.float32, got {call.dprob_tensor.dtype}") - if is_discrete and call.b_ptrs.numel() != call.padded_offsets.numel(): - raise ValueError(f"b_ptrs length mismatch: expected {call.padded_offsets.numel()}, " f"got {call.b_ptrs.numel()}") - if tuple(call.padded_offsets.shape) != (num_experts,): - raise ValueError(f"padded_offsets length mismatch: expected {num_experts}, got {call.padded_offsets.numel()}") - if tuple(call.alpha_tensor.shape) != (num_experts,): - raise ValueError(f"alpha_tensor must have shape {(num_experts,)}, got {tuple(call.alpha_tensor.shape)}") - if tuple(call.beta_tensor.shape) != (num_experts,): - raise ValueError(f"beta_tensor must have shape {(num_experts,)}, got {tuple(call.beta_tensor.shape)}") + offsets_shape = get_shape(call.padded_offsets) + if is_discrete and len(offsets_shape) == 1 and num_experts != offsets_shape[0]: + raise ValueError(f"b_ptrs length mismatch: expected {offsets_shape[0]}, " f"got {num_experts}") + if offsets_shape != (num_experts,): + raise ValueError(f"padded_offsets length mismatch: expected {num_experts}, got {offsets_shape}") + if get_shape(call.alpha_tensor) != (num_experts,): + raise ValueError(f"alpha_tensor must have shape {(num_experts,)}, got {get_shape(call.alpha_tensor)}") + if get_shape(call.beta_tensor) != (num_experts,): + raise ValueError(f"beta_tensor must have shape {(num_experts,)}, got {get_shape(call.beta_tensor)}") if is_discrete: - if call.b_ptrs.numel() != num_experts: - raise ValueError(f"b_ptrs length mismatch: expected {num_experts}, got {call.b_ptrs.numel()}") - if call.b_ptrs.device != call.a_tensor.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor ({call.a_tensor.device}), " f"got {call.b_ptrs.device}") - if not torch.cuda.is_available(): + if get_device(call.b_ptrs) != get_device(call.a_tensor): + raise ValueError(f"b_ptrs must be on the same device as a_tensor ({get_device(call.a_tensor)}), " f"got {get_device(call.b_ptrs)}") + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(call.a_tensor.device) + major, minor = get_compute_capability() capability = major * 10 + minor if capability < 100: raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability}") @@ -998,11 +1020,13 @@ def _normalize_dglu_call( def _dglu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + strides = get_strides(tensor) + shape = get_shape(tensor) return tuple( index for index, _ in sorted( - enumerate(tensor.stride()), - key=lambda item: (item[1], tensor.shape[item[0]]), + enumerate(strides), + key=lambda item: (item[1], shape[item[0]]), ) ) @@ -1010,34 +1034,50 @@ def _dglu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: def _dglu_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: if tensor is None: return (None, None, None, None) - shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + device = get_device(tensor) + shape = (None, *get_shape(tensor)[1:]) if dynamic_m else get_shape(tensor) return ( shape, _dglu_stride_order(tensor), - tensor.dtype, - (tensor.device.type, tensor.device.index), + _convert_to_cutlass_data_type(tensor.dtype), + (device.type, device.index), ) def _grouped_gemm_dglu_bf16_call(call: DgluCall) -> TupleDict: - import torch - - valid_m = call.a_tensor.shape[0] - n_weight = call.b_tensor.shape[0] if call.b_tensor is not None else call.n + framework = detect_framework(call.a_tensor) + valid_m = get_shape(call.a_tensor)[0] + n_weight = get_shape(call.b_tensor)[0] if call.b_tensor is not None else call.n two_n = 2 * n_weight - with _torch_stream_context(call.current_stream, call.a_tensor.device): - d_row_tensor = torch.empty_strided( - (valid_m, two_n, 1), - (two_n, 1, valid_m * two_n), - dtype=call.d_dtype, - device=call.a_tensor.device, - ) - dbias_tensor = ( - torch.zeros( - (call.num_experts, two_n, 1), - dtype=torch.bfloat16, + if framework == "torch": + import torch + + # _torch_stream_context is torch-only; other frameworks never enter it. + with _torch_stream_context(call.current_stream, call.a_tensor.device): + d_row_tensor = torch.empty_strided( + (valid_m, two_n, 1), + (two_n, 1, valid_m * two_n), + dtype=framework_dtype(call.d_dtype, "torch"), device=call.a_tensor.device, ) + dbias_tensor = ( + torch.zeros( + (call.num_experts, two_n, 1), + dtype=torch.bfloat16, + device=call.a_tensor.device, + ) + if call.generate_dbias + else None + ) + else: + import jax + import jax.numpy as jnp + + # n-major C-contiguous outputs; the extent-1 batch dim's stride is unobservable. + # The kernel writes into these buffers on the launch stream; materialize them first. + d_row_tensor = jax.block_until_ready(jnp.empty((valid_m, two_n, 1), dtype=framework_dtype(call.d_dtype, "jax"), device=call.a_tensor.device)) + dbias_tensor = ( + jax.block_until_ready(jnp.zeros((call.num_experts, two_n, 1), dtype=framework_dtype(cutlass.BFloat16, "jax"), device=call.a_tensor.device)) if call.generate_dbias else None ) @@ -1060,16 +1100,7 @@ def _grouped_gemm_dglu_bf16_call(call: DgluCall) -> TupleDict: _dglu_tensor_signature(call.prob_tensor, dynamic_m=True), _dglu_tensor_signature(call.dprob_tensor, dynamic_m=True), _dglu_tensor_signature(dbias_tensor), - ( - ( - tuple(call.b_ptrs.shape), - tuple(call.b_ptrs.stride()), - call.b_ptrs.dtype, - (call.b_ptrs.device.type, call.b_ptrs.device.index), - ) - if call.b_ptrs is not None - else None - ), + (_dglu_tensor_signature(call.b_ptrs) if call.b_ptrs is not None else None), call.acc_dtype, call.d_dtype, call.mma_tiler_mn, @@ -1079,7 +1110,7 @@ def _grouped_gemm_dglu_bf16_call(call: DgluCall) -> TupleDict: call.b_major, call.use_dynamic_sched, workspace_bytes, - (call.a_tensor.device.type, call.a_tensor.device.index), + ((get_device(call.a_tensor).type, get_device(call.a_tensor).index)), overlap_margin, ) @@ -1192,16 +1223,14 @@ def grouped_gemm_dglu_wrapper_sm100( current_stream: Optional[cuda.CUstream] = None, ) -> TupleDict: """Dispatch grouped GEMM dGLU once from an immutable normalized call.""" - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_dglu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if d_dtype is None: - d_dtype = torch.bfloat16 + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_dglu_wrapper_sm100; pass torch tensors or JAX arrays") + if framework == "jax" and b_tensor is not None: + raise ValueError(_JAX_DENSE_B_ERROR) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None call = DgluCall( a_tensor=a_tensor, c_tensor=c_tensor, @@ -1242,4 +1271,6 @@ def grouped_gemm_dglu_wrapper_sm100( normalized, backend = _normalize_dglu_call(call) if backend is GroupedGemmBackend.BF16: return _grouped_gemm_dglu_bf16_call(normalized) + if framework == "jax": + raise ValueError(_JAX_BLOCK_SCALED_ERROR) return _grouped_gemm_dglu_block_scaled_call(normalized) diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py new file mode 100644 index 000000000..de69ecc77 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the BF16 SM100 grouped GEMM dGLU +backward (discrete weight mode), built on :func:`cudnn.jax.call`. + +BF16 backend and discrete mode only: dense mode's expert-outermost strided B has +no row-major JAX equivalent, and the block-scaled backend's MMA-interleaved +scale-factor layouts cannot be presented as row-major JAX arrays. The per-expert +weight pointers travel as a regular device array whose *values* are raw addresses +— the referenced weight buffers are not visible to XLA, so the caller must keep +them alive (and unmoved) across every execution of the traced computation. +``dprob`` and ``dbias`` are kernel-accumulated outputs, so unlike the eager +wrapper they are not caller-provided buffers here: both are donated +zero-initialized outputs of the custom call. ``padded_offsets`` values cannot be +host-validated under tracing; malformed offsets are the caller's responsibility +here (the eager wrapper validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import call, gemm_operand_spec, zeros_init +from ..moe_utils import MoEWeightMode +from ..unfused.jax_api import _pointer_count, _prob_spec +from .moe_grouped_gemm_dglu_dbias import MoEGroupedGemmDgluDbiasBf16Kernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_output_dtypes = (cutlass.BFloat16, cutlass.Float16, cutlass.Float32) + +_JAX_BLOCK_SCALED_ERROR = ( + "the block-scaled grouped GEMM dGLU backend is not expressible as JAX arrays " + "(its scale-factor tensors use an MMA-interleaved layout with no row-major equivalent); " + "only the BF16 backend supports JAX inputs" +) + + +@cute.jit +def _grouped_dglu_bf16_adapter(stream, a, c, b_ptrs, padded_offsets, alpha, beta, prob, d, dprob, workspace, *, kernel, n, k, mac, linear_offset): + # Discrete-mode b is a raw pointer to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffer recasts for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d, + padded_offsets=padded_offsets, + alpha=alpha, + beta=beta, + prob=prob, + dprob=dprob, + linear_offset=cutlass.Float32(linear_offset), + dbias_tensor=None, + max_active_clusters=mac, + stream=stream, + ) + + +@cute.jit +def _grouped_dglu_bf16_dbias_adapter(stream, a, c, b_ptrs, padded_offsets, alpha, beta, prob, d, dprob, dbias, workspace, *, kernel, n, k, mac, linear_offset): + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d, + padded_offsets=padded_offsets, + alpha=alpha, + beta=beta, + prob=prob, + dprob=dprob, + linear_offset=cutlass.Float32(linear_offset), + dbias_tensor=dbias, + max_active_clusters=mac, + stream=stream, + ) + + +def grouped_gemm_dglu_jax_sm100( + a_tensor: Any, + c_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + b_ptrs: Any, + n: int, + prob_tensor: Any, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + generate_dbias: bool = False, + use_dynamic_sched: bool = False, +) -> Tuple[Any, Any, Optional[Any]]: + """BF16 grouped GEMM dGLU backward (discrete weights) as an XLA custom call. + + Same contract as the eager wrapper's BF16 discrete mode: A ``(m, k, 1)`` k-major + C-contiguous bfloat16, C ``(m, 2n, 1)`` n-major forward pre-activations, + ``padded_offsets (experts,)`` int32 cumulative 256-aligned row offsets, ``alpha`` + and ``beta`` ``(experts,)`` float32, ``prob (m, 1, 1)`` float32, and ``b_ptrs`` + holding per-expert ``(n, k)`` k-major bfloat16 weight base addresses (packed + little-endian uint8, 8 bytes per pointer — or int64 with x64 mode). ``n`` is the + per-expert weight N (half the pre-activation width). ``linear_offset`` defaults + per ``act_func`` (1.0 for ``"dgeglu"``, 0.0 for ``"dswiglu"``) and is a + compile-time constant of the traced call. Returns ``(d_row_tensor, dprob_tensor, + dbias_tensor)`` with ``dbias_tensor`` None unless ``generate_dbias``; rows + at/past ``padded_offsets[-1]`` come back zero-filled (the outputs are donated + zero-initialized buffers, matching the eager contract of a caller-zeroed + ``dprob``). + """ + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + if _convert_to_cutlass_data_type(a_tensor.dtype) is not cutlass.BFloat16: + raise ValueError(f"a_tensor must have dtype bfloat16, got {a_tensor.dtype}; " + _JAX_BLOCK_SCALED_ERROR) + if n is None or n <= 0 or n % 32 != 0: + raise ValueError(f"n must be positive and divisible by 32, got {n}") + two_n = 2 * n + if tuple(c_tensor.shape) != (m, two_n, 1): + raise ValueError(f"c_tensor must have shape ({m}, {two_n}, 1), got {tuple(c_tensor.shape)}") + c_dtype = _convert_to_cutlass_data_type(c_tensor.dtype) + if c_dtype not in _output_dtypes or d_dtype not in _output_dtypes: + raise ValueError(f"c_tensor/d_dtype must be BF16, FP16, or FP32, got {c_dtype}/{d_dtype}; " + _JAX_BLOCK_SCALED_ERROR) + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + if act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {act_func}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "dgeglu" else 0.0 + + expert_cnt = _pointer_count(b_ptrs) + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(beta_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(beta_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"beta_tensor must be ({expert_cnt},) float32, got {tuple(beta_tensor.shape)} {beta_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + + if not MoEGroupedGemmDgluDbiasBf16Kernel.can_implement( + cutlass.BFloat16, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + "k", + "n", + MoEGroupedGemmDgluDbiasBf16Kernel.FIX_PAD_SIZE, + act_func, + ): + raise ValueError("Unsupported BF16 grouped GEMM dGLU tile, cluster, alignment, or layout configuration") + + cache_key = ( + expert_cnt, + c_dtype, + d_dtype, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + vector_f32, + act_func, + use_dynamic_sched, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = MoEGroupedGemmDgluDbiasBf16Kernel( + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DISCRETE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + operand = gemm_operand_spec() + prob_spec = _prob_spec() + output_shape_dtype = [ + jax.ShapeDtypeStruct((m, two_n, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, 1, 1), jnp.float32), # dprob (kernel-accumulated) + ] + output_spec = [operand, prob_spec] + if generate_dbias: + output_shape_dtype.append(jax.ShapeDtypeStruct((expert_cnt, two_n, 1), framework_dtype(cutlass.BFloat16, "jax"))) + output_spec.append(operand) + output_shape_dtype.append(jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8)) + output_spec.append(None) + + results = call( + _grouped_dglu_bf16_dbias_adapter if generate_dbias else _grouped_dglu_bf16_adapter, + output_shape_dtype=tuple(output_shape_dtype), + input_spec=(operand, operand, None, None, None, None, prob_spec), + output_spec=tuple(output_spec), + # All outputs donated: d for the trailing-unit-dim layout spec (and defined + # bytes past the last offset); dprob/dbias because the kernel accumulates + # into them (atomic add) and expects zeroed buffers; the workspace because + # the helper kernel writes the per-expert TMA descriptors into it (XLA + # inputs are immutable). + initialized_outputs={index: zeros_init for index in range(len(output_shape_dtype))}, + kernel=kernel, + n=int(n), + k=int(k), + mac=mac, + linear_offset=float(linear_offset), + )(a_tensor, c_tensor, b_ptrs, padded_offsets, alpha_tensor, beta_tensor, prob_tensor) + + if generate_dbias: + d_row_tensor, dprob_tensor, dbias_tensor, _workspace = results + return d_row_tensor, dprob_tensor, dbias_tensor + d_row_tensor, dprob_tensor, _workspace = results + return d_row_tensor, dprob_tensor, None diff --git a/python/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.py b/python/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.py index 42652b1f2..fbb9016ae 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "GroupedGemmDsreluSm100", "grouped_gemm_dsrelu_wrapper_sm100", + "grouped_gemm_dsrelu_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "grouped_gemm_dsrelu_jax_sm100": + from .jax_api import grouped_gemm_dsrelu_jax_sm100 + + return grouped_gemm_dsrelu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py b/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py index 0ed234589..1ab5d423a 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py @@ -34,14 +34,26 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.runtime import from_dlpack, make_fake_stream -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.datatypes import _convert_to_cutlass_data_type, _convert_to_cutlass_data_type_or_none from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, + get_shape, + get_strides, + is_torch_tensor, +) def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: - import torch - - if tensor.dtype == torch.uint8: + if _convert_to_cutlass_data_type_or_none(tensor.dtype) is cutlass.Uint8: cute_tensor = from_dlpack(tensor, assumed_align=16, enable_tvm_ffi=True).mark_layout_dynamic(leading_dim=1) cute_tensor.element_type = cutlass.Float4E2M1FN return cute_tensor @@ -158,15 +170,27 @@ def __init__( :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing :param use_dsrelu_reuse: Reuse relu(C)^2 between d_srelu and dprob """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmDsreluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if sample_a is not None and framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmDsreluSm100; pass torch tensors or JAX arrays") + if framework == "jax": + if sample_b is not None: + raise ValueError( + "Dense weight mode (sample_b/sample_sfb) is not expressible as JAX arrays " + "(the expert-outermost strided B layout (n, k, l) has no row-major equivalent); " + "use discrete mode (num_experts, b_shape, b_dtype) with per-expert weight pointers" + ) + if _convert_to_cutlass_data_type_or_none(getattr(sample_a, "dtype", None)) in (cutlass.Float4E2M1FN, cutlass.Uint8) or ( + b_dtype is not None and _convert_to_cutlass_data_type_or_none(b_dtype) in (cutlass.Float4E2M1FN, cutlass.Uint8) + ): + raise ValueError( + "Packed fp4 A/B tensors (float4_e2m1fn / raw uint8) are not expressible as JAX arrays " + "(JAX has no packed fp4 dtype); use fp8 inputs from JAX, or torch tensors for fp4" + ) if acc_dtype is None: - import torch - - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() + self._framework = framework self._warn_experimental_api() self._logger.debug("Entering __init__") @@ -187,38 +211,38 @@ def __init__( self._sample_b_tensor = sample_b # ---- Common tensor descriptors ---- - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", interpret_uint8_as_fp4x2=False) - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - self.d_srelu_desc = self._make_tensor_desc(sample_d_srelu, name="sample_d_srelu") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") - self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") - - 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") - self.sfd_col_d_srelu_desc = self._make_tensor_desc(sample_sfd_col_d_srelu, name="sample_sfd_col_d_srelu") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", interpret_uint8_as_fp4x2=False, canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row", canonical=True) + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) + self.d_srelu_desc = self._make_tensor_desc(sample_d_srelu, name="sample_d_srelu", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob", canonical=True) + self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias", canonical=True) + + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.sfd_col_d_srelu_desc = self._make_tensor_desc(sample_sfd_col_d_srelu, name="sample_sfd_col_d_srelu", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) # ---- Mode-specific state ---- if self.weight_mode == MoEWeightMode.DENSE: - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", interpret_uint8_as_fp4x2=False) - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", interpret_uint8_as_fp4x2=False, canonical=True) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) self.expert_cnt = self.padded_offsets_desc.shape[0] else: self._value_error_if(num_experts == 0, "num_experts must be > 0") self.expert_cnt = num_experts self.b_shape = b_shape - self.b_dtype = b_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None self.b_major = b_major self._value_error_if( self.padded_offsets_desc.shape[0] != self.expert_cnt, @@ -226,7 +250,7 @@ def __init__( ) # ---- Configuration ---- - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -252,16 +276,63 @@ def __init__( self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") self._workspace = None + self._live_ptrs = None + self._compile_b_ptrs = None + self._compile_sfb_ptrs = None self._logger.debug("__init__ completed") + @staticmethod + def _sf_desc_is_physical(sf_desc) -> bool: + """True if the SF descriptor is the physical C-contiguous (L, MN', K', 32, 4, 4) + allocation rather than the torch-style permuted (32, 4, MN', 4, K', L) atom view. + + Frameworks that cannot express the permuted strided view (e.g. JAX, whose arrays + are row-major) pass the physical form. The kernel rebuilds the SF layout from the + GEMM shapes via ``tile_atom_to_shape_SF`` and consumes only the SF base pointer, + so the two forms are byte-identical. + """ + shape = sf_desc.shape + return not (len(shape) == 6 and shape[0] == 32 and shape[1] == 4 and shape[3] == 4) + + def _check_sf_shape(self, sf_desc, mn128: int, rest: int, l: int, name: str) -> None: + """Validate an SF tensor shape and strides, accepting the permuted atom view or (in + discrete weight mode) the physical C-contiguous form -- see ``_sf_desc_is_physical``. + + The kernel consumes only the SF base pointer and rebuilds the layout from the GEMM + shapes, so both forms must be exactly the C-contiguous physical allocation in memory: + strides are validated too (a shape-matching but differently-strided tensor would + silently produce wrong results). + """ + if sf_desc is None: + return + self._value_error_if(len(sf_desc.shape) != 6, f"{name} tensor must be 6-D, got shape {sf_desc.shape}") + if not self._sf_desc_is_physical(sf_desc): + self._check_tensor_shape(sf_desc, (32, 4, mn128, 4, rest, l), name) + _ = self._check_tensor_stride( + sf_desc, + stride=[canonicalize_unit_dim_strides((32, 4, mn128, 4, rest, l), (16, 4, rest * 512, 1, 512, mn128 * rest * 512))], + name=name, + extra_error_msg=f"{name} atom view must be the (3, 4, 1, 5, 2, 0) permutation of a C-contiguous (L, MN', K', 32, 4, 4) allocation", + ) + return + self._value_error_if( + self.weight_mode != MoEWeightMode.DISCRETE, + f"{name} physical (L, MN', K', 32, 4, 4) form is only supported in discrete weight mode; " "provide the permuted (32, 4, MN', 4, K', L) atom view", + ) + self._check_tensor_shape(sf_desc, (l, mn128, rest, 32, 4, 4), name) + _ = self._check_tensor_stride( + sf_desc, + stride=[canonicalize_unit_dim_strides((l, mn128, rest, 32, 4, 4), (mn128 * rest * 512, rest * 512, 512, 16, 4, 1))], + name=name, + extra_error_msg=f"{name} in the physical (L, MN', K', 32, 4, 4) form must be C-contiguous", + ) + def check_support(self) -> bool: """Check if the kernel configuration is supported. :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") # ---- SFD group validation ---- @@ -299,19 +370,15 @@ def check_support(self) -> bool: self._check_tensor_shape(self.d_srelu_desc, (tensor_m, n_out, 1), "D_srelu") rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") + self._check_sf_shape(self.sfa_desc, ceil_div(tensor_m, 128), rest_k, 1, "SFA") if self.weight_mode == MoEWeightMode.DENSE: self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") rest_n_out = ceil_div(ceil_div(n_out, self.sf_vec_size), 4) - self._check_tensor_shape( - self.sfd_row_desc, - (32, 4, ceil_div(tensor_m, 128), 4, rest_n_out, 1), - "SFD_row", - ) + self._check_sf_shape(self.sfd_row_desc, ceil_div(tensor_m, 128), rest_n_out, 1, "SFD_row") rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n_out, 128), 4, rest_m, 1), "SFD_col") - self._check_tensor_shape(self.sfd_col_d_srelu_desc, (32, 4, ceil_div(n_out, 128), 4, rest_m, 1), "SFD_col_d_srelu") + self._check_sf_shape(self.sfd_col_desc, ceil_div(n_out, 128), rest_m, 1, "SFD_col") + self._check_sf_shape(self.sfd_col_d_srelu_desc, ceil_div(n_out, 128), rest_m, 1, "SFD_col_d_srelu") self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") @@ -361,10 +428,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -383,7 +450,7 @@ def check_support(self) -> bool: self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD", ) if self.weight_mode == MoEWeightMode.DENSE: @@ -411,7 +478,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype is cutlass.Float8E4M3FN and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) self._value_error_if( @@ -421,31 +488,31 @@ def check_support(self) -> bool: self._check_dtype( self.acc_dtype, - dtype=torch.float32, + dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32", ) self._check_dtype( self.prob_desc, - dtype=torch.float32, + dtype=cutlass.Float32, name="Prob", extra_error_msg="Prob must be float32", ) self._check_dtype( self.dprob_desc, - dtype=torch.float32, + dtype=cutlass.Float32, name="Dprob", extra_error_msg="Dprob must be float32", ) self._check_dtype( self.dbias_desc, - dtype=torch.bfloat16, + dtype=cutlass.BFloat16, name="Dbias", extra_error_msg="dbias must be bfloat16", ) self.c_dtype = self._check_dtype( self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], + dtype=[cutlass.Float32, cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2], name="C", ) if self._is_fp8(self.c_dtype) and self.vector_f32: @@ -454,7 +521,7 @@ def check_support(self) -> bool: if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_row_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D_row", extra_error_msg="D_row must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -462,8 +529,8 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_row_desc, dtype=[ - torch.float8_e4m3fn, - torch.float8_e5m2, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, ], name="D_row", extra_error_msg="D_row must be fp8 dtype when ab_dtype is fp8", @@ -484,7 +551,7 @@ def check_support(self) -> bool: ) # ---- SFD generation logic ---- - kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype == torch.float8_e8m0fnu and self._is_fp8(self.d_dtype) + kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype is cutlass.Float8E8M0FNU and self._is_fp8(self.d_dtype) self._value_error_if( kernel_generate_sfd and not self._user_requested_sfd, "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=torch.float8_e8m0fnu", @@ -598,18 +665,17 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): # ---- Disabled configurations ---- self._not_implemented_error_if( - self.dbias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, + self.dbias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype is cutlass.Float32, "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. " "Please use sf_vec_size 32 or d_dtype bf16 instead", ) # ---- SM100+ check ---- - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmDsrelu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") + raise RuntimeError(f"GroupedGemmDsrelu requires SM100+ compute capability, " f"but found SM{compute_capability}") self._is_supported = True self._logger.debug("check_support completed successfully") @@ -617,8 +683,6 @@ def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): def compile(self) -> None: """Compile the kernel.""" - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -658,7 +722,9 @@ def compile(self) -> None: self._use_full_dynamic_mnkl = os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" workspace_bytes = gemm_dsrelu.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) if self.weight_mode == MoEWeightMode.DENSE: self._compile_dense(gemm_dsrelu, max_active_clusters, fake_stream) @@ -669,8 +735,6 @@ def compile(self) -> None: def _compile_dense(self, gemm_dsrelu, max_active_clusters, fake_stream) -> None: """Compile for dense (contiguous) weight mode.""" - import torch - self._logger.debug("Compiling grouped_gemm_dsrelu kernel") use_full_dynamic = self._use_full_dynamic_mnkl @@ -881,8 +945,8 @@ def _compile_dense(self, gemm_dsrelu, max_active_clusters, fake_stream) -> None: _compiled_kernel = cute.compile( gemm_dsrelu, - a=_reinterpret_raw_grouped_fp4_tensor(self._sample_a_tensor) if self.a_desc.dtype == torch.uint8 else a_cute_fake, - b=_reinterpret_raw_grouped_fp4_tensor(self._sample_b_tensor) if self.b_desc.dtype == torch.uint8 else b_cute_fake, + a=_reinterpret_raw_grouped_fp4_tensor(self._sample_a_tensor) if self.a_desc.dtype is cutlass.Uint8 else a_cute_fake, + b=_reinterpret_raw_grouped_fp4_tensor(self._sample_b_tensor) if self.b_desc.dtype is cutlass.Uint8 else b_cute_fake, sfb=sfb_cute_fake, n=cutlass.Int32(0), k=cutlass.Int32(0), @@ -963,8 +1027,6 @@ def tensor_api( def _compile_discrete(self, gemm_dsrelu, max_active_clusters, fake_stream) -> None: """Compile for discrete (per-expert pointer) weight mode.""" - import torch - if len(self.b_shape) == 2: n, k = self.b_shape else: @@ -1011,44 +1073,82 @@ def _compile_discrete(self, gemm_dsrelu, max_active_clusters, fake_stream) -> No tensor_m_128 = cute.sym_int() stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[5] = stride_tensor_m_128 - sfa_tensor = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - assumed_align=16, - ) + if self._sf_desc_is_physical(self.sfa_desc): + # Physical C-contiguous (1, M', K', 32, 4, 4) form (e.g. JAX): the kernel rebuilds + # the SF layout from the GEMM shapes and consumes only the base pointer. The + # extent-1 L dim's (M-dependent) stride is symbolic; the ABI cannot observe it. + rest_k_sfa = self.sfa_desc.shape[2] + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=(1, tensor_m_128, rest_k_sfa, 32, 4, 4), + stride=(stride_tensor_m_128, rest_k_sfa * 512, 512, 16, 4, 1), + assumed_align=16, + ) + else: + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[5] = stride_tensor_m_128 + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), + assumed_align=16, + ) sfd_row_tensor = None if self.sfd_row_desc is not None: stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), - assumed_align=16, - ) + if self._sf_desc_is_physical(self.sfd_row_desc): + rest_n_sfd = self.sfd_row_desc.shape[2] + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(1, tensor_m_128, rest_n_sfd, 32, 4, 4), + stride=(stride_sfd_m, rest_n_sfd * 512, 512, 16, 4, 1), + assumed_align=16, + ) + else: + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + assumed_align=16, + ) sfd_col_tensor = None sfd_col_d_srelu_tensor = None if self.sfd_col_desc is not None: rest_m = cute.sym_int(divisibility=1) stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), - assumed_align=16, - ) - if self.sfd_col_d_srelu_desc is not None: - sfd_col_d_srelu_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_d_srelu_desc.dtype, - shape=(32, 4, self.sfd_col_d_srelu_desc.shape[2], 4, rest_m, 1), + if self._sf_desc_is_physical(self.sfd_col_desc): + # Physical (1, N', M_rest, 32, 4, 4): both outer strides are M-dependent. + n_out_128 = self.sfd_col_desc.shape[1] + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(1, n_out_128, rest_m, 32, 4, 4), + stride=(stride_sfd_n, stride_rest_m, 512, 16, 4, 1), + assumed_align=16, + ) + if self.sfd_col_d_srelu_desc is not None: + sfd_col_d_srelu_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_d_srelu_desc.dtype, + shape=(1, n_out_128, rest_m, 32, 4, 4), + stride=(stride_sfd_n, stride_rest_m, 512, 16, 4, 1), + assumed_align=16, + ) + else: + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), assumed_align=16, ) + if self.sfd_col_d_srelu_desc is not None: + sfd_col_d_srelu_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_d_srelu_desc.dtype, + shape=(32, 4, self.sfd_col_d_srelu_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + assumed_align=16, + ) amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) @@ -1067,10 +1167,17 @@ def _compile_discrete(self, gemm_dsrelu, max_active_clusters, fake_stream) -> No ) dbias_tensor = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + # Compile-time placeholders for the pointer-array arguments: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + self._compile_sfb_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + b_ptrs_placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + b_ptrs_placeholder.element_type = cutlass.Int64 + b_ptrs_cute = b_ptrs_placeholder.iterator + sfb_ptrs_placeholder = from_dlpack(self._compile_sfb_ptrs, assumed_align=8) + sfb_ptrs_placeholder.element_type = cutlass.Int64 + sfb_ptrs_cute = sfb_ptrs_placeholder.iterator workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator @@ -1136,8 +1243,8 @@ def tensor_api( stream: cuda.CUstream, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + b_ptrs_addr = int(get_data_ptr(b_ptrs_device)) + sfb_ptrs_addr = int(get_data_ptr(sfb_ptrs_device)) _compiled_kernel( a_tensor, @@ -1220,7 +1327,10 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if a_tensor.shape[0] == 0: self._logger.debug("execute: valid_m is zero, skipping kernel execution") @@ -1260,6 +1370,11 @@ def execute( stream=current_stream, ) else: + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # pointer arrays referenced until the next execute so their buffers outlive + # the asynchronous launch. + self._live_ptrs = (b_ptrs, sfb_ptrs) self._compiled_kernel( a_tensor=a_tensor, b_ptrs_device=b_ptrs, @@ -1366,17 +1481,13 @@ def grouped_gemm_dsrelu_wrapper_sm100( TupleDict with keys: d_row_tensor, d_col_tensor, dprob_tensor, dbias_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor """ - from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_dsrelu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_dsrelu_wrapper_sm100; pass torch tensors or JAX arrays") - if acc_dtype is None: - acc_dtype = torch.float32 - if d_dtype is None: - d_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None is_dense = b_tensor is not None is_discrete = b_ptrs is not None @@ -1386,20 +1497,34 @@ def grouped_gemm_dsrelu_wrapper_sm100( if not is_dense and not is_discrete: raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") - valid_m, k_physical, _ = a_tensor.shape + if framework == "jax": + if is_dense: + raise ValueError( + "Dense weight mode (b_tensor/sfb_tensor) is not expressible as JAX arrays " + "(the expert-outermost strided B layout (n, k, l) has no row-major equivalent); " + "use discrete mode (b_ptrs/sfb_ptrs) with per-expert weight pointers" + ) + if _convert_to_cutlass_data_type(a_tensor.dtype) in (cutlass.Float4E2M1FN, cutlass.Uint8) or b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8): + raise ValueError( + "Packed fp4 A/B tensors (float4_e2m1fn / raw uint8) are not expressible as JAX arrays " + "(JAX has no packed fp4 dtype); use fp8 inputs from JAX, or torch tensors for fp4" + ) + if framework == "torch": + import torch + + valid_m, k_physical, _ = get_shape(a_tensor) if is_dense: weight_mode = MoEWeightMode.DENSE n_weight, _, l = b_tensor.shape else: weight_mode = MoEWeightMode.DISCRETE - _require_pointer_tensor(b_ptrs, "b_ptrs") - num_experts = b_ptrs.shape[0] - _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) + num_experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + _validate_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) if n is None or b_dtype is None: raise ValueError("n and b_dtype are required for discrete mode") n_weight = n - k_logical = k_physical * 2 if b_dtype in (torch.float4_e2m1fn_x2, torch.uint8) else k_physical + k_logical = k_physical * 2 if b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) else k_physical b_shape = (n_weight, k_logical) l = num_experts @@ -1407,13 +1532,31 @@ def grouped_gemm_dsrelu_wrapper_sm100( _logger.debug("grouped_gemm_dsrelu_wrapper_sm100: Creating output tensors") - if cd_major == "n": - d_row_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) - d_col_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) - d_srelu_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) - else: + if cd_major != "n": raise ValueError(f"cd_major must be 'n', got {cd_major}") + if framework == "jax": + import jax.numpy as jnp + + def _jax_alloc(builder): + import jax + + # The kernel writes into these buffers on the launch stream, outside XLA's + # tracking; materialize them before their pointers are taken. + return jax.block_until_ready(builder()) + + if framework == "torch": + d_torch_dtype = framework_dtype(d_dtype, "torch") + d_row_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_torch_dtype, device=a_tensor.device) + d_col_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_torch_dtype, device=a_tensor.device) + d_srelu_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_torch_dtype, device=a_tensor.device) + else: + # n-major C-contiguous; the extent-1 batch dim's stride is unobservable by the kernel. + d_jax_dtype = framework_dtype(d_dtype, "jax") + d_row_tensor = _jax_alloc(lambda: jnp.empty((valid_m, n_out, 1), dtype=d_jax_dtype, device=a_tensor.device)) + d_col_tensor = _jax_alloc(lambda: jnp.empty((valid_m, n_out, 1), dtype=d_jax_dtype, device=a_tensor.device)) + d_srelu_tensor = _jax_alloc(lambda: jnp.empty((valid_m, n_out, 1), dtype=d_jax_dtype, device=a_tensor.device)) + sfd_row_tensor = None sfd_col_tensor = None sfd_col_d_srelu_tensor = None @@ -1421,31 +1564,48 @@ def grouped_gemm_dsrelu_wrapper_sm100( dbias_tensor = None if dprob_tensor is None: - dprob_tensor = torch.zeros((valid_m, 1, 1), dtype=torch.float32, device=a_tensor.device) - - if a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [torch.float8_e8m0fnu, torch.float8_e4m3fn]: + if framework == "torch": + dprob_tensor = torch.zeros((valid_m, 1, 1), dtype=torch.float32, device=a_tensor.device) + else: + dprob_tensor = _jax_alloc(lambda: jnp.zeros((valid_m, 1, 1), dtype=jnp.float32, device=a_tensor.device)) + + if _convert_to_cutlass_data_type(a_tensor.dtype) in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ) and _convert_to_cutlass_data_type( + sfa_tensor.dtype + ) in (cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN): _logger.debug("grouped_gemm_dsrelu_wrapper_sm100: Detected fp8 config, constructing sfd tensors") sf_dtype = sfa_tensor.dtype - mma_permute_order = (3, 4, 1, 5, 2, 0) - sf_k_row = ceil_div(n_out, sf_vec_size) mma_shape_row = (1, ceil_div(valid_m, 128), ceil_div(sf_k_row, 4), 32, 4, 4) - sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - sf_k_col = ceil_div(valid_m, sf_vec_size) mma_shape_col = (1, ceil_div(n_out, 128), ceil_div(sf_k_col, 4), 32, 4, 4) - sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - sfd_col_d_srelu_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - - if d_dtype in [torch.bfloat16, torch.float16]: + if framework == "torch": + mma_permute_order = (3, 4, 1, 5, 2, 0) + sfd_row_tensor = torch.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + sfd_col_d_srelu_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) + else: + # Physical C-contiguous atom allocations: JAX cannot express the permuted view; + # the kernel rebuilds the SF layout from the GEMM shapes and consumes only the + # SF base pointer, so the physical form is byte-identical. + sfd_row_tensor = _jax_alloc(lambda: jnp.empty(mma_shape_row, dtype=sf_dtype, device=a_tensor.device)) + sfd_col_tensor = _jax_alloc(lambda: jnp.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device)) + sfd_col_d_srelu_tensor = _jax_alloc(lambda: jnp.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device)) + + if d_dtype in (cutlass.BFloat16, cutlass.Float16): _logger.debug("grouped_gemm_dsrelu_wrapper_sm100: Constructing amax_tensor") - amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) + if framework == "torch": + amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) + else: + amax_tensor = _jax_alloc(lambda: jnp.full((l, 1), float("-inf"), dtype=jnp.float32, device=a_tensor.device)) if generate_dbias: - dbias_tensor = torch.zeros((l, n_out, 1), dtype=torch.bfloat16, device=a_tensor.device) + if framework == "torch": + dbias_tensor = torch.zeros((l, n_out, 1), dtype=torch.bfloat16, device=a_tensor.device) + else: + dbias_tensor = _jax_alloc(lambda: jnp.zeros((l, n_out, 1), dtype=framework_dtype(cutlass.BFloat16, "jax"), device=a_tensor.device)) if valid_m == 0: _logger.debug("grouped_gemm_dsrelu_wrapper_sm100: valid_m is zero, skipping kernel execution") @@ -1463,37 +1623,52 @@ def grouped_gemm_dsrelu_wrapper_sm100( # ---- Build cache key ---- def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(i for i, s in sorted(enumerate(tensor.stride()), key=lambda x: x[1])) + return tuple(i for i, s in sorted(enumerate(get_strides(tensor)), key=lambda x: x[1])) def tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - return tuple(tensor.shape), tuple(tensor.stride()), tensor.dtype + return get_shape(tensor), get_strides(tensor), _convert_to_cutlass_data_type(tensor.dtype) def dynamic_tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - return None, stride_order(tensor), tensor.dtype + return None, stride_order(tensor), _convert_to_cutlass_data_type(tensor.dtype) def dynamic_m_tensor_signature( tensor: Optional[torch.Tensor], static_shape_suffix: Optional[Tuple[int, ...]], dynamic_stride_dims: Tuple[int, ...] = () ) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(tensor.stride())) - return static_shape_suffix, stride_signature, tensor.dtype + stride_signature = tuple(None if i in dynamic_stride_dims else s for i, s in enumerate(get_strides(tensor))) + return static_shape_suffix, stride_signature, _convert_to_cutlass_data_type(tensor.dtype) + + def _sf_is_physical(tensor) -> bool: + shape = get_shape(tensor) + return not (len(shape) == 6 and shape[0] == 32 and shape[1] == 4 and shape[3] == 4) + + def dynamic_m_sf_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: + """M-independent signature of an SF tensor in either atom form (see _sf_desc_is_physical).""" + if tensor is None: + return None, None, None + shape = get_shape(tensor) + if not _sf_is_physical(tensor): + # torch-style permuted view: M' at dim 2, M-dependent stride at dim 5 + return dynamic_m_tensor_signature(tensor, (shape[4], 1), dynamic_stride_dims=(5,)) + # physical C-contiguous form (e.g. JAX): M' at dim 1, M-dependent stride at dim 0 + static_shape = (shape[0], None, *shape[2:]) + return dynamic_m_tensor_signature(tensor, static_shape, dynamic_stride_dims=(0,)) def dynamic_sfd_col_tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - static_shape = ( - tensor.shape[0], - tensor.shape[1], - tensor.shape[2], - tensor.shape[3], - tensor.shape[5], - ) - return dynamic_m_tensor_signature(tensor, static_shape, dynamic_stride_dims=(2, 5)) + shape = get_shape(tensor) + if not _sf_is_physical(tensor): + static_shape = (shape[0], shape[1], shape[2], shape[3], shape[5]) + return dynamic_m_tensor_signature(tensor, static_shape, dynamic_stride_dims=(2, 5)) + # physical C-contiguous form (e.g. JAX): M_rest at dim 2, M-dependent strides at dims 0-1 + static_shape = (shape[0], shape[1], None, *shape[3:]) + return dynamic_m_tensor_signature(tensor, static_shape, dynamic_stride_dims=(0, 1)) use_full_dynamic = is_dense and os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" @@ -1543,11 +1718,11 @@ def dynamic_sfd_col_tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Op else: cache_key = ( weight_mode, - *dynamic_m_tensor_signature(a_tensor, tuple(a_tensor.shape[1:]), dynamic_stride_dims=(2,)), + *dynamic_m_tensor_signature(a_tensor, get_shape(a_tensor)[1:], dynamic_stride_dims=(2,)), b_shape, b_dtype, - *dynamic_m_tensor_signature(c_tensor, tuple(c_tensor.shape[1:]), dynamic_stride_dims=(2,)), - *dynamic_m_tensor_signature(sfa_tensor, (sfa_tensor.shape[4], 1) if sfa_tensor is not None else None, dynamic_stride_dims=(5,)), + *dynamic_m_tensor_signature(c_tensor, get_shape(c_tensor)[1:], dynamic_stride_dims=(2,)), + *dynamic_m_sf_signature(sfa_tensor), *tensor_signature(alpha_tensor), *dynamic_m_tensor_signature(prob_tensor, (1, 1)), *dynamic_m_tensor_signature(dprob_tensor, (1, 1)), @@ -1555,15 +1730,9 @@ def dynamic_sfd_col_tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Op *dynamic_m_tensor_signature(d_srelu_tensor, (n_out, 1), dynamic_stride_dims=(2,)), *dynamic_sfd_col_tensor_signature(sfd_col_d_srelu_tensor), *tensor_signature(norm_const_tensor), - tuple(b_ptrs.shape), - tuple(b_ptrs.stride()), - b_ptrs.dtype, - tuple(sfb_ptrs.shape), - tuple(sfb_ptrs.stride()), - sfb_ptrs.dtype, - tuple(padded_offsets.shape), - tuple(padded_offsets.stride()), - padded_offsets.dtype, + *tensor_signature(b_ptrs), + *tensor_signature(sfb_ptrs), + *tensor_signature(padded_offsets), acc_dtype, d_dtype, cd_major, diff --git a/python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py b/python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py new file mode 100644 index 000000000..4fc9ff9ed --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the grouped GEMM dSReLU backward +(discrete weight mode, FP8/blockscaled), built on :func:`cudnn.jax.call`. + +Discrete mode only (dense mode's expert-outermost strided B/SFB layouts have no +row-major JAX equivalent) and fp8 inputs only (JAX has no packed fp4 dtype) -- +the same contract the eager wrapper enforces for JAX arrays. Scale-factor +tensors travel in the physical C-contiguous atom shape ``(L, MN', K', 32, 4, 4)``: +the kernel rebuilds each SF layout from the GEMM shapes via +``tile_atom_to_shape_SF`` and consumes only the SF base pointers, so the +permuted torch atom view is never needed. The per-expert weight/SFB pointers +travel as regular device arrays whose *values* are raw addresses -- the +referenced buffers are not visible to XLA, so the caller must keep them alive +(and unmoved) across every execution of the traced computation. +``padded_offsets`` values cannot be host-validated under tracing; malformed +offsets are the caller's responsibility here (the eager wrapper validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.api_base import ceil_div +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import TensorSpec, call, gemm_operand_spec, zeros_init +from ..moe_utils import MoEWeightMode +from ..unfused.jax_api import _pointer_count +from .moe_blockscaled_grouped_gemm_dsrelu_quant import BlockScaledMoEGroupedGemmQuantBwdKernel, EpilogueType + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_fp8_dtypes = (cutlass.Float8E4M3FN, cutlass.Float8E5M2) +_c_dtypes = (cutlass.Float32, cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2) + + +def _prob_spec() -> TensorSpec: + # (m, 1, 1) with m innermost: explicit ranks because trailing unit dims make + # leading-dim inference ambiguous + return TensorSpec(layout=(0, 1, 2)) + + +def _sf_physical_spec() -> TensorSpec: + # Physical C-contiguous (L, MN', K', 32, 4, 4) atom form; explicit ranks because + # the extent-1 L dim makes leading-dim inference ambiguous. The kernel rebuilds + # the SF layout from the GEMM shapes and consumes only the base pointer. + return TensorSpec(layout=(5, 4, 3, 2, 1, 0)) + + +def _as_e8m0_array(scale: Any, name: str) -> Any: + """Present a uint8 E8M0-bit-pattern array as float8_e8m0fnu (free bitcast, jit-safe).""" + dtype = _convert_to_cutlass_data_type(scale.dtype) + if dtype is cutlass.Uint8: + import ml_dtypes + + return scale.view(ml_dtypes.float8_e8m0fnu) + if dtype is not cutlass.Float8E8M0FNU: + raise ValueError(f"{name} must be float8_e8m0fnu (or uint8 bit patterns), got {scale.dtype}; fp8 inputs require e8m0 scale factors with sf_vec_size=32") + return scale + + +@cute.jit +def _grouped_dsrelu_adapter( + stream, + a, + c, + sfa, + b_ptrs, + sfb_ptrs, + padded_offsets, + alpha, + prob, + norm_const, + d_row, + d_col, + d_srelu, + sfd_row, + sfd_col, + sfd_col_d_srelu, + dprob, + workspace, + *, + kernel, + n, + k, + b_stride, + b_major_mode, + mac, +): + # Discrete-mode b/sfb are raw pointers to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffers recast for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + sfb_arg = cute.recast_ptr(sfb_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + sfb=sfb_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(b_stride), + b_major_mode=b_major_mode, + workspace_ptr=workspace.iterator, + c=c, + d=d_row, + d_col=d_col, + sfa=sfa, + sfd_row_tensor=sfd_row, + sfd_col_tensor=sfd_col, + amax_tensor=None, # amax is only produced for bf16/fp16 D, unreachable with fp8 inputs + norm_const_tensor=norm_const, + padded_offsets=padded_offsets, + alpha=alpha, + prob=prob, + dprob=dprob, + dbias_tensor=None, + d_srelu=d_srelu, + sfd_col_d_srelu_tensor=sfd_col_d_srelu, + max_active_clusters=mac, + stream=stream, + ) + + +@cute.jit +def _grouped_dsrelu_dbias_adapter( + stream, + a, + c, + sfa, + b_ptrs, + sfb_ptrs, + padded_offsets, + alpha, + prob, + norm_const, + d_row, + d_col, + d_srelu, + sfd_row, + sfd_col, + sfd_col_d_srelu, + dprob, + dbias, + workspace, + *, + kernel, + n, + k, + b_stride, + b_major_mode, + mac, +): + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + sfb_arg = cute.recast_ptr(sfb_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + sfb=sfb_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(b_stride), + b_major_mode=b_major_mode, + workspace_ptr=workspace.iterator, + c=c, + d=d_row, + d_col=d_col, + sfa=sfa, + sfd_row_tensor=sfd_row, + sfd_col_tensor=sfd_col, + amax_tensor=None, + norm_const_tensor=norm_const, + padded_offsets=padded_offsets, + alpha=alpha, + prob=prob, + dprob=dprob, + dbias_tensor=dbias, + d_srelu=d_srelu, + sfd_col_d_srelu_tensor=sfd_col_d_srelu, + max_active_clusters=mac, + stream=stream, + ) + + +def grouped_gemm_dsrelu_jax_sm100( + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + b_ptrs: Any, + sfb_ptrs: Any, + n: int, + norm_const_tensor: Any, + b_dtype: Any = None, + b_major: str = "k", + d_dtype: Any = cutlass.Float8E4M3FN, + acc_dtype: Any = cutlass.Float32, + generate_dbias: bool = False, + 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, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, +) -> Tuple[Any, ...]: + """Grouped GEMM dSReLU backward (discrete FP8 weights) as an XLA custom call. + + Same contract as the eager wrapper's JAX (discrete, fp8) mode: ``a (m, k, 1)`` + k-major C-contiguous fp8, ``c (m, n, 1)`` n-major forward activations, + ``sfa`` in the physical atom shape ``(1, m/128, K', 32, 4, 4)`` as + ``float8_e8m0fnu`` (or uint8 bit patterns), ``padded_offsets (experts,)`` + int32 cumulative 256-aligned row offsets, ``alpha (experts,)`` float32, + ``prob (m, 1, 1)`` float32, ``norm_const (1,)`` float32, and + ``b_ptrs``/``sfb_ptrs`` holding per-expert weight/SFB base addresses (packed + little-endian uint8, 8 bytes per pointer -- or int64 with x64 mode). + + Returns, in the eager wrapper's key order, + ``(d_row, d_col, d_srelu, dprob, dbias, amax, sfd_row, sfd_col, sfd_col_d_srelu)`` + with ``dbias`` None unless ``generate_dbias`` and ``amax`` always None (fp8 + inputs force fp8 D, which produces SFD outputs instead of amax). SF outputs + come back in the physical atom form; rows at/past ``padded_offsets[-1]`` come + back zero-filled (the outputs are donated zero-initialized buffers). + """ + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + ab_dtype = _convert_to_cutlass_data_type(a_tensor.dtype) + if ab_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) or b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8): + raise ValueError( + "Packed fp4 A/B tensors (float4_e2m1fn / raw uint8) are not expressible as JAX arrays " + "(JAX has no packed fp4 dtype); use fp8 inputs from JAX, or torch tensors for fp4" + ) + if ab_dtype not in _fp8_dtypes: + raise ValueError(f"a_tensor must be fp8 (float8_e4m3fn/float8_e5m2), got {a_tensor.dtype}") + if b_dtype is not None and b_dtype is not ab_dtype: + raise ValueError(f"b_dtype ({b_dtype}) must match a_tensor dtype ({ab_dtype})") + if n is None or n <= 0: + raise ValueError(f"n must be > 0, got {n}") + if b_major not in ("k", "n"): + raise ValueError(f"b_major must be 'k' or 'n', got {b_major}") + if d_dtype not in _fp8_dtypes: + raise ValueError(f"d_dtype must be fp8 (float8_e4m3fn/float8_e5m2) when a/b are fp8, got {d_dtype}") + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + if sf_vec_size != 32: + raise ValueError(f"sf_vec_size must be 32 for fp8 inputs, got {sf_vec_size}") + + c_dtype = _convert_to_cutlass_data_type(c_tensor.dtype) + if tuple(c_tensor.shape) != (m, n, 1): + raise ValueError(f"c_tensor must have shape ({m}, {n}, 1), got {tuple(c_tensor.shape)}") + if c_dtype not in _c_dtypes: + raise ValueError(f"c_tensor must be fp32, fp16, bf16, or fp8, got {c_tensor.dtype}") + if c_dtype in _fp8_dtypes and vector_f32: + raise ValueError("Invalid configuration: fp8 c_dtype and vector_f32 is not supported. Please use vector_f32=False or c_dtype=bfloat16 instead") + + sfa_tensor = _as_e8m0_array(sfa_tensor, "sfa_tensor") + rest_k = ceil_div(ceil_div(k, sf_vec_size), 4) + if tuple(sfa_tensor.shape) != (1, ceil_div(m, 128), rest_k, 32, 4, 4): + raise ValueError( + f"sfa_tensor must be the physical atom shape (1, {ceil_div(m, 128)}, {rest_k}, 32, 4, 4) " + f"for (m, k)=({m}, {k}) with sf_vec_size={sf_vec_size}, got {tuple(sfa_tensor.shape)}" + ) + + expert_cnt = _pointer_count(b_ptrs) + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if _pointer_count(sfb_ptrs, "sfb_ptrs") != expert_cnt: + raise ValueError(f"sfb_ptrs must hold {expert_cnt} pointers to match b_ptrs") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required: fp8 inputs with e8m0 scale factors and fp8 D always generate SFD outputs") + if tuple(norm_const_tensor.shape) != (1,) or _convert_to_cutlass_data_type(norm_const_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"norm_const_tensor must be (1,) float32, got {tuple(norm_const_tensor.shape)} {norm_const_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + + if not BlockScaledMoEGroupedGemmQuantBwdKernel.can_implement( + ab_dtype, + cutlass.Float8E8M0FNU, + sf_vec_size, + acc_dtype, + d_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + b_major, + "n", + BlockScaledMoEGroupedGemmQuantBwdKernel.FIX_PAD_SIZE, + ): + raise ValueError("Unsupported grouped GEMM dSReLU tile, cluster, alignment, or layout configuration") + + cache_key = ( + expert_cnt, + ab_dtype, + c_dtype, + d_dtype, + acc_dtype, + b_major, + generate_dbias, + tuple(mma_tiler_mn), + cluster_shape_mn, + sf_vec_size, + vector_f32, + discrete_col_sfd, + use_dynamic_sched, + use_dsrelu_reuse, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = BlockScaledMoEGroupedGemmQuantBwdKernel( + sf_vec_size=sf_vec_size, + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=True, # fp8 inputs + e8m0 SF + fp8 D always generate SFD + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DISCRETE, + 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, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + d_jax_dtype = framework_dtype(d_dtype, "jax") + sf_jax_dtype = framework_dtype(cutlass.Float8E8M0FNU, "jax") + sfd_row_shape = (1, ceil_div(m, 128), ceil_div(ceil_div(n, sf_vec_size), 4), 32, 4, 4) + sfd_col_shape = (1, ceil_div(n, 128), ceil_div(ceil_div(m, sf_vec_size), 4), 32, 4, 4) + + output_shape_dtype = [ + jax.ShapeDtypeStruct((m, n, 1), d_jax_dtype), # d_row + jax.ShapeDtypeStruct((m, n, 1), d_jax_dtype), # d_col + jax.ShapeDtypeStruct((m, n, 1), d_jax_dtype), # d_srelu + jax.ShapeDtypeStruct(sfd_row_shape, sf_jax_dtype), + jax.ShapeDtypeStruct(sfd_col_shape, sf_jax_dtype), + jax.ShapeDtypeStruct(sfd_col_shape, sf_jax_dtype), # sfd_col_d_srelu + jax.ShapeDtypeStruct((m, 1, 1), jnp.float32), # dprob (atomic-add accumulator) + ] + operand = gemm_operand_spec() + sf = _sf_physical_spec() + output_spec = [operand, operand, operand, sf, sf, sf, _prob_spec()] + if generate_dbias: + # dbias (experts, n, 1) with n innermost shares the GEMM-operand stride ranks. + output_shape_dtype.append(jax.ShapeDtypeStruct((expert_cnt, n, 1), framework_dtype(cutlass.BFloat16, "jax"))) + output_spec.append(operand) + output_shape_dtype.append(jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8)) + output_spec.append(None) + + results = call( + _grouped_dsrelu_dbias_adapter if generate_dbias else _grouped_dsrelu_adapter, + output_shape_dtype=tuple(output_shape_dtype), + input_spec=(operand, operand, sf, None, None, None, None, _prob_spec(), None), + output_spec=tuple(output_spec), + # ALL outputs donated: d/sfd for the trailing-/leading-unit-dim layout specs + # (and defined bytes past the last padded offset); dprob and dbias because the + # kernel accumulates them with atomic adds; the workspace because the helper + # kernel writes the per-expert TMA descriptors into it (XLA inputs are immutable). + initialized_outputs={i: zeros_init for i in range(len(output_shape_dtype))}, + kernel=kernel, + n=int(n), + k=int(k), + b_stride=int(k if b_major == "k" else n), + b_major_mode=OperandMajorMode.K if b_major == "k" else OperandMajorMode.MN, + mac=mac, + )(a_tensor, c_tensor, sfa_tensor, b_ptrs, sfb_ptrs, padded_offsets, alpha_tensor, prob_tensor, norm_const_tensor) + + d_row, d_col, d_srelu, sfd_row, sfd_col, sfd_col_d_srelu, dprob = results[:7] + dbias = results[7] if generate_dbias else None + # Eager wrapper key order: (d_row, d_col, d_srelu, dprob, dbias, amax, sfd_row, + # sfd_col, sfd_col_d_srelu); amax is always None (fp8 D produces SFD, not amax). + return d_row, d_col, d_srelu, dprob, dbias, None, sfd_row, sfd_col, sfd_col_d_srelu diff --git a/python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py b/python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py index bc7bbbdcd..752f717c9 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py @@ -21,6 +21,13 @@ from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.tensor_adapter import ( + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, +) class GroupedGemmDswigluSm100(APIBase): @@ -101,41 +108,45 @@ def __init__( :param discrete_col_sfd: Boolean, True to generate discrete col-major scale factor tensor :param epilogue_op: Optional epilogue operation. Valid values: None, "none", "identity", "relu", "srelu" """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmDswigluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if sample_a is not None and framework != "torch": + if framework == "jax": + raise ValueError( + "GroupedGemmDswigluSm100 only supports dense weight mode, whose expert-outermost strided " + "B layout (n, k, l) is not expressible as JAX arrays (row-major only); " + "use torch tensors for this backward API" + ) + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmDswigluSm100; pass torch tensors") if acc_dtype is None: - import torch - - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() + self._framework = framework self._warn_experimental_api() self._logger.debug("Entering __init__") # Store sample tensor descriptors - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - 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.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row", canonical=True) + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob", canonical=True) # Optional quantization outputs - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") - self.norm_const_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), 1, "norm_const") + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) + self.norm_const_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const") # Configuration - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -171,8 +182,6 @@ def check_support(self) -> bool: :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -227,10 +236,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -238,7 +247,7 @@ def check_support(self) -> bool: self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD_row/SFD_col", ) self._check_dtype(self.sfb_desc, dtype=self.sf_dtype, name="SFB", extra_error_msg="SFB must have the same dtype as SFA") @@ -247,17 +256,17 @@ def check_support(self) -> bool: 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 self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32: + if self.sf_dtype is cutlass.Float8E4M3FN and self.sf_vec_size == 32: raise ValueError(f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported") if self._is_fp8(self.ab_dtype) and self.sf_vec_size == 16: raise ValueError(f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported") - self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator", extra_error_msg="Accumulator must be float32") - self._check_dtype(self.prob_desc, dtype=torch.float32, name="Prob", extra_error_msg="Prob must be float32") - self._check_dtype(self.dprob_desc, dtype=torch.float32, name="Dprob", extra_error_msg="Dprob must be float32") + self._check_dtype(self.acc_dtype, dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32") + self._check_dtype(self.prob_desc, dtype=cutlass.Float32, name="Prob", extra_error_msg="Prob must be float32") + self._check_dtype(self.dprob_desc, dtype=cutlass.Float32, name="Dprob", extra_error_msg="Dprob must be float32") self.c_dtype = self._check_dtype( self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], + dtype=[cutlass.Float32, cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2], name="C", ) if self._is_fp8(self.c_dtype) and self.vector_f32: @@ -268,7 +277,7 @@ def check_support(self) -> bool: if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_row_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D_row", extra_error_msg="D_row must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -276,8 +285,8 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_row_desc, dtype=[ - torch.float8_e4m3fn, - torch.float8_e5m2, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, ], name="D_row", extra_error_msg="D_row must be fp8 dtype when ab_dtype is fp8", @@ -338,13 +347,12 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): raise ValueError(f"expert_cnt must be <= 1024, got {self.expert_cnt}") # Check environment - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmDswiglu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") + raise RuntimeError(f"GroupedGemmDswiglu requires SM100+ compute capability, " f"but found SM{compute_capability}") self._is_supported = True self._logger.debug("check_support completed successfully") @@ -644,7 +652,10 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if self._compiled_kernel is None: raise RuntimeError("Kernel not compiled; call compile() first") @@ -742,16 +753,19 @@ def grouped_gemm_dswiglu_wrapper_sm100( - **sfd_row_tensor** (torch.Tensor or None): Row-wise scale factors for D - **sfd_col_tensor** (torch.Tensor or None): Column-wise scale factors for D """ - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_dswiglu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_dswiglu_wrapper_sm100; pass torch tensors or JAX arrays") + if framework == "jax": + raise ValueError( + "grouped_gemm_dswiglu_wrapper_sm100 only supports dense weight mode, whose expert-outermost strided " + "B layout (n, k, l) is not expressible as JAX arrays (row-major only); " + "use torch tensors for this backward API" + ) import torch - if acc_dtype is None: - acc_dtype = torch.float32 - if d_dtype is None: - d_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 valid_m = a_tensor.shape[0] n, _, l = b_tensor.shape @@ -797,13 +811,14 @@ def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: # Allocate M-dependent output tensors fresh every call (M varies across MoE steps). # Only M-independent tensors (amax, beta) are cached to avoid repeated allocation. _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: Allocating M-dependent output tensors") - d_row_tensor = torch.empty_strided((valid_m, n * 2, 1), (n * 2, 1, valid_m * n * 2), dtype=d_dtype, device=a_tensor.device) - d_col_tensor = torch.empty_strided((valid_m, n * 2, 1), (n * 2, 1, valid_m * n * 2), dtype=d_dtype, device=a_tensor.device) + d_torch_dtype = framework_dtype(d_dtype, "torch") + d_row_tensor = torch.empty_strided((valid_m, n * 2, 1), (n * 2, 1, valid_m * n * 2), dtype=d_torch_dtype, device=a_tensor.device) + d_col_tensor = torch.empty_strided((valid_m, n * 2, 1), (n * 2, 1, valid_m * n * 2), dtype=d_torch_dtype, device=a_tensor.device) dprob_tensor = dprob_tensor_buf.zero_() if dprob_tensor_buf is not None else torch.zeros((valid_m, 1, 1), dtype=torch.float32, device=a_tensor.device) if valid_m == 0: amax_tensor = None - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): amax_tensor = torch.full((l, 2, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: valid_m is zero, skipping kernel execution") return TupleDict( @@ -817,7 +832,10 @@ def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: sfd_row_tensor = None sfd_col_tensor = None - if a_tensor.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] and sfa_tensor.dtype in [torch.float8_e8m0fnu, torch.float8_e4m3fn]: + if _convert_to_cutlass_data_type(a_tensor.dtype) in (cutlass.Float8E4M3FN, cutlass.Float8E5M2) and _convert_to_cutlass_data_type(sfa_tensor.dtype) in ( + cutlass.Float8E8M0FNU, + cutlass.Float8E4M3FN, + ): _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: Detected fp8 a_dtype and sfa_dtype, constructing sfd_row_tensor and sfd_col_tensor") sf_dtype = sfa_tensor.dtype mma_permute_order = (3, 4, 1, 5, 2, 0) @@ -852,7 +870,7 @@ def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: cached_amax_tensor = None amax_tensor = None - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): _logger.debug("grouped_gemm_dswiglu_wrapper_sm100: Detected bf16/float16 d_dtype, constructing amax_tensor") cached_amax_tensor = torch.empty((l, 2, 1), dtype=torch.float32, device=a_tensor.device) amax_tensor = amax_tensor_buf if amax_tensor_buf is not None else cached_amax_tensor diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/__init__.py b/python/cudnn/gemm/cutedsl/grouped/glu/__init__.py index e27d2024c..3497f7e6e 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "GroupedGemmGluSm100", "grouped_gemm_glu_wrapper_sm100", + "grouped_gemm_glu_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "grouped_gemm_glu_jax_sm100": + from .jax_api import grouped_gemm_glu_jax_sm100 + + return grouped_gemm_glu_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py index 89feccf0a..dabdf9bbf 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py @@ -17,21 +17,27 @@ from cudnn.api_base import APIBase, TensorDesc from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _pointer_values, _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + get_compute_capability, + get_data_ptr, + get_device, + get_version, + is_torch_tensor, + to_host_list, +) from ..moe_utils import MoEWeightMode from .moe_grouped_gemm_glu_bias import MoEGroupedGemmGluBiasBf16Kernel -_OUTPUT_DTYPES = None - def _output_dtypes(): - global _OUTPUT_DTYPES - if _OUTPUT_DTYPES is None: - import torch - - _OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] - return _OUTPUT_DTYPES + return [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] class GroupedGemmGluBf16API(APIBase): @@ -60,12 +66,11 @@ def __init__( b_major: str = "k", use_dynamic_sched: bool = False, ) -> None: - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() + self._framework = detect_framework(sample_a) if sample_b is not None and num_experts is None: self.weight_mode = MoEWeightMode.DENSE @@ -76,20 +81,20 @@ def __init__( else: raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) self._sample_offsets_ref = weakref.ref(sample_padded_offsets) - self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_offsets_version = get_version(sample_padded_offsets) self._sample_data_ptrs = { - name: tensor.data_ptr() + name: get_data_ptr(tensor) for name, tensor in ( ("sample_a", sample_a), ("sample_b", sample_b), @@ -105,8 +110,8 @@ def __init__( self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) self.b_shape = tuple(b_shape) if b_shape is not None else None - self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype - self.acc_dtype = acc_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = tuple(mma_tiler_mn) self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) @@ -119,7 +124,8 @@ def __init__( self._has_bias = self.bias_desc is not None self._kernel = MoEGroupedGemmGluBiasBf16Kernel self._workspace: Optional[torch.Tensor] = None - self._compile_b_ptrs: Optional[torch.Tensor] = None + self._live_b_ptrs = None + self._compile_b_ptrs = None self._validated_offsets: dict[int, tuple] = {} self._validated_pointer_values: dict[int, tuple] = {} self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) @@ -141,12 +147,12 @@ def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: @staticmethod def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(int(value) for value in tensor.detach().cpu().tolist()) + return tuple(int(value) for value in to_host_list(tensor)) @staticmethod def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: cached = cache.get(id(tensor)) - return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + return bool(cached and cached[0]() is tensor and cached[1] == get_version(tensor) and cached[2] == extra) @staticmethod def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: @@ -155,7 +161,7 @@ def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) - def discard(_reference, *, cache=cache, key=key): cache.pop(key, None) - cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + cache[key] = (weakref.ref(tensor, discard), get_version(tensor), extra) @staticmethod def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None: @@ -182,23 +188,27 @@ def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> Non def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): return - pointer_values = self._copy_values_to_host(b_ptrs) + pointer_values = _pointer_values(b_ptrs) if any(value == 0 or value % 16 != 0 for value in pointer_values): raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) @staticmethod def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: - if tensor.data_ptr() % 16 != 0: + if get_data_ptr(tensor) % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") @staticmethod def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: - if tensor.data_ptr() % 8 != 0: + if get_data_ptr(tensor) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") - @staticmethod - def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + def _record_pointer_stream(self, b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # array referenced until the next execute so its buffer outlives the launch. + self._live_b_ptrs = b_ptrs + return import torch handle = int(current_stream) @@ -213,8 +223,6 @@ def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) b_ptrs.record_stream(launch_stream) def check_support(self) -> bool: - import torch - if self.a_desc.ndim != 3: raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") tensor_m, k, one = self.a_desc.shape @@ -255,17 +263,17 @@ def check_support(self) -> bool: self._expect_stride(self.d_desc, (n_out, 1, tensor_m * n_out), "sample_d") self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") self._expect_stride(self.alpha_desc, (1,), "sample_alpha") - self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") + self._expect_stride(self.prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_prob") - self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + self._check_dtype(self.a_desc, cutlass.BFloat16, "sample_a") if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") - self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.b_desc, cutlass.BFloat16, "sample_b") + self._check_dtype(self.b_dtype, cutlass.BFloat16, "b_dtype") self._check_dtype(self.c_desc, _output_dtypes(), "sample_c") self._check_dtype(self.d_desc, _output_dtypes(), "sample_d") - self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") - self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") - self._check_dtype(self.prob_desc, torch.float32, "sample_prob") + self._check_dtype(self.padded_offsets_desc, cutlass.Int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, cutlass.Float32, "sample_alpha") + self._check_dtype(self.prob_desc, cutlass.Float32, "sample_prob") device = self.a_desc.device for desc, name in ( @@ -289,7 +297,7 @@ def check_support(self) -> bool: if data_ptr % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") - if self.acc_dtype != torch.float32: + if self.acc_dtype is not cutlass.Float32: raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") if self.m_aligned != MoEGroupedGemmGluBiasBf16Kernel.FIX_PAD_SIZE: raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") @@ -308,7 +316,7 @@ def check_support(self) -> bool: tensor_m=tensor_m, ) sample_offsets = self._sample_offsets_ref() - if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + if sample_offsets is not None and get_version(sample_offsets) == self._sample_offsets_version: self._remember_validation( self._validated_offsets, sample_offsets, @@ -318,7 +326,7 @@ def check_support(self) -> bool: self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) if not self._kernel.can_implement( - _convert_to_cutlass_data_type(torch.bfloat16), + cutlass.BFloat16, _convert_to_cutlass_data_type(self.c_desc.dtype), _convert_to_cutlass_data_type(self.d_desc.dtype), _convert_to_cutlass_data_type(self.acc_dtype), @@ -336,19 +344,17 @@ def check_support(self) -> bool: ): raise ValueError("Unsupported BF16 grouped GEMM GLU tile, cluster, alignment, " "or layout configuration") - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(self.a_desc.device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmGluSm100 requires SM100+, found SM{compute_capability} " f"on {self.a_desc.device}") + raise RuntimeError(f"GroupedGemmGluSm100 requires SM100+, found SM{compute_capability}") self._is_supported = True return True def compile(self) -> None: - import torch - self._ensure_support_checked() if self._compiled_kernel is not None: return @@ -373,8 +379,10 @@ def compile(self) -> None: raise ValueError("max_active_clusters must be > 0 after applying " "CUDNNFE_CLUSTER_OVERLAP_MARGIN") workspace_bytes = kernel.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) - if self._workspace.data_ptr() % 128 != 0: + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) + if get_data_ptr(self._workspace) % 128 != 0: raise RuntimeError("workspace allocation must be 128-byte aligned") workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -410,9 +418,14 @@ def compile(self) -> None: b_stride = cutlass.Int64(0) b_major_mode = OperandMajorMode.K else: - self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + # Compile-time placeholder for the pointer-array argument: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) self._validate_pointer_array_alignment(self._compile_b_ptrs) - b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + placeholder.element_type = cutlass.Int64 + b_fake = placeholder.iterator n, k = self.b_shape[:2] n_value = cutlass.Int32(n) k_value = cutlass.Int32(k) @@ -457,7 +470,7 @@ def tensor_api( stream: cuda.CUstream, linear_offset: float, ) -> None: - b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(get_data_ptr(b_ptrs)) raw_compiled( a_tensor, b_arg, @@ -485,7 +498,7 @@ def _validate_live_tensor( *, dynamic_m: bool = False, ) -> TensorDesc: - desc = self._make_tensor_desc(tensor, name=name) + desc = self._make_tensor_desc(tensor, name=name, canonical=True) if desc.dtype != sample.dtype: raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") if desc.device != sample.device: @@ -513,7 +526,10 @@ def execute( linear_offset: float = 0.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if self._compiled_kernel is None: raise RuntimeError("Kernel not compiled; call compile() first") if prob_tensor is None: @@ -537,7 +553,7 @@ def execute( self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") self._expect_stride(c_desc, (n, 1, tensor_m * n), "c_tensor") self._expect_stride(d_desc, (n_out, 1, tensor_m * n_out), "d_tensor") - self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") + self._expect_stride(prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "prob_tensor") self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) for tensor, name in ( @@ -566,10 +582,10 @@ def execute( else: if b_tensor is not None or b_ptrs is None: raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") - _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) - if b_ptrs.device != self.a_desc.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {b_ptrs.device}") - if b_ptrs.data_ptr() % 8 != 0: + _validate_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if get_device(b_ptrs) != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {get_device(b_ptrs)}") + if get_data_ptr(b_ptrs) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") self._validate_pointer_values_once(b_ptrs) self._record_pointer_stream(b_ptrs, current_stream) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py index 2e39685e9..fd9e48c4f 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py @@ -20,7 +20,7 @@ from __future__ import annotations from .moe_blockscaled_grouped_gemm_glu_bias import BlockScaledMoEGroupedGemmGluBiasKernel -from ..backend_utils import _torch_stream_context, rubin_single_group_offsets_kwarg +from ..backend_utils import rubin_single_group_offsets_kwarg from ..moe_utils import MoEWeightMode from cuda.bindings import driver as cuda import os @@ -162,6 +162,13 @@ def __init__( :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ + from cudnn.tensor_adapter import detect_framework + + if sample_a is not None and detect_framework(sample_a) != "torch": + raise ValueError( + "GroupedGemmGluBlockScaledAPI supports torch tensors only: the block-scaled " + "scale-factor tensors use an MMA-interleaved layout that is not expressible as JAX arrays" + ) import torch if acc_dtype is None: diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/api.py b/python/cudnn/gemm/cutedsl/grouped/glu/api.py index af218df6c..f38ecfa66 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/api.py @@ -32,26 +32,47 @@ import os from typing import Any, Tuple, Optional, overload +import cutlass + from cudnn.api_base import APIBase, TupleDict, ceil_div, get_device_type +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + cuda_is_available, + detect_framework, + framework_dtype, + get_compute_capability, + get_device, + get_shape, + get_strides, +) -_BLOCK_SCALED_DTYPE_PAIRS = None +_JAX_DENSE_B_ERROR = ( + "Dense weight mode (b_tensor) is not expressible as JAX arrays " + "(the expert-outermost strided B layout has no row-major equivalent); " + "use discrete mode (b_ptrs) with per-expert weight pointers" +) +_JAX_BIAS_ERROR = ( + "bias_tensor is not expressible as a JAX array (its (n, experts) column-major layout has no row-major equivalent); " "omit bias for JAX inputs" +) +_JAX_BLOCK_SCALED_ERROR = ( + "The block-scaled grouped GEMM GLU backend is not expressible as JAX arrays " + "(its scale-factor tensors use an MMA-interleaved layout with no row-major equivalent); " + "only the BF16 backend supports JAX inputs" +) def _block_scaled_dtype_pairs(): - global _BLOCK_SCALED_DTYPE_PAIRS - if _BLOCK_SCALED_DTYPE_PAIRS is None: - import torch - - _BLOCK_SCALED_DTYPE_PAIRS = { - (dtype, dtype) - for dtype in ( - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ) - } - return _BLOCK_SCALED_DTYPE_PAIRS + # Canonical (cutlass) dtype vocabulary; select_grouped_gemm_backend canonicalizes + # the caller's dtypes so torch/jax/numpy/str dtypes all compare against these. + return { + (dtype, dtype) + for dtype in ( + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + ) + } from ._bf16_api import GroupedGemmGluBf16API @@ -173,14 +194,11 @@ def __init__( self._pending_init_kwargs = dict(locals()) self._pending_init_kwargs.pop("self") self._pending_init_kwargs.pop("__class__", None) - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmGluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if sample_a is not None and framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmGluSm100; pass torch tensors or JAX arrays") if acc_dtype is None: - import torch - - self._pending_init_kwargs["acc_dtype"] = torch.float32 + self._pending_init_kwargs["acc_dtype"] = cutlass.Float32 self._implementation = None def check_support(self) -> bool: @@ -233,8 +251,14 @@ def check_support(self) -> bool: use_dynamic_sched=kwargs["use_dynamic_sched"], ) else: + if detect_framework(kwargs["sample_a"]) == "jax": + raise ValueError(_JAX_BLOCK_SCALED_ERROR) block_kwargs = dict(kwargs) block_kwargs.pop("generate_c", None) + # The block-scaled implementation is torch-native: hand it torch dtypes. + block_kwargs["acc_dtype"] = framework_dtype(block_kwargs["acc_dtype"], "torch") + if block_kwargs.get("b_dtype") is not None: + block_kwargs["b_dtype"] = framework_dtype(block_kwargs["b_dtype"], "torch") self._implementation = GroupedGemmGluBlockScaledAPI(**block_kwargs) self._kernel = self._implementation._kernel self.weight_mode = self._implementation.weight_mode @@ -478,14 +502,16 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: b_ptrs = call.b_ptrs sfb_ptrs = call.sfb_ptrs n = call.n - b_dtype = call.b_dtype b_major = call.b_major norm_const_tensor = call.norm_const_tensor prob_tensor = call.prob_tensor - acc_dtype = call.acc_dtype - c_dtype = call.c_dtype - d_dtype = call.d_dtype cd_major = call.cd_major + # The block-scaled path is torch-native (torch-only allocations and kernels); + # the normalized call carries canonical (cutlass) dtypes, so map them back. + acc_dtype = framework_dtype(call.acc_dtype, "torch") + c_dtype = framework_dtype(call.c_dtype, "torch") + d_dtype = framework_dtype(call.d_dtype, "torch") + b_dtype = framework_dtype(call.b_dtype, "torch") if call.b_dtype is not None else None mma_tiler_mn = call.mma_tiler_mn cluster_shape_mn = call.cluster_shape_mn sf_vec_size = call.sf_vec_size @@ -827,17 +853,15 @@ def dynamic_m_tensor_signature( def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: - import torch - - from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor + from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor - if call.acc_dtype is None or call.c_dtype is None or call.d_dtype is None: - call = replace( - call, - acc_dtype=call.acc_dtype if call.acc_dtype is not None else torch.float32, - c_dtype=call.c_dtype if call.c_dtype is not None else torch.bfloat16, - d_dtype=call.d_dtype if call.d_dtype is not None else torch.bfloat16, - ) + call = replace( + call, + acc_dtype=_convert_to_cutlass_data_type(call.acc_dtype) if call.acc_dtype is not None else cutlass.Float32, + c_dtype=_convert_to_cutlass_data_type(call.c_dtype) if call.c_dtype is not None else cutlass.BFloat16, + d_dtype=_convert_to_cutlass_data_type(call.d_dtype) if call.d_dtype is not None else cutlass.BFloat16, + b_dtype=_convert_to_cutlass_data_type(call.b_dtype) if call.b_dtype is not None else None, + ) is_dense = call.b_tensor is not None is_discrete = call.b_ptrs is not None @@ -845,14 +869,16 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: raise ValueError("Provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs), not both") if not is_dense and not is_discrete: raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") - if call.a_tensor.ndim != 3 or call.a_tensor.shape[2] != 1: - raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(call.a_tensor.shape)}") + a_shape = get_shape(call.a_tensor) + if len(a_shape) != 3 or a_shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {a_shape}") - valid_m, k, _ = call.a_tensor.shape + valid_m, k, _ = a_shape if is_dense: - if call.b_tensor.ndim != 3: - raise ValueError(f"b_tensor must have shape (n, k, experts), got " f"{tuple(call.b_tensor.shape)}") - n_full, b_k, num_experts = call.b_tensor.shape + b_full_shape = get_shape(call.b_tensor) + if len(b_full_shape) != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got " f"{b_full_shape}") + n_full, b_k, num_experts = b_full_shape if b_k != k: raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") defining_b_dtype = call.b_tensor.dtype @@ -861,8 +887,7 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: if call.n is not None or call.b_dtype is not None: raise ValueError("Dense mode forbids n and b_dtype") else: - _require_pointer_tensor(call.b_ptrs, "b_ptrs") - num_experts = call.b_ptrs.numel() + num_experts = _validate_pointer_tensor(call.b_ptrs, "b_ptrs") if call.n is None or call.b_dtype is None: raise ValueError("n and b_dtype are required for discrete mode") n_full = call.n @@ -922,31 +947,32 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: raise ValueError(f"cd_major must be 'n', got {call.cd_major}") if call.act_func not in ("swiglu", "geglu"): raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {call.act_func}") - if call.c_dtype not in (torch.bfloat16, torch.float16, torch.float32): - raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {call.c_dtype}") - if call.d_dtype not in (torch.bfloat16, torch.float16, torch.float32): - raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {call.d_dtype}") + if normalized.c_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): + raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {normalized.c_dtype}") + if normalized.d_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): + raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {normalized.d_dtype}") if call.m_aligned != 256: raise ValueError(f"m_aligned must be 256, got {call.m_aligned}") if valid_m % 256 != 0: raise ValueError(f"a_tensor M dimension must be 256-aligned, got {valid_m}") if n_full <= 0 or n_full % 64 != 0: raise ValueError(f"N must be positive and divisible by 64, got {n_full}") - if tuple(call.prob_tensor.shape) != (valid_m, 1, 1): - raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got " f"{tuple(call.prob_tensor.shape)}") - if call.bias_tensor is not None and tuple(call.bias_tensor.shape) != ( + if get_shape(call.prob_tensor) != (valid_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got " f"{get_shape(call.prob_tensor)}") + if call.bias_tensor is not None and get_shape(call.bias_tensor) != ( n_full, num_experts, ): - raise ValueError(f"bias_tensor must have shape {(n_full, num_experts)}, got " f"{tuple(call.bias_tensor.shape)}") + raise ValueError(f"bias_tensor must have shape {(n_full, num_experts)}, got " f"{get_shape(call.bias_tensor)}") if is_discrete: - if call.b_ptrs.device != call.a_tensor.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({call.a_tensor.device}), got {call.b_ptrs.device}") - if call.b_ptrs.numel() != call.padded_offsets.numel(): - raise ValueError(f"b_ptrs length mismatch: expected {call.padded_offsets.numel()}, " f"got {call.b_ptrs.numel()}") - if not torch.cuda.is_available(): + if get_device(call.b_ptrs) != get_device(call.a_tensor): + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({get_device(call.a_tensor)}), got {get_device(call.b_ptrs)}") + offsets_shape = get_shape(call.padded_offsets) + if len(offsets_shape) == 1 and num_experts != offsets_shape[0]: + raise ValueError(f"b_ptrs length mismatch: expected {offsets_shape[0]}, " f"got {num_experts}") + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(call.a_tensor.device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: raise RuntimeError(f"GroupedGemmGluSm100 requires SM100+, found SM{compute_capability}") @@ -954,11 +980,13 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: def _glu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + strides = get_strides(tensor) + shape = get_shape(tensor) return tuple( index for index, _ in sorted( - enumerate(tensor.stride()), - key=lambda item: (item[1], tensor.shape[item[0]]), + enumerate(strides), + key=lambda item: (item[1], shape[item[0]]), ) ) @@ -966,37 +994,44 @@ def _glu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: def _glu_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: if tensor is None: return (None, None, None, None) - shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + device = get_device(tensor) + shape = (None, *get_shape(tensor)[1:]) if dynamic_m else get_shape(tensor) return ( shape, _glu_stride_order(tensor), - tensor.dtype, - (tensor.device.type, tensor.device.index), + _convert_to_cutlass_data_type(tensor.dtype), + (device.type, device.index), ) def _grouped_gemm_glu_bf16_call(call: GluCall) -> TupleDict: - import torch - - valid_m, k, _ = call.a_tensor.shape + framework = detect_framework(call.a_tensor) + valid_m, k, _ = get_shape(call.a_tensor) if call.weight_mode == MoEWeightMode.DENSE: - n_full = call.b_tensor.shape[0] + n_full = get_shape(call.b_tensor)[0] else: n_full = call.n n_out = n_full // 2 - c_tensor = torch.empty_strided( - (valid_m, n_full, 1), - (n_full, 1, valid_m * n_full), - dtype=call.c_dtype, - device=call.a_tensor.device, - ) - d_tensor = torch.empty_strided( - (valid_m, n_out, 1), - (n_out, 1, valid_m * n_out), - dtype=call.d_dtype, - device=call.a_tensor.device, - ) + def _allocate_output(shape, stride, dtype): + if framework == "torch": + import torch + + return torch.empty_strided( + shape, + stride, + dtype=framework_dtype(dtype, "torch"), + device=call.a_tensor.device, + ) + import jax + import jax.numpy as jnp + + # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. + # The kernel writes into this buffer on the launch stream; materialize it first. + return jax.block_until_ready(jnp.empty(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device)) + + c_tensor = _allocate_output((valid_m, n_full, 1), (n_full, 1, valid_m * n_full), call.c_dtype) + d_tensor = _allocate_output((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), call.d_dtype) overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) workspace_bytes = (128 * call.num_experts if call.weight_mode == MoEWeightMode.DISCRETE else 0) + (4 if call.use_dynamic_sched else 0) @@ -1014,16 +1049,7 @@ def _grouped_gemm_glu_bf16_call(call: GluCall) -> TupleDict: _glu_tensor_signature(call.bias_tensor), _glu_tensor_signature(call.padded_offsets), _glu_tensor_signature(call.prob_tensor, dynamic_m=True), - ( - ( - tuple(call.b_ptrs.shape), - tuple(call.b_ptrs.stride()), - call.b_ptrs.dtype, - (call.b_ptrs.device.type, call.b_ptrs.device.index), - ) - if call.b_ptrs is not None - else None - ), + (_glu_tensor_signature(call.b_ptrs) if call.b_ptrs is not None else None), call.acc_dtype, call.c_dtype, call.d_dtype, @@ -1035,7 +1061,7 @@ def _grouped_gemm_glu_bf16_call(call: GluCall) -> TupleDict: call.b_major, call.use_dynamic_sched, workspace_bytes, - (call.a_tensor.device.type, call.a_tensor.device.index), + ((get_device(call.a_tensor).type, get_device(call.a_tensor).index)), overlap_margin, ) @@ -1148,18 +1174,18 @@ def grouped_gemm_glu_wrapper_sm100( generate_c: bool = False, ) -> TupleDict: """Dispatch grouped GEMM GLU once from an immutable normalized call.""" - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_glu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if c_dtype is None: - c_dtype = torch.bfloat16 - if d_dtype is None: - d_dtype = torch.bfloat16 + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_glu_wrapper_sm100; pass torch tensors or JAX arrays") + if framework == "jax": + if b_tensor is not None: + raise ValueError(_JAX_DENSE_B_ERROR) + if bias_tensor is not None: + raise ValueError(_JAX_BIAS_ERROR) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + c_dtype = _convert_to_cutlass_data_type(c_dtype) if c_dtype is not None else cutlass.BFloat16 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None _reject_unsupported_rubin_glu_tune_params( get_device_type() == "rubin", geglu_alpha, @@ -1204,6 +1230,8 @@ def grouped_gemm_glu_wrapper_sm100( call, backend = _normalize_glu_call(call) if backend is GroupedGemmBackend.BF16: return _grouped_gemm_glu_bf16_call(call) + if framework == "jax": + raise ValueError(_JAX_BLOCK_SCALED_ERROR) return _grouped_gemm_glu_block_scaled_call(call) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py new file mode 100644 index 000000000..be72b5673 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py @@ -0,0 +1,221 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the BF16 SM100 grouped GEMM GLU +forward (discrete weight mode), built on :func:`cudnn.jax.call`. + +BF16 backend and discrete mode only: dense mode's expert-outermost strided B has +no row-major JAX equivalent, the (n, experts) column-major bias layout is likewise +inexpressible (``bias`` is a compile-time ``None`` inside the adapter), and the +block-scaled backend's MMA-interleaved scale-factor layouts cannot be presented as +row-major JAX arrays. The per-expert weight pointers travel as a regular device +array whose *values* are raw addresses — the referenced weight buffers are not +visible to XLA, so the caller must keep them alive (and unmoved) across every +execution of the traced computation. ``padded_offsets`` values cannot be +host-validated under tracing; malformed offsets are the caller's responsibility +here (the eager wrapper validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import call, gemm_operand_spec, zeros_init +from ..moe_utils import MoEWeightMode +from ..unfused.jax_api import _pointer_count, _prob_spec +from .moe_grouped_gemm_glu_bias import MoEGroupedGemmGluBiasBf16Kernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_output_dtypes = (cutlass.BFloat16, cutlass.Float16, cutlass.Float32) + +_JAX_BLOCK_SCALED_ERROR = ( + "the block-scaled grouped GEMM GLU backend is not expressible as JAX arrays " + "(its scale-factor tensors use an MMA-interleaved layout with no row-major equivalent); " + "only the BF16 backend supports JAX inputs" +) + + +@cute.jit +def _grouped_glu_bf16_adapter(stream, a, b_ptrs, padded_offsets, alpha, prob, d, c, workspace, *, kernel, n, k, mac, linear_offset): + # Discrete-mode b is a raw pointer to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffer recasts for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d, + padded_offsets=padded_offsets, + alpha=alpha, + prob=prob, + bias=None, + max_active_clusters=mac, + stream=stream, + linear_offset=cutlass.Float32(linear_offset), + ) + + +def grouped_gemm_glu_jax_sm100( + a_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_ptrs: Any, + n: int, + prob_tensor: Any, + c_dtype: Any = cutlass.BFloat16, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + generate_c: bool = False, + use_dynamic_sched: bool = False, +) -> Tuple[Any, Optional[Any]]: + """BF16 grouped GEMM GLU forward (discrete weights) as an XLA custom call. + + Same contract as the eager wrapper's BF16 discrete mode: A ``(m, k, 1)`` k-major + C-contiguous bfloat16, ``padded_offsets (experts,)`` int32 cumulative 256-aligned + row offsets, ``alpha (experts,)`` float32, ``prob (m, 1, 1)`` float32, and + ``b_ptrs`` holding per-expert ``(n, k)`` k-major bfloat16 weight base addresses + (packed little-endian uint8, 8 bytes per pointer — or int64 with x64 mode). + ``n`` is the full weight N before the GLU split; ``d`` comes back ``(m, n // 2, 1)``. + ``linear_offset`` defaults per ``act_func`` (1.0 for ``"geglu"``, 0.0 for + ``"swiglu"``) and is a compile-time constant of the traced call. Returns + ``(d_tensor, c_tensor)`` with ``c_tensor`` None unless ``generate_c``; rows + at/past ``padded_offsets[-1]`` come back zero-filled (the outputs are donated + zero-initialized buffers). + """ + c_dtype = _convert_to_cutlass_data_type(c_dtype) + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + if _convert_to_cutlass_data_type(a_tensor.dtype) is not cutlass.BFloat16: + raise ValueError(f"a_tensor must have dtype bfloat16, got {a_tensor.dtype}; " + _JAX_BLOCK_SCALED_ERROR) + if n is None or n <= 0 or n % 64 != 0: + raise ValueError(f"n must be positive and divisible by 64 for paired GLU blocks, got {n}") + if c_dtype not in _output_dtypes or d_dtype not in _output_dtypes: + raise ValueError(f"c_dtype/d_dtype must be BF16, FP16, or FP32, got {c_dtype}/{d_dtype}; " + _JAX_BLOCK_SCALED_ERROR) + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + if act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {act_func}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "geglu" else 0.0 + + expert_cnt = _pointer_count(b_ptrs) + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + + if not MoEGroupedGemmGluBiasBf16Kernel.can_implement( + cutlass.BFloat16, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + "k", + "n", + MoEGroupedGemmGluBiasBf16Kernel.FIX_PAD_SIZE, + ): + raise ValueError("Unsupported BF16 grouped GEMM GLU tile, cluster, alignment, or layout configuration") + + cache_key = ( + expert_cnt, + c_dtype, + d_dtype, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + vector_f32, + act_func, + generate_c, + use_dynamic_sched, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = MoEGroupedGemmGluBiasBf16Kernel( + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DISCRETE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + enable_bias=False, + generate_c=generate_c, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + n_out = n // 2 + operand = gemm_operand_spec() + d_tensor, c_tensor, _workspace = call( + _grouped_glu_bf16_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n_out, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, n, 1), framework_dtype(c_dtype, "jax")), + jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8), + ), + input_spec=(operand, None, None, None, _prob_spec()), + output_spec=(operand, operand, None), + # All three donated: c/d for the trailing-unit-dim layout spec (and defined + # bytes past the last offset); the workspace because the helper kernel writes + # the per-expert TMA descriptors into it (XLA inputs are immutable). + initialized_outputs={0: zeros_init, 1: zeros_init, 2: zeros_init}, + kernel=kernel, + n=int(n), + k=int(k), + mac=mac, + linear_offset=float(linear_offset), + )(a_tensor, b_ptrs, padded_offsets, alpha_tensor, prob_tensor) + + return d_tensor, (c_tensor if generate_c else None) diff --git a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py index 4914c2702..c1b973238 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.py @@ -7,7 +7,7 @@ import logging import os -from typing import Optional, Tuple +from typing import Any, Optional, Tuple from cuda.bindings import driver as cuda import cutlass @@ -22,6 +22,25 @@ from .hadamard_utils import HADAMARD_SIZE, hadamard_matrix from .moe_blockscaled_grouped_gemm_glu_hadamard import BlockScaledMoEGroupedGemmGluHadamardKernel +# The GLU + Hadamard fusion is block-scaled only: its mandatory scale-factor inputs +# (sfa/sfb) use an MMA-interleaved 6-D layout with no row-major equivalent, so they +# are not expressible as JAX arrays and the API stays torch-only. +_JAX_ERROR = ( + "grouped GEMM GLU hadamard is not supported for JAX arrays: the block-scaled " + "scale-factor tensors (sfa/sfb) use an MMA-interleaved layout that is not expressible as JAX arrays; " + "pass torch tensors" +) + + +def _require_torch_inputs(sample: Any, api_name: str) -> None: + from cudnn.tensor_adapter import detect_framework + + framework = detect_framework(sample) + if framework == "jax": + raise ValueError(_JAX_ERROR) + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for {api_name}; pass torch tensors") + def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: import torch @@ -65,10 +84,8 @@ def __init__( use_dynamic_sched: bool = False, use_tmem_post_rht_amax: bool = False, ): - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmGluHadamardSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + if sample_a is not None: + _require_torch_inputs(sample_a, "GroupedGemmGluHadamardSm100") import torch if acc_dtype is None: @@ -625,10 +642,9 @@ def grouped_gemm_glu_hadamard_wrapper_sm100( ) -> TupleDict: """High-level wrapper for grouped GEMM GLU + Hadamard forward fusion.""" from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor - from cudnn.tensor_adapter import is_torch_tensor - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_glu_hadamard_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + if a_tensor is not None: + _require_torch_inputs(a_tensor, "grouped_gemm_glu_hadamard_wrapper_sm100") import torch if acc_dtype is None: diff --git a/python/cudnn/gemm/cutedsl/grouped/quant/api.py b/python/cudnn/gemm/cutedsl/grouped/quant/api.py index ad16f6a25..f229770c0 100644 --- a/python/cudnn/gemm/cutedsl/grouped/quant/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/quant/api.py @@ -20,6 +20,16 @@ from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, get_device_type, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, + get_device, +) from .grouped_gemm_quant import ( BlockScaledMoEGroupedGemmQuantKernel, @@ -29,6 +39,12 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.runtime import from_dlpack +_JAX_SF_LAYOUT_ERROR = ( + "the block scale-factor tensors (sfa/sfb and the sfd outputs) are MMA-tiled " + "(32, 4, m//128, 4, rest_k, l) strided views that are not expressible as JAX arrays " + "(a row-major JAX array of that shape has different memory); pass torch tensors" +) + def _get_rubin_kernel(): from .moe_blockscaled_grouped_gemm_quant_rubin import ( @@ -120,15 +136,15 @@ def __init__( :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmQuantSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if framework == "jax": + raise ValueError(f"GroupedGemmQuantSm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmQuantSm100; pass torch tensors") if acc_dtype is None: - import torch - - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() + self._framework = framework self._warn_experimental_api() self._logger.debug("Entering __init__") @@ -145,14 +161,14 @@ def __init__( else: raise ValueError("Provide either (sample_b, sample_sfb) for dense mode " "or (num_experts, b_shape, b_dtype) for discrete mode, but not both.") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) self._has_d_col = sample_d_col is not None - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) if self.d_col_desc is None: self.d_col_desc = TensorDesc( dtype=self.d_desc.dtype, @@ -162,38 +178,38 @@ def __init__( device=self.d_desc.device, name="sample_d_col", ) - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) self.row_scale_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_row_scale, name="sample_row_scale"), + self._make_tensor_desc(sample_row_scale, name="sample_row_scale", canonical=True), 1, "row_scale", ) - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias", canonical=True) if self.weight_mode == MoEWeightMode.DENSE: - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) self.expert_cnt = self.padded_offsets_desc.shape[0] else: self._value_error_if(num_experts == 0, "num_experts must be > 0") self.expert_cnt = num_experts self.b_shape = b_shape - self.b_dtype = b_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) self.b_major = b_major self._value_error_if( self.padded_offsets_desc.shape[0] != self.expert_cnt, f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", ) - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -228,8 +244,6 @@ def check_support(self) -> bool: :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -333,10 +347,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -349,7 +363,7 @@ def check_support(self) -> bool: ) self._check_dtype( self.bias_desc, - dtype=[torch.bfloat16, torch.float16, torch.float32], + dtype=[cutlass.BFloat16, cutlass.Float16, cutlass.Float32], name="bias", extra_error_msg="bias must be fp16, bfloat16, or float32", ) @@ -360,14 +374,14 @@ def check_support(self) -> bool: ) self._check_dtype( self.bias_desc, - dtype=[torch.bfloat16, torch.float16], + dtype=[cutlass.BFloat16, cutlass.Float16], name="bias", extra_error_msg="bias must be fp16 or bfloat16 in discrete mode", ) self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD_row/SFD_col", ) if self.weight_mode == MoEWeightMode.DENSE: @@ -395,7 +409,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype in [cutlass.Float8E4M3FN] and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) self._value_error_if( @@ -405,14 +419,14 @@ def check_support(self) -> bool: self._check_dtype( self.acc_dtype, - dtype=torch.float32, + dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32", ) if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D", extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -420,11 +434,11 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_desc, dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="D", ) @@ -437,7 +451,7 @@ def check_support(self) -> bool: if not self._is_rubin_kernel: self._check_dtype( self.row_scale_desc, - dtype=torch.float32, + dtype=cutlass.Float32, name="row_scale", extra_error_msg="row_scale must be float32", ) @@ -532,13 +546,12 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): "Invalid configuration: fp8 ab_dtype and sf_vec_size 32 with mma_tiler_mn[1] == 128 and fp8 d_dtype is not supported. " "Please use mma_tiler_mn[1] == 256 instead", ) - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmQuant requires SM100+ compute capability, but found SM{compute_capability} on device {device}") + raise RuntimeError(f"GroupedGemmQuant requires SM100+ compute capability, but found SM{compute_capability}") self._is_supported = True self._logger.debug("check_support completed successfully") @@ -546,8 +559,6 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): def compile(self) -> None: """Compile the kernel.""" - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -590,7 +601,9 @@ def compile(self) -> None: fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) workspace_bytes = gemm_quant.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) if self.weight_mode == MoEWeightMode.DENSE: self._compile_dense(gemm_quant, max_active_clusters, fake_stream) @@ -886,8 +899,6 @@ def tensor_api( def _compile_discrete(self, gemm_quant, max_active_clusters, fake_stream) -> None: """Compile for discrete (per-expert pointer) weight mode.""" - import torch - if len(self.b_shape) == 2: n, k = self.b_shape else: @@ -971,10 +982,17 @@ def _compile_discrete(self, gemm_quant, max_active_clusters, fake_stream) -> Non ) bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + # Compile-time placeholders for the pointer-array arguments: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + self._compile_sfb_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + b_ptrs_placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + b_ptrs_placeholder.element_type = cutlass.Int64 + b_ptrs_cute = b_ptrs_placeholder.iterator + sfb_ptrs_placeholder = from_dlpack(self._compile_sfb_ptrs, assumed_align=8) + sfb_ptrs_placeholder.element_type = cutlass.Int64 + sfb_ptrs_cute = sfb_ptrs_placeholder.iterator workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator self._logger.debug("Compiling discrete grouped_gemm_quant kernel") @@ -1033,8 +1051,8 @@ def tensor_api( stream: cuda.CUstream, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + b_ptrs_addr = int(get_data_ptr(b_ptrs_device)) + sfb_ptrs_addr = int(get_data_ptr(sfb_ptrs_device)) if self._is_rubin_kernel: _compiled_kernel( a_tensor, @@ -1135,7 +1153,10 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if a_tensor.shape[0] == 0: self._logger.debug("execute: valid_m is zero, skipping kernel execution") @@ -1320,17 +1341,18 @@ def grouped_gemm_quant_wrapper_sm100( # Integer indexing d = result[0] # d_tensor """ - from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor - from cudnn.tensor_adapter import is_torch_tensor + from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_quant_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(a_tensor) + if framework == "jax": + raise ValueError(f"grouped_gemm_quant_wrapper_sm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_quant_wrapper_sm100; pass torch tensors") import torch - if acc_dtype is None: - acc_dtype = torch.float32 - if d_dtype is None: - d_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None is_dense = b_tensor is not None is_discrete = b_ptrs is not None @@ -1348,30 +1370,29 @@ def grouped_gemm_quant_wrapper_sm100( raise ValueError(f"bias_tensor must have shape {(n_out, l)}, got {tuple(bias_tensor.shape)}") else: weight_mode = MoEWeightMode.DISCRETE - _require_pointer_tensor(b_ptrs, "b_ptrs") - num_experts = b_ptrs.shape[0] - _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) + num_experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + _validate_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) if n is None or b_dtype is None: raise ValueError("n and b_dtype are required for discrete mode") - k_logical = k_physical * 2 if b_dtype in (torch.float4_e2m1fn_x2, torch.uint8) else k_physical + k_logical = k_physical * 2 if b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) else k_physical b_shape = (n, k_logical) n_out = n l = num_experts if bias_tensor is not None and tuple(bias_tensor.shape) != (n_out, num_experts): raise ValueError(f"bias_tensor must have shape {(n_out, num_experts)}, got {tuple(bias_tensor.shape)}") - is_fp8_input_config = a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [ - torch.float8_e8m0fnu, - torch.float8_e4m3fn, - ] - is_low_precision_output_config = d_dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, - ] + is_fp8_input_config = _convert_to_cutlass_data_type(a_tensor.dtype) in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ) and _convert_to_cutlass_data_type(sfa_tensor.dtype) in ( + cutlass.Float8E8M0FNU, + cutlass.Float8E4M3FN, + ) + is_low_precision_output_config = d_dtype in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, + ) _logger.debug("grouped_gemm_quant_wrapper_sm100: Creating output tensors") @@ -1379,12 +1400,12 @@ def grouped_gemm_quant_wrapper_sm100( expected_shape = (valid_m, n_out, 1) expected_stride = (n_out, 1, valid_m * n_out) if d_tensor is None: - d_tensor = torch.empty_strided(expected_shape, expected_stride, dtype=d_dtype, device=a_tensor.device) + d_tensor = torch.empty_strided(expected_shape, expected_stride, dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) elif ( tuple(d_tensor.shape) != expected_shape or tuple(d_tensor.stride()) != expected_stride - or d_tensor.dtype != d_dtype - or d_tensor.device != a_tensor.device + or _convert_to_cutlass_data_type(d_tensor.dtype) != d_dtype + or get_device(d_tensor) != get_device(a_tensor) ): raise ValueError( f"d_tensor must have shape {expected_shape}, stride {expected_stride}, " @@ -1392,7 +1413,7 @@ def grouped_gemm_quant_wrapper_sm100( f"stride {tuple(d_tensor.stride())}, dtype {d_tensor.dtype}, device {d_tensor.device}." ) d_col_tensor = ( - torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) + torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) if is_low_precision_output_config else None ) @@ -1441,7 +1462,7 @@ def grouped_gemm_quant_wrapper_sm100( ) sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): _logger.debug("grouped_gemm_quant_wrapper_sm100: Detected bf16/float16 d_dtype, constructing amax_tensor") amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) @@ -1449,7 +1470,7 @@ def grouped_gemm_quant_wrapper_sm100( if row_scale_tensor is not None: if device_type == "rubin": raise NotImplementedError("Rubin grouped GEMM quant does not support row_scale fusion") - if row_scale_tensor.dtype != torch.float32: + if _convert_to_cutlass_data_type(row_scale_tensor.dtype) is not cutlass.Float32: raise ValueError(f"row_scale_tensor must be float32, got {row_scale_tensor.dtype}") if tuple(row_scale_tensor.shape) != (valid_m,): raise ValueError(f"row_scale_tensor must have shape {(valid_m,)}, got {tuple(row_scale_tensor.shape)}") diff --git a/python/cudnn/gemm/cutedsl/grouped/srelu/api.py b/python/cudnn/gemm/cutedsl/grouped/srelu/api.py index cca3056c6..75afccef1 100644 --- a/python/cudnn/gemm/cutedsl/grouped/srelu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/srelu/api.py @@ -20,6 +20,15 @@ from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, +) from .moe_blockscaled_grouped_gemm_srelu_quant import ( BlockScaledMoEGroupedGemmQuantKernel, @@ -29,6 +38,12 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.runtime import from_dlpack +_JAX_SF_LAYOUT_ERROR = ( + "the block scale-factor tensors (sfa/sfb and the sfd outputs) are MMA-tiled " + "(32, 4, m//128, 4, rest_k, l) strided views that are not expressible as JAX arrays " + "(a row-major JAX array of that shape has different memory); pass torch tensors" +) + def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: import torch @@ -119,15 +134,15 @@ def __init__( :param b_major: Major dimension for B tensor, one of "k" or "n" :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmSreluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if framework == "jax": + raise ValueError(f"GroupedGemmSreluSm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmSreluSm100; pass torch tensors") if acc_dtype is None: - import torch - - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() + self._framework = framework self._warn_experimental_api() self._logger.debug("Entering __init__") @@ -147,15 +162,15 @@ def __init__( self._sample_a_tensor = sample_a self._sample_b_tensor = sample_b - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", interpret_uint8_as_fp4x2=False) - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", interpret_uint8_as_fp4x2=False, canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) self._has_d_col = sample_d_col is not None - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) if self.d_col_desc is None: self.d_col_desc = TensorDesc( dtype=self.d_desc.dtype, @@ -165,33 +180,33 @@ def __init__( device=self.d_desc.device, name="sample_d_col", ) - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias", canonical=True) if self.weight_mode == MoEWeightMode.DENSE: - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", interpret_uint8_as_fp4x2=False) - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", interpret_uint8_as_fp4x2=False, canonical=True) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) self.expert_cnt = self.padded_offsets_desc.shape[0] else: self._value_error_if(num_experts == 0, "num_experts must be > 0") self.expert_cnt = num_experts self.b_shape = b_shape - self.b_dtype = b_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) self.b_major = b_major self._value_error_if( self.padded_offsets_desc.shape[0] != self.expert_cnt, f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", ) - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -221,8 +236,6 @@ def check_support(self) -> bool: :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -325,10 +338,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -341,7 +354,7 @@ def check_support(self) -> bool: ) self._check_dtype( self.bias_desc, - dtype=[torch.bfloat16, torch.float16, torch.float32], + dtype=[cutlass.BFloat16, cutlass.Float16, cutlass.Float32], name="bias", extra_error_msg="bias must be fp16, bfloat16, or float32", ) @@ -352,14 +365,14 @@ def check_support(self) -> bool: ) self._check_dtype( self.bias_desc, - dtype=[torch.bfloat16, torch.float16], + dtype=[cutlass.BFloat16, cutlass.Float16], name="bias", extra_error_msg="bias must be fp16 or bfloat16 in discrete mode", ) self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD_row/SFD_col", ) if self.weight_mode == MoEWeightMode.DENSE: @@ -387,7 +400,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype in [cutlass.Float8E4M3FN] and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) self._value_error_if( @@ -397,19 +410,19 @@ def check_support(self) -> bool: self._check_dtype( self.acc_dtype, - dtype=torch.float32, + dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32", ) self.c_dtype = self._check_dtype( self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], + dtype=[cutlass.Float32, cutlass.Float16, cutlass.BFloat16, cutlass.Float8E4M3FN, cutlass.Float8E5M2], name="C", ) if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D", extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -417,11 +430,11 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_desc, dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="D", ) @@ -433,7 +446,7 @@ def check_support(self) -> bool: ) self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, + self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype is cutlass.Float32, "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. Please use sf_vec_size 32 or d_dtype bf16 instead", ) @@ -527,13 +540,12 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): "Invalid configuration: fp8 ab_dtype and sf_vec_size 32 with mma_tiler_mn[1] == 128 and fp8 d_dtype is not supported. " "Please use mma_tiler_mn[1] == 256 instead", ) - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmSrelu requires SM100+ compute capability, but found SM{compute_capability} on device {device}") + raise RuntimeError(f"GroupedGemmSrelu requires SM100+ compute capability, but found SM{compute_capability}") self._is_supported = True self._logger.debug("check_support completed successfully") @@ -541,8 +553,6 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): def compile(self) -> None: """Compile the kernel.""" - import torch - self._logger.debug("Entering compile") self._ensure_support_checked() if self._compiled_kernel is not None: @@ -581,7 +591,9 @@ def compile(self) -> None: fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) workspace_bytes = gemm_srelu.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) if self.weight_mode == MoEWeightMode.DENSE: self._compile_dense(gemm_srelu, max_active_clusters, fake_stream) @@ -592,8 +604,6 @@ def compile(self) -> None: def _compile_dense(self, gemm_srelu, max_active_clusters, fake_stream) -> None: """Compile for dense (contiguous) weight mode.""" - import torch - fake_workspace_ptr = cute.runtime.nullptr( dtype=cutlass.Uint8, assumed_align=128, @@ -775,8 +785,8 @@ def _compile_dense(self, gemm_srelu, max_active_clusters, fake_stream) -> None: _compiled_kernel = cute.compile( gemm_srelu, - a=_reinterpret_raw_grouped_fp4_tensor(self._sample_a_tensor) if self.a_desc.dtype == torch.uint8 else a_cute_fake, - b=_reinterpret_raw_grouped_fp4_tensor(self._sample_b_tensor) if self.b_desc.dtype == torch.uint8 else b_cute_fake, + a=_reinterpret_raw_grouped_fp4_tensor(self._sample_a_tensor) if self.a_desc.dtype is cutlass.Uint8 else a_cute_fake, + b=_reinterpret_raw_grouped_fp4_tensor(self._sample_b_tensor) if self.b_desc.dtype is cutlass.Uint8 else b_cute_fake, sfb=sfb_cute_fake, n=cutlass.Int32(0), k=cutlass.Int32(0), @@ -848,8 +858,6 @@ def tensor_api( def _compile_discrete(self, gemm_srelu, max_active_clusters, fake_stream) -> None: """Compile for discrete (per-expert pointer) weight mode.""" - import torch - if len(self.b_shape) == 2: n, k = self.b_shape else: @@ -928,10 +936,17 @@ def _compile_discrete(self, gemm_srelu, max_active_clusters, fake_stream) -> Non ) bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + # Compile-time placeholders for the pointer-array arguments: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + self._compile_sfb_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) + b_ptrs_placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + b_ptrs_placeholder.element_type = cutlass.Int64 + b_ptrs_cute = b_ptrs_placeholder.iterator + sfb_ptrs_placeholder = from_dlpack(self._compile_sfb_ptrs, assumed_align=8) + sfb_ptrs_placeholder.element_type = cutlass.Int64 + sfb_ptrs_cute = sfb_ptrs_placeholder.iterator workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator self._logger.debug("Compiling discrete grouped_gemm_srelu kernel") @@ -987,8 +1002,8 @@ def tensor_api( stream: cuda.CUstream, ) -> None: norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + b_ptrs_addr = int(get_data_ptr(b_ptrs_device)) + sfb_ptrs_addr = int(get_data_ptr(sfb_ptrs_device)) _compiled_kernel( a_tensor, b_ptrs_addr, @@ -1061,7 +1076,10 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if a_tensor.shape[0] == 0: self._logger.debug("execute: valid_m is zero, skipping kernel execution") @@ -1230,19 +1248,19 @@ def grouped_gemm_srelu_wrapper_sm100( c = result[0] # c_tensor d = result[1] # d_tensor """ - from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor - from cudnn.tensor_adapter import is_torch_tensor + from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_srelu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(a_tensor) + if framework == "jax": + raise ValueError(f"grouped_gemm_srelu_wrapper_sm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_srelu_wrapper_sm100; pass torch tensors") import torch - if acc_dtype is None: - acc_dtype = torch.float32 - if c_dtype is None: - c_dtype = torch.bfloat16 - if d_dtype is None: - d_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + c_dtype = _convert_to_cutlass_data_type(c_dtype) if c_dtype is not None else cutlass.BFloat16 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None is_dense = b_tensor is not None is_discrete = b_ptrs is not None @@ -1260,38 +1278,37 @@ def grouped_gemm_srelu_wrapper_sm100( raise ValueError(f"bias_tensor must have shape {(n_out, l)}, got {tuple(bias_tensor.shape)}") else: weight_mode = MoEWeightMode.DISCRETE - _require_pointer_tensor(b_ptrs, "b_ptrs") - num_experts = b_ptrs.shape[0] - _require_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) + num_experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + _validate_pointer_tensor(sfb_ptrs, "sfb_ptrs", num_experts) if n is None or b_dtype is None: raise ValueError("n and b_dtype are required for discrete mode") - k_logical = k_physical * 2 if b_dtype in (torch.float4_e2m1fn_x2, torch.uint8) else k_physical + k_logical = k_physical * 2 if b_dtype in (cutlass.Float4E2M1FN, cutlass.Uint8) else k_physical b_shape = (n, k_logical) n_out = n l = num_experts if bias_tensor is not None and tuple(bias_tensor.shape) != (n_out, num_experts): raise ValueError(f"bias_tensor must have shape {(n_out, num_experts)}, got {tuple(bias_tensor.shape)}") - is_fp8_input_config = a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [ - torch.float8_e8m0fnu, - torch.float8_e4m3fn, - ] - is_low_precision_output_config = d_dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, - ] + is_fp8_input_config = _convert_to_cutlass_data_type(a_tensor.dtype) in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ) and _convert_to_cutlass_data_type(sfa_tensor.dtype) in ( + cutlass.Float8E8M0FNU, + cutlass.Float8E4M3FN, + ) + is_low_precision_output_config = d_dtype in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, + ) _logger.debug("grouped_gemm_srelu_wrapper_sm100: Creating output tensors") if cd_major == "n": - c_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=c_dtype, device=a_tensor.device) - d_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(c_dtype, "torch"), device=a_tensor.device) + d_tensor = torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) d_col_tensor = ( - torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=d_dtype, device=a_tensor.device) + torch.empty_strided((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device) if is_low_precision_output_config else None ) @@ -1340,7 +1357,7 @@ def grouped_gemm_srelu_wrapper_sm100( ) sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): _logger.debug("grouped_gemm_srelu_wrapper_sm100: Detected bf16/float16 d_dtype, constructing amax_tensor") amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) diff --git a/python/cudnn/gemm/cutedsl/grouped/swiglu/api.py b/python/cudnn/gemm/cutedsl/grouped/swiglu/api.py index c41717d74..ccb11c4ea 100644 --- a/python/cudnn/gemm/cutedsl/grouped/swiglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/swiglu/api.py @@ -22,6 +22,19 @@ from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.tensor_adapter import ( + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, +) + +_JAX_SF_LAYOUT_ERROR = ( + "the block scale-factor tensors (sfa/sfb and the sfd outputs) are MMA-tiled " + "(32, 4, m//128, 4, rest_k, l) strided views that are not expressible as JAX arrays " + "(a row-major JAX array of that shape has different memory); pass torch tensors" +) class GroupedGemmSwigluSm100(APIBase): @@ -96,46 +109,46 @@ def __init__( :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Boolean, True to generate discrete col-major scale factor tensor. Only applies when already output scale factor tensors are provided. """ - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmSwigluSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(sample_a) + if framework == "jax": + raise ValueError(f"GroupedGemmSwigluSm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmSwigluSm100; pass torch tensors") if acc_dtype is None: - import torch - - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() + self._framework = framework self._warn_experimental_api() self._logger.debug("Entering __init__") # Store sample tensor descriptors - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - 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.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) # Optional quantization outputs - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - 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") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col", canonical=True) + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row", canonical=True) + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col", canonical=True) + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + self._make_tensor_desc(sample_norm_const, name="sample_norm_const", canonical=True), 1, "norm_const", ) - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) # expert_cnt derived from padded_offsets shape self.expert_cnt = self.padded_offsets_desc.shape[0] # Configuration - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = mma_tiler_mn self.use_2cta_instrs = mma_tiler_mn[0] == 256 if cluster_shape_mn is None: @@ -159,8 +172,6 @@ def check_support(self) -> bool: :return: True if supported, raises exception otherwise """ - import torch - self._logger.debug("Entering check_support") all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) @@ -235,10 +246,10 @@ def check_support(self) -> bool: self.ab_dtype = self._check_dtype( self.a_desc, dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, ], name="A/B", ) @@ -251,7 +262,7 @@ def check_support(self) -> bool: self.sf_dtype = self._check_dtype( self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + dtype=[cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN], name="SFA/SFB/SFD_row/SFD_col", ) self._check_dtype( @@ -278,7 +289,7 @@ def check_support(self) -> bool: f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", ) self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + self.sf_dtype in [cutlass.Float8E4M3FN] and self.sf_vec_size == 32, f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) self._value_error_if( @@ -288,19 +299,19 @@ def check_support(self) -> bool: self._check_dtype( self.acc_dtype, - dtype=torch.float32, + dtype=cutlass.Float32, name="Accumulator", extra_error_msg="Accumulator must be float32", ) self.c_dtype = self._check_dtype( self.c_desc, dtype=[ - torch.float32, - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, ], name="C", ) @@ -308,7 +319,7 @@ def check_support(self) -> bool: if self._is_fp4x2(self.ab_dtype): self.d_dtype = self._check_dtype( self.d_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], + dtype=[cutlass.Float16, cutlass.BFloat16, cutlass.Float32], name="D", extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", ) @@ -316,12 +327,12 @@ def check_support(self) -> bool: self.d_dtype = self._check_dtype( self.d_desc, dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, - ], # torch.float32 fails non-deterministicly + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + cutlass.Float4E2M1FN, + ], # float32 fails non-deterministicly name="D", ) self._check_dtype( @@ -332,7 +343,7 @@ def check_support(self) -> bool: ) self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, # Fails to compile + self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype is cutlass.Float32, # Fails to compile f"Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. Please use sf_vec_size 32 or d_dtype bf16 instead", ) @@ -416,18 +427,17 @@ def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): "Please use mma_tiler_mn[1] == 256 instead", ) self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [torch.float16, torch.bfloat16]), + self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [cutlass.Float16, cutlass.BFloat16]), f"Invalid configuration: for fp4 ab_dtype, c_dtype must be float16 or bfloat16, got {self.c_dtype}", ) # Check environment - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmSwiglu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") + raise RuntimeError(f"GroupedGemmSwiglu requires SM100+ compute capability, " f"but found SM{compute_capability}") self._is_supported = True self._logger.debug("check_support completed successfully") @@ -725,7 +735,10 @@ def execute( :param current_stream: CUDA stream """ self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if a_tensor.shape[0] == 0: self._logger.debug("execute: valid_m is zero, skipping kernel execution") @@ -834,18 +847,16 @@ def grouped_gemm_swiglu_wrapper_sm100( # Integer indexing c = result[0] # c_tensor """ - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_swiglu_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") + framework = detect_framework(a_tensor) + if framework == "jax": + raise ValueError(f"grouped_gemm_swiglu_wrapper_sm100 does not support JAX arrays: {_JAX_SF_LAYOUT_ERROR}") + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_swiglu_wrapper_sm100; pass torch tensors") import torch - if acc_dtype is None: - acc_dtype = torch.float32 - if c_dtype is None: - c_dtype = torch.bfloat16 - if d_dtype is None: - d_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + c_dtype = _convert_to_cutlass_data_type(c_dtype) if c_dtype is not None else cutlass.BFloat16 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 valid_m, k, _ = a_tensor.shape n, _, l = b_tensor.shape n_out = n // 2 # After SwiGLU @@ -854,17 +865,17 @@ def grouped_gemm_swiglu_wrapper_sm100( if cd_major == "n": # 1, m, n, permute (1, 2, 0) -> (m, n, 1) - c_tensor = torch.empty_strided((valid_m, n, 1), (n, 1, valid_m * n), dtype=c_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided((valid_m, n, 1), (n, 1, valid_m * n), dtype=framework_dtype(c_dtype, "torch"), device=a_tensor.device) d_tensor = torch.empty_strided( (valid_m, n_out, 1), (n_out, 1, valid_m * n_out), - dtype=d_dtype, + dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device, ) d_col_tensor = torch.empty_strided( (valid_m, n_out, 1), (n_out, 1, valid_m * n_out), - dtype=d_dtype, + dtype=framework_dtype(d_dtype, "torch"), device=a_tensor.device, ) else: @@ -874,10 +885,12 @@ def grouped_gemm_swiglu_wrapper_sm100( sfd_col_tensor = None amax_tensor = None - if a_tensor.dtype in [ - torch.float8_e4m3fn, - torch.float8_e5m2, - ] and sfa_tensor.dtype in [torch.float8_e8m0fnu, torch.float8_e4m3fn]: + if _convert_to_cutlass_data_type(a_tensor.dtype) in ( + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + ) and _convert_to_cutlass_data_type( + sfa_tensor.dtype + ) in (cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN): _logger.debug("grouped_gemm_swiglu_wrapper_sm100: Detected fp8 a_dtype and sfa_dtype, constructing sfd_row_tensor and sfd_col_tensor") sf_dtype = sfa_tensor.dtype @@ -908,7 +921,7 @@ def grouped_gemm_swiglu_wrapper_sm100( sfd_col_tensor = torch.empty(mma_shape_col, dtype=sf_dtype, device=a_tensor.device).permute(mma_permute_order) if valid_m == 0: - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) _logger.debug("grouped_gemm_swiglu_wrapper_sm100: valid_m is zero, skipping kernel execution") @@ -983,7 +996,7 @@ def stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: _logger.debug("group_gemm_swiglu_wrapper_sm100: No previously cached GroupedGemmSwigluSm100 object found, creating new GroupedGemmSwigluSm100 object") # Allocate amax_tensor once here; cache-hit calls reuse this buffer so # the FillFunctor (torch.full) only fires during warmup, not every step. - if d_dtype in [torch.bfloat16, torch.float16]: + if d_dtype in (cutlass.BFloat16, cutlass.Float16): amax_tensor = torch.full((l, 1), float("-inf"), dtype=torch.float32, device=a_tensor.device) grouped_gemm_swiglu = GroupedGemmSwigluSm100( sample_a=a_tensor, diff --git a/python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py b/python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py index 00ff3e3df..713e28fac 100644 --- a/python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py @@ -3,4 +3,14 @@ from .api import GroupedGemmSm100, grouped_gemm_wrapper_sm100 -__all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100"] +__all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100", "grouped_gemm_jax_sm100"] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "grouped_gemm_jax_sm100": + from .jax_api import grouped_gemm_jax_sm100 + + return grouped_gemm_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py b/python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py index 6d22afa87..9a018a093 100644 --- a/python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py @@ -17,21 +17,68 @@ from cudnn.api_base import APIBase, TensorDesc from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + get_compute_capability, + get_data_ptr, + get_device, + get_shape, + get_version, + is_torch_tensor, + to_host_list, +) from ..moe_utils import MoEWeightMode from .moe_grouped_gemm import MoEGroupedGemmBf16Kernel -_OUTPUT_DTYPES = None - def _output_dtypes(): - global _OUTPUT_DTYPES - if _OUTPUT_DTYPES is None: - import torch - - _OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] - return _OUTPUT_DTYPES + return [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] + + +def _validate_pointer_tensor(ptrs, name: str, expected_len: int | None = None) -> int: + """Framework-neutral check of a device pointer-array tensor; returns the pointer count. + + torch: a contiguous 1-D CUDA int64 tensor (unchanged contract). JAX: a 1-D int64 + array (requires jax x64 mode) or, since JAX truncates int64 without x64 mode, a + 1-D uint8 array of length 8*count holding the packed little-endian pointers. + """ + import math + + shape = get_shape(ptrs) + if len(shape) != 1: + raise ValueError(f"{name} must be 1-D, got shape={shape}") + numel = math.prod(shape) + dtype = _convert_to_cutlass_data_type(ptrs.dtype) + if dtype is cutlass.Int64: + count = numel + elif dtype is cutlass.Uint8 and not is_torch_tensor(ptrs): + if numel % 8 != 0: + raise ValueError(f"{name} packed uint8 length must be a multiple of 8, got {numel}") + count = numel // 8 + else: + raise ValueError(f"{name} must be int64 (or, for JAX without x64 mode, packed uint8), got {ptrs.dtype}") + if expected_len is not None and count != expected_len: + raise ValueError(f"{name} length mismatch: expected {expected_len}, got {count}") + if is_torch_tensor(ptrs): + if not ptrs.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor, got device={ptrs.device}") + if not ptrs.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + return count + + +def _pointer_values(ptrs) -> tuple: + """Host copy of the pointer values, decoding the packed-uint8 JAX form if needed.""" + if not is_torch_tensor(ptrs) and _convert_to_cutlass_data_type(ptrs.dtype) is cutlass.Uint8: + import numpy as np + + return tuple(int(v) for v in np.asarray(ptrs).view(np.int64)) + return tuple(int(v) for v in to_host_list(ptrs)) class GroupedGemmBf16API(APIBase): @@ -59,12 +106,11 @@ def __init__( b_major: str = "k", use_dynamic_sched: bool = False, ) -> None: - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() + self._framework = detect_framework(sample_a) if sample_b is not None and num_experts is None: self.weight_mode = MoEWeightMode.DENSE @@ -75,19 +121,19 @@ def __init__( else: raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d", canonical=True) + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets", canonical=True) + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias", canonical=True) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob", canonical=True) self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) self._sample_offsets_ref = weakref.ref(sample_padded_offsets) - self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_offsets_version = get_version(sample_padded_offsets) self._sample_data_ptrs = { - name: tensor.data_ptr() + name: get_data_ptr(tensor) for name, tensor in ( ("sample_a", sample_a), ("sample_b", sample_b), @@ -103,8 +149,8 @@ def __init__( self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) self.b_shape = tuple(b_shape) if b_shape is not None else None - self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype - self.acc_dtype = acc_dtype + self.b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = tuple(mma_tiler_mn) self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) @@ -116,7 +162,8 @@ def __init__( self._has_bias = self.bias_desc is not None self._kernel = MoEGroupedGemmBf16Kernel self._workspace: Optional[torch.Tensor] = None - self._compile_b_ptrs: Optional[torch.Tensor] = None + self._live_b_ptrs = None + self._compile_b_ptrs = None self._validated_offsets: dict[int, tuple] = {} self._validated_pointer_values: dict[int, tuple] = {} self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) @@ -138,12 +185,12 @@ def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: @staticmethod def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(int(value) for value in tensor.detach().cpu().tolist()) + return tuple(int(value) for value in to_host_list(tensor)) @staticmethod def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: cached = cache.get(id(tensor)) - return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + return bool(cached and cached[0]() is tensor and cached[1] == get_version(tensor) and cached[2] == extra) @staticmethod def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: @@ -154,7 +201,7 @@ def discard(_reference, *, cache=cache, key=key): cache[key] = ( weakref.ref(tensor, discard), - int(tensor._version), + get_version(tensor), extra, ) @@ -169,7 +216,7 @@ def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> Non def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): return - pointer_values = self._copy_values_to_host(b_ptrs) + pointer_values = _pointer_values(b_ptrs) if any(value == 0 or value % 16 != 0 for value in pointer_values): raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) @@ -190,16 +237,20 @@ def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tenso @staticmethod def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: - if tensor.data_ptr() % 16 != 0: + if get_data_ptr(tensor) % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") @staticmethod def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: - if tensor.data_ptr() % 8 != 0: + if get_data_ptr(tensor) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") - @staticmethod - def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + def _record_pointer_stream(self, b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + if not is_torch_tensor(b_ptrs): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # array referenced until the next execute so its buffer outlives the launch. + self._live_b_ptrs = b_ptrs + return import torch handle = int(current_stream) @@ -214,8 +265,6 @@ def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) b_ptrs.record_stream(launch_stream) def check_support(self) -> bool: - import torch - if self.a_desc.ndim != 3: raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") tensor_m, k, one = self.a_desc.shape @@ -253,17 +302,17 @@ def check_support(self) -> bool: self._expect_stride(self.d_desc, (n, 1, tensor_m * n), "sample_d") self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") self._expect_stride(self.alpha_desc, (1,), "sample_alpha") - self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") + self._expect_stride(self.prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "sample_prob") - self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + self._check_dtype(self.a_desc, cutlass.BFloat16, "sample_a") if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") - self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.b_desc, cutlass.BFloat16, "sample_b") + self._check_dtype(self.b_dtype, cutlass.BFloat16, "b_dtype") self._check_dtype(self.c_desc, _output_dtypes(), "sample_c") self._check_dtype(self.d_desc, _output_dtypes(), "sample_d") - self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") - self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") - self._check_dtype(self.prob_desc, torch.float32, "sample_prob") + self._check_dtype(self.padded_offsets_desc, cutlass.Int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, cutlass.Float32, "sample_alpha") + self._check_dtype(self.prob_desc, cutlass.Float32, "sample_prob") device = self.a_desc.device for desc, name in ( @@ -287,7 +336,7 @@ def check_support(self) -> bool: if data_ptr % 16 != 0: raise ValueError(f"{name} data pointer must be 16-byte aligned") - if self.acc_dtype != torch.float32: + if self.acc_dtype is not cutlass.Float32: raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") if self.m_aligned != MoEGroupedGemmBf16Kernel.FIX_PAD_SIZE: raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") @@ -304,7 +353,7 @@ def check_support(self) -> bool: tensor_m=tensor_m, ) sample_offsets = self._sample_offsets_ref() - if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + if sample_offsets is not None and get_version(sample_offsets) == self._sample_offsets_version: self._remember_validation( self._validated_offsets, sample_offsets, @@ -314,7 +363,7 @@ def check_support(self) -> bool: self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) if not self._kernel.can_implement( - _convert_to_cutlass_data_type(torch.bfloat16), + cutlass.BFloat16, _convert_to_cutlass_data_type(self.c_desc.dtype), _convert_to_cutlass_data_type(self.d_desc.dtype), _convert_to_cutlass_data_type(self.acc_dtype), @@ -332,19 +381,17 @@ def check_support(self) -> bool: ): raise ValueError("Unsupported BF16 grouped GEMM tile, cluster, alignment, or layout configuration") - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(self.a_desc.device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: - raise RuntimeError(f"GroupedGemmSm100 requires SM100+, found SM{compute_capability} " f"on {self.a_desc.device}") + raise RuntimeError(f"GroupedGemmSm100 requires SM100+, found SM{compute_capability}") self._is_supported = True return True def compile(self) -> None: - import torch - self._ensure_support_checked() if self._compiled_kernel is not None: return @@ -368,8 +415,10 @@ def compile(self) -> None: raise ValueError("max_active_clusters must be > 0 after applying " "CUDNNFE_CLUSTER_OVERLAP_MARGIN") workspace_bytes = kernel.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) - if self._workspace.data_ptr() % 128 != 0: + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, workspace_bytes, self.a_desc.device) + if get_data_ptr(self._workspace) % 128 != 0: raise RuntimeError("workspace allocation must be 128-byte aligned") workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -408,9 +457,14 @@ def compile(self) -> None: k_value = cutlass.Int32(0) b_stride = cutlass.Int64(0) else: - self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + # Compile-time placeholder for the pointer-array argument: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_b_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) self._validate_pointer_array_alignment(self._compile_b_ptrs) - b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + placeholder = from_dlpack(self._compile_b_ptrs, assumed_align=8) + placeholder.element_type = cutlass.Int64 + b_fake = placeholder.iterator n, k = self.b_shape[:2] n_value = cutlass.Int32(n) k_value = cutlass.Int32(k) @@ -452,7 +506,7 @@ def tensor_api( prob_tensor: torch.Tensor, stream: cuda.CUstream, ) -> None: - b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(get_data_ptr(b_ptrs)) raw_compiled( a_tensor, b_arg, @@ -479,7 +533,7 @@ def _validate_live_tensor( *, dynamic_m: bool = False, ) -> TensorDesc: - desc = self._make_tensor_desc(tensor, name=name) + desc = self._make_tensor_desc(tensor, name=name, canonical=True) if desc.dtype != sample.dtype: raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") if desc.device != sample.device: @@ -506,7 +560,10 @@ def execute( prob_tensor: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if self._compiled_kernel is None: raise RuntimeError("Kernel not compiled; call compile() first") if prob_tensor is None: @@ -529,7 +586,7 @@ def execute( self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") self._expect_stride(c_desc, (n, 1, tensor_m * n), "c_tensor") self._expect_stride(d_desc, (n, 1, tensor_m * n), "d_tensor") - self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") + self._expect_stride(prob_desc, canonicalize_unit_dim_strides((tensor_m, 1, 1), (1, 1, 1)), "prob_tensor") self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) for tensor, name in ( @@ -558,9 +615,9 @@ def execute( else: if b_tensor is not None or b_ptrs is None: raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") - _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) - if b_ptrs.device != self.a_desc.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {b_ptrs.device}") + _validate_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if get_device(b_ptrs) != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {get_device(b_ptrs)}") self._validate_pointer_array_alignment(b_ptrs) self._validate_pointer_values_once(b_ptrs) self._record_pointer_stream(b_ptrs, current_stream) diff --git a/python/cudnn/gemm/cutedsl/grouped/unfused/api.py b/python/cudnn/gemm/cutedsl/grouped/unfused/api.py index 57e0e3f22..53b009eeb 100644 --- a/python/cudnn/gemm/cutedsl/grouped/unfused/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/unfused/api.py @@ -10,11 +10,22 @@ from cuda.bindings import driver as cuda -from cudnn.api_base import APIBase, TupleDict -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor - +import cutlass -from ._bf16_api import GroupedGemmBf16API +from cudnn.api_base import APIBase, TupleDict +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import ( + canonicalize_unit_dim_strides, + cuda_is_available, + detect_framework, + get_compute_capability, + get_data_ptr, + get_device, + get_shape, + get_strides, +) + +from ._bf16_api import GroupedGemmBf16API, _validate_pointer_tensor from ..moe_utils import MoEWeightMode __all__ = ["GroupedGemmBf16API"] @@ -49,14 +60,8 @@ def __init__( self._pending_init_kwargs = dict(locals()) self._pending_init_kwargs.pop("self") self._pending_init_kwargs.pop("__class__", None) - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") if acc_dtype is None: - import torch - - self._pending_init_kwargs["acc_dtype"] = torch.float32 + self._pending_init_kwargs["acc_dtype"] = cutlass.Float32 self._implementation = None def check_support(self) -> bool: @@ -113,11 +118,13 @@ def execute( def _stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + strides = get_strides(tensor) + shape = get_shape(tensor) return tuple( index for index, _ in sorted( - enumerate(tensor.stride()), - key=lambda item: (item[1], tensor.shape[item[0]]), + enumerate(strides), + key=lambda item: (item[1], shape[item[0]]), ) ) @@ -125,12 +132,13 @@ def _stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: def _tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: if tensor is None: return (None, None, None, None) - shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + device = get_device(tensor) + shape = (None, *get_shape(tensor)[1:]) if dynamic_m else get_shape(tensor) return ( shape, _stride_order(tensor), - tensor.dtype, - (tensor.device.type, tensor.device.index), + _convert_to_cutlass_data_type(tensor.dtype), + (device.type, device.index), ) @@ -143,11 +151,14 @@ def _validate_output( dtype: torch.dtype, device: torch.device, ) -> None: - if tuple(tensor.shape) != shape or tuple(tensor.stride()) != stride or tensor.dtype != dtype or tensor.device != device: + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + expected_stride = canonicalize_unit_dim_strides(shape, stride) + if tensor_shape != shape or tensor_stride != expected_stride or _convert_to_cutlass_data_type(tensor.dtype) != dtype or get_device(tensor) != device: raise ValueError( f"{name} must have shape {shape}, stride {stride}, dtype {dtype}, " - f"device {device}; got shape {tuple(tensor.shape)}, stride " - f"{tuple(tensor.stride())}, dtype {tensor.dtype}, device {tensor.device}" + f"device {device}; got shape {tensor_shape}, stride " + f"{tensor_stride}, dtype {tensor.dtype}, device {get_device(tensor)}" ) @@ -163,71 +174,76 @@ def _normalize_call( d_dtype: torch.dtype, cd_major: str, m_aligned: int, + framework: str, ) -> tuple[bool, int, int, int]: - import torch - is_dense = b_tensor is not None is_discrete = b_ptrs is not None if is_dense and is_discrete: raise ValueError("Provide either b_tensor or b_ptrs, not both") if not is_dense and not is_discrete: raise ValueError("Must provide either b_tensor or b_ptrs") + if framework == "jax" and is_dense: + raise ValueError( + "Dense weight mode (b_tensor) is not expressible as JAX arrays " + "(the expert-outermost strided B layout has no row-major equivalent); " + "use discrete mode (b_ptrs) with per-expert weight pointers" + ) - if a_tensor.dtype != torch.bfloat16: + if _convert_to_cutlass_data_type(a_tensor.dtype) is not cutlass.BFloat16: raise ValueError(f"a_tensor must have dtype torch.bfloat16, got {a_tensor.dtype}") - if a_tensor.ndim != 3 or a_tensor.shape[2] != 1: - raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + if len(get_shape(a_tensor)) != 3 or get_shape(a_tensor)[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {get_shape(a_tensor)}") if prob_tensor is None: raise ValueError("prob_tensor is required") if cd_major != "n": raise ValueError(f"cd_major must be 'n', got {cd_major}") - if c_dtype not in (torch.bfloat16, torch.float16, torch.float32): + if c_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {c_dtype}") - if d_dtype not in (torch.bfloat16, torch.float16, torch.float32): + if d_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {d_dtype}") if m_aligned != 256: raise ValueError(f"m_aligned must be 256, got {m_aligned}") - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(a_tensor.device) + major, minor = get_compute_capability() compute_capability = major * 10 + minor if compute_capability < 100: raise RuntimeError(f"GroupedGemmSm100 requires SM100+, found SM{compute_capability}") - tensor_m, k, _ = a_tensor.shape + tensor_m, k, _ = get_shape(a_tensor) if tensor_m % 256 != 0: raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}") - if tuple(prob_tensor.shape) != (tensor_m, 1, 1): - raise ValueError(f"prob_tensor must have shape {(tensor_m, 1, 1)}, got " f"{tuple(prob_tensor.shape)}") - if prob_tensor.dtype != torch.float32: + if get_shape(prob_tensor) != (tensor_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(tensor_m, 1, 1)}, got " f"{get_shape(prob_tensor)}") + if _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: raise ValueError(f"prob_tensor must have dtype torch.float32, got {prob_tensor.dtype}") if is_dense: if n is not None: raise ValueError("Dense mode forbids n") if b_dtype is not None: raise ValueError("Dense mode forbids b_dtype") - if b_tensor.dtype != torch.bfloat16: + if _convert_to_cutlass_data_type(b_tensor.dtype) is not cutlass.BFloat16: raise ValueError(f"b_tensor must have dtype torch.bfloat16, got {b_tensor.dtype}") - if b_tensor.ndim != 3: - raise ValueError(f"b_tensor must have shape (n, k, experts), got {tuple(b_tensor.shape)}") - n, b_k, experts = b_tensor.shape + if len(get_shape(b_tensor)) != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got {get_shape(b_tensor)}") + n, b_k, experts = get_shape(b_tensor) if b_k != k: raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") else: - _require_pointer_tensor(b_ptrs, "b_ptrs") - if b_ptrs.device != a_tensor.device: - raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({a_tensor.device}), got {b_ptrs.device}") - if b_ptrs.data_ptr() % 8 != 0: + experts = _validate_pointer_tensor(b_ptrs, "b_ptrs") + if get_device(b_ptrs) != get_device(a_tensor): + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({get_device(a_tensor)}), got {get_device(b_ptrs)}") + if get_data_ptr(b_ptrs) % 8 != 0: raise ValueError("b_ptrs data pointer must be 8-byte aligned") - if padded_offsets.ndim == 1 and b_ptrs.numel() != padded_offsets.numel(): - raise ValueError(f"b_ptrs length mismatch: expected {padded_offsets.numel()}, " f"got {b_ptrs.numel()}") + offsets_shape = get_shape(padded_offsets) + if len(offsets_shape) == 1 and experts != offsets_shape[0]: + raise ValueError(f"b_ptrs length mismatch: expected {offsets_shape[0]}, " f"got {experts}") if n is None or b_dtype is None: raise ValueError("Discrete mode requires n and b_dtype") - if b_dtype != torch.bfloat16: + if b_dtype is not cutlass.BFloat16: raise ValueError(f"b_dtype must be torch.bfloat16 for the BF16 backend, got {b_dtype}") if n <= 0: raise ValueError(f"n must be > 0, got {n}") - experts = b_ptrs.numel() return is_dense, tensor_m, n, experts @@ -257,18 +273,17 @@ def grouped_gemm_wrapper_sm100( use_dynamic_sched: bool = False, current_stream: Optional[cuda.CUstream] = None, ) -> TupleDict: - from cudnn.tensor_adapter import is_torch_tensor - - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if c_dtype is None: - c_dtype = torch.bfloat16 - if d_dtype is None: - d_dtype = torch.bfloat16 + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_wrapper_sm100; pass torch tensors or JAX arrays") + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + c_dtype = _convert_to_cutlass_data_type(c_dtype) if c_dtype is not None else cutlass.BFloat16 + d_dtype = _convert_to_cutlass_data_type(d_dtype) if d_dtype is not None else cutlass.BFloat16 + b_dtype = _convert_to_cutlass_data_type(b_dtype) if b_dtype is not None else None + if framework == "jax" and bias_tensor is not None: + raise ValueError( + "bias_tensor is not expressible as a JAX array (its (n, experts) column-major layout has no row-major equivalent); " "omit bias for JAX inputs" + ) is_dense, tensor_m, n_out, expert_cnt = _normalize_call( a_tensor, padded_offsets, @@ -281,16 +296,32 @@ def grouped_gemm_wrapper_sm100( d_dtype, cd_major, m_aligned, + framework, ) expected_shape = (tensor_m, n_out, 1) expected_stride = (n_out, 1, tensor_m * n_out) + + def _allocate_output(dtype): + from cudnn.tensor_adapter import framework_dtype + + if framework == "torch": + import torch + + return torch.empty_strided( + expected_shape, + expected_stride, + dtype=framework_dtype(dtype, "torch"), + device=a_tensor.device, + ) + import jax + import jax.numpy as jnp + + # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. + # The kernel writes into this buffer on the launch stream; materialize it first. + return jax.block_until_ready(jnp.empty(expected_shape, dtype=framework_dtype(dtype, "jax"), device=a_tensor.device)) + if c_tensor is None: - internal_c = torch.empty_strided( - expected_shape, - expected_stride, - dtype=c_dtype, - device=a_tensor.device, - ) + internal_c = _allocate_output(c_dtype) else: _validate_output( c_tensor, @@ -298,16 +329,11 @@ def grouped_gemm_wrapper_sm100( shape=expected_shape, stride=expected_stride, dtype=c_dtype, - device=a_tensor.device, + device=get_device(a_tensor), ) internal_c = c_tensor if d_tensor is None: - d_tensor = torch.empty_strided( - expected_shape, - expected_stride, - dtype=d_dtype, - device=a_tensor.device, - ) + d_tensor = _allocate_output(d_dtype) else: _validate_output( d_tensor, @@ -315,14 +341,15 @@ def grouped_gemm_wrapper_sm100( shape=expected_shape, stride=expected_stride, dtype=d_dtype, - device=a_tensor.device, + device=get_device(a_tensor), ) overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) workspace_bytes = (128 * expert_cnt if not is_dense else 0) + (4 if use_dynamic_sched else 0) + a_device = get_device(a_tensor) workspace_signature = ( (max(workspace_bytes, 1),), - (a_tensor.device.type, a_tensor.device.index), + (a_device.type, a_device.index), ) cache_key = ( "bf16", diff --git a/python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py b/python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py new file mode 100644 index 000000000..259549402 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the unfused BF16 grouped GEMM +(discrete weight mode), built on :func:`cudnn.jax.call`. + +Discrete mode only (dense mode's expert-outermost strided B has no row-major JAX +equivalent) and no bias (its (n, experts) column-major layout is likewise +inexpressible; ``bias`` is a compile-time ``None`` inside the adapter). The +per-expert weight pointers travel as a regular device array whose *values* are +raw addresses — the referenced weight buffers are not visible to XLA, so the +caller must keep them alive (and unmoved) across every execution of the traced +computation. ``padded_offsets`` values cannot be host-validated under tracing; +malformed offsets are the caller's responsibility here (the eager wrapper +validates them). +""" + +import os +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils +from cutlass.cute.nvgpu import OperandMajorMode + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import TensorSpec, call, gemm_operand_spec, zeros_init +from ..moe_utils import MoEWeightMode +from .moe_grouped_gemm import MoEGroupedGemmBf16Kernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_output_dtypes = (cutlass.BFloat16, cutlass.Float16, cutlass.Float32) + + +def _prob_spec() -> TensorSpec: + # (m, 1, 1) with m innermost: explicit ranks because trailing unit dims make + # leading-dim inference ambiguous + return TensorSpec(layout=(0, 1, 2)) + + +@cute.jit +def _grouped_bf16_adapter(stream, a, b_ptrs, padded_offsets, alpha, prob, d, c, workspace, *, kernel, n, k, mac): + # Discrete-mode b is a raw pointer to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffer recasts for free. + b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) + kernel( + a=a, + b=b_arg, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace.iterator, + c=c, + d=d, + padded_offsets=padded_offsets, + alpha=alpha, + bias=None, + prob=prob, + max_active_clusters=mac, + stream=stream, + ) + + +def _pointer_count(b_ptrs: Any, name: str = "b_ptrs") -> int: + """Tracing-safe pointer-array shape/dtype check; returns the pointer count.""" + shape = tuple(b_ptrs.shape) + if len(shape) != 1: + raise ValueError(f"{name} must be 1-D, got shape={shape}") + dtype = _convert_to_cutlass_data_type(b_ptrs.dtype) + if dtype is cutlass.Int64: + return shape[0] + if dtype is cutlass.Uint8: + if shape[0] % 8 != 0: + raise ValueError(f"{name} packed uint8 length must be a multiple of 8, got {shape[0]}") + return shape[0] // 8 + raise ValueError(f"{name} must be int64 (or, without x64 mode, packed uint8), got {b_ptrs.dtype}") + + +def grouped_gemm_jax_sm100( + a_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_ptrs: Any, + n: int, + prob_tensor: Any, + c_dtype: Any = cutlass.BFloat16, + d_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + generate_c: bool = False, + use_dynamic_sched: bool = False, +) -> Tuple[Any, Optional[Any]]: + """Unfused BF16 grouped GEMM (discrete weights) as an XLA custom call. + + Same contract as the eager wrapper's discrete mode: A ``(m, k, 1)`` k-major + C-contiguous bfloat16, ``padded_offsets (experts,)`` int32 cumulative 256-aligned + row offsets, ``alpha (experts,)`` float32, ``prob (m, 1, 1)`` float32, and + ``b_ptrs`` holding per-expert ``(n, k)`` k-major bfloat16 weight base addresses + (packed little-endian uint8, 8 bytes per pointer — or int64 with x64 mode). + Returns ``(d_tensor, c_tensor)`` with ``c_tensor`` None unless ``generate_c``; + rows at/past ``padded_offsets[-1]`` come back zero-filled (the outputs are + donated zero-initialized buffers). + """ + c_dtype = _convert_to_cutlass_data_type(c_dtype) + d_dtype = _convert_to_cutlass_data_type(d_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + + if len(a_tensor.shape) != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + m, k, _ = a_tensor.shape + if m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {m}") + if _convert_to_cutlass_data_type(a_tensor.dtype) is not cutlass.BFloat16: + raise ValueError(f"a_tensor must have dtype bfloat16, got {a_tensor.dtype}") + if n is None or n <= 0: + raise ValueError(f"n must be > 0, got {n}") + if c_dtype not in _output_dtypes or d_dtype not in _output_dtypes: + raise ValueError(f"c_dtype/d_dtype must be BF16, FP16, or FP32, got {c_dtype}/{d_dtype}") + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + + expert_cnt = _pointer_count(b_ptrs) + if expert_cnt <= 0 or expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {expert_cnt}") + if tuple(padded_offsets.shape) != (expert_cnt,): + raise ValueError(f"padded_offsets must have shape ({expert_cnt},), got {tuple(padded_offsets.shape)}") + if _convert_to_cutlass_data_type(padded_offsets.dtype) is not cutlass.Int32: + raise ValueError(f"padded_offsets must have dtype int32, got {padded_offsets.dtype}") + if tuple(alpha_tensor.shape) != (expert_cnt,) or _convert_to_cutlass_data_type(alpha_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"alpha_tensor must be ({expert_cnt},) float32, got {tuple(alpha_tensor.shape)} {alpha_tensor.dtype}") + if tuple(prob_tensor.shape) != (m, 1, 1) or _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: + raise ValueError(f"prob_tensor must be ({m}, 1, 1) float32, got {tuple(prob_tensor.shape)} {prob_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + + if not MoEGroupedGemmBf16Kernel.can_implement( + cutlass.BFloat16, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + k, + expert_cnt, + "k", + "k", + "n", + MoEGroupedGemmBf16Kernel.FIX_PAD_SIZE, + ): + raise ValueError("Unsupported BF16 grouped GEMM tile, cluster, alignment, or layout configuration") + + cache_key = ( + expert_cnt, + c_dtype, + d_dtype, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + vector_f32, + generate_c, + use_dynamic_sched, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = MoEGroupedGemmBf16Kernel( + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_c=generate_c, + enable_bias=False, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DISCRETE, + use_dynamic_sched=use_dynamic_sched, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + operand = gemm_operand_spec() + d_tensor, c_tensor, _workspace = call( + _grouped_bf16_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n, 1), framework_dtype(d_dtype, "jax")), + jax.ShapeDtypeStruct((m, n, 1), framework_dtype(c_dtype, "jax")), + jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8), + ), + input_spec=(operand, None, None, None, _prob_spec()), + output_spec=(operand, operand, None), + # All three donated: c/d for the trailing-unit-dim layout spec (and defined + # bytes past the last offset); the workspace because the helper kernel writes + # the per-expert TMA descriptors into it (XLA inputs are immutable). + initialized_outputs={0: zeros_init, 1: zeros_init, 2: zeros_init}, + kernel=kernel, + n=int(n), + k=int(k), + mac=mac, + )(a_tensor, b_ptrs, padded_offsets, alpha_tensor, prob_tensor) + + return d_tensor, (c_tensor if generate_c else None) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py index b68d8926d..85d717865 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/__init__.py @@ -9,4 +9,15 @@ __all__ = [ "GroupedGemmWgradSm100", "grouped_gemm_wgrad_wrapper_sm100", + "grouped_gemm_wgrad_jax_sm100", ] + + +def __getattr__(name): + # Lazy: the jax entry point imports jax/cutlass.jax, which must not be pulled in + # for torch-only users. + if name == "grouped_gemm_wgrad_jax_sm100": + from .jax_api import grouped_gemm_wgrad_jax_sm100 + + return grouped_gemm_wgrad_jax_sm100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py index 8f33a166d..e8008274b 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py @@ -16,22 +16,31 @@ from cudnn.api_base import APIBase, TensorDesc from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _pointer_values, _validate_pointer_tensor +from cudnn.tensor_adapter import ( + allocate_byte_workspace, + canonicalize_unit_dim_strides, + cuda_is_available, + default_stream, + detect_framework, + framework_dtype, + get_compute_capability, + get_data_ptr, + get_device, + get_shape, + get_strides, + get_version, + is_torch_tensor, + to_host_list, +) from ..backend_utils import _torch_stream_context from ..moe_utils import MoEWeightMode, WGradInputOrder from .moe_grouped_gemm_wgrad import MoEGroupedGemmWgradBF16Kernel -_OUTPUT_DTYPES = None - def _output_dtypes(): - global _OUTPUT_DTYPES - if _OUTPUT_DTYPES is None: - import torch - - _OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] - return _OUTPUT_DTYPES + return [cutlass.BFloat16, cutlass.Float16, cutlass.Float32] class GroupedGemmWgradBf16API(APIBase): @@ -58,12 +67,11 @@ def __init__( accumulate_on_output: bool = False, input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, ) -> None: - import torch - if acc_dtype is None: - acc_dtype = torch.float32 + acc_dtype = cutlass.Float32 super().__init__() self._warn_experimental_api() + self._framework = detect_framework(sample_a) self.input_order = WGradInputOrder(input_order) if sample_wgrad is not None and num_experts is None: self.weight_mode = MoEWeightMode.DENSE @@ -74,14 +82,14 @@ def __init__( else: raise ValueError("Provide either sample_wgrad for dense mode or " "(num_experts, wgrad_shape, wgrad_dtype) for discrete mode, but not both") - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.offsets_desc = self._make_tensor_desc(sample_offsets, name="sample_offsets") - self.wgrad_desc = self._make_tensor_desc(sample_wgrad, name="sample_wgrad") - self.single_expert_wgrad_desc = self._make_tensor_desc(sample_wgrad_expert, name="sample_wgrad_expert") + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) + self.offsets_desc = self._make_tensor_desc(sample_offsets, name="sample_offsets", canonical=True) + self.wgrad_desc = self._make_tensor_desc(sample_wgrad, name="sample_wgrad", canonical=True) + self.single_expert_wgrad_desc = self._make_tensor_desc(sample_wgrad_expert, name="sample_wgrad_expert", canonical=True) self.expert_cnt = self.wgrad_desc.shape[0] if self.weight_mode == MoEWeightMode.DENSE and self.wgrad_desc.ndim == 3 else int(num_experts or 0) self.wgrad_shape = self.wgrad_desc.shape[1:] if self.weight_mode == MoEWeightMode.DENSE and self.wgrad_desc.ndim == 3 else tuple(wgrad_shape or ()) - self.wgrad_dtype = self.wgrad_desc.dtype if self.weight_mode == MoEWeightMode.DENSE else wgrad_dtype + self.wgrad_dtype = self.wgrad_desc.dtype if self.weight_mode == MoEWeightMode.DENSE else _convert_to_cutlass_data_type(wgrad_dtype) if self.weight_mode == MoEWeightMode.DISCRETE and self.single_expert_wgrad_desc is None: self.single_expert_wgrad_desc = TensorDesc( dtype=self.wgrad_dtype, @@ -92,7 +100,7 @@ def __init__( name="single_expert_wgrad", ) - self.acc_dtype = acc_dtype + self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) self.mma_tiler_mn = tuple(mma_tiler_mn) self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) @@ -108,13 +116,14 @@ def __init__( self._workspace: Optional[torch.Tensor] = None self._compile_wgrad_ptrs: Optional[torch.Tensor] = None self._single_expert_placeholder: Optional[torch.Tensor] = None + self._live_wgrad_ptrs = None self._validated_offsets: dict[int, tuple] = {} self._validated_pointer_values: dict[int, tuple] = {} self._sample_offset_values = self._copy_values_to_host(sample_offsets) self._sample_offsets_ref = weakref.ref(sample_offsets) - self._sample_offsets_version = int(sample_offsets._version) + self._sample_offsets_version = get_version(sample_offsets) self._sample_data_ptrs = { - name: tensor.data_ptr() + name: get_data_ptr(tensor) for name, tensor in ( ("sample_a", sample_a), ("sample_b", sample_b), @@ -130,12 +139,12 @@ def __init__( @staticmethod def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(int(value) for value in tensor.detach().cpu().tolist()) + return tuple(int(value) for value in to_host_list(tensor)) @staticmethod def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: cached = cache.get(id(tensor)) - return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + return bool(cached and cached[0]() is tensor and cached[1] == get_version(tensor) and cached[2] == extra) @staticmethod def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: @@ -144,7 +153,7 @@ def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) - def discard(_reference, *, cache=cache, key=key): cache.pop(key, None) - cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + cache[key] = (weakref.ref(tensor, discard), get_version(tensor), extra) @staticmethod def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tokens_sum: int) -> Tuple[int, ...]: @@ -175,23 +184,27 @@ def _validate_offsets_once(self, offsets: torch.Tensor, *, tokens_sum: int) -> N def _validate_pointer_values_once(self, pointers: torch.Tensor) -> None: if self._is_validation_cached(self._validated_pointer_values, pointers, self.expert_cnt): return - values = self._copy_values_to_host(pointers) + values = _pointer_values(pointers) if any(value == 0 or value % 16 != 0 for value in values): raise ValueError("wgrad_ptrs entries must be non-null and 16-byte aligned") self._remember_validation(self._validated_pointer_values, pointers, self.expert_cnt) @staticmethod def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: - if tensor.data_ptr() % 8 != 0: + if get_data_ptr(tensor) % 8 != 0: raise ValueError("wgrad_ptrs data pointer must be 8-byte aligned") @staticmethod def _validate_data_alignment(tensor: torch.Tensor, name: str, alignment: int = 16) -> None: - if tensor.data_ptr() % alignment != 0: + if get_data_ptr(tensor) % alignment != 0: raise ValueError(f"{name} data pointer must be {alignment}-byte aligned") - @staticmethod - def _record_pointer_stream(pointers: torch.Tensor, current_stream: cuda.CUstream) -> None: + def _record_pointer_stream(self, pointers: torch.Tensor, current_stream: cuda.CUstream) -> None: + if not is_torch_tensor(pointers): + # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the + # array referenced until the next execute so its buffer outlives the launch. + self._live_wgrad_ptrs = pointers + return import torch handle = int(current_stream) @@ -208,18 +221,20 @@ def _record_pointer_stream(pointers: torch.Tensor, current_stream: cuda.CUstream @staticmethod def _infer_a_major(desc: TensorDesc) -> str: m, tokens = desc.shape - if desc.stride == (tokens, 1): + stride = canonicalize_unit_dim_strides(desc.shape, desc.stride) + if stride == canonicalize_unit_dim_strides(desc.shape, (tokens, 1)): return "k" - if desc.stride == (1, m): + if stride == canonicalize_unit_dim_strides(desc.shape, (1, m)): return "m" raise ValueError(f"A tensor must use a supported K-major or M-major layout, got stride {desc.stride}") @staticmethod def _infer_b_major(desc: TensorDesc) -> str: tokens, n = desc.shape - if desc.stride == (1, tokens): + stride = canonicalize_unit_dim_strides(desc.shape, desc.stride) + if stride == canonicalize_unit_dim_strides(desc.shape, (1, tokens)): return "k" - if desc.stride == (n, 1): + if stride == canonicalize_unit_dim_strides(desc.shape, (n, 1)): return "n" raise ValueError(f"B tensor must use a supported K-major or N-major layout, got stride {desc.stride}") @@ -229,8 +244,6 @@ def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: raise ValueError(f"{name} must be on {device}, got {desc.device}") def check_support(self) -> bool: - import torch - if self.a_desc.ndim != 2: raise ValueError(f"sample_a must be rank-2, got {self.a_desc.shape}") if self.b_desc.ndim != 2: @@ -241,11 +254,11 @@ def check_support(self) -> bool: raise ValueError(f"sample_a and sample_b token dimensions must match, got {tokens_sum} and {tokens_b}") self.a_major = self._infer_a_major(self.a_desc) self.b_major = self._infer_b_major(self.b_desc) - self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") - self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") - self._check_dtype(self.offsets_desc, torch.int32, "sample_offsets") + self._check_dtype(self.a_desc, cutlass.BFloat16, "sample_a") + self._check_dtype(self.b_desc, cutlass.BFloat16, "sample_b") + self._check_dtype(self.offsets_desc, cutlass.Int32, "sample_offsets") self._check_dtype(self.wgrad_dtype, _output_dtypes(), "wgrad_dtype") - if self.acc_dtype != torch.float32: + if self.acc_dtype is not cutlass.Float32: raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") if any(control is not None for control in self._scale_controls): raise ValueError("BF16 wgrad forbids scale and global-scale tensors") @@ -261,7 +274,7 @@ def check_support(self) -> bool: raise ValueError(f"sample_wgrad must be rank-3, got {self.wgrad_desc.shape}") if self.wgrad_desc.shape != (self.expert_cnt, m, n): raise ValueError(f"sample_wgrad shape mismatch: expected {(self.expert_cnt, m, n)}, got {self.wgrad_desc.shape}") - if self.wgrad_desc.stride != (m * n, n, 1): + if self.wgrad_desc.stride != canonicalize_unit_dim_strides(self.wgrad_desc.shape, (m * n, n, 1)): raise ValueError("sample_wgrad must be contiguous in expert/M/N order") output_desc = self.wgrad_desc else: @@ -270,7 +283,7 @@ def check_support(self) -> bool: output_desc = self.single_expert_wgrad_desc if output_desc.shape not in ((m, n), (m, n, 1)): raise ValueError(f"sample_wgrad_expert shape mismatch: expected {(m, n)}, got {output_desc.shape}") - expected_stride = (n, 1) if output_desc.ndim == 2 else (n, 1, 1) + expected_stride = (n, 1) if output_desc.ndim == 2 else canonicalize_unit_dim_strides((m, n, 1), (n, 1, 1)) if output_desc.stride != expected_stride: raise ValueError("sample_wgrad_expert must be contiguous in M/N order") self._check_dtype(output_desc, self.wgrad_dtype, "sample_wgrad_expert") @@ -293,7 +306,7 @@ def check_support(self) -> bool: tokens_sum=tokens_sum, ) sample_offsets = self._sample_offsets_ref() - if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + if sample_offsets is not None and get_version(sample_offsets) == self._sample_offsets_version: self._remember_validation( self._validated_offsets, sample_offsets, @@ -303,7 +316,7 @@ def check_support(self) -> bool: self._validate_offsets_once(sample_offsets, tokens_sum=tokens_sum) if not self._kernel.can_implement( - _convert_to_cutlass_data_type(torch.bfloat16), + cutlass.BFloat16, _convert_to_cutlass_data_type(self.wgrad_dtype), _convert_to_cutlass_data_type(self.acc_dtype), self.use_2cta_instrs, @@ -319,18 +332,38 @@ def check_support(self) -> bool: self.input_order, ): raise ValueError("Unsupported BF16 grouped GEMM wgrad configuration: check mma_tiler, cluster, and alignment") - if not torch.cuda.is_available(): + if not cuda_is_available(): raise RuntimeError("CUDA is not available") - major, minor = torch.cuda.get_device_capability(self.a_desc.device) + major, minor = get_compute_capability() capability = major * 10 + minor if capability < 100: - raise RuntimeError(f"GroupedGemmWgradSm100 requires SM100+, found SM{capability} on {self.a_desc.device}") + raise RuntimeError(f"GroupedGemmWgradSm100 requires SM100+, found SM{capability}") self._is_supported = True return True - def compile(self) -> None: - import torch + def _allocate_single_expert_placeholder(self) -> None: + """Allocate the real (never read) discrete-mode single-expert template tensor.""" + desc = self.single_expert_wgrad_desc + if self._framework == "torch": + import torch + + self._single_expert_placeholder = torch.empty_strided( + desc.shape, + desc.stride, + dtype=framework_dtype(desc.dtype, "torch"), + device=desc.device, + ) + return + import jax + import jax.numpy as jnp + + if canonicalize_unit_dim_strides(desc.shape, desc.stride) != canonicalize_unit_dim_strides( + desc.shape, TensorDesc._compute_contiguous_stride(desc.shape) + ): + raise ValueError(f"single expert placeholder layout {desc.stride} is not expressible as a C-contiguous JAX array") + self._single_expert_placeholder = jax.block_until_ready(jnp.empty(desc.shape, dtype=framework_dtype(desc.dtype, "jax"))) + def compile(self) -> None: self._ensure_support_checked() if self._compiled_kernel is not None: return @@ -348,11 +381,9 @@ def compile(self) -> None: max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin if max_active_clusters <= 0: raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") - self._workspace = torch.empty( - max(kernel.get_workspace_bytes(), 1), - dtype=torch.uint8, - device=self.a_desc.device, - ) + # Internal scratch in the caller's framework allocator; kernels write through its + # raw pointer and it is never surfaced as a framework array. + self._workspace = allocate_byte_workspace(self._framework, kernel.get_workspace_bytes(), self.a_desc.device) self._validate_data_alignment(self._workspace, "workspace", 128) workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -377,9 +408,14 @@ def compile(self) -> None: out_fake = self._make_fake_cute_tensor_from_desc(self.wgrad_desc, assumed_align=16) single_expert_fake = None else: - self._compile_wgrad_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + # Compile-time placeholder for the pointer-array argument: real device bytes + # (fake tensors have dummy iterators) allocated in the caller's framework, + # retyped to Int64 via the element_type override. + self._compile_wgrad_ptrs = allocate_byte_workspace(self._framework, 8 * self.expert_cnt, self.a_desc.device) self._validate_pointer_array_alignment(self._compile_wgrad_ptrs) - out_fake = from_dlpack(self._compile_wgrad_ptrs, assumed_align=8).iterator + placeholder = from_dlpack(self._compile_wgrad_ptrs, assumed_align=8) + placeholder.element_type = cutlass.Int64 + out_fake = placeholder.iterator single_expert_fake = self._make_fake_cute_tensor_from_desc(self.single_expert_wgrad_desc, assumed_align=16) raw_compiled = cute.compile( kernel, @@ -395,12 +431,7 @@ def compile(self) -> None: ) cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) if self.weight_mode == MoEWeightMode.DISCRETE: - self._single_expert_placeholder = torch.empty_strided( - self.single_expert_wgrad_desc.shape, - self.single_expert_wgrad_desc.stride, - dtype=self.single_expert_wgrad_desc.dtype, - device=self.single_expert_wgrad_desc.device, - ) + self._allocate_single_expert_placeholder() self._validate_data_alignment(self._single_expert_placeholder, "single expert placeholder") cached_single_expert = from_dlpack( self._single_expert_placeholder, @@ -411,7 +442,7 @@ def compile(self) -> None: cached_single_expert = None def tensor_api(a_tensor, b_tensor, output, offsets, stream) -> None: - out_arg = output if self.weight_mode == MoEWeightMode.DENSE else int(output.data_ptr()) + out_arg = output if self.weight_mode == MoEWeightMode.DENSE else int(get_data_ptr(output)) raw_compiled( a_tensor, b_tensor, @@ -432,7 +463,7 @@ def _validate_live_input( *, token_axis: int, ) -> TensorDesc: - desc = self._make_tensor_desc(tensor, name=name) + desc = self._make_tensor_desc(tensor, name=name, canonical=True) if desc.dtype != sample.dtype: raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") if desc.device != sample.device: @@ -447,14 +478,17 @@ def _validate_live_input( return desc def _validate_live_output(self, tensor: torch.Tensor) -> None: - desc = self._make_tensor_desc(tensor, name="wgrad_tensor") + desc = self._make_tensor_desc(tensor, name="wgrad_tensor", canonical=True) expected = (self.expert_cnt, *self.wgrad_shape) if desc.shape != expected: raise ValueError(f"wgrad_tensor shape mismatch: expected {expected}, got {desc.shape}") - if desc.stride != ( - self.wgrad_shape[0] * self.wgrad_shape[1], - self.wgrad_shape[1], - 1, + if desc.stride != canonicalize_unit_dim_strides( + expected, + ( + self.wgrad_shape[0] * self.wgrad_shape[1], + self.wgrad_shape[1], + 1, + ), ): raise ValueError("wgrad_tensor must be contiguous in expert/M/N order") if desc.dtype != self.wgrad_dtype: @@ -462,6 +496,24 @@ def _validate_live_output(self, tensor: torch.Tensor) -> None: if desc.device != self.a_desc.device: raise ValueError(f"wgrad_tensor device mismatch: expected {self.a_desc.device}, got {desc.device}") + def _generate_wgrad_ptrs(self, wgrad_tensor: torch.Tensor, current_stream: cuda.CUstream): + """Derive the per-expert pointer array from a rank-3 wgrad tensor.""" + element_bits = _convert_to_cutlass_data_type(self.wgrad_dtype).width + stride_bytes = get_strides(wgrad_tensor)[0] * element_bits // 8 + base_ptr = get_data_ptr(wgrad_tensor) + values = [base_ptr + index * stride_bytes for index in range(self.expert_cnt)] + if is_torch_tensor(wgrad_tensor): + import torch + + with _torch_stream_context(current_stream, wgrad_tensor.device): + return torch.tensor(values, dtype=torch.int64, device=wgrad_tensor.device) + import jax + import jax.numpy as jnp + import numpy as np + + # Packed little-endian uint8 pointer bytes: JAX truncates int64 without x64 mode. + return jax.block_until_ready(jnp.asarray(np.asarray(values, dtype=np.int64).view(np.uint8))) + def execute( self, a_tensor: torch.Tensor, @@ -475,9 +527,10 @@ def execute( global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: - import torch - - current_stream = self._get_default_stream(current_stream) + if current_stream is None: + # torch inputs stay ordered with the caller's current torch stream; + # other frameworks (e.g. JAX) default to the CUDA legacy default stream. + current_stream = default_stream(detect_framework(a_tensor)) if self._compiled_kernel is None: raise RuntimeError("Kernel not compiled; call compile() first") forbidden = ( @@ -494,8 +547,8 @@ def execute( tokens_sum = a_desc.shape[1] if b_desc.shape[0] != tokens_sum: raise ValueError("a_tensor and b_tensor token dimensions must match") - offsets_desc = self._make_tensor_desc(offsets_tensor, name="offsets_tensor") - if offsets_desc.shape != (self.expert_cnt,) or offsets_desc.stride != (1,) or offsets_desc.dtype != torch.int32: + offsets_desc = self._make_tensor_desc(offsets_tensor, name="offsets_tensor", canonical=True) + if offsets_desc.shape != (self.expert_cnt,) or offsets_desc.stride != (1,) or offsets_desc.dtype is not cutlass.Int32: raise ValueError("offsets_tensor must be a contiguous rank-1 int32 tensor with one entry per expert") if offsets_desc.device != self.a_desc.device: raise ValueError(f"offsets_tensor device mismatch: expected {self.a_desc.device}, got {offsets_desc.device}") @@ -518,16 +571,10 @@ def execute( if wgrad_ptrs is None: if wgrad_tensor is None: raise ValueError("Discrete execution requires wgrad_tensor or wgrad_ptrs") - stride_bytes = wgrad_tensor.stride(0) * wgrad_tensor.element_size() - with _torch_stream_context(current_stream, wgrad_tensor.device): - wgrad_ptrs = torch.tensor( - [wgrad_tensor.data_ptr() + index * stride_bytes for index in range(self.expert_cnt)], - dtype=torch.int64, - device=wgrad_tensor.device, - ) - _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) - if wgrad_ptrs.device != self.a_desc.device: - raise ValueError(f"wgrad_ptrs must be on {self.a_desc.device}, got {wgrad_ptrs.device}") + wgrad_ptrs = self._generate_wgrad_ptrs(wgrad_tensor, current_stream) + _validate_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + if get_device(wgrad_ptrs) != self.a_desc.device: + raise ValueError(f"wgrad_ptrs must be on {self.a_desc.device}, got {get_device(wgrad_ptrs)}") self._validate_pointer_array_alignment(wgrad_ptrs) if not generated_wgrad_ptrs: self._validate_pointer_values_once(wgrad_ptrs) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py index eb3a958a4..1c0094169 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py @@ -14,7 +14,8 @@ from cudnn.api_base import APIBase, TensorDesc, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor +from cudnn.tensor_adapter import is_torch_tensor from .moe_blockscaled_grouped_gemm_wgrad import BlockScaledMoEGroupedGemmWgradKernel from ..moe_utils import MoEWeightMode, WGradInputOrder @@ -65,10 +66,20 @@ def __init__( accumulate_on_output: bool = False, input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, ): + if sample_a is not None and not is_torch_tensor(sample_a): + raise ValueError( + "The block-scaled wgrad backend supports torch tensors only: its B operand " + "(and fp4-packed A/B operands) require K-major (token-innermost) layouts that " + "are not expressible as row-major JAX arrays" + ) import torch - if acc_dtype is None: - acc_dtype = torch.float32 + from cudnn.tensor_adapter import framework_dtype + + # This backend is torch-internal: normalize loose dtype parameters (which the + # type-erased wrapper/facade may pass as cutlass or numpy dtypes) to torch dtypes. + acc_dtype = framework_dtype(acc_dtype, "torch") if acc_dtype is not None else torch.float32 + wgrad_dtype = framework_dtype(wgrad_dtype, "torch") if wgrad_dtype is not None else None super().__init__() self._warn_experimental_api() self.input_order = WGradInputOrder(input_order) @@ -575,7 +586,7 @@ def execute( expert_stride_bytes = wgrad_tensor.stride(0) * wgrad_tensor.element_size() ptrs = [wgrad_tensor.data_ptr() + i * expert_stride_bytes for i in range(wgrad_tensor.shape[0])] wgrad_ptrs = torch.tensor(ptrs, dtype=torch.int64, device=wgrad_tensor.device) - _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + _validate_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) self._compiled_kernel( a_tensor, b_tensor, diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index b8bd6705d..6ddf55588 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -10,8 +10,19 @@ from cuda.bindings import driver as cuda +import cutlass + from cudnn.api_base import APIBase, TupleDict, get_device_type -from cudnn.gemm.cutedsl.discrete_grouped.discrete_kernel_utils import _require_pointer_tensor +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.gemm.cutedsl.grouped.unfused._bf16_api import _validate_pointer_tensor +from cudnn.tensor_adapter import ( + canonicalize_unit_dim_strides, + detect_framework, + framework_dtype, + get_device, + get_shape, + get_strides, +) from ..backend_utils import ( GroupedGemmBackend, @@ -21,24 +32,17 @@ ) from ..moe_utils import WGradInputOrder -_BLOCK_SCALED_DTYPE_PAIRS = None - def _block_scaled_dtype_pairs(): - global _BLOCK_SCALED_DTYPE_PAIRS - if _BLOCK_SCALED_DTYPE_PAIRS is None: - import torch - - _BLOCK_SCALED_DTYPE_PAIRS = { - (dtype, dtype) - for dtype in ( - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ) - } - return _BLOCK_SCALED_DTYPE_PAIRS + return { + (dtype, dtype) + for dtype in ( + cutlass.Float4E2M1FN, + cutlass.Uint8, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + ) + } _cache_of_GroupedGemmWgradSm100Objects = {} @@ -51,6 +55,12 @@ def _block_scaled_dtype_pairs(): _is_supported_rubin_quantization, ) +_BLOCK_SCALED_JAX_ERROR = ( + "the block-scaled wgrad backend is not expressible as JAX arrays " + "(its B operand requires a K-major, token-innermost layout and fp4 operands are K-packed, " + "neither of which has a row-major equivalent); use torch tensors, or bfloat16 operands for the BF16 backend" +) + class GroupedGemmWgradSm100(APIBase): """Stable public facade that selects the WGrad backend during support checking.""" @@ -106,14 +116,8 @@ def __init__( self._pending_init_kwargs = dict(locals()) self._pending_init_kwargs.pop("self") self._pending_init_kwargs.pop("__class__", None) - from cudnn.tensor_adapter import is_torch_tensor - - if sample_a is not None and not is_torch_tensor(sample_a): - raise ValueError("GroupedGemmWgradSm100 currently supports torch tensors only; JAX support is not yet implemented for this API") if acc_dtype is None: - import torch - - self._pending_init_kwargs["acc_dtype"] = torch.float32 + self._pending_init_kwargs["acc_dtype"] = cutlass.Float32 self._implementation = None def check_support(self) -> bool: @@ -136,6 +140,8 @@ def check_support(self) -> bool: if backend is GroupedGemmBackend.BF16: self._implementation = GroupedGemmWgradBf16API(**kwargs) else: + if detect_framework(kwargs["sample_a"]) == "jax": + raise ValueError(_BLOCK_SCALED_JAX_ERROR) self._implementation = GroupedGemmWgradBlockScaledAPI(**kwargs) self._kernel = self._implementation._kernel self.weight_mode = self._implementation.weight_mode @@ -218,10 +224,12 @@ def execute( def _wgrad_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_dims: tuple[int, ...] = (), exact_stride: bool): if tensor is None: return None - shape = tuple(None if index in dynamic_dims else int(value) for index, value in enumerate(tensor.shape)) - stride = tuple(int(value) for value in tensor.stride()) - layout = stride if exact_stride else tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], tensor.shape[item[0]]))) - return (shape, layout, tensor.dtype, tensor.device) + tensor_shape = get_shape(tensor) + tensor_stride = canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)) + shape = tuple(None if index in dynamic_dims else int(value) for index, value in enumerate(tensor_shape)) + stride = tuple(int(value) for value in tensor_stride) + layout = stride if exact_stride else tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], tensor_shape[item[0]]))) + return (shape, layout, _convert_to_cutlass_data_type(tensor.dtype), get_device(tensor)) def grouped_gemm_wgrad_wrapper_sm100( @@ -245,32 +253,28 @@ def grouped_gemm_wgrad_wrapper_sm100( current_stream: Optional[cuda.CUstream] = None, ) -> TupleDict: """Compile and execute grouped GEMM wgrad through the selected backend API.""" - from cudnn.tensor_adapter import is_torch_tensor + framework = detect_framework(a_tensor) + if framework not in ("torch", "jax"): + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_wgrad_wrapper_sm100; pass torch tensors or JAX arrays") - if a_tensor is not None and not is_torch_tensor(a_tensor): - raise ValueError("grouped_gemm_wgrad_wrapper_sm100 currently supports torch tensors only; JAX support is not yet implemented for this API") - import torch - - if acc_dtype is None: - acc_dtype = torch.float32 - if wgrad_dtype is None: - wgrad_dtype = torch.bfloat16 + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) if acc_dtype is not None else cutlass.Float32 + wgrad_dtype = _convert_to_cutlass_data_type(wgrad_dtype) if wgrad_dtype is not None else cutlass.BFloat16 if output_mode not in ("dense", "discrete"): raise ValueError(f'output_mode must be "dense" or "discrete", got {output_mode}') - if a_tensor.ndim != 2 or b_tensor.ndim != 2: + if len(get_shape(a_tensor)) != 2 or len(get_shape(b_tensor)) != 2: raise ValueError("a_tensor and b_tensor must both be rank-2") - hidden, tokens_sum = a_tensor.shape - tokens_b, intermediate = b_tensor.shape + hidden, tokens_sum = get_shape(a_tensor) + tokens_b, intermediate = get_shape(b_tensor) if tokens_sum != tokens_b: raise ValueError(f"a_tensor and b_tensor token dimensions must match, got {tokens_sum} and {tokens_b}") - if offsets_tensor.ndim != 1: - raise ValueError(f"offsets_tensor must be rank-1, got shape {tuple(offsets_tensor.shape)}") + if len(get_shape(offsets_tensor)) != 1: + raise ValueError(f"offsets_tensor must be rank-1, got shape {get_shape(offsets_tensor)}") input_order = WGradInputOrder(input_order) - expert_cnt = offsets_tensor.numel() + expert_cnt = get_shape(offsets_tensor)[0] if output_mode == "dense" and wgrad_ptrs is not None: raise ValueError("dense output_mode forbids wgrad_ptrs") if wgrad_ptrs is not None: - _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", expert_cnt) + _validate_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", expert_cnt) backend = select_grouped_gemm_backend( operation="grouped_gemm_wgrad_sm100", a_dtype=a_tensor.dtype, @@ -284,10 +288,24 @@ def grouped_gemm_wgrad_wrapper_sm100( ), block_scaled_dtype_pairs=_block_scaled_dtype_pairs(), ) + if framework == "jax" and backend is GroupedGemmBackend.BLOCK_SCALED: + raise ValueError(_BLOCK_SCALED_JAX_ERROR) if wgrad_tensor is None and wgrad_ptrs is None: - allocator = torch.zeros if accumulate_on_output else torch.empty - with _torch_stream_context(current_stream, a_tensor.device): - wgrad_tensor = allocator((expert_cnt, hidden, intermediate), dtype=wgrad_dtype, device=a_tensor.device) + wgrad_shape = (expert_cnt, hidden, intermediate) + if framework == "torch": + import torch + + allocator = torch.zeros if accumulate_on_output else torch.empty + with _torch_stream_context(current_stream, a_tensor.device): + wgrad_tensor = allocator(wgrad_shape, dtype=framework_dtype(wgrad_dtype, "torch"), device=a_tensor.device) + else: + import jax + import jax.numpy as jnp + + # C-contiguous expert/M/N-order output; the kernel writes into this buffer + # on the launch stream, so materialize it before its pointer is taken. + allocator = jnp.zeros if accumulate_on_output else jnp.empty + wgrad_tensor = jax.block_until_ready(allocator(wgrad_shape, dtype=framework_dtype(wgrad_dtype, "jax"), device=a_tensor.device)) cache_key = backend_cache_key( backend, get_device_type(), @@ -312,6 +330,19 @@ def grouped_gemm_wgrad_wrapper_sm100( ) op = _cache_of_GroupedGemmWgradSm100Objects.get(cache_key) if op is None: + + def _sample_wgrad_expert(): + if wgrad_tensor is not None: + return wgrad_tensor[0] + if framework == "torch": + import torch + + return torch.empty((hidden, intermediate), dtype=framework_dtype(wgrad_dtype, "torch"), device=a_tensor.device) + import jax + import jax.numpy as jnp + + return jax.block_until_ready(jnp.empty((hidden, intermediate), dtype=framework_dtype(wgrad_dtype, "jax"), device=a_tensor.device)) + common = dict( sample_a=a_tensor, sample_b=b_tensor, @@ -331,9 +362,7 @@ def grouped_gemm_wgrad_wrapper_sm100( common["sample_wgrad"] = wgrad_tensor else: common.update( - sample_wgrad_expert=( - wgrad_tensor[0] if wgrad_tensor is not None else torch.empty((hidden, intermediate), dtype=wgrad_dtype, device=a_tensor.device) - ), + sample_wgrad_expert=_sample_wgrad_expert(), num_experts=expert_cnt, wgrad_shape=(hidden, intermediate), wgrad_dtype=wgrad_dtype, diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.py new file mode 100644 index 000000000..17c5e08d4 --- /dev/null +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX-native (XLA custom call) entry point for the BF16 grouped GEMM wgrad +(discrete output mode), built on :func:`cudnn.jax.call`. + +BF16 backend, discrete (pointer-array) output mode only. The block-scaled +backend stays rejected exactly as in the eager wrapper (its B operand requires a +K-major, token-innermost layout and fp4 operands are K-packed — neither has a +row-major JAX equivalent), and only bfloat16 operands reach this entry point. + +Output pointer situation: the per-expert weight gradients are *not* XLA outputs. +``wgrad_ptrs`` is a regular input array whose values are raw device addresses of +caller-owned ``(m, n)`` row-major buffers; the kernel writes through them, so +under jit those buffers live outside XLA's buffer management. The caller must +keep them alive (and unmoved) across every execution of the traced computation, +and must order reads of them on the returned token (``jax.block_until_ready``, +or any data dependency on it). The token is the kernel's never-read discrete-mode +single-expert template output — an ``(m, n)`` zero-filled donated buffer whose +only job is to give the custom call an XLA-managed result (preventing dead-code +elimination) and to carry completion ordering; its contents are meaningless. + +``offsets_tensor`` values cannot be host-validated under tracing; malformed +per-expert offsets (non-cumulative or not 256-aligned) are the caller's +responsibility here (the eager wrapper validates them). Only the total token +count — a static shape — is checked for 256-alignment. +""" + +import os +from typing import Any, Optional, Tuple, Union + +import jax +import jax.numpy as jnp + +import cutlass +import cutlass.cute as cute +import cutlass.utils + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.tensor_adapter import framework_dtype +from cudnn.jax import call, zeros_init +from ..moe_utils import MoEWeightMode, WGradInputOrder +from ..unfused.jax_api import _pointer_count +from .moe_grouped_gemm_wgrad import MoEGroupedGemmWgradBF16Kernel + +# cache_key -> (kernel instance, max_active_clusters, workspace_bytes); reusing the +# instance keeps cutlass_call's compile cache warm (its FunctionSpec keys on the +# constexpr kwargs). +_kernel_cache: dict = {} + +_output_dtypes = (cutlass.BFloat16, cutlass.Float16, cutlass.Float32) + +_BLOCK_SCALED_JAX_ERROR = ( + "only the BF16 wgrad backend is supported for JAX (the block-scaled wgrad backend is " + "not expressible as JAX arrays: its B operand requires a K-major, token-innermost " + "layout and fp4 operands are K-packed, neither of which has a row-major equivalent)" +) + + +@cute.jit +def _wgrad_bf16_adapter(stream, a, b, offsets, wgrad_ptrs, wgrad_template, workspace, *, kernel, mac): + # Discrete-mode out is a raw pointer to the device int64[] of per-expert base + # addresses; the packed uint8 (or int64) input buffer recasts for free. + out_arg = cute.recast_ptr(wgrad_ptrs.iterator, dtype=cutlass.Int64) + kernel( + mat_a=a, + mat_b=b, + out=out_arg, + offs=offsets, + workspace=workspace, + max_active_clusters=mac, + stream=stream, + # Never-read (m, n) layout/dtype template for the per-expert TMA descriptors; + # doubles as the XLA-managed token output returned to the caller. + out_single_expert=wgrad_template, + ) + + +def grouped_gemm_wgrad_jax_sm100( + a_tensor: Any, + b_tensor: Any, + offsets_tensor: Any, + wgrad_ptrs: Any, + wgrad_dtype: Any = cutlass.BFloat16, + acc_dtype: Any = cutlass.Float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + accumulate_on_output: bool = False, + input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, +) -> Any: + """BF16 grouped GEMM wgrad (discrete per-expert output pointers) as an XLA custom call. + + Same contract as the eager wrapper's discrete mode with explicit ``wgrad_ptrs``: + ``a_tensor (m, tokens_sum)`` K-major and ``b_tensor (tokens_sum, n)`` N-major (both + plain C-contiguous bfloat16 JAX arrays), ``offsets_tensor (experts,)`` int32 + cumulative 256-aligned token end-offsets, and ``wgrad_ptrs`` holding the per-expert + ``(m, n)`` row-major ``wgrad_dtype`` output base addresses (packed little-endian + uint8, 8 bytes per pointer — or int64 with x64 mode). The kernel writes each + expert's weight gradient through those caller-owned buffers (with + ``accumulate_on_output`` it TMA-reduces into them, so pre-zero/pre-seed them); + empty experts are zero-filled. + + Returns an opaque ``(m, n)`` token array: block on it (or thread it through a data + dependency) before reading the external per-expert buffers. Its values are + unspecified. Dense (single 3-D wgrad tensor) output is available through the eager + wrapper only. + """ + wgrad_dtype = _convert_to_cutlass_data_type(wgrad_dtype) + acc_dtype = _convert_to_cutlass_data_type(acc_dtype) + input_order = WGradInputOrder(input_order) + + if len(a_tensor.shape) != 2: + raise ValueError(f"a_tensor must have shape (m, tokens_sum), got {tuple(a_tensor.shape)}") + if len(b_tensor.shape) != 2: + raise ValueError(f"b_tensor must have shape (tokens_sum, n), got {tuple(b_tensor.shape)}") + m, tokens_sum = a_tensor.shape + tokens_b, n = b_tensor.shape + if tokens_b != tokens_sum: + raise ValueError(f"a_tensor and b_tensor token dimensions must match, got {tokens_sum} and {tokens_b}") + for name, tensor in (("a_tensor", a_tensor), ("b_tensor", b_tensor)): + if _convert_to_cutlass_data_type(tensor.dtype) is not cutlass.BFloat16: + raise ValueError(f"{name} must have dtype bfloat16, got {tensor.dtype}; {_BLOCK_SCALED_JAX_ERROR}") + if tokens_sum % MoEGroupedGemmWgradBF16Kernel.FIX_PAD_SIZE != 0: + raise ValueError(f"total token count must be {MoEGroupedGemmWgradBF16Kernel.FIX_PAD_SIZE}-aligned, got {tokens_sum}") + if wgrad_dtype not in _output_dtypes: + raise ValueError(f"wgrad_dtype must be BF16, FP16, or FP32, got {wgrad_dtype}") + if acc_dtype is not cutlass.Float32: + raise ValueError(f"acc_dtype must be float32, got {acc_dtype}") + + expert_cnt = _pointer_count(wgrad_ptrs, "wgrad_ptrs") + if expert_cnt <= 0: + raise ValueError(f"expert count must be > 0, got {expert_cnt}") + if tuple(offsets_tensor.shape) != (expert_cnt,): + raise ValueError(f"offsets_tensor must have shape ({expert_cnt},), got {tuple(offsets_tensor.shape)}") + if _convert_to_cutlass_data_type(offsets_tensor.dtype) is not cutlass.Int32: + raise ValueError(f"offsets_tensor must have dtype int32, got {offsets_tensor.dtype}") + + use_2cta_instrs = mma_tiler_mn[0] == 256 + cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if use_2cta_instrs else (1, 1))) + + # Per-expert token counts live in device memory (tracing-safe: no host reads), so + # feed can_implement a synthetic 256-aligned split with the correct static sum; + # it exercises every shape/tile/alignment rule that does not depend on the split. + synthetic_group_k = [int(tokens_sum)] + [0] * (expert_cnt - 1) + if not MoEGroupedGemmWgradBF16Kernel.can_implement( + cutlass.BFloat16, + wgrad_dtype, + acc_dtype, + use_2cta_instrs, + tuple(mma_tiler_mn), + cluster_shape_mn, + m, + n, + synthetic_group_k, + expert_cnt, + "k", # C-contiguous (m, tokens_sum) a_tensor is K-major + "n", # C-contiguous (tokens_sum, n) b_tensor is N-major + MoEWeightMode.DISCRETE, + input_order, + ): + raise ValueError("Unsupported BF16 grouped GEMM wgrad tile, cluster, alignment, or layout configuration") + + cache_key = ( + expert_cnt, + wgrad_dtype, + acc_dtype, + tuple(mma_tiler_mn), + cluster_shape_mn, + accumulate_on_output, + input_order, + ) + entry = _kernel_cache.get(cache_key) + if entry is None: + kernel = MoEGroupedGemmWgradBF16Kernel( + acc_dtype=acc_dtype, + use_2cta_instrs=use_2cta_instrs, + mma_tiler_mn=tuple(mma_tiler_mn), + cluster_shape_mn=cluster_shape_mn, + accumulate_on_output=accumulate_on_output, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DISCRETE, + input_order=input_order, + ) + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - overlap_margin + if mac <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + entry = (kernel, mac, max(kernel.get_workspace_bytes(), 1)) + _kernel_cache[cache_key] = entry + kernel, mac, workspace_bytes = entry + + token, _workspace = call( + _wgrad_bf16_adapter, + output_shape_dtype=( + jax.ShapeDtypeStruct((m, n), framework_dtype(wgrad_dtype, "jax")), + jax.ShapeDtypeStruct((workspace_bytes,), jnp.uint8), + ), + # Both donated: the template so the returned token has defined (zero) bytes + # (the kernel never writes it); the workspace because the helper kernel writes + # the per-expert TMA descriptors into it (XLA inputs are immutable). + initialized_outputs={0: zeros_init, 1: zeros_init}, + kernel=kernel, + mac=mac, + )(a_tensor, b_tensor, offsets_tensor, wgrad_ptrs) + + return token diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py new file mode 100644 index 000000000..87bbf6cca --- /dev/null +++ b/python/cudnn/jax/__init__.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""JAX integration for the cuDNN frontend CuTeDSL APIs. + +Built on CuTeDSL's native JAX bridge (``cutlass.jax.cutlass_call``): kernels run on +XLA's compute stream as FFI custom calls, outputs are XLA-managed, and calls compose +with ``jax.jit`` (and CUDA graph capture). :func:`cudnn.jax.call` is a thin wrapper +adding the conveniences the cuDNN kernels need — pre-initialized accumulator outputs +and TensorSpec presets for the layouts the GEMM fusions use. + +Requires jax >= 0.5 and the CuTeDSL JAX extensions (shipped with nvidia-cutlass-dsl). +""" + +from .call import ( + call, + gemm_operand_spec, + row_major_desc, + sf_atom_spec, + zeros_init, + neg_inf_init, +) +from cutlass.jax import TensorSpec + +__all__ = [ + "call", + "row_major_desc", + "TensorSpec", + "gemm_operand_spec", + "sf_atom_spec", + "zeros_init", + "neg_inf_init", +] diff --git a/python/cudnn/jax/call.py b/python/cudnn/jax/call.py new file mode 100644 index 000000000..c35d2375a --- /dev/null +++ b/python/cudnn/jax/call.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""cudnn.jax.call: cutlass.jax.cutlass_call with cuDNN conveniences.""" + +from typing import Any, Callable, Mapping, Optional, Sequence + +import jax +import jax.numpy as jnp + +import cutlass.jax +from cutlass.jax import TensorSpec, cutlass_call + +if not cutlass.jax.is_available(): # pragma: no cover - guarded import surface + raise ImportError( + "cudnn.jax requires the CuTeDSL JAX extensions (cutlass.jax), which need jax >= 0.5; " "install/upgrade jax (`pip install --group jax` from a checkout)" + ) + + +def row_major_desc(shape, dtype, name: str): + """Metadata-only TensorDesc for a C-contiguous (row-major) JAX buffer. + + Built from aval metadata so this works for jax.jit tracers as well as concrete + arrays (tracers expose .shape/.dtype but no device or DLPack); used to reuse the + class APIs' check_support for validation. + """ + from cudnn.api_base import TensorDesc + from cudnn.datatypes import _convert_to_cutlass_data_type + from cudnn.tensor_adapter import Device + + shape = tuple(shape) + strides, acc = [1] * len(shape), 1 + for i in range(len(shape) - 1, -1, -1): + strides[i] = acc + acc *= shape[i] + stride = tuple(strides) + return TensorDesc( + dtype=_convert_to_cutlass_data_type(dtype), + shape=shape, + stride=stride, + stride_order=TensorDesc._compute_stride_order(shape, stride), + device=Device("cuda", 0), + name=name, + ) + + +def zeros_init(shape_dtype: jax.ShapeDtypeStruct) -> jax.Array: + """Zero-filled initializer for accumulator outputs (e.g. atomic-max amax, atomic-add dprob).""" + return jnp.zeros(shape_dtype.shape, shape_dtype.dtype) + + +def neg_inf_init(shape_dtype: jax.ShapeDtypeStruct) -> jax.Array: + """-inf-filled initializer for max-accumulator outputs.""" + return jnp.full(shape_dtype.shape, -float("inf"), shape_dtype.dtype) + + +def gemm_operand_spec() -> TensorSpec: + """Spec for the (MN, K, 1)-shaped k-/n-major GEMM operands and outputs. + + The trailing unit batch dim makes leading-dim inference ambiguous for a + C-contiguous buffer, so the minor-to-major stride ranks are declared + explicitly: K/N innermost (rank 0), MN next (rank 1), L outermost (rank 2). + """ + return TensorSpec(layout=(1, 0, 2)) + + +def sf_atom_spec() -> TensorSpec: + """Spec presenting a physical C-contiguous (L, MN', K', 32, 4, 4) scale-factor + buffer to the kernel in the logical MMA atom view (32, 4, MN', 4, K', L). + + ``mode`` remaps dimensions without materializing a transpose, so the kernel sees + exactly the layout the torch path compiles from its permuted view. + """ + return TensorSpec(mode=(3, 4, 1, 5, 2, 0)) + + +def call( + fn: Callable[..., None], + *, + output_shape_dtype: Any, + input_spec: Optional[Sequence[Optional[TensorSpec]]] = None, + output_spec: Optional[Sequence[Optional[TensorSpec]]] = None, + initialized_outputs: Optional[Mapping[int, Callable[[jax.ShapeDtypeStruct], jax.Array]]] = None, + input_output_aliases: Optional[dict[int, int]] = None, + allow_cuda_graph: bool = True, + compile_options: Optional[str] = None, + use_static_tensors: bool = False, + **kwargs: Any, +) -> Callable[..., Any]: + """Invoke a ``@cute.jit`` kernel adapter from JAX; see :func:`cutlass.jax.cutlass_call`. + + Same contract as ``cutlass_call`` plus: + + initialized_outputs: ``{output_index: init_fn}`` for outputs the kernel + *accumulates into* rather than fully writes (atomic max/add). For each entry, + ``init_fn(ShapeDtypeStruct) -> jax.Array`` produces the pre-initialized buffer + (e.g. :func:`zeros_init`), which is appended as a trailing input and donated to + that output via ``input_output_aliases`` — the bridge drops aliased inputs from + the kernel's argument list, so ``fn``'s signature stays exactly the kernel's. + """ + output_leaves = jax.tree.leaves( + output_shape_dtype, + is_leaf=lambda x: hasattr(x, "shape") and hasattr(x, "dtype"), + ) + initialized_outputs = dict(initialized_outputs or {}) + input_output_aliases = dict(input_output_aliases or {}) + + def wrapper(*arrays: Any) -> Any: + inits = [] + aliases = dict(input_output_aliases) + extra_specs = [] + for offset, (out_index, init_fn) in enumerate(sorted(initialized_outputs.items())): + inits.append(init_fn(output_leaves[out_index])) + aliases[len(arrays) + offset] = out_index + extra_specs.append(output_spec[out_index] if output_spec is not None else None) + + full_input_spec = input_spec + if inits and input_spec is not None: + full_input_spec = tuple(input_spec) + tuple(extra_specs) + + return cutlass_call( + fn, + output_shape_dtype=output_shape_dtype, + input_spec=full_input_spec, + output_spec=output_spec, + input_output_aliases=aliases, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + **kwargs, + )(*arrays, *inits) + + return wrapper diff --git a/python/cudnn/tensor_adapter.py b/python/cudnn/tensor_adapter.py index 9917864aa..5602a8a68 100644 --- a/python/cudnn/tensor_adapter.py +++ b/python/cudnn/tensor_adapter.py @@ -161,6 +161,61 @@ def canonicalize_unit_dim_strides(shape: Tuple[int, ...], stride: Tuple[int, ... return tuple(numel if dim == 1 else s for dim, s in zip(shape, stride)) +def get_data_ptr(tensor: Any) -> int: + """Device data pointer of the tensor, in the caller's framework. + + JAX note: the pointer is only valid while the array is alive and not donated; + callers must hold a reference for the duration of any kernel that uses it. + """ + if is_torch_tensor(tensor): + return tensor.data_ptr() + if is_jax_array(tensor): + return tensor.unsafe_buffer_pointer() + data_ptr = getattr(tensor, "data_ptr", None) + if callable(data_ptr): + return data_ptr() + raise ValueError(f"Cannot extract a device pointer from {type(tensor)!r}") + + +def get_version(tensor: Any) -> int: + """Mutation counter for validation caching: torch's ._version, 0 for immutable arrays (JAX).""" + return int(getattr(tensor, "_version", 0)) + + +def to_host_list(tensor: Any) -> list: + """Copy a small device tensor to host and return its values as a flat Python list.""" + if is_torch_tensor(tensor): + return tensor.detach().cpu().flatten().tolist() + import numpy as np + + return np.asarray(tensor).flatten().tolist() + + +def allocate_byte_workspace(framework: str, nbytes: int, device: Any) -> Any: + """Allocate an internal uint8 workspace buffer in the caller's framework allocator. + + The buffer is written by kernels through its raw pointer and never surfaced as a + framework array, so allocating it as a (zero-initialized, for JAX) framework tensor + is safe; the caller must keep a reference alive for the compiled kernel's lifetime. + """ + nbytes = max(int(nbytes), 1) + if framework == "torch": + import torch + + return torch.empty(nbytes, dtype=torch.uint8, device=device) + if framework == "jax": + import jax + import jax.numpy as jnp + + if isinstance(device, Device): + # Canonical descriptor device -> the corresponding jax device + device = jax.devices("gpu")[device.index or 0] if device.type == "cuda" else None + buffer = jnp.zeros((nbytes,), dtype=jnp.uint8, device=device) + # Materialize before anyone reads its pointer + return jax.block_until_ready(buffer) + raise ValueError(f"Cannot allocate a workspace for framework '{framework}'") + + def pad_to_ndim(tensor: Any, ndim: int) -> Any: """Append size-1 dims up to ndim; works for any framework tensor exposing reshape.""" shape = get_shape(tensor) diff --git a/test/python/conftest.py b/test/python/conftest.py index 26a9443ec..90c6ca55c 100644 --- a/test/python/conftest.py +++ b/test/python/conftest.py @@ -12,6 +12,12 @@ "expandable_segments:True,garbage_collection_threshold:0.6", ) +# The JAX interop tests (fe_api/**/test_*_jax.py) initialize XLA in the same pytest +# process as the torch suites; XLA's default 75%-of-GPU preallocation starves later +# torch tests of memory (CUDA_ERROR_OUT_OF_MEMORY at kernel-compile time). Must be set +# before jax initializes its backend. +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + import sys import time import traceback diff --git a/test/python/fe_api/gemm/test_cutedsl_jax_guards.py b/test/python/fe_api/gemm/test_cutedsl_jax_guards.py deleted file mode 100644 index 197e604ed..000000000 --- a/test/python/fe_api/gemm/test_cutedsl_jax_guards.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -The grouped/discrete-grouped/proj_rope CuTeDSL APIs are torch-only for now: they must -reject JAX arrays with a clear error at the public entry points (instead of failing -deep inside pointer-array/workspace machinery), and their modules must import without -torch. Real JAX support for these APIs is future work. -""" - -import inspect - -import numpy as np -import pytest - -jax = pytest.importorskip("jax") -import jax.numpy as jnp - -TORCH_ONLY_WRAPPERS = [ - "grouped_gemm_wrapper_sm100", - "grouped_gemm_swiglu_wrapper_sm100", - "grouped_gemm_dswiglu_wrapper_sm100", - "grouped_gemm_srelu_wrapper_sm100", - "grouped_gemm_dsrelu_wrapper_sm100", - "grouped_gemm_glu_wrapper_sm100", - "grouped_gemm_dglu_wrapper_sm100", - "grouped_gemm_glu_hadamard_wrapper_sm100", - "grouped_gemm_quant_wrapper_sm100", - "grouped_gemm_wgrad_wrapper_sm100", - "discrete_grouped_gemm_swiglu_wrapper_sm100", - "discrete_grouped_gemm_dswiglu_wrapper_sm100", - "gemm_proj_rope_mxfp8_wrapper_sm100", -] - - -@pytest.mark.L0 -@pytest.mark.parametrize("wrapper_name", TORCH_ONLY_WRAPPERS) -def test_torch_only_wrapper_rejects_jax(wrapper_name): - import cudnn - - try: - wrapper = getattr(cudnn, wrapper_name) - except (AttributeError, ImportError): - pytest.skip(f"{wrapper_name} is not exported in this build") - - dummy = jnp.asarray(np.zeros((16, 16, 1), dtype=np.float32)) - # Fill every required parameter with the dummy JAX array; the framework guard - # runs on the first tensor argument before any validation. - kwargs = {} - for name, param in inspect.signature(wrapper).parameters.items(): - if param.default is inspect.Parameter.empty and param.kind in ( - inspect.Parameter.POSITIONAL_ONLY, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - ): - kwargs[name] = dummy - - with pytest.raises(ValueError, match="torch tensors only"): - wrapper(**kwargs) diff --git a/test/python/fe_api/gemm/test_gemm_amax.py b/test/python/fe_api/gemm/test_gemm_amax.py index ecaa38151..dc5fd2651 100644 --- a/test/python/fe_api/gemm/test_gemm_amax.py +++ b/test/python/fe_api/gemm/test_gemm_amax.py @@ -298,3 +298,38 @@ def _test_gemm_amax_wrapper( pytest.skip(f"Unsupported testcase: {e}") check_ref_gemm_amax(a_ref, b_ref, sfa_ref, sfb_ref, c_torch, amax_torch, skip_ref=cfg["skip_ref"]) + + +@pytest.mark.L0 +def test_gemm_amax_rejects_noncontiguous_scale_factors(): + """SF tensors are consumed by base pointer only (the kernel rebuilds the layout from + the GEMM shapes), so a shape-matching but differently-strided tensor must be rejected + rather than silently producing wrong results.""" + try: + from cudnn import gemm_amax_wrapper_sm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: + pytest.skip("requires SM100+") + + m, n, k, sf_vec_size = 512, 256, 256, 32 + a = torch.randn(m, k, 1, device="cuda").to(torch.float8_e5m2) + b = torch.randn(n, k, 1, device="cuda").to(torch.float8_e5m2) + sf_dtype = torch.float8_e8m0fnu + sfa = torch.ones(1, m // 128, k // (4 * sf_vec_size), 32, 4, 4, device="cuda", dtype=torch.uint8).view(sf_dtype) + sfb = torch.ones(1, n // 128, k // (4 * sf_vec_size), 32, 4, 4, device="cuda", dtype=torch.uint8).view(sf_dtype) + + # Valid: the physical form and its (3, 4, 1, 5, 2, 0)-permuted atom view + gemm_amax_wrapper_sm100(a, b, sfa, sfb, sf_vec_size=sf_vec_size) + gemm_amax_wrapper_sm100(a, b, sfa.permute(3, 4, 1, 5, 2, 0), sfb.permute(3, 4, 1, 5, 2, 0), sf_vec_size=sf_vec_size) + + # Non-contiguous tensor whose shape matches the physical form + bad_physical = torch.ones(1, k // (4 * sf_vec_size), m // 128, 32, 4, 4, device="cuda", dtype=torch.uint8).view(sf_dtype).permute(0, 2, 1, 3, 4, 5) + assert tuple(bad_physical.shape) == tuple(sfa.shape) and not bad_physical.is_contiguous() + with pytest.raises(ValueError, match="stride"): + gemm_amax_wrapper_sm100(a, b, bad_physical, sfb, sf_vec_size=sf_vec_size) + + # C-contiguous allocation in the atom-view shape (not a permutation of the physical form) + bad_atom = torch.ones(32, 4, m // 128, 4, k // (4 * sf_vec_size), 1, device="cuda", dtype=torch.uint8).view(sf_dtype) + with pytest.raises(ValueError, match="stride"): + gemm_amax_wrapper_sm100(a, b, bad_atom, sfb, sf_vec_size=sf_vec_size) diff --git a/test/python/fe_api/gemm/test_gemm_amax_jax.py b/test/python/fe_api/gemm/test_gemm_amax_jax.py index 062a2da36..6cb636820 100644 --- a/test/python/fe_api/gemm/test_gemm_amax_jax.py +++ b/test/python/fe_api/gemm/test_gemm_amax_jax.py @@ -242,12 +242,20 @@ def test_gemm_amax_jax_wrapper_errors(): @pytest.mark.L0 -def test_gemm_amax_jax_ffi_sm100(): - """XLA custom-call entry point (jax-tvm-ffi): eager, jitted, cached, and composed.""" - pytest.importorskip("jax_tvm_ffi") +def test_gemm_amax_jax_jit_sm100(): + """XLA custom-call entry point (cudnn.jax.call): eager, jitted, cached, and composed.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") skip_unless_sm100() from cudnn import gemm_amax_jax_sm100 + # cudnn.jax is reachable from a bare `import cudnn` (lazy submodule export) + import cudnn + + assert cudnn.jax.TensorSpec is cutlass.jax.TensorSpec + m, n, k = 512, 256, 256 sf_vec_size = 32 rng = np.random.default_rng(3) diff --git a/test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py b/test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py new file mode 100644 index 000000000..027434e33 --- /dev/null +++ b/test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the fused projection GEMM + RoPE + MXFP8 quantize wrapper. + +JAX contract: w_out_in=True only (the [in, out] weight layout reaches the kernel +through a transposed strided view, which has no row-major JAX equivalent). Outputs +are checked bit-identical against the torch wrapper run on identical input bytes +(both paths share one compiled kernel). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 +from fe_api.gemm.test_gemm_proj_rope_mxfp8_utils import BLOCK, HEAD_DIM, NUM_HEADS, Q_LORA, Q_OUT, QK_ROPE + + +def to_torch(x, dtype): + return torch.from_numpy(np.ascontiguousarray(x).view(np.uint8)).view(dtype).reshape(x.shape).cuda() + + +def run_both(kwargs_np, jax_dtypes, torch_dtypes): + """Run the wrapper with torch tensors and JAX arrays built from the same bytes.""" + import cudnn + + torch_kwargs = {name: to_torch(arr, torch_dtypes[name]) for name, arr in kwargs_np.items()} + result_t = cudnn.gemm_proj_rope_mxfp8_wrapper_sm100(**torch_kwargs, w_out_in=True) + torch.cuda.synchronize() + + jax_kwargs = {name: jnp.asarray(arr.view(jax_dtypes[name])) for name, arr in kwargs_np.items()} + jax.block_until_ready(tuple(jax_kwargs.values())) + result_j = cudnn.gemm_proj_rope_mxfp8_wrapper_sm100(**jax_kwargs, w_out_in=True) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + for key in ("out_fp8_row", "out_scales_row", "out_fp8_col", "out_scales_col"): + np.testing.assert_array_equal( + np.asarray(result_j[key]).view(np.uint8), + result_t[key].view(torch.uint8).cpu().numpy(), + err_msg=f"proj_rope {key}: JAX output differs from torch output on identical input bytes", + ) + + +@pytest.mark.L0 +def test_gemm_proj_rope_mxfp8_bf16in_jax_matches_torch(): + skip_unless_sm100() + tokens = 256 + rng = np.random.default_rng(0) + kwargs_np = { + "x": (rng.standard_normal((tokens, Q_LORA), dtype=np.float32) * 0.5).astype(ml_dtypes.bfloat16), + "w": (rng.standard_normal((Q_OUT, Q_LORA), dtype=np.float32) * 0.02).astype(ml_dtypes.bfloat16), + "cos": rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16), + "sin": rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16), + } + dtypes_j = {"x": ml_dtypes.bfloat16, "w": ml_dtypes.bfloat16, "cos": ml_dtypes.bfloat16, "sin": ml_dtypes.bfloat16} + dtypes_t = {"x": torch.bfloat16, "w": torch.bfloat16, "cos": torch.bfloat16, "sin": torch.bfloat16} + run_both(kwargs_np, dtypes_j, dtypes_t) + + +@pytest.mark.L0 +def test_gemm_proj_rope_mxfp8_mxfp8in_jax_matches_torch(): + skip_unless_sm100() + tokens = 256 + rng = np.random.default_rng(1) + kwargs_np = { + "x": (rng.standard_normal((tokens, Q_LORA), dtype=np.float32) * 0.5).astype(ml_dtypes.float8_e4m3fn), + "w": (rng.standard_normal((Q_OUT, Q_LORA), dtype=np.float32) * 0.02).astype(ml_dtypes.float8_e4m3fn), + "cos": rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16), + "sin": rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16), + # E8M0 biased exponents around 1.0 (byte 127) + "x_scale": rng.integers(125, 130, size=(tokens, Q_LORA // BLOCK)).astype(np.uint8), + "w_scale": rng.integers(125, 130, size=(Q_OUT, Q_LORA // BLOCK)).astype(np.uint8), + } + dtypes_j = { + "x": ml_dtypes.float8_e4m3fn, + "w": ml_dtypes.float8_e4m3fn, + "cos": ml_dtypes.bfloat16, + "sin": ml_dtypes.bfloat16, + "x_scale": np.uint8, + "w_scale": np.uint8, + } + dtypes_t = { + "x": torch.float8_e4m3fn, + "w": torch.float8_e4m3fn, + "cos": torch.bfloat16, + "sin": torch.bfloat16, + "x_scale": torch.uint8, + "w_scale": torch.uint8, + } + run_both(kwargs_np, dtypes_j, dtypes_t) + + +@pytest.mark.L0 +@pytest.mark.parametrize("input_path", ["bf16in", "mxfp8in"]) +def test_gemm_proj_rope_mxfp8_jax_jit_matches_eager(input_path): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + import cudnn + from cudnn import gemm_proj_rope_mxfp8_jax_sm100 + + tokens = 256 + rng = np.random.default_rng(4) + cos = jnp.asarray(rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16)) + sin = jnp.asarray(rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16)) + if input_path == "bf16in": + x = jnp.asarray((rng.standard_normal((tokens, Q_LORA), dtype=np.float32) * 0.5).astype(ml_dtypes.bfloat16)) + w = jnp.asarray((rng.standard_normal((Q_OUT, Q_LORA), dtype=np.float32) * 0.02).astype(ml_dtypes.bfloat16)) + scales = {} + else: + x = jnp.asarray((rng.standard_normal((tokens, Q_LORA), dtype=np.float32) * 0.5).astype(ml_dtypes.float8_e4m3fn)) + w = jnp.asarray((rng.standard_normal((Q_OUT, Q_LORA), dtype=np.float32) * 0.02).astype(ml_dtypes.float8_e4m3fn)) + scales = { + "x_scale": jnp.asarray(rng.integers(125, 130, size=(tokens, Q_LORA // BLOCK)).astype(np.uint8)), + "w_scale": jnp.asarray(rng.integers(125, 130, size=(Q_OUT, Q_LORA // BLOCK)).astype(np.uint8)), + } + jax.block_until_ready((x, w, cos, sin, *scales.values())) + + result_eager = cudnn.gemm_proj_rope_mxfp8_wrapper_sm100(x, w, cos, sin, **scales, w_out_in=True) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = [np.asarray(result_eager[key]).view(np.uint8) for key in ("out_fp8_row", "out_scales_row", "out_fp8_col", "out_scales_col")] + + def check(outputs): + jax.block_until_ready(outputs) + for got, want, key in zip(outputs, expected, ("out_fp8_row", "out_scales_row", "out_fp8_col", "out_scales_col")): + np.testing.assert_array_equal( + np.asarray(got).view(np.uint8), + want, + err_msg=f"proj_rope {key}: jit output differs from eager wrapper output on identical input bytes", + ) + + # Eager custom call + check(gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, **scales)) + + # Under jax.jit, twice (compiled-kernel / registration cache) + jitted = jax.jit(lambda *args: gemm_proj_rope_mxfp8_jax_sm100(*args)) + check(jitted(x, w, cos, sin, *scales.values())) + check(jitted(x, w, cos, sin, *scales.values())) + + +@pytest.mark.L0 +def test_gemm_proj_rope_mxfp8_jax_errors(): + skip_unless_sm100() + import cudnn + + tokens = 256 + rng = np.random.default_rng(2) + x = jnp.asarray((rng.standard_normal((tokens, Q_LORA), dtype=np.float32) * 0.5).astype(ml_dtypes.bfloat16)) + w_in_out = jnp.asarray((rng.standard_normal((Q_LORA, Q_OUT), dtype=np.float32) * 0.02).astype(ml_dtypes.bfloat16)) + cos = jnp.asarray(rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16)) + sin = jnp.asarray(rng.standard_normal((tokens, QK_ROPE), dtype=np.float32).astype(ml_dtypes.bfloat16)) + + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + cudnn.gemm_proj_rope_mxfp8_wrapper_sm100(x, w_in_out, cos, sin, w_out_in=False) + + with pytest.raises(ValueError, match="Unsupported tensor framework"): + cudnn.gemm_proj_rope_mxfp8_wrapper_sm100(np.asarray(x), np.asarray(w_in_out), np.asarray(cos), np.asarray(sin)) diff --git a/test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py b/test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py index 8b89a6998..9fff9fb89 100644 --- a/test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py +++ b/test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py @@ -137,3 +137,47 @@ def test_gemm_srelu_dsrelu_jax_errors(): with pytest.raises(ValueError, match="Unsupported tensor framework"): cudnn.gemm_srelu_wrapper_sm100(a_np, b_np, sfa_np, sfb_np, prob_np, sf_vec_size=32) + + +@pytest.mark.L0 +def test_gemm_srelu_dsrelu_jax_jit_matches_eager(): + """The XLA custom-call entry points must agree bit-for-bit with the eager wrappers.""" + skip_unless_sm100() + import cudnn + from cudnn import gemm_dsrelu_jax_sm100, gemm_srelu_jax_sm100 + + m, n, k, l = 256, 256, 512, 1 + sf_vec_size = 32 + rng = np.random.default_rng(3) + a_np, b_np, sfa_np, sfb_np, prob_np = make_inputs(m, n, k, l, sf_vec_size, rng) + + a_j, b_j, sfa_j, sfb_j, prob_j = (jnp.asarray(x) for x in (a_np, b_np, sfa_np, sfb_np, prob_np)) + jax.block_until_ready((a_j, b_j, sfa_j, sfb_j, prob_j)) + + eager = cudnn.gemm_srelu_wrapper_sm100( + a_tensor=a_j, b_tensor=b_j, sfa_tensor=sfa_j, sfb_tensor=sfb_j, prob_tensor=prob_j, c_dtype="bfloat16", d_dtype="bfloat16", sf_vec_size=sf_vec_size + ) + device_sync() + + jitted = jax.jit(lambda a, b, sfa, sfb, prob: gemm_srelu_jax_sm100(a, b, sfa, sfb, prob, c_dtype="bfloat16", d_dtype="bfloat16", sf_vec_size=sf_vec_size)) + for _ in range(2): # repeat: donation safety + c_jit, d_jit = jitted(a_j, b_j, sfa_j, sfb_j, prob_j) + jax.block_until_ready((c_jit, d_jit)) # XLA-stream ordered; no manual sync needed + np.testing.assert_array_equal(np.asarray(c_jit).view(np.uint8), np.asarray(eager["c_tensor"]).view(np.uint8)) + np.testing.assert_array_equal(np.asarray(d_jit).view(np.uint8), np.asarray(eager["d_tensor"]).view(np.uint8)) + + # backward + c_in_np = rng.standard_normal((m, n, l), dtype=np.float32).astype(ml_dtypes.bfloat16) + c_in_j = jnp.asarray(c_in_np) + jax.block_until_ready(c_in_j) + eager_b = cudnn.gemm_dsrelu_wrapper_sm100( + a_tensor=a_j, b_tensor=b_j, c_tensor=c_in_j, sfa_tensor=sfa_j, sfb_tensor=sfb_j, prob_tensor=prob_j, d_dtype="bfloat16", sf_vec_size=sf_vec_size + ) + device_sync() + + jitted_b = jax.jit(lambda a, b, c, sfa, sfb, prob: gemm_dsrelu_jax_sm100(a, b, c, sfa, sfb, prob, d_dtype="bfloat16", sf_vec_size=sf_vec_size)) + d_jit, dprob_jit = jitted_b(a_j, b_j, c_in_j, sfa_j, sfb_j, prob_j) + jax.block_until_ready((d_jit, dprob_jit)) + np.testing.assert_array_equal(np.asarray(d_jit).view(np.uint8), np.asarray(eager_b["d_tensor"]).view(np.uint8)) + # dprob accumulates via FP32 atomic adds (ordering-nondeterministic): tight tolerance + np.testing.assert_allclose(np.asarray(dprob_jit), np.asarray(eager_b["dprob_tensor"]), rtol=1e-5, atol=1e-5) diff --git a/test/python/fe_api/gemm/test_gemm_swiglu_jax.py b/test/python/fe_api/gemm/test_gemm_swiglu_jax.py index ad9aa1b9b..56652d990 100644 --- a/test/python/fe_api/gemm/test_gemm_swiglu_jax.py +++ b/test/python/fe_api/gemm/test_gemm_swiglu_jax.py @@ -104,9 +104,12 @@ def test_gemm_swiglu_jax_wrapper_quant_fp8(): @pytest.mark.L0 -def test_gemm_swiglu_jax_ffi_sm100(): +def test_gemm_swiglu_jax_jit_sm100(): """XLA custom-call entry point: jitted, repeated (donation safety), and alpha attr.""" - pytest.importorskip("jax_tvm_ffi") + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") skip_unless_sm100() from cudnn import gemm_swiglu_jax_sm100 @@ -127,15 +130,24 @@ def test_gemm_swiglu_jax_ffi_sm100(): np.testing.assert_allclose(np.asarray(ab12)[:, :, 0], ab12_ref, atol=0.02, rtol=0.02) np.testing.assert_allclose(np.asarray(c).astype(np.float32)[:, :, 0], c_ref, atol=0.05, rtol=0.05) - # The quantized (blockscaled) config is eager-wrapper-only for now + # Quantized (blockscaled MXFP8) config through the same entry point, under jit sf_vec_size = 32 - a2_np, _ = make_ab_fp8(m, k, ml_dtypes.float8_e4m3fn, rng) - b2_np, _ = make_ab_fp8(n, k, ml_dtypes.float8_e4m3fn, rng) - sfa_np, _ = make_sf_physical(m, k, sf_vec_size, ml_dtypes.float8_e8m0fnu, rng) - sfb_np, _ = make_sf_physical(n, k, sf_vec_size, ml_dtypes.float8_e8m0fnu, rng) + a2_np, a2_ref = make_ab_fp8(m, k, ml_dtypes.float8_e4m3fn, rng) + b2_np, b2_ref = make_ab_fp8(n, k, ml_dtypes.float8_e4m3fn, rng) + sfa_np, sfa_expanded = make_sf_physical(m, k, sf_vec_size, ml_dtypes.float8_e8m0fnu, rng) + sfb_np, sfb_expanded = make_sf_physical(n, k, sf_vec_size, ml_dtypes.float8_e8m0fnu, rng) a2, b2, sfa, sfb = (jax.device_put(x) for x in (a2_np, b2_np, sfa_np, sfb_np)) - with pytest.raises(NotImplementedError, match="non-quantized kernel only"): - gemm_swiglu_jax_sm100(a2, b2, sfa_tensor=sfa, sfb_tensor=sfb, sf_vec_size=sf_vec_size) + + quant = jax.jit( + lambda a, b, sfa, sfb: gemm_swiglu_jax_sm100( + a, b, ab12_dtype=jnp.bfloat16, c_dtype=jnp.bfloat16, sfa_tensor=sfa, sfb_tensor=sfb, sf_vec_size=sf_vec_size + ) + ) + ab12q, cq = quant(a2, b2, sfa, sfb) + jax.block_until_ready((ab12q, cq)) + ab12q_ref = (a2_ref * sfa_expanded) @ (b2_ref * sfb_expanded).T + np.testing.assert_allclose(np.asarray(ab12q).astype(np.float32)[:, :, 0], ab12q_ref, atol=0.5, rtol=0.05) + np.testing.assert_allclose(np.asarray(cq).astype(np.float32)[:, :, 0], swiglu_block_ref(ab12q_ref, n), atol=1.0, rtol=0.05) @pytest.mark.L0 diff --git a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py new file mode 100644 index 000000000..857d55c19 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the discrete-weight grouped GEMM dSwiGLU (backward) wrapper. + +JAX contract mirrors the SwiGLU forward: FP8 inputs, per-expert B/SFB as packed-uint8 +pointer arrays, scale-factor tensors (SFA input and the sfd_row/sfd_col outputs) in the +physical C-contiguous atom shape (1, MN', K', 32, 4, 4) (the kernel rebuilds every SF +layout from the A/D shapes and consumes only the base pointers), plus the +backward-specific C (forward activations), beta, prob, and zero-initialized dprob +inputs. d_row/d_col are checked bit-identical against the torch wrapper on identical +input bytes; dprob accumulates through floating-point atomics whose ordering is not +deterministic, so it is compared with a tight tolerance instead. + +Rejected for JAX: packed-fp4 inputs (JAX has no packed fp4 dtype, and uint8 container +arrays are rejected at the kernel entry). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import ceil_div, device_sync, skip_unless_sm100 + +M, N, K, EXPERTS = 1024, 512, 512, 4 +N_OUT = 2 * N +SF_VEC_SIZE = 32 + + +def make_problem(): + rng = np.random.default_rng(20260811) + a_np = rng.integers(-2, 3, (M, K, 1)).astype(np.float32).astype(ml_dtypes.float8_e4m3fn) # gradient input, k-major + b_np = [rng.integers(-2, 3, (N, K)).astype(np.float32).astype(ml_dtypes.float8_e4m3fn) for _ in range(EXPERTS)] # k-major + c_np = (rng.standard_normal((M, N_OUT, 1), dtype=np.float32) * 0.5).astype(ml_dtypes.bfloat16) # forward activations + rest_k = ceil_div(ceil_div(K, SF_VEC_SIZE), 4) + sfa_u8 = rng.integers(126, 130, (1, ceil_div(M, 128), rest_k, 32, 4, 4), dtype=np.uint8) + sfb_u8 = [rng.integers(126, 130, (1, ceil_div(N, 128), rest_k, 32, 4, 4), dtype=np.uint8) for _ in range(EXPERTS)] + offsets_np = np.arange(M // EXPERTS, M + 1, M // EXPERTS, dtype=np.int32) + alpha_np = rng.uniform(-1.5, 1.5, (EXPERTS,)).astype(np.float32) + beta_np = rng.uniform(-1.5, 1.5, (EXPERTS,)).astype(np.float32) + prob_np = rng.uniform(-1.0, 1.0, (M, 1, 1)).astype(np.float32) + norm_const_np = np.array([0.01], dtype=np.float32) + return a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np + + +MMA_PERMUTE_ORDER = (3, 4, 1, 5, 2, 0) + + +def run_torch(a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np): + """Reference run through the established torch contract on identical bytes.""" + from cudnn import discrete_grouped_gemm_dswiglu_wrapper_sm100 + + a_t = torch.from_numpy(a_np.view(np.uint8)).view(torch.float8_e4m3fn).reshape(a_np.shape).cuda() + b_t = [torch.from_numpy(b.view(np.uint8)).view(torch.float8_e4m3fn).reshape(b.shape).cuda() for b in b_np] + c_t = torch.from_numpy(c_np.view(np.uint8)).view(torch.bfloat16).reshape(c_np.shape).cuda() + sfa_t = torch.from_numpy(sfa_u8).cuda().view(torch.float8_e8m0fnu).permute(MMA_PERMUTE_ORDER) # torch atom view + sfb_t = [torch.from_numpy(sfb).cuda().view(torch.float8_e8m0fnu) for sfb in sfb_u8] + dprob_t = torch.zeros((M, 1, 1), dtype=torch.float32, device="cuda") + result = discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_t, + b_ptrs=torch.tensor([b.data_ptr() for b in b_t], dtype=torch.int64, device="cuda"), + c_tensor=c_t, + sfa_tensor=sfa_t, + sfb_ptrs=torch.tensor([sfb.data_ptr() for sfb in sfb_t], dtype=torch.int64, device="cuda"), + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + beta_tensor=torch.from_numpy(beta_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + dprob_tensor=dprob_t, + norm_const_tensor=torch.from_numpy(norm_const_np).cuda(), + n=N, + b_dtype=torch.float8_e4m3fn, + d_dtype=torch.float8_e4m3fn, + sf_vec_size=SF_VEC_SIZE, + act_func="dswiglu", + ) + torch.cuda.synchronize() + return result, (a_t, b_t, c_t, sfa_t, sfb_t) + + +def packed_ptrs(arrays): + values = np.array([array.unsafe_buffer_pointer() for array in arrays], dtype=np.int64) + return jax.block_until_ready(jnp.asarray(values.view(np.uint8))) + + +def as_bytes(array_or_tensor): + """Raw little-endian memory bytes as a uint8 numpy array (exact comparison incl. fp8).""" + if isinstance(array_or_tensor, torch.Tensor): + # The wrapper outputs are (m, n, 1) with an extent-1 batch dim of arbitrary + # stride; squeeze it so the row-major bytes match the JAX C-contiguous bytes. + data = array_or_tensor.squeeze(-1).contiguous().view(torch.uint8).cpu().numpy().tobytes() + else: + data = np.asarray(array_or_tensor).tobytes() + return np.frombuffer(data, dtype=np.uint8) + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_dswiglu_jax_fp8_matches_torch(): + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_dswiglu_wrapper_sm100 + + a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np = make_problem() + result_t, _torch_keepalive = run_torch(a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np) + + a_j = jnp.asarray(a_np) + b_j = [jnp.asarray(b) for b in b_np] + c_j = jnp.asarray(c_np) + sfa_j = jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)) # physical atom shape + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + offsets_j, alpha_j, beta_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, beta_np, prob_np, norm_const_np)) + dprob_j = jnp.zeros((M, 1, 1), dtype=jnp.float32) # output accumulator, must be zero-initialized + jax.block_until_ready((a_j, c_j, sfa_j, offsets_j, alpha_j, beta_j, prob_j, norm_const_j, dprob_j, *b_j, *sfb_j)) + + # The per-expert weight/scale arrays must stay alive while the kernel runs. + result_j = discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_j, + b_ptrs=packed_ptrs(b_j), + c_tensor=c_j, + sfa_tensor=sfa_j, + sfb_ptrs=packed_ptrs(sfb_j), + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + norm_const_tensor=norm_const_j, + n=N, + b_dtype="float8_e4m3fn", + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="dswiglu", + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + # d_row/d_col are deterministic per-tile kernel outputs; compare raw bytes. + for key in ("d_row_tensor", "d_col_tensor"): + np.testing.assert_array_equal( + as_bytes(result_j[key]), + as_bytes(result_t[key]), + err_msg=f"dswiglu {key}: JAX output differs from torch output on identical input bytes", + ) + # dprob accumulates through floating-point atomics; ordering is not deterministic, + # so compare with a tight tolerance instead of bit equality. + np.testing.assert_allclose( + np.asarray(result_j["dprob_tensor"]), + result_t["dprob_tensor"].float().cpu().numpy(), + rtol=2e-5, + atol=1e-4, + err_msg="dswiglu dprob_tensor: JAX output differs from torch output beyond atomic-ordering tolerance", + ) + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_dswiglu_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper. + + d_row/d_col are deterministic per-tile kernel outputs and are compared bitwise; + dprob accumulates through floating-point atomics whose ordering is not + deterministic across runs, so it is compared with the same tight tolerance the + torch-parity test uses. + """ + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_dswiglu_jax_sm100, discrete_grouped_gemm_dswiglu_wrapper_sm100 + + a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np = make_problem() + + a_j = jnp.asarray(a_np) + b_j = [jnp.asarray(b) for b in b_np] + c_j = jnp.asarray(c_np) + sfa_j = jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)) # physical atom shape + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + offsets_j, alpha_j, beta_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, beta_np, prob_np, norm_const_np)) + dprob_j = jnp.zeros((M, 1, 1), dtype=jnp.float32) # eager wrapper takes dprob as a zeroed input buffer + jax.block_until_ready((a_j, c_j, sfa_j, offsets_j, alpha_j, beta_j, prob_j, norm_const_j, dprob_j, *b_j, *sfb_j)) + b_ptrs_j = packed_ptrs(b_j) + sfb_ptrs_j = packed_ptrs(sfb_j) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row and d_row/d_col are fully overwritten per-row kernel outputs, so + # full-buffer bitwise comparison against the jit path is well-defined. + result_eager = discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_j, + b_ptrs=b_ptrs_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + sfb_ptrs=sfb_ptrs_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + norm_const_tensor=norm_const_j, + n=N, + b_dtype="float8_e4m3fn", + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="dswiglu", + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_row_tensor", "d_col_tensor")} + dprob_expected = np.asarray(result_eager["dprob_tensor"]) + + def run_jit_entry(a, c, sfa, offsets, alpha, beta, prob, norm_const, b_ptrs, sfb_ptrs): + # dprob is a donated zero-initialized output of the custom call (the jit + # entry has no caller-provided dprob buffer, unlike the eager wrapper). + return discrete_grouped_gemm_dswiglu_jax_sm100( + a_tensor=a, + b_ptrs=b_ptrs, + c_tensor=c, + sfa_tensor=sfa, + sfb_ptrs=sfb_ptrs, + padded_offsets=offsets, + alpha_tensor=alpha, + beta_tensor=beta, + prob_tensor=prob, + norm_const_tensor=norm_const, + n=N, + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="dswiglu", + ) + + def check(result): + jax.block_until_ready(tuple(value for value in result.values() if value is not None)) + for key in ("d_row_tensor", "d_col_tensor"): + np.testing.assert_array_equal( + np.asarray(result[key]).view(np.uint8), + expected[key], + err_msg=f"dswiglu {key}: jit output differs from eager wrapper output on identical input bytes", + ) + np.testing.assert_allclose( + np.asarray(result["dprob_tensor"]), + dprob_expected, + rtol=2e-5, + atol=1e-4, + err_msg="dswiglu dprob_tensor: jit output differs from eager wrapper output beyond atomic-ordering tolerance", + ) + assert result["amax_tensor"] is None # fp8 d_dtype: no amax, matching the eager wrapper + assert result["dbias_tensor"] is None # generate_dbias not requested + + # Eager custom call + check(run_jit_entry(a_j, c_j, sfa_j, offsets_j, alpha_j, beta_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit(run_jit_entry) + check(jitted(a_j, c_j, sfa_j, offsets_j, alpha_j, beta_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + check(jitted(a_j, c_j, sfa_j, offsets_j, alpha_j, beta_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_dswiglu_jax_errors(): + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_dswiglu_wrapper_sm100 + + a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np = make_problem() + b_j = [jnp.asarray(b) for b in b_np] + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + jax.block_until_ready((*b_j, *sfb_j)) + + # Packed fp4 has no JAX dtype; uint8 container inputs are rejected. + a_fp4_j = jnp.asarray(np.zeros((M, K // 2, 1), dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_fp4_j, + b_ptrs=packed_ptrs(b_j), + c_tensor=jnp.asarray(c_np), + sfa_tensor=jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)), + sfb_ptrs=packed_ptrs(sfb_j), + padded_offsets=jnp.asarray(offsets_np), + alpha_tensor=jnp.asarray(alpha_np), + beta_tensor=jnp.asarray(beta_np), + prob_tensor=jnp.asarray(prob_np), + dprob_tensor=jnp.zeros((M, 1, 1), dtype=jnp.float32), + n=N, + b_dtype="uint8", + sf_vec_size=SF_VEC_SIZE, + act_func="dswiglu", + ) diff --git a/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py new file mode 100644 index 000000000..7890df2f8 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the discrete-weight grouped GEMM SwiGLU (forward) wrapper. + +JAX contract: FP8 inputs (ml_dtypes float8 dtypes map 1:1 onto the kernel's dtypes); +per-expert B/SFB are passed as packed-uint8 pointer arrays (8 bytes per pointer, since +JAX truncates int64 without x64 mode); scale-factor tensors (SFA input and the +sfd_row/sfd_col outputs) use the physical C-contiguous atom shape (1, MN', K', 32, 4, 4) +-- the kernel provably rebuilds every SF layout from the A/D shapes and consumes only +the SF base pointers, so the permuted torch atom view and the physical allocation are +byte-identical. Outputs are checked bit-identical against the torch wrapper run on +identical input bytes. + +Rejected for JAX: bias (column-major (n, experts) layout) and packed-fp4 inputs (JAX +has no packed fp4 dtype, and uint8 container arrays are rejected at the kernel entry). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import ceil_div, device_sync, skip_unless_sm100 + +M, N, K, EXPERTS = 1024, 512, 512, 4 +SF_VEC_SIZE = 32 + + +def make_problem(): + rng = np.random.default_rng(20260810) + a_np = rng.integers(-2, 3, (M, K, 1)).astype(np.float32).astype(ml_dtypes.float8_e4m3fn) # k-major + b_np = [rng.integers(-2, 3, (N, K)).astype(np.float32).astype(ml_dtypes.float8_e4m3fn) for _ in range(EXPERTS)] # k-major + rest_k = ceil_div(ceil_div(K, SF_VEC_SIZE), 4) + # e8m0 scale factors in the physical C-contiguous atom shape (1, MN', K', 32, 4, 4); + # byte value 127 is 1.0, so keep exponents near it. + sfa_u8 = rng.integers(126, 130, (1, ceil_div(M, 128), rest_k, 32, 4, 4), dtype=np.uint8) + sfb_u8 = [rng.integers(126, 130, (1, ceil_div(N, 128), rest_k, 32, 4, 4), dtype=np.uint8) for _ in range(EXPERTS)] + offsets_np = np.arange(M // EXPERTS, M + 1, M // EXPERTS, dtype=np.int32) + alpha_np = rng.uniform(-1.5, 1.5, (EXPERTS,)).astype(np.float32) + prob_np = rng.uniform(-1.0, 1.0, (M, 1, 1)).astype(np.float32) + norm_const_np = np.array([0.01], dtype=np.float32) + return a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np + + +MMA_PERMUTE_ORDER = (3, 4, 1, 5, 2, 0) + + +def run_torch(a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np): + """Reference run through the established torch contract on identical bytes.""" + from cudnn import discrete_grouped_gemm_swiglu_wrapper_sm100 + + a_t = torch.from_numpy(a_np.view(np.uint8)).view(torch.float8_e4m3fn).reshape(a_np.shape).cuda() + b_t = [torch.from_numpy(b.view(np.uint8)).view(torch.float8_e4m3fn).reshape(b.shape).cuda() for b in b_np] + sfa_t = torch.from_numpy(sfa_u8).cuda().view(torch.float8_e8m0fnu).permute(MMA_PERMUTE_ORDER) # torch atom view + sfb_t = [torch.from_numpy(sfb).cuda().view(torch.float8_e8m0fnu) for sfb in sfb_u8] + result = discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_t, + b_ptrs=torch.tensor([b.data_ptr() for b in b_t], dtype=torch.int64, device="cuda"), + sfa_tensor=sfa_t, + sfb_ptrs=torch.tensor([sfb.data_ptr() for sfb in sfb_t], dtype=torch.int64, device="cuda"), + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + norm_const_tensor=torch.from_numpy(norm_const_np).cuda(), + n=N, + b_dtype=torch.float8_e4m3fn, + d_dtype=torch.float8_e4m3fn, + sf_vec_size=SF_VEC_SIZE, + act_func="swiglu", + ) + torch.cuda.synchronize() + return result, (a_t, b_t, sfa_t, sfb_t) + + +def packed_ptrs(arrays): + values = np.array([array.unsafe_buffer_pointer() for array in arrays], dtype=np.int64) + return jax.block_until_ready(jnp.asarray(values.view(np.uint8))) + + +def as_bytes(array_or_tensor): + """Raw little-endian memory bytes as a uint8 numpy array (exact comparison incl. fp8).""" + if isinstance(array_or_tensor, torch.Tensor): + # The wrapper outputs are (m, n, 1) with an extent-1 batch dim of arbitrary + # stride; squeeze it so the row-major bytes match the JAX C-contiguous bytes. + data = array_or_tensor.squeeze(-1).contiguous().view(torch.uint8).cpu().numpy().tobytes() + else: + data = np.asarray(array_or_tensor).tobytes() + return np.frombuffer(data, dtype=np.uint8) + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_swiglu_jax_fp8_matches_torch(): + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_swiglu_wrapper_sm100 + + a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np = make_problem() + result_t, _torch_keepalive = run_torch(a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np) + + a_j = jnp.asarray(a_np) + b_j = [jnp.asarray(b) for b in b_np] + sfa_j = jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)) # physical atom shape + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + offsets_j, alpha_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np, norm_const_np)) + jax.block_until_ready((a_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, *b_j, *sfb_j)) + + # The per-expert weight/scale arrays must stay alive while the kernel runs. + result_j = discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_j, + b_ptrs=packed_ptrs(b_j), + sfa_tensor=sfa_j, + sfb_ptrs=packed_ptrs(sfb_j), + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + norm_const_tensor=norm_const_j, + n=N, + b_dtype="float8_e4m3fn", + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="swiglu", + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + # c/d/d_col are deterministic per-tile kernel outputs; compare raw bytes (fp8 + # outputs are compared exactly). amax is None for fp8 d_dtype; the sfd outputs + # may contain unwritten padding regions and are covered transitively through + # the quantized d/d_col values. + for key in ("c_tensor", "d_tensor", "d_col_tensor"): + np.testing.assert_array_equal( + as_bytes(result_j[key]), + as_bytes(result_t[key]), + err_msg=f"swiglu {key}: JAX output differs from torch output on identical input bytes", + ) + assert result_j["amax_tensor"] is None and result_t["amax_tensor"] is None + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_swiglu_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_swiglu_jax_sm100, discrete_grouped_gemm_swiglu_wrapper_sm100 + + a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np = make_problem() + + a_j = jnp.asarray(a_np) + b_j = [jnp.asarray(b) for b in b_np] + sfa_j = jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)) # physical atom shape + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + offsets_j, alpha_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np, norm_const_np)) + jax.block_until_ready((a_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, *b_j, *sfb_j)) + b_ptrs_j = packed_ptrs(b_j) + sfb_ptrs_j = packed_ptrs(sfb_j) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row and c/d/d_col are fully overwritten per-row kernel outputs, so full-buffer + # bitwise comparison against the jit path is well-defined. + result_eager = discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_j, + b_ptrs=b_ptrs_j, + sfa_tensor=sfa_j, + sfb_ptrs=sfb_ptrs_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + norm_const_tensor=norm_const_j, + n=N, + b_dtype="float8_e4m3fn", + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="swiglu", + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("c_tensor", "d_tensor", "d_col_tensor")} + + def run_jit_entry(a, sfa, offsets, alpha, prob, norm_const, b_ptrs, sfb_ptrs): + return discrete_grouped_gemm_swiglu_jax_sm100( + a_tensor=a, + b_ptrs=b_ptrs, + sfa_tensor=sfa, + sfb_ptrs=sfb_ptrs, + padded_offsets=offsets, + alpha_tensor=alpha, + prob_tensor=prob, + norm_const_tensor=norm_const, + n=N, + d_dtype="float8_e4m3fn", + sf_vec_size=SF_VEC_SIZE, + act_func="swiglu", + ) + + def check(result): + jax.block_until_ready(tuple(value for value in result.values() if value is not None)) + for key in ("c_tensor", "d_tensor", "d_col_tensor"): + np.testing.assert_array_equal( + np.asarray(result[key]).view(np.uint8), + expected[key], + err_msg=f"swiglu {key}: jit output differs from eager wrapper output on identical input bytes", + ) + assert result["amax_tensor"] is None # fp8 d_dtype: no amax, matching the eager wrapper + + # Eager custom call + check(run_jit_entry(a_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit(run_jit_entry) + check(jitted(a_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + check(jitted(a_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, b_ptrs_j, sfb_ptrs_j)) + + +@pytest.mark.L0 +def test_discrete_grouped_gemm_swiglu_jax_errors(): + skip_unless_sm100() + from cudnn import discrete_grouped_gemm_swiglu_wrapper_sm100 + + a_np, b_np, sfa_u8, sfb_u8, offsets_np, alpha_np, prob_np, norm_const_np = make_problem() + a_j = jnp.asarray(a_np) + b_j = [jnp.asarray(b) for b in b_np] + sfa_j = jnp.asarray(sfa_u8.view(ml_dtypes.float8_e8m0fnu)) + sfb_j = [jnp.asarray(sfb.view(ml_dtypes.float8_e8m0fnu)) for sfb in sfb_u8] + offsets_j, alpha_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np, norm_const_np)) + jax.block_until_ready((a_j, sfa_j, *b_j, *sfb_j)) + + common = dict( + b_ptrs=packed_ptrs(b_j), + sfb_ptrs=packed_ptrs(sfb_j), + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + n=N, + sf_vec_size=SF_VEC_SIZE, + act_func="swiglu", + ) + + # Column-major bias layout is not expressible as a JAX array. + bias_j = jnp.zeros((N, EXPERTS), dtype=ml_dtypes.bfloat16) + with pytest.raises(ValueError, match="bias_tensor is not expressible"): + discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=sfa_j, + norm_const_tensor=norm_const_j, + b_dtype="float8_e4m3fn", + bias_tensor=bias_j, + **common, + ) + + # Packed fp4 has no JAX dtype; uint8 container inputs are rejected. + a_fp4_j = jnp.asarray(np.zeros((M, K // 2, 1), dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_fp4_j, + sfa_tensor=sfa_j, + b_dtype="uint8", + **common, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py new file mode 100644 index 000000000..8bf3be24b --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM dGLU backward wrapper. + +JAX contract: BF16 backend, discrete weight mode only (dense mode's expert-outermost +strided B is not expressible as row-major JAX arrays, and the block-scaled backend's +MMA-interleaved scale-factor layouts are likewise inexpressible), with b_ptrs built +from per-expert weight pointers as a packed uint8 array. dprob is a caller-provided +zero-initialized (m, 1, 1) f32 buffer that the kernel writes; dbias is wrapper-allocated. +Outputs are checked bit-identical against the torch wrapper run on identical input bytes. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 + + +def make_problem(m=512, n_weight=128, k=128, experts=2): + rng = np.random.default_rng(20260809) + two_n = 2 * n_weight + a_np = (rng.standard_normal((m, k, 1), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + c_np = (rng.standard_normal((m, two_n, 1), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + b_storage_np = (rng.standard_normal((experts, n_weight, k), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + group_m = m // experts + offsets_np = np.arange(group_m, m + 1, group_m, dtype=np.int32) + alpha_np = np.array([0.75, -1.25][:experts], dtype=np.float32) + beta_np = np.array([1.5, 0.5][:experts], dtype=np.float32) + prob_np = np.linspace(0.25, 0.875, m, dtype=np.float32).reshape(m, 1, 1) + return a_np, c_np, b_storage_np, offsets_np, alpha_np, beta_np, prob_np + + +def _packed_jax_ptrs(weights): + ptr_values = np.array([w.unsafe_buffer_pointer() for w in weights], dtype=np.int64) + return jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + + +@pytest.mark.L0 +def test_grouped_gemm_dglu_jax_discrete_matches_torch(): + skip_unless_sm100() + from cudnn import grouped_gemm_dglu_wrapper_sm100 + + m, n_weight, k, experts = 512, 128, 128, 2 + two_n = 2 * n_weight + a_np, c_np, b_storage_np, offsets_np, alpha_np, beta_np, prob_np = make_problem(m, n_weight, k, experts) + + # ---- torch run (discrete BF16 mode) ---- + a_t = torch.from_numpy(a_np.view(np.uint8)).view(torch.bfloat16).reshape(m, k, 1).cuda() + c_t = torch.from_numpy(c_np.view(np.uint8)).view(torch.bfloat16).reshape(m, two_n, 1).cuda() + b_storage_t = torch.from_numpy(b_storage_np.view(np.uint8)).view(torch.bfloat16).reshape(experts, n_weight, k).cuda() + b_ptrs_t = torch.tensor([b_storage_t[i].data_ptr() for i in range(experts)], dtype=torch.int64, device="cuda") + dprob_t = torch.zeros((m, 1, 1), dtype=torch.float32, device="cuda") + result_t = grouped_gemm_dglu_wrapper_sm100( + a_tensor=a_t, + c_tensor=c_t, + sfa_tensor=None, + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + beta_tensor=torch.from_numpy(beta_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + dprob_tensor=dprob_t, + b_ptrs=b_ptrs_t, + n=n_weight, + b_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + generate_dbias=True, + ) + torch.cuda.synchronize() + + # ---- jax run on identical bytes ---- + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n_weight, k) k-major + offsets_j, alpha_j, beta_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, beta_np, prob_np)) + # Kernel-written output buffer: zero-initialized, materialized before its pointer is used. + dprob_j = jax.block_until_ready(jnp.zeros((m, 1, 1), dtype=jnp.float32)) + jax.block_until_ready((a_j, c_j, offsets_j, alpha_j, beta_j, prob_j, *b_experts_j)) + + # Packed uint8 pointer array (8 little-endian bytes per pointer): JAX truncates + # int64 without x64 mode. The weight arrays must stay alive while the kernel runs. + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + + result_j = grouped_gemm_dglu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + b_ptrs=b_ptrs_j, + n=n_weight, + b_dtype="bfloat16", + d_dtype="bfloat16", + generate_dbias=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + for key in ("d_row_tensor", "dprob_tensor", "dbias_tensor"): + np.testing.assert_array_equal( + np.asarray(result_j[key]).astype(np.float32), + result_t[key].float().cpu().numpy(), + err_msg=f"grouped dGLU {key}: JAX output differs from torch output on identical input bytes", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_dglu_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import grouped_gemm_dglu_jax_sm100, grouped_gemm_dglu_wrapper_sm100 + + m, n_weight, k, experts = 512, 128, 128, 2 + a_np, c_np, b_storage_np, offsets_np, alpha_np, beta_np, prob_np = make_problem(m, n_weight, k, experts) + + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n_weight, k) k-major + offsets_j, alpha_j, beta_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, beta_np, prob_np)) + # Eager path only: kernel-written output buffer, zero-initialized and materialized + # before its pointer is used (the jit entry allocates dprob as a donated output). + dprob_j = jax.block_until_ready(jnp.zeros((m, 1, 1), dtype=jnp.float32)) + jax.block_until_ready((a_j, c_j, offsets_j, alpha_j, beta_j, prob_j, *b_experts_j)) + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row, so full-buffer bitwise comparison against the jit path is well-defined. + result_eager = grouped_gemm_dglu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + b_ptrs=b_ptrs_j, + n=n_weight, + b_dtype="bfloat16", + generate_dbias=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_row_tensor", "dprob_tensor", "dbias_tensor")} + + def check(d_row_tensor, dprob_tensor, dbias_tensor): + jax.block_until_ready((d_row_tensor, dprob_tensor, dbias_tensor)) + for got, key in ((d_row_tensor, "d_row_tensor"), (dprob_tensor, "dprob_tensor"), (dbias_tensor, "dbias_tensor")): + np.testing.assert_array_equal( + np.asarray(got).view(np.uint8), + expected[key], + err_msg=f"grouped dGLU {key}: jit output differs from eager wrapper output on identical input bytes", + ) + + # Eager custom call + check(*grouped_gemm_dglu_jax_sm100(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, n_weight, prob_j, generate_dbias=True)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit( + lambda a, c, offsets, alpha, beta, ptrs, prob: grouped_gemm_dglu_jax_sm100(a, c, offsets, alpha, beta, ptrs, n_weight, prob, generate_dbias=True), + ) + check(*jitted(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, prob_j)) + check(*jitted(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, prob_j)) + + # generate_dbias=False returns (d_row, dprob, None) + d_row_only, dprob_only, dbias_none = grouped_gemm_dglu_jax_sm100(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, n_weight, prob_j) + jax.block_until_ready((d_row_only, dprob_only)) + assert dbias_none is None + np.testing.assert_array_equal(np.asarray(d_row_only).view(np.uint8), expected["d_row_tensor"]) + np.testing.assert_array_equal(np.asarray(dprob_only).view(np.uint8), expected["dprob_tensor"]) + + +@pytest.mark.L0 +def test_grouped_gemm_dglu_jax_errors(): + skip_unless_sm100() + from cudnn import grouped_gemm_dglu_wrapper_sm100 + + m, n_weight, k, experts = 512, 128, 128, 2 + a_np, c_np, b_storage_np, offsets_np, alpha_np, beta_np, prob_np = make_problem(m, n_weight, k, experts) + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + offsets_j, alpha_j, beta_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, beta_np, prob_np)) + dprob_j = jnp.zeros((m, 1, 1), dtype=jnp.float32) + + # Dense weight mode: (n, k, experts) expert-outermost strides are inexpressible. + b_dense_j = jnp.asarray(b_storage_np) # (experts, n, k): not the dense-mode layout + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_dglu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + b_tensor=b_dense_j, + ) + + # Block-scaled backend (fp4/fp8 dtypes): MMA-interleaved SF layouts are inexpressible. + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] + jax.block_until_ready(b_experts_j) + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + a_fp4_j = jnp.asarray(np.zeros((m, k, 1), dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_dglu_wrapper_sm100( + a_tensor=a_fp4_j, + c_tensor=c_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + dprob_tensor=dprob_j, + b_ptrs=b_ptrs_j, + n=n_weight, + b_dtype="uint8", + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py new file mode 100644 index 000000000..3dc704094 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM dSReLU backward wrapper. + +JAX contract: discrete weight mode only (dense mode's expert-outermost strided B +layout is not expressible as row-major JAX arrays) with fp8 inputs (JAX has no +packed fp4 dtype). Pointer arrays are passed as packed uint8 (8 bytes per pointer) +because JAX truncates int64 without x64 mode. Scale-factor tensors are passed in +the physical C-contiguous atom shape (L, MN', K', 32, 4, 4) -- the permuted torch +view of the same bytes is not expressible in JAX, and the kernel rebuilds the SF +layout from the GEMM shapes, consuming only the SF base pointer. + +Outputs are checked bit-identical against the torch wrapper run on identical input +bytes, except dprob, which the kernel accumulates with atomic float adds (ordering +is not deterministic), checked with a tight allclose instead. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 + +MMA_PERMUTE_ORDER = (3, 4, 1, 5, 2, 0) + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def make_problem(m=512, n=256, k=256, experts=2, sf_vec_size=32): + rng = np.random.default_rng(20260809) + rk = ceil_div(ceil_div(k, sf_vec_size), 4) + a_np = (rng.integers(-4, 5, size=(m, k, 1)) * 0.25).astype(ml_dtypes.float8_e4m3fn) + b_np = (rng.integers(-4, 5, size=(experts, n, k)) * 0.25).astype(ml_dtypes.float8_e4m3fn) # per-expert (n, k) k-major + c_np = (rng.standard_normal((m, n, 1), dtype=np.float32) * 0.5).astype(ml_dtypes.bfloat16) + # SF tensors in physical C-contiguous atom form; e8m0 holds powers of two exactly. + sfa_np = (2.0 ** rng.integers(-1, 2, size=(1, ceil_div(m, 128), rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu) + sfb_np = (2.0 ** rng.integers(-1, 2, size=(experts, 1, ceil_div(n, 128), rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu) + offsets_np = np.arange(m // experts, m + 1, m // experts, dtype=np.int32) + alpha_np = np.array([0.75, 1.25][:experts], dtype=np.float32) + prob_np = np.linspace(0.25, 1.0, m, dtype=np.float32).reshape(m, 1, 1) + norm_const_np = np.ones(1, dtype=np.float32) + return a_np, b_np, c_np, sfa_np, sfb_np, offsets_np, alpha_np, prob_np, norm_const_np + + +def _torch_from_bytes(arr, torch_dtype, shape): + return torch.from_numpy(np.ascontiguousarray(arr).view(np.uint8)).view(torch_dtype).reshape(shape).cuda() + + +def _u8(tensor_or_array): + """Byte view of an fp8/bf16 tensor or array for exact comparisons.""" + if isinstance(tensor_or_array, torch.Tensor): + return tensor_or_array.view(torch.uint8).cpu().numpy() + return np.asarray(tensor_or_array).view(np.uint8) + + +@pytest.mark.L0 +def test_grouped_gemm_dsrelu_jax_discrete_fp8_matches_torch(): + skip_unless_sm100() + from cudnn import grouped_gemm_dsrelu_wrapper_sm100 + + m, n, k, experts, sf_vec_size = 512, 256, 256, 2, 32 + a_np, b_np, c_np, sfa_np, sfb_np, offsets_np, alpha_np, prob_np, norm_const_np = make_problem(m, n, k, experts, sf_vec_size) + + # ---- torch run (discrete mode, permuted SF atom views over the same bytes) ---- + a_t = _torch_from_bytes(a_np, torch.float8_e4m3fn, (m, k, 1)) + c_t = _torch_from_bytes(c_np, torch.bfloat16, (m, n, 1)) + sfa_t = _torch_from_bytes(sfa_np, torch.float8_e8m0fnu, sfa_np.shape).permute(MMA_PERMUTE_ORDER) + b_t = _torch_from_bytes(b_np, torch.float8_e4m3fn, b_np.shape) + sfb_t = _torch_from_bytes(sfb_np, torch.float8_e8m0fnu, sfb_np.shape) + b_ptrs_t = torch.tensor([b_t[i].data_ptr() for i in range(experts)], dtype=torch.int64, device="cuda") + sfb_ptrs_t = torch.tensor([sfb_t[i].data_ptr() for i in range(experts)], dtype=torch.int64, device="cuda") + + result_t = grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=a_t, + c_tensor=c_t, + sfa_tensor=sfa_t, + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + b_ptrs=b_ptrs_t, + sfb_ptrs=sfb_ptrs_t, + n=n, + b_dtype=torch.float8_e4m3fn, + b_major="k", + norm_const_tensor=torch.from_numpy(norm_const_np).cuda(), + d_dtype=torch.float8_e4m3fn, + sf_vec_size=sf_vec_size, + ) + torch.cuda.synchronize() + + # ---- jax run on identical bytes (SF tensors in the physical atom form) ---- + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + sfa_j = jnp.asarray(sfa_np) + b_experts_j = [jnp.asarray(b_np[i]) for i in range(experts)] + sfb_experts_j = [jnp.asarray(sfb_np[i]) for i in range(experts)] + offsets_j = jnp.asarray(offsets_np) + alpha_j = jnp.asarray(alpha_np) + prob_j = jnp.asarray(prob_np) + norm_const_j = jnp.asarray(norm_const_np) + jax.block_until_ready((a_j, c_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, *b_experts_j, *sfb_experts_j)) + + # Packed uint8 pointer arrays (8 little-endian bytes per pointer): JAX truncates + # int64 without x64 mode. The weight/SF arrays must stay alive while the kernel runs. + b_ptr_values = np.array([w.unsafe_buffer_pointer() for w in b_experts_j], dtype=np.int64) + sfb_ptr_values = np.array([w.unsafe_buffer_pointer() for w in sfb_experts_j], dtype=np.int64) + b_ptrs_j = jax.block_until_ready(jnp.asarray(b_ptr_values.view(np.uint8))) + sfb_ptrs_j = jax.block_until_ready(jnp.asarray(sfb_ptr_values.view(np.uint8))) + + result_j = grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + sfb_ptrs=sfb_ptrs_j, + n=n, + b_dtype="float8_e4m3fn", + b_major="k", + norm_const_tensor=norm_const_j, + d_dtype="float8_e4m3fn", + sf_vec_size=sf_vec_size, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + # Sanity: the kernels actually ran and produced non-trivial values. + assert np.count_nonzero(result_t["dprob_tensor"].cpu().numpy()) > 0 + assert np.count_nonzero(_u8(result_t["d_row_tensor"])) > 0 + + # Elementwise-quantized outputs must be bit-identical on identical input bytes. + for key in ("d_row_tensor", "d_col_tensor", "d_srelu_tensor"): + np.testing.assert_array_equal( + _u8(result_j[key]), + _u8(result_t[key]), + err_msg=f"dsrelu discrete {key}: JAX output differs from torch output on identical input bytes", + ) + + # SF outputs: JAX holds the physical atom form; the torch view of the same logical + # tensor is its (3, 4, 1, 5, 2, 0) permutation. + for key in ("sfd_row_tensor", "sfd_col_tensor", "sfd_col_d_srelu_tensor"): + np.testing.assert_array_equal( + np.transpose(_u8(result_j[key]), MMA_PERMUTE_ORDER), + _u8(result_t[key]), + err_msg=f"dsrelu discrete {key}: JAX output differs from torch output on identical input bytes", + ) + + # dprob is accumulated with atomic float adds; ordering is nondeterministic, so + # compare with a tight tolerance rather than bit-identically. + np.testing.assert_allclose( + np.asarray(result_j["dprob_tensor"]), + result_t["dprob_tensor"].cpu().numpy(), + rtol=1e-4, + atol=1e-4, + err_msg="dsrelu discrete dprob_tensor: JAX output differs from torch output beyond atomic-add tolerance", + ) + + assert result_j["amax_tensor"] is None and result_t["amax_tensor"] is None + assert result_j["dbias_tensor"] is None and result_t["dbias_tensor"] is None + + +@pytest.mark.L0 +def test_grouped_gemm_dsrelu_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import grouped_gemm_dsrelu_jax_sm100, grouped_gemm_dsrelu_wrapper_sm100 + + m, n, k, experts, sf_vec_size = 512, 256, 256, 2, 32 + a_np, b_np, c_np, sfa_np, sfb_np, offsets_np, alpha_np, prob_np, norm_const_np = make_problem(m, n, k, experts, sf_vec_size) + + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + sfa_j = jnp.asarray(sfa_np) + b_experts_j = [jnp.asarray(b_np[i]) for i in range(experts)] # per-expert (n, k) k-major + sfb_experts_j = [jnp.asarray(sfb_np[i]) for i in range(experts)] + offsets_j, alpha_j, prob_j, norm_const_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np, norm_const_np)) + jax.block_until_ready((a_j, c_j, sfa_j, offsets_j, alpha_j, prob_j, norm_const_j, *b_experts_j, *sfb_experts_j)) + b_ptr_values = np.array([w.unsafe_buffer_pointer() for w in b_experts_j], dtype=np.int64) + sfb_ptr_values = np.array([w.unsafe_buffer_pointer() for w in sfb_experts_j], dtype=np.int64) + b_ptrs_j = jax.block_until_ready(jnp.asarray(b_ptr_values.view(np.uint8))) + sfb_ptrs_j = jax.block_until_ready(jnp.asarray(sfb_ptr_values.view(np.uint8))) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row, so full-buffer bitwise comparison against the jit path is well-defined. + result_eager = grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + sfb_ptrs=sfb_ptrs_j, + n=n, + b_dtype="float8_e4m3fn", + b_major="k", + norm_const_tensor=norm_const_j, + d_dtype="float8_e4m3fn", + sf_vec_size=sf_vec_size, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + exact_keys = ("d_row_tensor", "d_col_tensor", "d_srelu_tensor", "sfd_row_tensor", "sfd_col_tensor", "sfd_col_d_srelu_tensor") + assert np.count_nonzero(_u8(result_eager["d_row_tensor"])) > 0 + expected = {key: _u8(result_eager[key]) for key in exact_keys} + expected_dprob = np.asarray(result_eager["dprob_tensor"]) + + def check(results): + # jit entry returns the eager wrapper's key order as a tuple. + d_row, d_col, d_srelu, dprob, dbias, amax, sfd_row, sfd_col, sfd_col_d_srelu = results + jax.block_until_ready((d_row, d_col, d_srelu, dprob, sfd_row, sfd_col, sfd_col_d_srelu)) + assert dbias is None and amax is None + got = { + "d_row_tensor": d_row, + "d_col_tensor": d_col, + "d_srelu_tensor": d_srelu, + "sfd_row_tensor": sfd_row, + "sfd_col_tensor": sfd_col, + "sfd_col_d_srelu_tensor": sfd_col_d_srelu, + } + for key in exact_keys: + np.testing.assert_array_equal( + _u8(got[key]), + expected[key], + err_msg=f"dsrelu jit {key}: output differs from eager wrapper output on identical input bytes", + ) + # dprob is accumulated with atomic float adds; ordering is nondeterministic. + np.testing.assert_allclose( + np.asarray(dprob), + expected_dprob, + rtol=1e-4, + atol=1e-4, + err_msg="dsrelu jit dprob: output differs from eager wrapper output beyond atomic-add tolerance", + ) + + def run(a, c, sfa, offsets, alpha, prob, b_ptrs, sfb_ptrs, norm_const): + return grouped_gemm_dsrelu_jax_sm100( + a_tensor=a, + c_tensor=c, + sfa_tensor=sfa, + padded_offsets=offsets, + alpha_tensor=alpha, + prob_tensor=prob, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + n=n, + norm_const_tensor=norm_const, + b_dtype="float8_e4m3fn", + b_major="k", + d_dtype="float8_e4m3fn", + sf_vec_size=sf_vec_size, + ) + + args = (a_j, c_j, sfa_j, offsets_j, alpha_j, prob_j, b_ptrs_j, sfb_ptrs_j, norm_const_j) + + # Eager custom call + check(run(*args)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit(run) + check(jitted(*args)) + check(jitted(*args)) + + +@pytest.mark.L0 +def test_grouped_gemm_dsrelu_jax_errors(): + skip_unless_sm100() + from cudnn import grouped_gemm_dsrelu_wrapper_sm100 + + m, n, k, experts, sf_vec_size = 512, 256, 256, 2, 32 + a_np, b_np, c_np, sfa_np, sfb_np, offsets_np, alpha_np, prob_np, norm_const_np = make_problem(m, n, k, experts, sf_vec_size) + a_j = jnp.asarray(a_np) + c_j = jnp.asarray(c_np) + sfa_j = jnp.asarray(sfa_np) + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + + # Dense weight mode: the expert-outermost strided (n, k, l) B layout is inexpressible. + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=a_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_tensor=jnp.asarray(np.zeros((n, k, experts), dtype=ml_dtypes.float8_e4m3fn)), + sfb_tensor=jnp.asarray(sfb_np), + sf_vec_size=sf_vec_size, + ) + + # Packed fp4 (raw uint8) inputs: JAX has no packed fp4 dtype. + a_u8_j = jnp.asarray(np.zeros((m, k // 2, 1), dtype=np.uint8)) + b_ptrs_j = jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)) + sfb_ptrs_j = jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=a_u8_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + sfb_ptrs=sfb_ptrs_j, + n=n, + b_dtype="uint8", + sf_vec_size=16, + ) + + # Plain numpy arrays are neither torch nor JAX device tensors. + with pytest.raises(ValueError, match="Unsupported tensor framework"): + grouped_gemm_dsrelu_wrapper_sm100( + a_tensor=np.asarray(a_np), + c_tensor=np.asarray(c_np), + sfa_tensor=np.asarray(sfa_np), + padded_offsets=offsets_np, + alpha_tensor=alpha_np, + prob_tensor=prob_np, + b_ptrs=np.zeros(8 * experts, dtype=np.uint8), + sfb_ptrs=np.zeros(8 * experts, dtype=np.uint8), + n=n, + b_dtype="float8_e4m3fn", + sf_vec_size=sf_vec_size, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py new file mode 100644 index 000000000..6441d1702 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM dSwiGLU backward wrapper. + +JAX contract: this backward API only supports dense weight mode, whose +expert-outermost strided B layout (n, k, l) has no row-major JAX equivalent, +so JAX inputs are rejected with a clear error (and unknown frameworks with +an "Unsupported tensor framework" error). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import skip_unless_sm100 + + +def _make_jax_inputs(m=256, n=128, k=128, l=2): + rng = np.random.default_rng(20260809) + a_j = jnp.asarray((rng.integers(-4, 5, size=(m, k, 1)) * 0.25).astype(ml_dtypes.float8_e4m3fn)) + b_j = jnp.asarray((rng.integers(-4, 5, size=(n, k, l)) * 0.25).astype(ml_dtypes.float8_e4m3fn)) + c_j = jnp.asarray(rng.standard_normal((m, n * 2, 1), dtype=np.float32).astype(ml_dtypes.bfloat16)) + rk = ((k + 31) // 32 + 3) // 4 + sfa_j = jnp.asarray((2.0 ** rng.integers(-2, 3, size=(1, (m + 127) // 128, rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu)) + sfb_j = jnp.asarray((2.0 ** rng.integers(-2, 3, size=(l, (n + 127) // 128, rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu)) + offsets_j = jnp.asarray(np.arange(m // l, m + 1, m // l, dtype=np.int32)) + alpha_j = jnp.asarray(np.ones(l, dtype=np.float32)) + beta_j = jnp.asarray(np.ones(l, dtype=np.float32)) + prob_j = jnp.asarray(np.ones((m, 1, 1), dtype=np.float32)) + return a_j, b_j, c_j, sfa_j, sfb_j, offsets_j, alpha_j, beta_j, prob_j + + +@pytest.mark.L0 +def test_grouped_gemm_dswiglu_jax_rejected(): + """JAX inputs are rejected: the dense-only B layout is not expressible as JAX arrays.""" + skip_unless_sm100() + from cudnn import grouped_gemm_dswiglu_wrapper_sm100 + + a_j, b_j, c_j, sfa_j, sfb_j, offsets_j, alpha_j, beta_j, prob_j = _make_jax_inputs() + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + c_tensor=c_j, + sfa_tensor=sfa_j, + sfb_tensor=sfb_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + beta_tensor=beta_j, + prob_tensor=prob_j, + sf_vec_size=32, + ) + + +@pytest.mark.L0 +def test_grouped_gemm_dswiglu_jax_api_class_rejected(): + """Constructing the API class directly with JAX samples raises the same clear error.""" + skip_unless_sm100() + from cudnn import GroupedGemmDswigluSm100 + + a_j, b_j, c_j, sfa_j, sfb_j, offsets_j, alpha_j, beta_j, prob_j = _make_jax_inputs() + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + GroupedGemmDswigluSm100( + sample_a=a_j, + sample_b=b_j, + sample_c=c_j, + sample_d_row=c_j, + sample_d_col=c_j, + sample_sfa=sfa_j, + sample_sfb=sfb_j, + sample_padded_offsets=offsets_j, + sample_alpha=alpha_j, + sample_beta=beta_j, + sample_prob=prob_j, + sample_dprob=prob_j, + sf_vec_size=32, + ) + + +@pytest.mark.L0 +def test_grouped_gemm_dswiglu_unknown_framework_rejected(): + """Plain numpy arrays are neither torch nor JAX device tensors: clear unsupported error.""" + skip_unless_sm100() + from cudnn import grouped_gemm_dswiglu_wrapper_sm100 + + m, n, k, l = 256, 128, 128, 2 + a_np = np.zeros((m, k, 1), dtype=np.uint8) + with pytest.raises(ValueError, match="Unsupported tensor framework"): + grouped_gemm_dswiglu_wrapper_sm100( + a_tensor=a_np, + b_tensor=np.zeros((n, k, l), dtype=np.uint8), + c_tensor=np.zeros((m, n * 2, 1), dtype=np.float32), + sfa_tensor=None, + sfb_tensor=None, + padded_offsets=np.arange(m // l, m + 1, m // l, dtype=np.int32), + alpha_tensor=np.ones(l, dtype=np.float32), + beta_tensor=np.ones(l, dtype=np.float32), + prob_tensor=np.ones((m, 1, 1), dtype=np.float32), + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.py new file mode 100644 index 000000000..d0830ef81 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM GLU + Hadamard wrapper. + +The GLU + Hadamard fusion is block-scaled only: its mandatory scale-factor inputs +(sfa/sfb) use an MMA-interleaved 6-D layout with no row-major equivalent, so no +configuration of this API is expressible as JAX arrays. The wrapper and the API +class must therefore reject JAX inputs with a clear error. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import skip_unless_sm100 + + +@pytest.mark.L0 +def test_grouped_gemm_glu_hadamard_jax_rejected(): + skip_unless_sm100() + from cudnn import grouped_gemm_glu_hadamard_wrapper_sm100 + + m, k, experts = 512, 128, 2 + a_j = jnp.asarray(np.zeros((m, k // 2, 1), dtype=np.uint8)) # packed fp4 container + offsets_j = jnp.asarray(np.array([256, 512], dtype=np.int32)) + alpha_j = jnp.asarray(np.ones(experts, dtype=np.float32)) + prob_j = jnp.asarray(np.ones((m, 1, 1), dtype=np.float32)) + sfa_j = jnp.asarray(np.zeros((1,), dtype=np.uint8)) # never inspected: rejection happens first + + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_glu_hadamard_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)), + sfb_ptrs=jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)), + n=256, + b_dtype="uint8", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_glu_hadamard_class_jax_rejected(): + skip_unless_sm100() + from cudnn.gemm.cutedsl.grouped.glu_hadamard.api import GroupedGemmGluHadamardSm100 + + a_j = jnp.asarray(np.zeros((512, 64, 1), dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + GroupedGemmGluHadamardSm100( + sample_a=a_j, + sample_c=None, + sample_d=None, + sample_sfa=None, + sample_padded_offsets=None, + sample_alpha=None, + sample_prob=None, + num_experts=2, + b_shape=(256, 128), + b_dtype="uint8", + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py new file mode 100644 index 000000000..b87a9471a --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM GLU forward wrapper. + +JAX contract: BF16 backend, discrete weight mode only (dense mode's expert-outermost +strided B and the column-major bias layout are not expressible as row-major JAX arrays, +and the block-scaled backend's MMA-interleaved scale-factor layouts are likewise +inexpressible), with b_ptrs built from per-expert weight pointers -- as a packed uint8 +array (8 bytes per pointer) since JAX truncates int64 without x64 mode. Outputs are +checked bit-identical against the torch wrapper run on identical input bytes (both +paths share one compiled kernel). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 + + +def make_problem(m=512, n_full=256, k=128, experts=2): + rng = np.random.default_rng(20260808) + a_np = (rng.standard_normal((m, k, 1), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + b_storage_np = (rng.standard_normal((experts, n_full, k), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + group_m = m // experts + offsets_np = np.arange(group_m, m + 1, group_m, dtype=np.int32) + alpha_np = np.array([0.75, -1.25][:experts], dtype=np.float32) + prob_np = np.linspace(0.25, 0.875, m, dtype=np.float32).reshape(m, 1, 1) + return a_np, b_storage_np, offsets_np, alpha_np, prob_np + + +def _packed_jax_ptrs(weights): + ptr_values = np.array([w.unsafe_buffer_pointer() for w in weights], dtype=np.int64) + return jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + + +@pytest.mark.L0 +@pytest.mark.parametrize("act_func", ["swiglu", "geglu"]) +def test_grouped_gemm_glu_jax_discrete_matches_torch(act_func): + skip_unless_sm100() + from cudnn import grouped_gemm_glu_wrapper_sm100 + + m, n_full, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n_full, k, experts) + + # ---- torch run (discrete BF16 mode, no bias) ---- + a_t = torch.from_numpy(a_np.view(np.uint8)).view(torch.bfloat16).reshape(m, k, 1).cuda() + b_storage_t = torch.from_numpy(b_storage_np.view(np.uint8)).view(torch.bfloat16).reshape(experts, n_full, k).cuda() + b_ptrs_t = torch.tensor([b_storage_t[i].data_ptr() for i in range(experts)], dtype=torch.int64, device="cuda") + result_t = grouped_gemm_glu_wrapper_sm100( + a_tensor=a_t, + sfa_tensor=None, + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + b_ptrs=b_ptrs_t, + n=n_full, + b_dtype=torch.bfloat16, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + act_func=act_func, + generate_c=True, + ) + torch.cuda.synchronize() + + # ---- jax run on identical bytes ---- + a_j = jnp.asarray(a_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n_full, k) k-major + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + jax.block_until_ready((a_j, offsets_j, alpha_j, prob_j, *b_experts_j)) + + # Packed uint8 pointer array (8 little-endian bytes per pointer): JAX truncates + # int64 without x64 mode. The weight arrays must stay alive while the kernel runs. + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + + result_j = grouped_gemm_glu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n_full, + b_dtype="bfloat16", + c_dtype="bfloat16", + d_dtype="bfloat16", + act_func=act_func, + generate_c=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + for key in ("c_tensor", "d_tensor"): + np.testing.assert_array_equal( + np.asarray(result_j[key]).astype(np.float32), + result_t[key].float().cpu().numpy(), + err_msg=f"grouped GLU {key} ({act_func}): JAX output differs from torch output on identical input bytes", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_glu_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import grouped_gemm_glu_jax_sm100, grouped_gemm_glu_wrapper_sm100 + + m, n_full, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n_full, k, experts) + + a_j = jnp.asarray(a_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n_full, k) k-major + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + jax.block_until_ready((a_j, offsets_j, alpha_j, prob_j, *b_experts_j)) + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row, so full-buffer bitwise comparison against the jit path is well-defined. + result_eager = grouped_gemm_glu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n_full, + b_dtype="bfloat16", + generate_c=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_tensor", "c_tensor")} + + def check(d_tensor, c_tensor): + jax.block_until_ready((d_tensor, c_tensor)) + for got, key in ((d_tensor, "d_tensor"), (c_tensor, "c_tensor")): + np.testing.assert_array_equal( + np.asarray(got).view(np.uint8), + expected[key], + err_msg=f"grouped GLU {key}: jit output differs from eager wrapper output on identical input bytes", + ) + + # Eager custom call + check(*grouped_gemm_glu_jax_sm100(a_j, offsets_j, alpha_j, b_ptrs_j, n_full, prob_j, generate_c=True)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit( + lambda a, offsets, alpha, ptrs, prob: grouped_gemm_glu_jax_sm100(a, offsets, alpha, ptrs, n_full, prob, generate_c=True), + ) + check(*jitted(a_j, offsets_j, alpha_j, b_ptrs_j, prob_j)) + check(*jitted(a_j, offsets_j, alpha_j, b_ptrs_j, prob_j)) + + # generate_c=False returns (d, None) + d_only, c_none = grouped_gemm_glu_jax_sm100(a_j, offsets_j, alpha_j, b_ptrs_j, n_full, prob_j) + jax.block_until_ready(d_only) + assert c_none is None + np.testing.assert_array_equal(np.asarray(d_only).view(np.uint8), expected["d_tensor"]) + + +@pytest.mark.L0 +def test_grouped_gemm_glu_jax_errors(): + skip_unless_sm100() + from cudnn import grouped_gemm_glu_wrapper_sm100 + + m, n_full, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n_full, k, experts) + a_j = jnp.asarray(a_np) + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + + # Dense weight mode: (n, k, experts) expert-outermost strides are inexpressible. + b_dense_j = jnp.asarray(b_storage_np) # (experts, n, k): not the dense-mode layout + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_glu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_tensor=b_dense_j, + ) + + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] + jax.block_until_ready(b_experts_j) + b_ptrs_j = _packed_jax_ptrs(b_experts_j) + + # Column-major (n, experts) bias is inexpressible. + bias_j = jnp.zeros((n_full, experts), dtype=ml_dtypes.bfloat16) + with pytest.raises(ValueError, match="bias_tensor is not expressible"): + grouped_gemm_glu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n_full, + b_dtype="bfloat16", + bias_tensor=bias_j, + ) + + # Block-scaled backend (fp4/fp8 dtypes): MMA-interleaved SF layouts are inexpressible. + a_fp4_j = jnp.asarray(np.zeros((m, k, 1), dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_glu_wrapper_sm100( + a_tensor=a_fp4_j, + sfa_tensor=None, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n_full, + b_dtype="uint8", + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py new file mode 100644 index 000000000..dd23f8ad1 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the unfused SM100 grouped GEMM wrapper. + +JAX contract: discrete weight mode only (dense mode's expert-outermost strided B and the +column-major bias layout are not expressible as row-major JAX arrays), with b_ptrs built +from per-expert weight pointers — as a packed uint8 array (8 bytes per pointer) since JAX +truncates int64 without x64 mode. Outputs are checked bit-identical against the torch +wrapper run on identical input bytes (both paths share one compiled kernel). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 + + +def make_problem(m=512, n=256, k=128, experts=2): + rng = np.random.default_rng(20260716) + a_np = (rng.standard_normal((m, k, 1), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + b_storage_np = (rng.standard_normal((experts, n, k), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) + group_m = m // experts + offsets_np = np.arange(group_m, m + 1, group_m, dtype=np.int32) + alpha_np = np.array([0.75, -1.25][:experts], dtype=np.float32) + prob_np = np.linspace(0.25, 0.875, m, dtype=np.float32).reshape(m, 1, 1) + return a_np, b_storage_np, offsets_np, alpha_np, prob_np + + +@pytest.mark.L0 +def test_grouped_gemm_jax_discrete_matches_torch(): + skip_unless_sm100() + from cudnn import grouped_gemm_wrapper_sm100 + + m, n, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n, k, experts) + + # ---- torch run (discrete mode, no bias) ---- + a_t = torch.from_numpy(a_np.view(np.uint8)).view(torch.bfloat16).reshape(m, k, 1).cuda() + b_storage_t = torch.from_numpy(b_storage_np.view(np.uint8)).view(torch.bfloat16).reshape(experts, n, k).cuda() + b_ptrs_t = torch.tensor([b_storage_t[i].data_ptr() for i in range(experts)], dtype=torch.int64, device="cuda") + result_t = grouped_gemm_wrapper_sm100( + a_tensor=a_t, + padded_offsets=torch.from_numpy(offsets_np).cuda(), + alpha_tensor=torch.from_numpy(alpha_np).cuda(), + prob_tensor=torch.from_numpy(prob_np).cuda(), + b_ptrs=b_ptrs_t, + n=n, + b_dtype=torch.bfloat16, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + generate_c=True, + ) + torch.cuda.synchronize() + + # ---- jax run on identical bytes ---- + a_j = jnp.asarray(a_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n, k) k-major + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + jax.block_until_ready((a_j, offsets_j, alpha_j, prob_j, *b_experts_j)) + + # Packed uint8 pointer array (8 little-endian bytes per pointer): JAX truncates + # int64 without x64 mode. The weight arrays must stay alive while the kernel runs. + ptr_values = np.array([w.unsafe_buffer_pointer() for w in b_experts_j], dtype=np.int64) + b_ptrs_j = jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + + result_j = grouped_gemm_wrapper_sm100( + a_tensor=a_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n, + b_dtype="bfloat16", + c_dtype="bfloat16", + d_dtype="bfloat16", + generate_c=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + for key in ("c_tensor", "d_tensor"): + np.testing.assert_array_equal( + np.asarray(result_j[key]).astype(np.float32), + result_t[key].float().cpu().numpy(), + err_msg=f"unfused grouped {key}: JAX output differs from torch output on identical input bytes", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper.""" + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import grouped_gemm_jax_sm100, grouped_gemm_wrapper_sm100 + + m, n, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n, k, experts) + + a_j = jnp.asarray(a_np) + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] # per-expert (n, k) k-major + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + jax.block_until_ready((a_j, offsets_j, alpha_j, prob_j, *b_experts_j)) + ptr_values = np.array([w.unsafe_buffer_pointer() for w in b_experts_j], dtype=np.int64) + b_ptrs_j = jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + + # Eager wrapper baseline on the same bytes; the last padded offset covers every + # row, so full-buffer bitwise comparison against the jit path is well-defined. + result_eager = grouped_gemm_wrapper_sm100( + a_tensor=a_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n, + b_dtype="bfloat16", + generate_c=True, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_tensor", "c_tensor")} + + def check(d_tensor, c_tensor): + jax.block_until_ready((d_tensor, c_tensor)) + for got, key in ((d_tensor, "d_tensor"), (c_tensor, "c_tensor")): + np.testing.assert_array_equal( + np.asarray(got).view(np.uint8), + expected[key], + err_msg=f"unfused grouped {key}: jit output differs from eager wrapper output on identical input bytes", + ) + + # Eager custom call + check(*grouped_gemm_jax_sm100(a_j, offsets_j, alpha_j, b_ptrs_j, n, prob_j, generate_c=True)) + + # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. + jitted = jax.jit( + lambda a, offsets, alpha, ptrs, prob: grouped_gemm_jax_sm100(a, offsets, alpha, ptrs, n, prob, generate_c=True), + ) + check(*jitted(a_j, offsets_j, alpha_j, b_ptrs_j, prob_j)) + check(*jitted(a_j, offsets_j, alpha_j, b_ptrs_j, prob_j)) + + # generate_c=False returns (d, None) + d_only, c_none = grouped_gemm_jax_sm100(a_j, offsets_j, alpha_j, b_ptrs_j, n, prob_j) + jax.block_until_ready(d_only) + assert c_none is None + np.testing.assert_array_equal(np.asarray(d_only).view(np.uint8), expected["d_tensor"]) + + +@pytest.mark.L0 +def test_grouped_gemm_jax_errors(): + skip_unless_sm100() + from cudnn import grouped_gemm_wrapper_sm100 + + m, n, k, experts = 512, 256, 128, 2 + a_np, b_storage_np, offsets_np, alpha_np, prob_np = make_problem(m, n, k, experts) + a_j = jnp.asarray(a_np) + offsets_j, alpha_j, prob_j = (jnp.asarray(x) for x in (offsets_np, alpha_np, prob_np)) + b_dense_j = jnp.asarray(b_storage_np) # (experts, n, k): not the dense-mode layout + + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_wrapper_sm100( + a_tensor=a_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_tensor=b_dense_j, + ) + + b_experts_j = [jnp.asarray(b_storage_np[i]) for i in range(experts)] + jax.block_until_ready(b_experts_j) + ptr_values = np.array([w.unsafe_buffer_pointer() for w in b_experts_j], dtype=np.int64) + b_ptrs_j = jnp.asarray(ptr_values.view(np.uint8)) + bias_j = jnp.zeros((n, experts), dtype=ml_dtypes.bfloat16) + with pytest.raises(ValueError, match="bias_tensor is not expressible"): + grouped_gemm_wrapper_sm100( + a_tensor=a_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=b_ptrs_j, + n=n, + b_dtype="bfloat16", + bias_tensor=bias_j, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py new file mode 100644 index 000000000..c06059c24 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM Quant wrapper. + +JAX contract: every configuration of this kernel is block-scaled and consumes the +scale-factor tensor sfa (and, in dense mode, sfb; plus the sfd outputs for FP8 configs) +as MMA-tiled (32, 4, m//128, 4, rest_k, l) strided views built via torch .permute(). +Those layouts are not expressible as row-major JAX arrays, so JAX inputs are rejected +with a clear ValueError at both the wrapper and the API class — in dense AND discrete +(b_ptrs) weight modes, since sfa is required in both. Torch behavior is unchanged. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +import jax.numpy as jnp + + +def make_jax_inputs(m=256, n=256, k=128, experts=2): + """Plausible (physical-layout) inputs; rejection fires before shape validation.""" + rng = np.random.default_rng(20260809) + a_j = jnp.asarray(rng.integers(0, 255, (m, k, 1), dtype=np.uint8).view(ml_dtypes.float8_e4m3fn)) + rest_k = -(-(-(-k // 32) // 4)) # ceil_div(ceil_div(k, 32), 4) + sfa_j = jnp.asarray(rng.integers(0, 127, (32, 4, -(-m // 128), 4, rest_k, 1), dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + group_m = m // experts + offsets_j = jnp.asarray(np.arange(group_m, m + 1, group_m, dtype=np.int32)) + alpha_j = jnp.asarray(np.ones(experts, dtype=np.float32)) + return a_j, sfa_j, offsets_j, alpha_j + + +@pytest.mark.L0 +def test_grouped_gemm_quant_jax_discrete_rejected_with_clear_error(): + from cudnn import grouped_gemm_quant_wrapper_sm100 + + m, n, k, experts = 256, 256, 128, 2 + a_j, sfa_j, offsets_j, alpha_j = make_jax_inputs(m, n, k, experts) + # Discrete mode ships weights as pointer arrays, but sfa is still an MMA-tiled + # cute tensor argument, so the whole jax config is rejected up front. + ptrs_j = jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_quant_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + b_ptrs=ptrs_j, + sfb_ptrs=ptrs_j, + n=n, + b_dtype="float8_e4m3fn", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_quant_jax_api_class_rejected(): + from cudnn.gemm.cutedsl.grouped.quant.api import GroupedGemmQuantSm100 + + m, n, k, experts = 256, 256, 128, 2 + a_j, sfa_j, offsets_j, alpha_j = make_jax_inputs(m, n, k, experts) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + GroupedGemmQuantSm100( + sample_a=a_j, + sample_sfa=sfa_j, + sample_padded_offsets=offsets_j, + sample_alpha=alpha_j, + sample_d=None, + num_experts=experts, + b_shape=(n, k), + b_dtype="float8_e4m3fn", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_quant_unknown_framework_rejected(): + from cudnn import grouped_gemm_quant_wrapper_sm100 + + a_np = np.zeros((256, 128, 1), dtype=np.uint8) + with pytest.raises(ValueError, match="Unsupported tensor framework 'numpy'"): + grouped_gemm_quant_wrapper_sm100( + a_tensor=a_np, + sfa_tensor=None, + padded_offsets=None, + alpha_tensor=None, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py new file mode 100644 index 000000000..2519b1da2 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM SReLU wrapper. + +JAX contract: every configuration of this kernel is block-scaled and consumes the +scale-factor tensor sfa (and, in dense mode, sfb; plus the sfd outputs for FP8 configs) +as MMA-tiled (32, 4, m//128, 4, rest_k, l) strided views built via torch .permute(). +Those layouts are not expressible as row-major JAX arrays, so JAX inputs are rejected +with a clear ValueError at both the wrapper and the API class — in dense AND discrete +(b_ptrs) weight modes, since sfa is required in both. Torch behavior is unchanged. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +import jax.numpy as jnp + + +def make_jax_inputs(m=256, n=256, k=128, experts=2): + """Plausible (physical-layout) inputs; rejection fires before shape validation.""" + rng = np.random.default_rng(20260809) + a_j = jnp.asarray(rng.integers(0, 255, (m, k, 1), dtype=np.uint8).view(ml_dtypes.float8_e4m3fn)) + rest_k = -(-(-(-k // 32) // 4)) # ceil_div(ceil_div(k, 32), 4) + sfa_j = jnp.asarray(rng.integers(0, 127, (32, 4, -(-m // 128), 4, rest_k, 1), dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + group_m = m // experts + offsets_j = jnp.asarray(np.arange(group_m, m + 1, group_m, dtype=np.int32)) + alpha_j = jnp.asarray(np.ones(experts, dtype=np.float32)) + prob_j = jnp.asarray(np.ones((m, 1, 1), dtype=np.float32)) + return a_j, sfa_j, offsets_j, alpha_j, prob_j + + +@pytest.mark.L0 +def test_grouped_gemm_srelu_jax_discrete_rejected_with_clear_error(): + from cudnn import grouped_gemm_srelu_wrapper_sm100 + + m, n, k, experts = 256, 256, 128, 2 + a_j, sfa_j, offsets_j, alpha_j, prob_j = make_jax_inputs(m, n, k, experts) + # Discrete mode ships weights as pointer arrays, but sfa is still an MMA-tiled + # cute tensor argument, so the whole jax config is rejected up front. + ptrs_j = jnp.asarray(np.zeros(8 * experts, dtype=np.uint8)) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_srelu_wrapper_sm100( + a_tensor=a_j, + sfa_tensor=sfa_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + prob_tensor=prob_j, + b_ptrs=ptrs_j, + sfb_ptrs=ptrs_j, + n=n, + b_dtype="float8_e4m3fn", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_srelu_jax_api_class_rejected(): + from cudnn.gemm.cutedsl.grouped.srelu.api import GroupedGemmSreluSm100 + + m, n, k, experts = 256, 256, 128, 2 + a_j, sfa_j, offsets_j, alpha_j, prob_j = make_jax_inputs(m, n, k, experts) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + GroupedGemmSreluSm100( + sample_a=a_j, + sample_sfa=sfa_j, + sample_padded_offsets=offsets_j, + sample_alpha=alpha_j, + sample_prob=prob_j, + num_experts=experts, + b_shape=(n, k), + b_dtype="float8_e4m3fn", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_srelu_unknown_framework_rejected(): + from cudnn import grouped_gemm_srelu_wrapper_sm100 + + a_np = np.zeros((256, 128, 1), dtype=np.uint8) + with pytest.raises(ValueError, match="Unsupported tensor framework 'numpy'"): + grouped_gemm_srelu_wrapper_sm100( + a_tensor=a_np, + sfa_tensor=None, + padded_offsets=None, + alpha_tensor=None, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py new file mode 100644 index 000000000..56ae53cab --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM SwiGLU wrapper. + +JAX contract: every configuration of this kernel is block-scaled and consumes the +scale-factor tensors (sfa/sfb, and the sfd outputs for FP8 configs) as MMA-tiled +(32, 4, m//128, 4, rest_k, l) strided views built via torch .permute(). Those layouts +are not expressible as row-major JAX arrays, so JAX inputs are rejected with a clear +ValueError at both the wrapper and the API class; torch behavior is unchanged. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +import jax.numpy as jnp + + +def make_jax_inputs(m=256, n=256, k=128, experts=2): + """Plausible (physical-layout) inputs; rejection fires before shape validation.""" + rng = np.random.default_rng(20260809) + a_j = jnp.asarray(rng.integers(0, 255, (m, k // 2, 1), dtype=np.uint8)) # fp4x2 container + b_j = jnp.asarray(rng.integers(0, 255, (n, k // 2, experts), dtype=np.uint8)) + rest_k = -(-(-(-k // 16) // 4)) # ceil_div(ceil_div(k, 16), 4) + sfa_j = jnp.asarray(rng.integers(0, 127, (32, 4, -(-m // 128), 4, rest_k, 1), dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + sfb_j = jnp.asarray(rng.integers(0, 127, (32, 4, -(-n // 128), 4, rest_k, experts), dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + group_m = m // experts + offsets_j = jnp.asarray(np.arange(group_m, m + 1, group_m, dtype=np.int32)) + alpha_j = jnp.asarray(np.ones(experts, dtype=np.float32)) + return a_j, b_j, sfa_j, sfb_j, offsets_j, alpha_j + + +@pytest.mark.L0 +def test_grouped_gemm_swiglu_jax_rejected_with_clear_error(): + from cudnn import grouped_gemm_swiglu_wrapper_sm100 + + a_j, b_j, sfa_j, sfb_j, offsets_j, alpha_j = make_jax_inputs() + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + sfa_tensor=sfa_j, + sfb_tensor=sfb_j, + padded_offsets=offsets_j, + alpha_tensor=alpha_j, + ) + + +@pytest.mark.L0 +def test_grouped_gemm_swiglu_jax_api_class_rejected(): + from cudnn.gemm.cutedsl.grouped.swiglu.api import GroupedGemmSwigluSm100 + + a_j, b_j, sfa_j, sfb_j, offsets_j, alpha_j = make_jax_inputs() + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + GroupedGemmSwigluSm100( + sample_a=a_j, + sample_b=b_j, + sample_c=None, + sample_d=None, + sample_sfa=sfa_j, + sample_sfb=sfb_j, + sample_padded_offsets=offsets_j, + sample_alpha=alpha_j, + sample_d_col=None, + ) + + +@pytest.mark.L0 +def test_grouped_gemm_swiglu_unknown_framework_rejected(): + from cudnn import grouped_gemm_swiglu_wrapper_sm100 + + a_np = np.zeros((256, 64, 1), dtype=np.uint8) + with pytest.raises(ValueError, match="Unsupported tensor framework 'numpy'"): + grouped_gemm_swiglu_wrapper_sm100( + a_tensor=a_np, + b_tensor=None, + sfa_tensor=None, + sfb_tensor=None, + padded_offsets=None, + alpha_tensor=None, + ) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py new file mode 100644 index 000000000..e0ce85d0f --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +JAX coverage for the SM100 grouped GEMM wgrad wrapper. + +JAX contract: BF16 backend only, with A K-major and B N-major (both plain C-contiguous +JAX arrays); dense outputs are C-contiguous (expert, M, N) arrays, and discrete outputs +take a packed-uint8 wgrad_ptrs array (8 bytes per pointer, since JAX truncates int64 +without x64 mode). Outputs are checked bit-identical against the torch wrapper run on +identical input bytes. The block-scaled backend is rejected (its B operand is K-major, +i.e. column-major, and fp4 operands are K-packed -- neither is expressible as a +row-major JAX array). +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +ml_dtypes = pytest.importorskip("ml_dtypes") +torch = pytest.importorskip("torch") +import jax.numpy as jnp + +from fe_api.gemm.test_gemm_amax_jax import device_sync, skip_unless_sm100 + + +def make_problem(m=128, n=128, group_k_list=(256, 256)): + rng = np.random.default_rng(20260809) + tokens = sum(group_k_list) + a_np = (rng.standard_normal((m, tokens), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) # K-major + b_np = (rng.standard_normal((tokens, n), dtype=np.float32) * 0.125).astype(ml_dtypes.bfloat16) # N-major + offsets_np = np.array([sum(group_k_list[: i + 1]) for i in range(len(group_k_list))], dtype=np.int32) + return a_np, b_np, offsets_np + + +def to_torch_bf16(x_np): + return torch.from_numpy(x_np.view(np.uint8)).view(torch.bfloat16).reshape(x_np.shape).cuda() + + +def reference_wgrad(a_np, b_np, offsets_np): + a32 = a_np.astype(np.float32) + b32 = b_np.astype(np.float32) + result = [] + begin = 0 + for end in offsets_np.tolist(): + result.append(a32[:, begin:end] @ b32[begin:end, :]) + begin = end + return np.stack(result) + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_jax_dense_matches_torch(): + skip_unless_sm100() + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + a_np, b_np, offsets_np = make_problem() + + kwargs = dict( + sfa_tensor=None, + sfb_tensor=None, + output_mode="dense", + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + result_t = grouped_gemm_wgrad_wrapper_sm100( + a_tensor=to_torch_bf16(a_np), + b_tensor=to_torch_bf16(b_np), + offsets_tensor=torch.from_numpy(offsets_np).cuda(), + **kwargs, + ) + torch.cuda.synchronize() + + a_j, b_j, offsets_j = (jnp.asarray(x) for x in (a_np, b_np, offsets_np)) + jax.block_until_ready((a_j, b_j, offsets_j)) + result_j = grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + offsets_tensor=offsets_j, + **kwargs, + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + + wgrad_j = np.asarray(result_j["wgrad_tensor"]).astype(np.float32) + wgrad_t = result_t["wgrad_tensor"].float().cpu().numpy() + np.testing.assert_array_equal(wgrad_j, wgrad_t, err_msg="wgrad dense: JAX output differs from torch output on identical input bytes") + # Value sanity against the fp32 grouped-mm oracle (upstream tolerance). + np.testing.assert_allclose(wgrad_j, reference_wgrad(a_np, b_np, offsets_np), rtol=3e-2, atol=8e-2) + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_jax_discrete_matches_torch(): + skip_unless_sm100() + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + a_np, b_np, offsets_np = make_problem() + m, n, experts = a_np.shape[0], b_np.shape[1], len(offsets_np) + + kwargs = dict( + sfa_tensor=None, + sfb_tensor=None, + output_mode="discrete", + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + # ---- torch run: discrete mode with auto-generated pointers ---- + result_t = grouped_gemm_wgrad_wrapper_sm100( + a_tensor=to_torch_bf16(a_np), + b_tensor=to_torch_bf16(b_np), + offsets_tensor=torch.from_numpy(offsets_np).cuda(), + **kwargs, + ) + torch.cuda.synchronize() + wgrad_t = result_t["wgrad_tensor"].float().cpu().numpy() + + a_j, b_j, offsets_j = (jnp.asarray(x) for x in (a_np, b_np, offsets_np)) + jax.block_until_ready((a_j, b_j, offsets_j)) + + # ---- jax run 1: auto-allocated wgrad tensor and auto-generated pointers ---- + result_j = grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + offsets_tensor=offsets_j, + **kwargs, + ) + device_sync() + np.testing.assert_array_equal( + np.asarray(result_j["wgrad_tensor"]).astype(np.float32), + wgrad_t, + err_msg="wgrad discrete: JAX output differs from torch output on identical input bytes", + ) + + # ---- jax run 2: explicit packed-uint8 wgrad_ptrs into per-expert buffers ---- + expert_outputs = [jax.block_until_ready(jnp.zeros((m, n), dtype=ml_dtypes.bfloat16)) for _ in range(experts)] + ptr_values = np.array([buf.unsafe_buffer_pointer() for buf in expert_outputs], dtype=np.int64) + wgrad_ptrs_j = jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + offsets_tensor=offsets_j, + wgrad_ptrs=wgrad_ptrs_j, + **kwargs, + ) + device_sync() + for expert in range(experts): + np.testing.assert_array_equal( + np.asarray(expert_outputs[expert]).astype(np.float32), + wgrad_t[expert], + err_msg=f"wgrad discrete expert {expert}: explicit-ptrs JAX output differs from torch output", + ) + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_jax_jit_matches_eager(): + """XLA custom-call entry point (cudnn.jax.call): bit-identical to the eager JAX wrapper. + + The jit entry is discrete-mode only: the per-expert outputs are caller-owned + external buffers reached through the wgrad_ptrs input array (not XLA outputs), so + each run gets fresh zero-filled buffers and the comparison reads them after + blocking on the returned token plus a device sync. + """ + import cutlass.jax + + if not cutlass.jax.is_available(): + pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") + skip_unless_sm100() + from cudnn import grouped_gemm_wgrad_jax_sm100, grouped_gemm_wgrad_wrapper_sm100 + + a_np, b_np, offsets_np = make_problem() + m, n, experts = a_np.shape[0], b_np.shape[1], len(offsets_np) + + a_j, b_j, offsets_j = (jnp.asarray(x) for x in (a_np, b_np, offsets_np)) + jax.block_until_ready((a_j, b_j, offsets_j)) + + # Eager JAX wrapper baseline (dense output) on the same bytes and kernel config. + result_eager = grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + sfa_tensor=None, + sfb_tensor=None, + offsets_tensor=offsets_j, + output_mode="dense", + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + device_sync() # eager JAX path runs on the CUDA legacy default stream + expected = np.asarray(result_eager["wgrad_tensor"]).view(np.uint8) # (experts, m, 2n) bytes + + def run_and_check(fn, label): + # Fresh zero-filled external buffers per run so each check observes that + # run's writes (the kernel fully overwrites; buffers are immutable JAX + # arrays mutated behind XLA's back through their raw addresses). + expert_outputs = [jax.block_until_ready(jnp.zeros((m, n), dtype=ml_dtypes.bfloat16)) for _ in range(experts)] + ptr_values = np.array([buf.unsafe_buffer_pointer() for buf in expert_outputs], dtype=np.int64) + wgrad_ptrs_j = jax.block_until_ready(jnp.asarray(ptr_values.view(np.uint8))) + token = fn(wgrad_ptrs_j) + assert token.shape == (m, n) + jax.block_until_ready(token) + device_sync() # external-buffer writes are outside XLA's dataflow + for expert in range(experts): + np.testing.assert_array_equal( + np.asarray(expert_outputs[expert]).view(np.uint8), + expected[expert], + err_msg=f"wgrad jit expert {expert} ({label}): output differs from eager wrapper output on identical input bytes", + ) + + kwargs = dict(mma_tiler_mn=(128, 128), cluster_shape_mn=(1, 1)) + + # Eager custom call + run_and_check(lambda ptrs: grouped_gemm_wgrad_jax_sm100(a_j, b_j, offsets_j, ptrs, **kwargs), "eager custom call") + + # Under jax.jit, twice (compiled-kernel / registration cache). + jitted = jax.jit(lambda a, b, offsets, ptrs: grouped_gemm_wgrad_jax_sm100(a, b, offsets, ptrs, **kwargs)) + run_and_check(lambda ptrs: jitted(a_j, b_j, offsets_j, ptrs), "jit call 1") + run_and_check(lambda ptrs: jitted(a_j, b_j, offsets_j, ptrs), "jit call 2") + + +@pytest.mark.L0 +def test_grouped_gemm_wgrad_jax_block_scaled_rejected(): + skip_unless_sm100() + import cudnn + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + rng = np.random.default_rng(0) + m, n, tokens = 128, 128, 512 + a_j = jnp.asarray(rng.integers(0, 100, (m, tokens), dtype=np.uint8).view(ml_dtypes.float8_e4m3fn)) + b_j = jnp.asarray(rng.integers(0, 100, (tokens, n), dtype=np.uint8).view(ml_dtypes.float8_e4m3fn)) + sfa_j = jnp.asarray(np.full((128, 32), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + sfb_j = jnp.asarray(np.full((128, 32), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) + offsets_j = jnp.asarray(np.array([256, 512], dtype=np.int32)) + + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_j, + b_tensor=b_j, + sfa_tensor=sfa_j, + sfb_tensor=sfb_j, + offsets_tensor=offsets_j, + output_mode="dense", + sf_vec_size=32, + ) + + # Class API path rejects too. + op = cudnn.GroupedGemmWgradSm100( + sample_a=a_j, + sample_b=b_j, + sample_sfa=sfa_j, + sample_sfb=sfb_j, + sample_offsets=offsets_j, + sample_wgrad=jnp.zeros((2, m, n), dtype=ml_dtypes.bfloat16), + sf_vec_size=32, + ) + with pytest.raises(ValueError, match="not expressible as JAX arrays"): + op.check_support()