Add cudnn.jax.call on CuTeDSL's native JAX bridge; jit entry points across the GEMM CuTeDSL APIs (stacked on #534)

AGENTS.md

@@ -44,7 +44,9 @@ Python (editable; compiles the pybind11 extension via CMake):
 
 ```bash
 pip install -e .              # core graph API only
-pip install -e ".[cutedsl]"   # + OSS CuTeDSL kernels (torch, nvidia-cutlass-dsl, cuda-python)
+pip install -e ".[cutedsl]"   # + OSS CuTeDSL kernels (nvidia-cutlass-dsl, cuda-python, tvm-ffi; framework-neutral)
+pip install --group torch      # + torch for the CuTeDSL APIs (torch, torch-c-dlpack-ext)
+pip install --group jax        # + jax for the CuTeDSL APIs (jax >= 0.5; XLA entry points via cutlass.jax)
 ```
 
 `setup.py` honors env vars: `CUDNN_PATH`, `CUDA_PATH` / `CUDAToolkit_ROOT`, `DEBUG=1` (debug build), `CMAKE_BUILD_PARALLEL_LEVEL`, `CMAKE_GENERATOR`.
@@ -61,7 +63,7 @@ cd test/python
 pytest                        # default is -m L0 (smoke level) per pytest.ini
 pytest -m L1                  # deeper levels: L0..L4
 pytest test_conv_fprop.py     # one file (still filtered by -m L0 — pass -m "L0 or L1" to widen)
-pytest fe_api/                # OSS kernel tests; require ".[cutedsl]" install + SM90/SM100 GPU
+pytest fe_api/                # OSS kernel tests; require ".[cutedsl]" + `--group torch` (and `--group jax` for the *_jax tests) + SM90/SM100 GPU
 ```
 
 Read [test/AGENTS.md](test/AGENTS.md) before touching tests — `test/python/conftest.py` has import-order and env-var requirements that are easy to break.

docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md

@@ -2,6 +2,12 @@
 
 **This is an experimental API and subject to change.**
 
+## JAX support
+
+Supports **JAX arrays** in FP8 configurations: b_ptrs/sfb_ptrs as packed-uint8 (or x64 int64) pointer arrays, SFA in the physical C-contiguous atom shape, SFD outputs allocated the same way (the kernel rebuilds all SF layouts from the GEMM shapes and reads only base pointers). Packed-fp4 inputs are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes.
+
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `discrete_grouped_gemm_dswiglu_jax_sm100` (built on `cudnn.jax.call`; k-major weights only): all outputs (d_row/d_col, SFD tensors, amax, `dprob` as a bridge-managed zero-initialized accumulator, optional `dbias`) are XLA-managed donated buffers — no manual synchronization. Under tracing the offsets *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation.
+
 ## Overview
 
 **Discrete Grouped GEMM + dGLU backward fusion**: A block-scaled grouped GEMM fused with a dSwiGLU/dGeGLU backward epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE workloads where each expert weight/scale lives in a separate allocation.

docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md

@@ -2,6 +2,12 @@
 
 **This is an experimental API and subject to change.**
 
+## JAX support
+
+Supports **JAX arrays** in FP8 configurations: b_ptrs/sfb_ptrs as packed-uint8 (or x64 int64) pointer arrays, SFA in the physical C-contiguous atom shape `(1, MN', K', 32, 4, 4)`, SFD outputs allocated the same way (the kernel rebuilds all SF layouts from the GEMM shapes and reads only base pointers). Column-major bias and packed-fp4 inputs are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes.
+
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `discrete_grouped_gemm_swiglu_jax_sm100` (built on `cudnn.jax.call`; k-major weights only): all outputs (c/d/d_col, SFD tensors, amax) are XLA-managed donated buffers — no manual synchronization. Under tracing the offsets *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation.
+
 ## Overview
 
 **Discrete Grouped GEMM + SwiGLU fusion**: A block-scaled grouped GEMM fused with a SwiGLU/GeGLU epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE workloads where each expert weight lives in a separate allocation.

docs/fe-oss-apis/gemm_fusions/gemm_amax.md

@@ -109,7 +109,7 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None)
 
 Two integration levels are available:
 
-1. **`gemm_amax_jax_sm100`** (recommended for jitted programs) — an XLA custom call via [jax-tvm-ffi]. The kernel runs on XLA's compute stream (correctly ordered with surrounding ops), outputs are XLA-managed fresh arrays, and the call composes with `jax.jit`. No manual synchronization is needed. Requires the `jax` dependency group (`pip install --group jax`, which brings `jax-tvm-ffi`; Python >= 3.11).
+1. **`gemm_amax_jax_sm100`** (recommended for jitted programs) — an XLA custom call built on `cudnn.jax.call` (CuTeDSL's native `cutlass.jax` bridge). The kernel runs on XLA's compute stream (correctly ordered with surrounding ops), outputs are XLA-managed fresh arrays, and the call composes with `jax.jit`. No manual synchronization is needed. Requires the `jax` dependency group (`pip install --group jax`; jax >= 0.5).
 
 ```python
 from cudnn import gemm_amax_jax_sm100
