Make the cutedsl extra framework-neutral: torch moves to its dependency group (stacked on #530)

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.

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.

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.

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.

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:

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

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`.

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.

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

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`.

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.

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`.

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.

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

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)
@@ -39,6 +43,13 @@ All Frontend OSS APIs come installed with the `nvidia-cudnn-frontend` package. H
 pip install nvidia-cudnn-frontend[cutedsl]
 ```
 
+The `cutedsl` extra is framework-neutral (nvidia-cutlass-dsl, cuda-python, apache-tvm-ffi). Install your tensor framework separately — from a checkout, the PEP 735 dependency groups pin the right companion packages:
+```bash
+pip install --group torch   # torch + torch-c-dlpack-ext
+pip install --group jax     # jax + 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

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]