From 6dcef0d097b2f056ca09f3c0035e428c0520880e Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 9 Aug 2026 21:09:44 -0700 Subject: [PATCH 1/2] 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/2] 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]