@@ -120,7 +120,7 @@ def quantized_matmul(a, b, sfa, sfb):
     return c, amax
 ```
 
-Calling it eagerly works but re-traces the `ffi_call` on every invocation; call it from inside a jitted function in hot loops.
+Calling it eagerly works but re-traces the custom call on every invocation; call it from inside a jitted function in hot loops.
 
 2. **The eager entry points below** (`gemm_amax_wrapper_sm100`, `GemmAmaxSm100`) also accept JAX arrays via DLPack. In hot loops prefer the **class API with pre-allocated output buffers** (~15 µs CPU per launch) over the wrapper — per-call `jnp` output allocation in the wrapper costs hundreds of µs of XLA dispatch.
 
@@ -254,7 +254,7 @@ Tuple unpacking order is: `(c_tensor, amax_tensor)`.
 
 - `L == 1`; `A`/`B` k-major; `C` n-major only
 - `SFA`/`SFB` in the physical atom shape `(L, MN', K', 32, 4, 4)` (see "Using JAX arrays")
-- Eager use only (no `jax.jit` over these entry points); synchronize before reading outputs
+- The wrapper/class entry points are eager-only (use `gemm_amax_jax_sm100` under `jax.jit`); synchronize before reading outputs
 
 ---
 

docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md

@@ -73,7 +73,7 @@ A (MxKxL), SFA                   B (NxKxL), SFB
 
 ## API Usage
 
-The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). A `jax.jit`-compatible XLA custom-call entry point is not yet available for this kernel (its signature carries optional None-typed parameters that the jax-tvm-ffi bridge cannot supply).
+The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point (`gemm_srelu_jax_sm100` / `gemm_dsrelu_jax_sm100`, built on `cudnn.jax.call`); the kernels' optional parameters are compile-time constants inside its adapter.
 
 ### High-level wrapper
 

docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md

@@ -2,6 +2,14 @@
 
 **This is an experimental API and subject to change.**
 
+## JAX support
+
+Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` (the `[in, out]` weight layout reaches the kernel through a transposed strided view, which has no row-major JAX equivalent and raises a clear error). The E8M0 scale inputs stay `uint8` as with torch. Outputs are allocated as C-contiguous `jnp` arrays. The wrapper is eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs.
+
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"): same contract as the wrapper with `w_out_in=True`, dispatching on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs.
+
+The API is compiled with `--enable-tvm-ffi`: raw framework tensors go straight to the compiled kernel (no per-call `from_dlpack` conversion), cutting per-launch CPU overhead roughly in half for torch callers as well.
+
 ## Overview
 
 **Fused projection GEMM + per-head YARN RoPE + dual-direction MXFP8 quantize**: a persistent dense GEMM on NVIDIA Blackwell GPUs (SM100+) that projects activations, applies the Megatron MLA-YARN rotary embedding to each attention head's trailing rotary features, and MXFP8 (E4M3, block=32) quantizes the result in **both** the rowwise (D-direction) and columnwise (S-direction) layouts. Implemented with CUTLASS/CUTE.

docs/fe-oss-apis/gemm_fusions/gemm_srelu.md

@@ -66,7 +66,7 @@ A (MxKxL), SFA                   B (NxKxL), SFB
 
 ## API Usage
 
-The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). A `jax.jit`-compatible XLA custom-call entry point is not yet available for this kernel (its signature carries optional None-typed parameters that the jax-tvm-ffi bridge cannot supply).
+The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, scale-factor tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs). For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point (`gemm_srelu_jax_sm100` / `gemm_dsrelu_jax_sm100`, built on `cudnn.jax.call`); the kernels' optional parameters are compile-time constants inside its adapter.
 
 ### High-level wrapper
 

docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md

@@ -68,7 +68,7 @@ Notes:
 
 The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is only imported when torch tensors/dtypes are passed, jax only when JAX arrays are passed). Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types. The JAX contract matches gemm_amax (see `gemm_amax.md` "Using JAX arrays"): A/B k-major `(M, K, 1)`/`(N, K, 1)`, outputs n-major only, batch `L == 1`, SF tensors accepted in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)`; the eager entry points run on the CUDA legacy default stream (synchronize before reading outputs).
 
-For jitted JAX programs, use **`gemm_swiglu_jax_sm100`** — an XLA custom call (via jax-tvm-ffi, `jax` dependency group) that runs on XLA's compute stream, returns fresh `(ab12, c)` arrays, and composes with `jax.jit`. It currently supports the standard (non-quantized) kernel only; use the eager wrapper for blockscaled MXFP8 inputs from JAX. `alpha` is a static (trace-time) parameter.
+For jitted JAX programs, use **`gemm_swiglu_jax_sm100`** — an XLA custom call (built on `cudnn.jax.call` / CuTeDSL's native `cutlass.jax` bridge, `jax` dependency group) that runs on XLA's compute stream, returns fresh `(ab12, c)` arrays, and composes with `jax.jit`. Both the standard and the blockscaled MXFP8 quantized kernels are supported. `alpha` is a static (trace-time) parameter.
 
 ```python
 from cudnn import gemm_swiglu_jax_sm100

