From 6dcef0d097b2f056ca09f3c0035e428c0520880e Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 9 Aug 2026 21:09:44 -0700 Subject: [PATCH 1/3] Extend JAX support to the grouped/discrete-grouped GEMM APIs Applies the dense-fusion type-erasure + JAX pattern to the grouped family, with real JAX eager support wherever the kernel's tensor layouts are expressible as row-major arrays and clear rejections where they are not. Per-API JAX support (all eager; discrete/pointer-array weight modes): - grouped_gemm (unfused): BF16 discrete mode - grouped_gemm_glu / dglu: BF16 backend, discrete mode, swiglu+geglu / dswiglu+dgeglu incl. generate_dbias and caller-provided dprob - grouped_gemm_dsrelu: discrete FP8 (scale factors in the physical C-contiguous atom shape -- the backward kernels provably rebuild SF layouts from the GEMM shapes and read only base pointers) - grouped_gemm_wgrad: BF16 backend, dense (experts, m, n) or discrete pointer outputs - discrete_grouped_gemm_swiglu / dswiglu: FP8 (SF physical atom shape for SFA and the SFD outputs) Rejected with clear "not expressible as JAX arrays" errors: - grouped swiglu/srelu/quant: their SFA scale factors are MMA-permuted strided cute tensor arguments in every mode (unlike amax/dsrelu, the kernel consumes the full layout, which has no row-major equivalent) - grouped dswiglu (dense-weight-mode only) and glu_hadamard (block-scaled only); the block-scaled glu/dglu/wgrad backends - dense-mode b_tensor (expert-outermost strides), column-major bias, and packed-fp4 inputs everywhere Mechanics shared across the family (unfused is the template): - b_ptrs/sfb_ptrs/wgrad_ptrs pointer arrays from JAX: int64 (jax x64 mode) or packed little-endian uint8 (8 bytes per pointer), since JAX truncates int64 without x64; framework-neutral validation + host decoding in unfused._bf16_api (_validate_pointer_tensor/_pointer_values); pointers come from jax.Array.unsafe_buffer_pointer() and the arrays must stay alive until kernel completion (record_stream is torch-only; the JAX path keeps live references instead) - internal workspaces via tensor_adapter.allocate_byte_workspace: allocated in the caller's framework allocator (torch.empty / jnp.zeros + block_until_ready), written through raw pointers, never surfaced as arrays; compile-time Int64 pointer placeholders are real bytes retyped via the from_dlpack element_type override (fake tensors have dummy iterators) - new tensor_adapter helpers: get_data_ptr (torch data_ptr / jax unsafe_buffer_pointer), get_version (0 for immutable arrays), to_host_list, allocate_byte_workspace - canonical (cutlass) dtype vocabulary and canonical TensorDescs throughout, incl. live-tensor validation; expected-stride literals with extent-1 dims wrapped in canonicalize_unit_dim_strides; select_grouped_gemm_backend accepts torch/jax/numpy/str dtypes - execute stream defaulting per framework; wrapper output allocation branches (torch empty_strided byte-identical; jnp.empty n-major C-contiguous + block_until_ready) Also: - discrete_grouped swiglu/dswiglu now set _interpret_uint8_as_fp4x2 before descriptor creation (the torch uint8-container path previously built descs with the flag unset and was silently broken) - test conftest sets XLA_PYTHON_CLIENT_PREALLOCATE=false: the JAX interop tests share the pytest process with the torch suites, and XLA's default 75%-of-GPU preallocation starved later torch kernel compiles (12 CUDA_ERROR_OUT_OF_MEMORY failures in full-suite runs) - per-API "JAX support" docs sections + overview matrix; the blanket torch-only guard test narrows to proj_rope (each grouped family now has its own JAX test file) proj_rope_mxfp8 (added after review): migrated both classes to the TVM-FFI compile path (--enable-tvm-ffi + fake stream) so raw DLPack tensors go straight to the compiled kernel -- the per-call from_dlpack(x.detach()) conversion loop is gone from the hot path (~10.8 us/launch CPU after, vs a per-call conversion protocol that cost 2-3 us per tensor across 8-10 tensors before). torch inputs keep cheap detach views for autograd safety; the uint8 E8M0 scale inputs keep a per-call element-type reinterpret (now tvm-ffi-enabled). JAX supported on both input paths with w_out_in=True (the [in, out] weight reaches the kernel through a transposed strided view -- torch-only, clear error); bit-identical torch-vs-JAX tests for the bf16 and mxfp8 paths. With proj_rope no longer torch-only, the blanket guard test file is removed (every API now has its own JAX test file). Tests: per-family JAX tests assert bit-identical outputs between torch and JAX wrapper runs on identical input bytes (both paths share one compiled kernel) for every supported config -- unfused, glu (swiglu+geglu), dglu (d_row/dprob/dbias), dsrelu (d_row/d_col/d_srelu + all three SFD outputs), wgrad (dense+discrete), discrete swiglu/dswiglu (fp8, byte-exact) -- with dprob-style atomic accumulators compared at tight tolerance; rejected configs assert their clear errors. Verified on SM100: full fe_api/gemm + fe_api/grouped_gemm + unfused suite run yields 64 failed / 1437 passed / 1027 skipped / 2 xfailed / 2 errors -- the failure list is byte-identical to the known pre-existing test_gemm_swiglu env-numerics failures, and the 2 collection errors are the pre-existing upstream test_grouped_gemm_{glu,dglu}.py missing-module imports. All grouped/discrete modules import with torch absent. Co-Authored-By: Claude Fable 5 --- .../discrete_grouped_gemm_dswiglu.md | 4 + .../discrete_grouped_gemm_swiglu.md | 4 + .../gemm_fusions/gemm_proj_rope_mxfp8.md | 6 + docs/fe-oss-apis/gemm_fusions/grouped_gemm.md | 11 + .../gemm_fusions/grouped_gemm_dglu.md | 4 + .../gemm_fusions/grouped_gemm_dsrelu.md | 4 + .../gemm_fusions/grouped_gemm_dswiglu.md | 4 + .../gemm_fusions/grouped_gemm_glu.md | 4 + .../gemm_fusions/grouped_gemm_glu_hadamard.md | 4 + .../gemm_fusions/grouped_gemm_quant.md | 4 + .../gemm_fusions/grouped_gemm_srelu.md | 4 + .../gemm_fusions/grouped_gemm_swiglu.md | 4 + .../gemm_fusions/grouped_gemm_wgrad.md | 4 + docs/fe-oss-apis/overview.md | 6 +- .../gemm/cutedsl/dense/proj_rope_mxfp8/api.py | 259 ++++++---- .../gemm_proj_rope_mxfp8_mxfp8in.py | 4 +- .../discrete_grouped/discrete_kernel_utils.py | 6 + .../cutedsl/discrete_grouped/dswiglu/api.py | 424 ++++++++++------ .../cutedsl/discrete_grouped/swiglu/api.py | 428 ++++++++++------ .../gemm/cutedsl/grouped/backend_utils.py | 21 +- .../gemm/cutedsl/grouped/dglu/_bf16_api.py | 156 +++--- .../cutedsl/grouped/dglu/_blockscaled_api.py | 7 + python/cudnn/gemm/cutedsl/grouped/dglu/api.py | 243 +++++---- .../cudnn/gemm/cutedsl/grouped/dsrelu/api.py | 480 ++++++++++++------ .../cudnn/gemm/cutedsl/grouped/dswiglu/api.py | 132 ++--- .../gemm/cutedsl/grouped/glu/_bf16_api.py | 140 ++--- .../cutedsl/grouped/glu/_blockscaled_api.py | 9 +- python/cudnn/gemm/cutedsl/grouped/glu/api.py | 232 +++++---- .../gemm/cutedsl/grouped/glu_hadamard/api.py | 32 +- .../cudnn/gemm/cutedsl/grouped/quant/api.py | 197 +++---- .../cudnn/gemm/cutedsl/grouped/srelu/api.py | 203 ++++---- .../cudnn/gemm/cutedsl/grouped/swiglu/api.py | 153 +++--- .../gemm/cutedsl/grouped/unfused/_bf16_api.py | 179 ++++--- .../cudnn/gemm/cutedsl/grouped/unfused/api.py | 169 +++--- .../gemm/cutedsl/grouped/wgrad/_bf16_api.py | 213 +++++--- .../cutedsl/grouped/wgrad/_blockscaled_api.py | 19 +- .../cudnn/gemm/cutedsl/grouped/wgrad/api.py | 129 +++-- python/cudnn/tensor_adapter.py | 55 ++ test/python/conftest.py | 6 + .../fe_api/gemm/test_cutedsl_jax_guards.py | 58 --- .../gemm/test_gemm_proj_rope_mxfp8_jax.py | 115 +++++ .../test_discrete_grouped_gemm_dswiglu_jax.py | 186 +++++++ .../test_discrete_grouped_gemm_swiglu_jax.py | 186 +++++++ .../test_grouped_gemm_dglu_jax.py | 160 ++++++ .../test_grouped_gemm_dsrelu_jax.py | 230 +++++++++ .../test_grouped_gemm_dswiglu_jax.py | 104 ++++ .../test_grouped_gemm_glu_hadamard_jax.py | 67 +++ .../grouped_gemm/test_grouped_gemm_glu_jax.py | 159 ++++++ .../grouped_gemm/test_grouped_gemm_jax.py | 130 +++++ .../test_grouped_gemm_quant_jax.py | 87 ++++ .../test_grouped_gemm_srelu_jax.py | 89 ++++ .../test_grouped_gemm_swiglu_jax.py | 84 +++ .../test_grouped_gemm_wgrad_jax.py | 190 +++++++ 53 files changed, 4252 insertions(+), 1556 deletions(-) delete mode 100644 test/python/fe_api/gemm/test_cutedsl_jax_guards.py create mode 100644 test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py create mode 100644 test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py 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..d5846bc24 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,10 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + ## 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..330ca661f 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,10 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + ## 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_proj_rope_mxfp8.md b/docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md index eed9aa9bd..54bccab92 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,12 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. + +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/grouped_gemm.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md index ec932bae3..0bd0fc890 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md @@ -53,6 +53,17 @@ 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. Eager use only — no `jax.jit`. + +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..4c8e3c722 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,10 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + ## 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..18dad9faf 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,10 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + ## 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..e3e636257 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,10 @@ **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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. + ## 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..10d19fd24 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,10 @@ 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. + ## 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..f83844504 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 amax and swiglu (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). 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). 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) 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/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/api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py index c9b7433f6..229a96b91 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,31 @@ 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, + ) + 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 +244,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 +290,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 +305,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 +326,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 +334,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 +380,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 +472,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 +484,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 +526,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 +557,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 +642,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 +732,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 +819,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 +944,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/swiglu/api.py b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py index 5b85f781c..4d5a78f6b 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,34 @@ 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, + ) + 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 +252,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 +299,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 +314,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 +331,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 +341,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 +351,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 +371,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 +379,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 +395,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 +484,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 +492,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 +505,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 +552,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 +583,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 +663,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 +751,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 +831,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 +975,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/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/_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/dsrelu/api.py b/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py index 0ed234589..7db24df94 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py @@ -34,14 +34,25 @@ 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, + 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 +169,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 +210,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 +249,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 +275,45 @@ 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, accepting the permuted atom view or (in discrete + weight mode) the physical C-contiguous form -- see ``_sf_desc_is_physical``.""" + 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) + 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) + 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 +351,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 +409,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 +431,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 +459,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 +469,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 +502,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 +510,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 +532,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 +646,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 +664,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 +703,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 +716,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 +926,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 +1008,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 +1054,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 +1148,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 +1224,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 +1308,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 +1351,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 +1462,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 +1478,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 +1513,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 +1545,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 +1604,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 +1699,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 +1711,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/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/_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_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/_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/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/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_proj_rope_mxfp8_jax.py b/test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py new file mode 100644 index 000000000..e4b5a02f4 --- /dev/null +++ b/test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py @@ -0,0 +1,115 @@ +# 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 +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/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..715d81f4e --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py @@ -0,0 +1,186 @@ +# 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_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..bb9dc5ed2 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py @@ -0,0 +1,186 @@ +# 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_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..763c43f8d --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py @@ -0,0 +1,160 @@ +# 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_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..2211cfcd2 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py @@ -0,0 +1,230 @@ +# 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_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..ddeb95532 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py @@ -0,0 +1,159 @@ +# 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_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..2d7538782 --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py @@ -0,0 +1,130 @@ +# 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_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..d9f0f207e --- /dev/null +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py @@ -0,0 +1,190 @@ +# 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_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() From d31d137ab7b36298eb658aa87ad25b461726983c Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 9 Aug 2026 22:33:30 -0700 Subject: [PATCH 2/3] Make the cutedsl extra framework-neutral: torch moves to its dependency group Remove torch and torch-c-dlpack-ext from the [cutedsl] optional extra. The CuTeDSL APIs are type-erased and torch-lazy, so torch is now opt-in exactly like jax, via the PEP 735 dependency groups introduced earlier: pip install -e ".[cutedsl]" # framework-neutral core pip install --group torch # torch + torch-c-dlpack-ext pip install --group jax # jax + jax-tvm-ffi (py3.11+) The cutedsl extra keeps nvidia-cutlass-dsl, cuda-python, and apache-tvm-ffi. Compatibility note: `pip install nvidia-cudnn-frontend[cutedsl]` no longer pulls torch. Users of the torch-only OSS APIs behind this extra (SDPA, BSA/DSA/NSA, and the torch-only grouped configurations) must install torch via the group (from a checkout) or directly (from the published wheel, since PEP 735 groups are not part of wheel metadata). AGENTS.md and the FE-OSS overview installation docs updated accordingly. Co-Authored-By: Claude Fable 5 --- AGENTS.md | 6 ++++-- docs/fe-oss-apis/overview.md | 7 +++++++ pyproject.toml | 4 ++-- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6b37a67d5..5d617348d 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, jax-tvm-ffi; py3.11+) ``` `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/overview.md b/docs/fe-oss-apis/overview.md index f83844504..7634a6375 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -43,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 + jax-tvm-ffi (Python >= 3.11) +``` +(For the published wheel, `pip install torch torch-c-dlpack-ext` or `pip install jax jax-tvm-ffi` 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..016023bf7 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] From 2ff2239a6a51fad21349276aaf5aaa95c9311f35 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 10 Aug 2026 23:22:42 -0700 Subject: [PATCH 3/3] Add cudnn.jax.call on CuTeDSL's native JAX bridge; jit entry points across the GEMM CuTeDSL APIs Replace the jax-tvm-ffi backend with cutlass.jax.cutlass_call wrapped as cudnn.jax.call, and add jax.jit-compatible XLA custom-call entry points for every JAX-reachable GEMM API: the four dense fusions (amax, swiglu incl. quantized, srelu, dsrelu), proj_rope_mxfp8 (both input paths), and the discrete-mode grouped family (unfused, glu, dglu, dsrelu, wgrad, discrete-grouped swiglu/dswiglu). Co-Authored-By: Claude Fable 5 --- AGENTS.md | 2 +- .../discrete_grouped_gemm_dswiglu.md | 4 +- .../discrete_grouped_gemm_swiglu.md | 4 +- docs/fe-oss-apis/gemm_fusions/gemm_amax.md | 6 +- docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md | 2 +- .../gemm_fusions/gemm_proj_rope_mxfp8.md | 4 +- docs/fe-oss-apis/gemm_fusions/gemm_srelu.md | 2 +- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 2 +- docs/fe-oss-apis/gemm_fusions/grouped_gemm.md | 3 +- .../gemm_fusions/grouped_gemm_dglu.md | 4 +- .../gemm_fusions/grouped_gemm_dsrelu.md | 4 +- .../gemm_fusions/grouped_gemm_glu.md | 4 +- .../gemm_fusions/grouped_gemm_wgrad.md | 4 +- docs/fe-oss-apis/overview.md | 10 +- pyproject.toml | 8 +- python/cudnn/__init__.py | 20 +- python/cudnn/gemm/cutedsl/_jax_ffi.py | 72 --- .../cudnn/gemm/cutedsl/dense/amax/__init__.py | 2 +- python/cudnn/gemm/cutedsl/dense/amax/api.py | 49 ++- .../cudnn/gemm/cutedsl/dense/amax/jax_api.py | 103 +++-- .../gemm/cutedsl/dense/dsrelu/__init__.py | 13 + .../cutedsl/dense/proj_rope_mxfp8/__init__.py | 11 + .../cutedsl/dense/proj_rope_mxfp8/jax_api.py | 146 +++++++ .../gemm/cutedsl/dense/srelu/__init__.py | 13 + .../cudnn/gemm/cutedsl/dense/srelu/jax_api.py | 291 ++++++++++++ .../gemm/cutedsl/dense/swiglu/__init__.py | 2 +- python/cudnn/gemm/cutedsl/dense/swiglu/api.py | 13 +- .../gemm/cutedsl/dense/swiglu/jax_api.py | 157 ++++--- .../gemm/cutedsl/discrete_grouped/__init__.py | 18 + .../discrete_grouped/dswiglu/__init__.py | 11 + .../cutedsl/discrete_grouped/dswiglu/api.py | 11 + .../discrete_grouped/dswiglu/jax_api.py | 399 +++++++++++++++++ .../discrete_grouped/swiglu/__init__.py | 11 + .../cutedsl/discrete_grouped/swiglu/api.py | 11 + .../discrete_grouped/swiglu/jax_api.py | 369 ++++++++++++++++ python/cudnn/gemm/cutedsl/grouped/__init__.py | 24 + .../gemm/cutedsl/grouped/dglu/__init__.py | 11 + .../gemm/cutedsl/grouped/dglu/jax_api.py | 268 ++++++++++++ .../gemm/cutedsl/grouped/dsrelu/__init__.py | 11 + .../cudnn/gemm/cutedsl/grouped/dsrelu/api.py | 23 +- .../gemm/cutedsl/grouped/dsrelu/jax_api.py | 413 ++++++++++++++++++ .../gemm/cutedsl/grouped/glu/__init__.py | 11 + .../cudnn/gemm/cutedsl/grouped/glu/jax_api.py | 221 ++++++++++ .../gemm/cutedsl/grouped/unfused/__init__.py | 12 +- .../gemm/cutedsl/grouped/unfused/jax_api.py | 220 ++++++++++ .../gemm/cutedsl/grouped/wgrad/__init__.py | 11 + .../gemm/cutedsl/grouped/wgrad/jax_api.py | 206 +++++++++ python/cudnn/jax/__init__.py | 33 ++ python/cudnn/jax/call.py | 133 ++++++ test/python/fe_api/gemm/test_gemm_amax.py | 35 ++ test/python/fe_api/gemm/test_gemm_amax_jax.py | 14 +- .../gemm/test_gemm_proj_rope_mxfp8_jax.py | 51 +++ .../fe_api/gemm/test_gemm_srelu_dsrelu_jax.py | 44 ++ .../fe_api/gemm/test_gemm_swiglu_jax.py | 30 +- .../test_discrete_grouped_gemm_dswiglu_jax.py | 101 +++++ .../test_discrete_grouped_gemm_swiglu_jax.py | 77 ++++ .../test_grouped_gemm_dglu_jax.py | 69 +++ .../test_grouped_gemm_dsrelu_jax.py | 107 +++++ .../grouped_gemm/test_grouped_gemm_glu_jax.py | 61 +++ .../grouped_gemm/test_grouped_gemm_jax.py | 61 +++ .../test_grouped_gemm_wgrad_jax.py | 65 +++ 61 files changed, 3863 insertions(+), 234 deletions(-) delete mode 100644 python/cudnn/gemm/cutedsl/_jax_ffi.py create mode 100644 python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py create mode 100644 python/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.py create mode 100644 python/cudnn/jax/__init__.py create mode 100644 python/cudnn/jax/call.py diff --git a/AGENTS.md b/AGENTS.md index 5d617348d..75b58e12f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,7 +46,7 @@ Python (editable; compiles the pybind11 extension via CMake): pip install -e . # core graph API only 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, jax-tvm-ffi; py3.11+) +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`. 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 d5846bc24..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 @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. +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 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 330ca661f..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 @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. +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 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 54bccab92..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 @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. +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. 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 0bd0fc890..9a303ec9a 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md @@ -60,7 +60,8 @@ The tensor parameters are type-erased: torch tensors and JAX arrays are both acc - **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. Eager use only — no `jax.jit`. +- 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. 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 4c8e3c722..b8a85fd2a 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. +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 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 18dad9faf..ac8b0e625 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. +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 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 e3e636257..eacd4a3e5 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md @@ -4,7 +4,9 @@ ## 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes. +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 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 10d19fd24..88514b2ba 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -13,7 +13,9 @@ 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. Eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. +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 diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 7634a6375..da2101001 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -4,9 +4,9 @@ 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 amax and swiglu (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). 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). +- **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) @@ -46,9 +46,9 @@ 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 + jax-tvm-ffi (Python >= 3.11) +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 jax-tvm-ffi` directly.) +(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}` diff --git a/pyproject.toml b/pyproject.toml index 016023bf7..75a9751ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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/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/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 229a96b91..bcd8e439d 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py @@ -215,6 +215,17 @@ def _check_sf_shape(self, desc, mn_div_128: int, rest: int, name: str) -> bool: (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: 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 4d5a78f6b..779e947d8 100644 --- a/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py +++ b/python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py @@ -214,6 +214,17 @@ def _check_sf_shape(self, desc, mn_div_128: int, rest: int, name: str) -> bool: (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: 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/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/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 7db24df94..1ab5d423a 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py @@ -39,6 +39,7 @@ 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, @@ -295,19 +296,37 @@ def _sf_desc_is_physical(sf_desc) -> bool: 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, accepting the permuted atom view or (in discrete - weight mode) the physical C-contiguous form -- see ``_sf_desc_is_physical``.""" + """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. 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/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/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/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/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/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/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 index e4b5a02f4..027434e33 100644 --- 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 @@ -96,6 +96,57 @@ def test_gemm_proj_rope_mxfp8_mxfp8in_jax_matches_torch(): 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() 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 index 715d81f4e..857d55c19 100644 --- 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 @@ -155,6 +155,107 @@ def test_discrete_grouped_gemm_dswiglu_jax_fp8_matches_torch(): ) +@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() 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 index bb9dc5ed2..7890df2f8 100644 --- 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 @@ -139,6 +139,83 @@ def test_discrete_grouped_gemm_swiglu_jax_fp8_matches_torch(): 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() 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 index 763c43f8d..8bf3be24b 100644 --- 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 @@ -112,6 +112,75 @@ def test_grouped_gemm_dglu_jax_discrete_matches_torch(): ) +@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() 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 index 2211cfcd2..3dc704094 100644 --- 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 @@ -168,6 +168,113 @@ def test_grouped_gemm_dsrelu_jax_discrete_fp8_matches_torch(): 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() 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 index ddeb95532..b87a9471a 100644 --- 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 @@ -103,6 +103,67 @@ def test_grouped_gemm_glu_jax_discrete_matches_torch(act_func): ) +@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() 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 index 2d7538782..dd23f8ad1 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py @@ -92,6 +92,67 @@ def test_grouped_gemm_jax_discrete_matches_torch(): ) +@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() 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 index d9f0f207e..e0ce85d0f 100644 --- 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 @@ -151,6 +151,71 @@ def test_grouped_gemm_wgrad_jax_discrete_matches_torch(): ) +@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()