docs/fe-oss-apis/gemm_fusions/grouped_gemm.md

@@ -53,6 +53,18 @@ alignment contract. The API records the pointer-array tensor on the launch
 stream. The caller must keep every pointed-to expert allocation alive and must
 not modify or free it until that stream completes.
 
+## Using JAX arrays
+
+The tensor parameters are type-erased: torch tensors and JAX arrays are both accepted (torch is imported only when torch tensors are passed, jax only when JAX arrays are passed). Because JAX arrays are always row-major, the JAX contract is narrower than torch's:
+
+- **Discrete weight mode only** (`b_ptrs`): dense mode's `b_tensor` uses an expert-outermost strided layout with no row-major equivalent, and `bias_tensor`'s `(n, experts)` column-major layout is likewise not expressible — both raise clear errors for JAX inputs. Each per-expert weight is a plain k-major `(n, k)` C-contiguous JAX array.
+- **`b_ptrs` from JAX**: build the pointer array from `weight.unsafe_buffer_pointer()` per expert. JAX truncates int64 without x64 mode, so pass the pointers either as an int64 array (with `jax_enable_x64`) or as a **packed uint8 array** (8 little-endian bytes per pointer): `jnp.asarray(np.array(ptrs, dtype=np.int64).view(np.uint8))`. The weight arrays (and `b_ptrs`) must stay alive and un-donated until the kernel completes.
+- A/offsets/alpha/prob are plain C-contiguous JAX arrays of the documented shapes; outputs are allocated as n-major C-contiguous `jnp` arrays. Dtype parameters accept torch dtypes, numpy/ml_dtypes dtypes, dtype name strings, or `cutlass` types.
+- The eager path launches on the **CUDA legacy default stream** (XLA does not track it): `jax.block_until_ready(...)` your inputs before calling, and synchronize the device (or the stream you passed) before reading the outputs.
+- For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_jax_sm100(a_tensor, padded_offsets, alpha_tensor, b_ptrs, n, prob_tensor, ...)` (built on `cudnn.jax.call`; discrete mode, no bias): outputs are fresh XLA-managed arrays with rows at/past `padded_offsets[-1]` zero-filled, and no manual synchronization is needed. Under tracing the `padded_offsets` *values* cannot be host-validated (shapes/dtypes still are), and the per-expert weight buffers behind `b_ptrs` must stay alive and unmoved across every execution of the traced computation.
+
+Internal workspaces are allocated in the caller's framework allocator (torch caching allocator or XLA's pool) and written through raw pointers; they are never surfaced as arrays.
+
 ## Wrapper API
 
 Dense mode:

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md

@@ -2,6 +2,12 @@
 
 **This is an experimental API and subject to change.**
 
+## JAX support
+
+Supports **JAX arrays** on the BF16 backend in discrete weight mode (dswiglu and dgeglu), including `generate_dbias=True` and caller-provided zero-initialized `dprob`. Dense `b_tensor` and the block-scaled backend (MMA-interleaved scale-factor layouts) are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes.
+
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_dglu_jax_sm100` (built on `cudnn.jax.call`; discrete mode): `dprob` and (with `generate_dbias=True`) `dbias` come back as bridge-managed zero-initialized accumulator outputs — no caller-zeroed buffers, no manual synchronization. Under tracing the `padded_offsets` *values* cannot be host-validated, and the per-expert weight buffers behind `b_ptrs` must stay alive and unmoved across every execution of the traced computation.
+
 ## Overview
 
 **Unified Grouped GEMM + dGLU fusion**: one public class and wrapper select a

docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md

@@ -2,6 +2,12 @@
 
 **This is an experimental API and subject to change.**
 
+## JAX support
+
+Supports **JAX arrays** in the discrete (b_ptrs) FP8 configurations: pointer arrays as int64 (jax x64 mode) or packed uint8 (8 bytes per pointer), scale-factor tensors in the physical C-contiguous atom shape `(L, MN', K', 32, 4, 4)` (the kernel rebuilds SF layouts from the GEMM shapes and reads only the base pointer), outputs allocated as C-contiguous `jnp` arrays. Dense weight mode and packed-fp4 A/B are not expressible as JAX arrays and raise clear errors. The wrapper is eager, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs; keep weight arrays alive until the kernel completes.
+
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `grouped_gemm_dsrelu_jax_sm100` (built on `cudnn.jax.call`; discrete FP8 mode, `sf_vec_size=32`): all outputs (d/SFD tensors, `dprob`, and with `generate_dbias=True` `dbias`) are XLA-managed donated zero-initialized buffers — no manual synchronization. Under tracing the `padded_offsets` *values* cannot be host-validated, and the weight/scale buffers behind the pointer arrays must stay alive and unmoved across every execution of the traced computation.
+
 ## Overview
 
 **Grouped GEMM + dsReLU backward fusion**: A grouped block-scaled GEMM fused with a probability-gradient backward epilogue on NVIDIA Blackwell GPUs (SM100+), designed for MoE-style workloads. The API supports dense contiguous weights and discrete per-expert weight allocations. Groups are contiguous in the `M` dimension and described by `padded_offsets`.

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.