Make the cutedsl extra framework-neutral: torch moves to its dependency group (stacked on #530) - #534
Closed
Anerudhan wants to merge 2 commits into
Closed
Make the cutedsl extra framework-neutral: torch moves to its dependency group (stacked on #530)#534Anerudhan wants to merge 2 commits into
Anerudhan wants to merge 2 commits into
Conversation
Anerudhan
added
mod-infra
Contributor
📝 WalkthroughWalkthroughCuTeDSL now separates core, Torch, and JAX dependencies. GEMM APIs use framework-neutral tensor adapters and CUTLASS dtypes. Selected APIs support JAX execution, while unsupported layouts raise explicit errors. Documentation and Torch-versus-JAX tests cover these contracts. ChangesCuTeDSL framework and JAX support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant CuTeDSL_API
participant tensor_adapter
participant TVM_FFI
participant CUDA
Caller->>CuTeDSL_API: pass Torch or JAX tensors
CuTeDSL_API->>tensor_adapter: detect framework and normalize metadata
tensor_adapter-->>CuTeDSL_API: dtype, shape, pointer, device, and stream
CuTeDSL_API->>TVM_FFI: submit framework-compatible kernel arguments
TVM_FFI->>CUDA: launch compiled kernel
CUDA-->>Caller: write framework-native outputs
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 17
🧹 Nitpick comments (8)
python/cudnn/gemm/cutedsl/dense/amax/jax_api.py (1)
28-30: 🚀 Performance & Scalability | 🔵 TrivialThe target registry grows without bound.
Each new
cache_keyregisters a new global XLA FFI target throughget_or_register_env_stream_targetand retains theGemmAmaxSm100object plus the compiled kernel in_registered_targets. XLA FFI target registration is process-global and cannot be undone. A long-running service that sweeps many shapes accumulates targets and compiled kernels for the process lifetime.Consider documenting the expected number of distinct configurations, or adding a registry size limit that raises a clear error instead of growing silently.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py` around lines 28 - 30, Add a clear registry-size policy for _registered_targets, either documenting the expected bounded number of distinct configurations or enforcing a maximum that raises an explicit error before registering another target. Apply the policy in the cache-key registration path using get_or_register_env_stream_target, while preserving reuse of existing entries.python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py (2)
578-585: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRaise instead of asserting on the dtype mismatch.
Line 583 validates a user-supplied argument with
assert. Python removes assert statements under-O. If a user runs with optimizations enabled and passes mismatchedxandwdtypes, the check disappears and the mismatch reaches kernel dispatch. Line 580 already uses an explicitraisefor the framework check; use the same form here.Lines 620 and 641 use the same pattern for
x_scale/w_scalevalidation.🛡️ Proposed fix
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}" + if x_cutlass_dtype != _convert_to_cutlass_data_type(w.dtype): + raise ValueError(f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py` around lines 578 - 585, Replace the assert in gemm_proj_rope_mxfp8_wrapper_sm100 that validates matching x and w dtypes with an explicit ValueError, preserving the existing message. Apply the same assert-to-explicit-raise change to the x_scale and w_scale validation checks around the corresponding validation blocks.
64-73: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the
torch.Tensorannotations withAny.The constructor now accepts JAX arrays, and
from __future__ import annotationskeeps these strings from being evaluated at import time. The annotations still statetorch.Tensor, which contradicts the runtime contract and differs from the sibling APIs in this layer (amax/api.py,srelu/api.py,dsrelu/api.py,swiglu/api.pyall useAny). Any runtime introspection throughtyping.get_type_hints()also raisesNameErrorbecausetorchis not imported in this module.The same applies to
GemmProjRopeMxfp8Mxfp8InSm100.__init__at lines 282-291 andgemm_proj_rope_mxfp8_wrapper_sm100at lines 554-559.♻️ Proposed annotation change
def __init__( self, - sample_x: torch.Tensor, - sample_w: torch.Tensor, - sample_cos: torch.Tensor, - sample_sin: torch.Tensor, - sample_out_fp8_row: torch.Tensor, - sample_out_scales_row: torch.Tensor, - sample_out_fp8_col: torch.Tensor, - sample_out_scales_col: torch.Tensor, + sample_x: Any, + sample_w: Any, + sample_cos: Any, + sample_sin: Any, + sample_out_fp8_row: Any, + sample_out_scales_row: Any, + sample_out_fp8_col: Any, + sample_out_scales_col: Any, w_out_in: bool = False, ):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py` around lines 64 - 73, Replace all torch.Tensor annotations with Any in the first constructor, GemmProjRopeMxfp8Mxfp8InSm100.__init__, and gemm_proj_rope_mxfp8_wrapper_sm100, including every affected parameter and return annotation. Preserve the existing signatures and behavior while matching the sibling APIs and allowing JAX arrays and runtime type-hint resolution without requiring torch.python/cudnn/gemm/cutedsl/dense/amax/api.py (1)
469-484: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winDocument the eager JAX ownership contract for every wrapper. These branches pass JAX arrays to kernels that write through DLPack.
jax.block_until_readyorders materialization but does not declare the writes to XLA. Keep these wrappers eager-only and explicitly forbid use underjax.jitor with donated buffers. Apply the explicit contract forgemm_amax_wrapper_sm100to the other four wrappers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/dense/amax/api.py` around lines 469 - 484, Document and enforce the eager JAX ownership contract in gemm_amax_wrapper_sm100 and the corresponding JAX branches of python/cudnn/gemm/cutedsl/dense/dsrelu/api.py:544-566, python/cudnn/gemm/cutedsl/dense/srelu/api.py:527-549, python/cudnn/gemm/cutedsl/dense/swiglu/api.py:658-677, and python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py:611-617. Keep these wrappers eager-only and explicitly reject invocation under jax.jit or with donated buffers; retain block_until_ready only for materialization ordering and apply the same explicit contract consistently across all five wrappers.python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py (1)
202-218: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftScale-factor layout classification is duplicated and divergent across three APIs. The permuted atom view and the physical C-contiguous allocation are one wire contract, but three copies now decide which form a descriptor uses, and two of them use different predicates. A change to the accepted layouts must be applied in every copy. Extract one helper, for example in
discrete_grouped/discrete_kernel_utils.py, that classifies the form and validates the expected shape, then call it from all three APIs.
python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py#L202-L218: replace_check_sf_shapewith a call to the shared helper; the body is identical to the copy inpython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py#L203-L218.python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py#L284-L311: replace_sf_desc_is_physicaland_check_sf_shapewith the shared helper, and keep the discrete-mode restriction as a caller-supplied option.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py` around lines 202 - 218, Extract the shared scale-factor layout classifier and shape validator into discrete_kernel_utils.py, covering both the permuted atom view and physical C-contiguous form with one consistent predicate. Replace _check_sf_shape in swiglu/api.py and the _sf_desc_is_physical/_check_sf_shape pair in dsrelu/api.py with calls to this helper, while passing the discrete-mode restriction from the dsrelu caller. Apply the same shared-helper change to dswiglu/api.py, preserving each API’s existing return and validation behavior.python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py (1)
585-589: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReuse the existing alignment helper.
Lines 588-589 repeat the body of
_validate_pointer_array_alignment(lines 201-204). Call the helper so the rule stays in one place.♻️ Proposed refactor
- if get_data_ptr(b_ptrs) % 8 != 0: - raise ValueError("b_ptrs data pointer must be 8-byte aligned") + self._validate_pointer_array_alignment(b_ptrs)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py` around lines 585 - 589, In the validation flow after _validate_pointer_tensor for b_ptrs, replace the duplicated get_data_ptr alignment check and error with a call to the existing _validate_pointer_array_alignment helper, preserving the 8-byte alignment validation and centralized error behavior.python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py (1)
756-764: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the framework rejection message.
Line 758 tells the caller that JAX arrays are accepted. Lines 759-764 then reject JAX unconditionally. Report a single, accurate constraint for this API.
♻️ Proposed simplification
- 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": + framework = detect_framework(a_tensor) + 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" ) + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for grouped_gemm_dswiglu_wrapper_sm100; pass torch tensors")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py` around lines 756 - 764, Update the framework validation in the grouped_gemm_dswiglu_wrapper_sm100 path so JAX is rejected by one accurate message rather than first being listed as supported and then rejected unconditionally. Keep torch as the accepted framework and state the dense-weight/layout constraint consistently in the resulting error.python/cudnn/gemm/cutedsl/grouped/glu/api.py (1)
1016-1034: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe
strideargument is unused on the JAX branch.
_allocate_outputacceptsstridebut only the torch branch uses it. The current shapes make the two layouts equivalent, because only the extent-1 batch dimension differs. A future non-unit trailing dimension would silently diverge. Add an assertion thatstridematches the C-contiguous layout, or drop the parameter and derive it in the torch branch.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py` around lines 1016 - 1034, Update _allocate_output so its JAX and torch allocation paths cannot silently diverge on layout: either validate that stride matches the expected C-contiguous layout before allocation, or remove the stride parameter and derive that layout directly in the torch branch. Preserve the existing output allocation behavior for c_tensor and d_tensor.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/fe-oss-apis/gemm_fusions/gemm_amax.md`:
- Around line 114-120: Add imports for the JAX symbols used in the
`quantized_matmul` example: import `jax` for `jax.jit` and `jax.numpy` as `jnp`
for `jnp.float32`, while preserving the existing `gemm_amax_jax_sm100` import.
In `@docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md`:
- Around line 73-80: Add the required JAX imports at the start of the documented
example so both jax.jit and jnp.float32/jnp.bfloat16 resolve before swiglu_mlp
calls gemm_swiglu_jax_sm100.
In `@python/cudnn/api_base.py`:
- Line 690: Resolve the undefined torch annotations in python/cudnn/api_base.py,
including _get_innermost_stride_dim, by adding a type-only torch import guarded
by TYPE_CHECKING while preserving lazy runtime imports, or replacing direct
torch annotations with framework-neutral types where appropriate.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/api.py`:
- Around line 167-183: Update the validation messages in the relevant dtype
checks to name CUTLASS dtype values consistently, including the unsupported
acc_dtype error and the message around line 250; ensure the expected values
match the actual CUTLASS types interpolated or compared. Remove the unnecessary
f-string prefix from the static ab12_dtype message to resolve Ruff F541.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py`:
- Around line 35-40: Reject non-default values for sf_vec_size, vector_f32, and
ab12_stages at the entry point before make_gemm, since this path ignores them
and rejects blockscaled inputs. Preserve the existing defaults; alternatively,
if these options must be supported, forward them through make_gemm and include
them in cache_key.
- Around line 56-59: Validate n in the input-shape guard near the existing L
check, rejecting odd values before invoking the non-quantized SwiGLU kernel.
Preserve the current L validation and ensure even n remains on the existing
execution path.
- Around line 69-80: Remove alpha from the cache_key tuple in the relevant JAX
API caching logic, while retaining all shape, dtype, tiling, and configuration
components. Do not change the docstring’s “static (trace-time)” wording;
continue passing alpha through the FFI attribute without using it to distinguish
compiled targets.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Line 110: Replace user-facing torch-qualified dtype names with
framework-neutral CUTLASS dtype names: update the acc_dtype documentation and
sf_dtype default in dswiglu/api.py at lines 110 and 386, update the acc_dtype,
c_dtype, and d_dtype defaults in swiglu/api.py at lines 943-945, and change the
validation message in _bf16_api.py at lines 300-301 from torch.float32 to
float32.
- Around line 822-834: The four pointer-retention paths use one slot, allowing
replacement before the prior kernel finishes. In
python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py:822-834 and
python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py:834-849, replace
retention of (b_ptrs, sfb_ptrs) with CUDA-event-gated retention; in
python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py:206-213, update
_record_pointer_stream similarly for self._live_b_ptrs; and in
python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py:1354-1358, update the discrete
branch. Record an event after each launch and release each retained array set
only after its event completes, or maintain a bounded event-keyed retention
list, preserving the existing torch record_stream behavior.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/api.py`:
- Around line 304-321: Resolve JAX devices through the adapter before
allocation, avoiding direct use of a potentially callable .device attribute. In
python/cudnn/gemm/cutedsl/grouped/unfused/api.py:304-321, update
_allocate_output before jnp.empty; in
python/cudnn/gemm/cutedsl/grouped/wgrad/api.py:299-308, resolve it before
allocator; in python/cudnn/gemm/cutedsl/grouped/wgrad/api.py:337-344, resolve it
before jnp.empty; and in
python/cudnn/gemm/cutedsl/grouped/dglu/api.py:1072-1080, resolve it once and
reuse it for both d_row_tensor and dbias_tensor allocations.
- Around line 262-264: The torch type annotations trigger Ruff F821 because
torch was removed from module scope. In each affected
file—python/cudnn/gemm/cutedsl/grouped/unfused/api.py (262-264), quant/api.py
(95), srelu/api.py (48), swiglu/api.py (80), unfused/_bf16_api.py (100),
dglu/_bf16_api.py (61), dglu/_blockscaled_api.py (127), dglu/api.py (104-105),
wgrad/api.py (246-247), and wgrad/_blockscaled_api.py (62)—import TYPE_CHECKING
and add one TYPE_CHECKING-guarded torch import, covering all torch.dtype and
torch.Tensor annotations without restoring the runtime import.
- Around line 192-193: Replace framework-specific dtype names in validation
error messages with framework-neutral names: in
python/cudnn/gemm/cutedsl/grouped/unfused/api.py lines 192-193, 219, 226, and
244 use “bfloat16” instead of “torch.bfloat16”; in
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py:339-340,
python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py:310-311, and
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:261-262 use “float32”
instead of “torch.float32”; and in
python/cudnn/gemm/cutedsl/grouped/dglu/api.py:999-1000 update the dprob_tensor
message similarly.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py`:
- Around line 344-364: Update _allocate_single_expert_placeholder to pass
desc.device to the JAX jnp.empty call, ensuring the placeholder is allocated on
the input device. Also update the related jnp.asarray path for
cached_single_expert to pass the same device, preserving correct placement for
multi-GPU inputs.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py`:
- Around line 25-33: Rename the ambiguous l parameter and all related references
to batch_size in make_inputs and each affected test case, including shape
construction, helper calls, and loops. Preserve the existing batch-dimension
behavior while eliminating Ruff E741 violations.
In `@test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py`:
- Line 163: Prefix the unused tuple-unpacked values with underscores at both
error-test setup sites: in
test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py lines
163-163, rename a_np and norm_const_np to _a_np and _norm_const_np; in
test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py lines 177-177,
rename b_np and norm_const_np to _b_np and _norm_const_np.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py`:
- Line 23: Rename the _make_jax_inputs parameter l to a descriptive expert-count
name such as num_experts, and update all local uses in the function and the
additional affected declaration accordingly, preserving the existing behavior.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py`:
- Around line 35-87: Gate all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py (lines 35-87),
all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py (lines 36-89),
and all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py (lines 36-84) by
importing and calling skip_unless_sm100() at the start of each test.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/dense/amax/api.py`:
- Around line 469-484: Document and enforce the eager JAX ownership contract in
gemm_amax_wrapper_sm100 and the corresponding JAX branches of
python/cudnn/gemm/cutedsl/dense/dsrelu/api.py:544-566,
python/cudnn/gemm/cutedsl/dense/srelu/api.py:527-549,
python/cudnn/gemm/cutedsl/dense/swiglu/api.py:658-677, and
python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py:611-617. Keep these
wrappers eager-only and explicitly reject invocation under jax.jit or with
donated buffers; retain block_until_ready only for materialization ordering and
apply the same explicit contract consistently across all five wrappers.
In `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py`:
- Around line 28-30: Add a clear registry-size policy for _registered_targets,
either documenting the expected bounded number of distinct configurations or
enforcing a maximum that raises an explicit error before registering another
target. Apply the policy in the cache-key registration path using
get_or_register_env_stream_target, while preserving reuse of existing entries.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 578-585: Replace the assert in gemm_proj_rope_mxfp8_wrapper_sm100
that validates matching x and w dtypes with an explicit ValueError, preserving
the existing message. Apply the same assert-to-explicit-raise change to the
x_scale and w_scale validation checks around the corresponding validation
blocks.
- Around line 64-73: Replace all torch.Tensor annotations with Any in the first
constructor, GemmProjRopeMxfp8Mxfp8InSm100.__init__, and
gemm_proj_rope_mxfp8_wrapper_sm100, including every affected parameter and
return annotation. Preserve the existing signatures and behavior while matching
the sibling APIs and allowing JAX arrays and runtime type-hint resolution
without requiring torch.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py`:
- Around line 202-218: Extract the shared scale-factor layout classifier and
shape validator into discrete_kernel_utils.py, covering both the permuted atom
view and physical C-contiguous form with one consistent predicate. Replace
_check_sf_shape in swiglu/api.py and the _sf_desc_is_physical/_check_sf_shape
pair in dsrelu/api.py with calls to this helper, while passing the discrete-mode
restriction from the dsrelu caller. Apply the same shared-helper change to
dswiglu/api.py, preserving each API’s existing return and validation behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py`:
- Around line 756-764: Update the framework validation in the
grouped_gemm_dswiglu_wrapper_sm100 path so JAX is rejected by one accurate
message rather than first being listed as supported and then rejected
unconditionally. Keep torch as the accepted framework and state the
dense-weight/layout constraint consistently in the resulting error.
In `@python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py`:
- Around line 585-589: In the validation flow after _validate_pointer_tensor for
b_ptrs, replace the duplicated get_data_ptr alignment check and error with a
call to the existing _validate_pointer_array_alignment helper, preserving the
8-byte alignment validation and centralized error behavior.
In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py`:
- Around line 1016-1034: Update _allocate_output so its JAX and torch allocation
paths cannot silently diverge on layout: either validate that stride matches the
expected C-contiguous layout before allocation, or remove the stride parameter
and derive that layout directly in the torch branch. Preserve the existing
output allocation behavior for c_tensor and d_tensor.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bf9180c3-d107-4f82-99d1-adf482bb162b
📒 Files selected for processing (80)
AGENTS.mddocs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/gemm_amax.mddocs/fe-oss-apis/gemm_fusions/gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.mddocs/fe-oss-apis/gemm_fusions/gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.mddocs/fe-oss-apis/overview.mdpyproject.tomlpython/cudnn/__init__.pypython/cudnn/api_base.pypython/cudnn/datatypes.pypython/cudnn/gemm/cutedsl/_jax_ffi.pypython/cudnn/gemm/cutedsl/dense/amax/__init__.pypython/cudnn/gemm/cutedsl/dense/amax/api.pypython/cudnn/gemm/cutedsl/dense/amax/jax_api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_bf16in.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pypython/cudnn/gemm/cutedsl/dense/srelu/api.pypython/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.pypython/cudnn/gemm/cutedsl/dense/swiglu/__init__.pypython/cudnn/gemm/cutedsl/dense/swiglu/api.pypython/cudnn/gemm/cutedsl/dense/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/grouped/backend_utils.pypython/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/dglu/api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/api.pypython/cudnn/gemm/cutedsl/grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/glu/api.pypython/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.pypython/cudnn/gemm/cutedsl/grouped/glu_hadamard/hadamard_utils.pypython/cudnn/gemm/cutedsl/grouped/moe_kernel_helpers.pypython/cudnn/gemm/cutedsl/grouped/quant/api.pypython/cudnn/gemm/cutedsl/grouped/srelu/api.pypython/cudnn/gemm/cutedsl/grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/unfused/api.pypython/cudnn/gemm/cutedsl/grouped/utils.pypython/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/api.pypython/cudnn/tensor_adapter.pytest/python/conftest.pytest/python/fe_api/gemm/test_gemm_amax_jax.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.pytest/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.pytest/python/fe_api/gemm/test_gemm_swiglu_jax.pytest/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.pytest/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py
💤 Files with no reviewable changes (2)
- python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py
- python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py
Comment on lines
+114
to
+120
| ```python | ||
| from cudnn import gemm_amax_jax_sm100 | ||
|
|
||
| @jax.jit | ||
| def quantized_matmul(a, b, sfa, sfb): | ||
| c, amax = gemm_amax_jax_sm100(a, b, sfa, sfb, c_dtype=jnp.float32, sf_vec_size=32) | ||
| return c, amax |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Import the JAX symbols used by this example.
Line 117 uses jax.jit. Line 119 uses jnp.float32. The snippet raises NameError when copied without imports.
Proposed fix
+import jax
+import jax.numpy as jnp
from cudnn import gemm_amax_jax_sm100📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| ```python | |
| from cudnn import gemm_amax_jax_sm100 | |
| @jax.jit | |
| def quantized_matmul(a, b, sfa, sfb): | |
| c, amax = gemm_amax_jax_sm100(a, b, sfa, sfb, c_dtype=jnp.float32, sf_vec_size=32) | |
| return c, amax |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/fe-oss-apis/gemm_fusions/gemm_amax.md` around lines 114 - 120, Add
imports for the JAX symbols used in the `quantized_matmul` example: import `jax`
for `jax.jit` and `jax.numpy` as `jnp` for `jnp.float32`, while preserving the
existing `gemm_amax_jax_sm100` import.
Comment on lines
+73
to
+80
| ```python | ||
| from cudnn import gemm_swiglu_jax_sm100 | ||
|
|
||
| @jax.jit | ||
| def swiglu_mlp(a, b): | ||
| ab12, c = gemm_swiglu_jax_sm100(a, b, alpha=1.0, ab12_dtype=jnp.float32, c_dtype=jnp.bfloat16) | ||
| return c | ||
| ``` |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Import JAX in the example.
The example uses jax.jit and jnp.float32 without defining jax or jnp. The documented code raises NameError before it calls gemm_swiglu_jax_sm100.
Proposed fix
+import jax
+import jax.numpy as jnp
from cudnn import gemm_swiglu_jax_sm100📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| ```python | |
| from cudnn import gemm_swiglu_jax_sm100 | |
| @jax.jit | |
| def swiglu_mlp(a, b): | |
| ab12, c = gemm_swiglu_jax_sm100(a, b, alpha=1.0, ab12_dtype=jnp.float32, c_dtype=jnp.bfloat16) | |
| return c | |
| ``` |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md` around lines 73 - 80, Add the
required JAX imports at the start of the documented example so both jax.jit and
jnp.float32/jnp.bfloat16 resolve before swiglu_mlp calls gemm_swiglu_jax_sm100.
| dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, TensorDesc) or _is_framework_tensor(tensor_or_dtype) else tensor_or_dtype | ||
| return _convert_to_cutlass_data_type_or_none(dtype) in {cutlass.Float16, cutlass.BFloat16} | ||
|
|
||
| def _get_innermost_stride_dim(self, tensor: torch.Tensor, name: str = "") -> int: |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- imports and module header ---'
sed -n '1,90p' python/cudnn/api_base.py
printf '%s\n' '--- target method and nearby definitions ---'
sed -n '660,725p' python/cudnn/api_base.py
printf '%s\n' '--- torch references in the file ---'
rg -n '\btorch\b|from __future__ import annotations|TYPE_CHECKING' python/cudnn/api_base.py
printf '%s\n' '--- lint configuration relevant to F821 ---'
rg -n 'ruff|F821|target-version|select|ignore' pyproject.toml setup.cfg tox.ini .ruff.toml ruff.toml 2>/dev/null || trueRepository: NVIDIA/cudnn-frontend
Length of output: 10628
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- available static-analysis tools ---'
command -v ruff || true
command -v pyflakes || true
printf '%s\n' '--- Ruff behavior for postponed annotations ---'
if command -v ruff >/dev/null 2>&1; then
printf '%s\n' 'from __future__ import annotations' 'def f(x: torch.Tensor) -> None: ...' |
ruff check --stdin-filename probe.py --select F821 -
else
printf '%s\n' 'ruff is not installed'
fi
printf '%s\n' '--- exact annotation definitions in api_base.py ---'
sed -n '560,625p' python/cudnn/api_base.py
sed -n '635,710p' python/cudnn/api_base.py
sed -n '945,965p' python/cudnn/api_base.pyRepository: NVIDIA/cudnn-frontend
Length of output: 557
🏁 Script executed:
#!/bin/bash
set +e
printf '%s\n' '--- Ruff diagnostics for api_base.py ---'
ruff check python/cudnn/api_base.py --select F821
printf 'ruff_exit=%s\n' "$?"
printf '%s\n' '--- call sites of _get_innermost_stride_dim ---'
rg -n -C 2 '_get_innermost_stride_dim' python
printf '%s\n' '--- get_strides definition and accepted input types ---'
rg -n -C 8 'def get_strides|def _get_innermost_stride_dim' python/cudnn
printf '%s\n' '--- focused diff context ---'
git diff --unified=3 -- python/cudnn/api_base.py | rg -n -C 8 '_get_innermost_stride_dim|torch import|_torch\(\)'Repository: NVIDIA/cudnn-frontend
Length of output: 16834
Add a type-only torch import or remove direct torch annotations. Ruff reports 23 F821 errors in python/cudnn/api_base.py, including line 690. Use if TYPE_CHECKING: import torch to preserve lazy runtime imports, or use framework-neutral annotations where appropriate.
🧰 Tools
🪛 Ruff (0.16.1)
[error] 690-690: Undefined name torch
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/api_base.py` at line 690, Resolve the undefined torch
annotations in python/cudnn/api_base.py, including _get_innermost_stride_dim, by
adding a type-only torch import guarded by TYPE_CHECKING while preserving lazy
runtime imports, or replacing direct torch annotations with framework-neutral
types where appropriate.
Source: Linters/SAST tools
Comment on lines
+167
to
+183
| self._not_implemented_error_if( | ||
| self._is_fp8(self.ab12_dtype), | ||
| f"ab12_dtype {{torch.float8_e5m2, torch.float8_e4m3fn}} is currently disabled", | ||
| ) | ||
| elif self.acc_dtype is cutlass.Float16: | ||
| self.ab12_dtype = self._check_dtype( | ||
| self.ab12_desc, | ||
| dtype=[cutlass.Float16, cutlass.BFloat16], | ||
| name="AB12 (for float16 acc_dtype)", | ||
| ) | ||
| self._check_dtype( | ||
| self.a_desc, | ||
| dtype=[cutlass.Float16, cutlass.Float8E4M3FN, cutlass.Float8E5M2], | ||
| name="A/B (for float16 acc_dtype)", | ||
| ) | ||
| else: | ||
| raise ValueError(f"Unsupported acc_dtype: expected one of {{torch.float32, torch.float16}}, got {self.acc_dtype}") |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the error messages to name CUTLASS dtypes.
The compared values are now CUTLASS types, but these messages still name torch dtypes. Line 183 interpolates self.acc_dtype, so the raised text reads "expected one of {torch.float32, torch.float16}, got Float32". That is contradictory and slows debugging. Line 169 also has an f prefix with no placeholder, which Ruff reports as F541.
Line 250 carries the same stale torch wording.
✏️ Proposed message fix
self._not_implemented_error_if(
self._is_fp8(self.ab12_dtype),
- f"ab12_dtype {{torch.float8_e5m2, torch.float8_e4m3fn}} is currently disabled",
+ "ab12_dtype {Float8E5M2, Float8E4M3FN} is currently disabled",
) else:
- raise ValueError(f"Unsupported acc_dtype: expected one of {{torch.float32, torch.float16}}, got {self.acc_dtype}")
+ raise ValueError(f"Unsupported acc_dtype: expected one of {{cutlass.Float32, cutlass.Float16}}, got {self.acc_dtype}")🧰 Tools
🪛 Ruff (0.16.1)
[error] 169-169: f-string without any placeholders
Remove extraneous f prefix
(F541)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/api.py` around lines 167 - 183, Update
the validation messages in the relevant dtype checks to name CUTLASS dtype
values consistently, including the unsupported acc_dtype error and the message
around line 250; ensure the expected values match the actual CUTLASS types
interpolated or compared. Remove the unnecessary f-string prefix from the static
ab12_dtype message to resolve Ruff F541.
Source: Linters/SAST tools
Comment on lines
+35
to
+40
| ### Quantize only arguments | ||
| sfa_tensor: Optional[Any] = None, | ||
| sfb_tensor: Optional[Any] = None, | ||
| sf_vec_size: int = 16, | ||
| vector_f32: bool = False, | ||
| ab12_stages: int = 4, |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject the parameters that this entry point ignores.
sf_vec_size, vector_f32, and ab12_stages are accepted but never used. make_gemm at lines 82-92 does not forward them to GemmSwigluSm100, so the constructor defaults apply. vector_f32 and ab12_stages reach cute.compile only through the blockscaled branch of GemmSwigluSm100._compile_kernel, and this entry point rejects blockscaled inputs at lines 61-67. A caller who sets ab12_stages=8 therefore gets the default value 4 with no error.
The parameters are also absent from cache_key, so a future change that makes them meaningful would return a stale compiled kernel.
Either forward them to make_gemm and add them to cache_key, or reject non-default values.
🛡️ Proposed guard
if sfa_tensor is not None or sfb_tensor is not None:
# The quantized kernel's compiled signature carries explicit None-typed
# parameters (amax/sfc/norm_const) that the XLA FFI bridge cannot supply.
raise NotImplementedError(
"gemm_swiglu_jax_sm100 currently supports the non-quantized kernel only; "
"use gemm_swiglu_wrapper_sm100 (eager) for blockscaled MXFP8 inputs from JAX"
)
+ if vector_f32 or ab12_stages != 4:
+ raise ValueError("vector_f32 and ab12_stages apply to the blockscaled kernel only, which gemm_swiglu_jax_sm100 does not support")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py` around lines 35 - 40,
Reject non-default values for sf_vec_size, vector_f32, and ab12_stages at the
entry point before make_gemm, since this path ignores them and rejects
blockscaled inputs. Preserve the existing defaults; alternatively, if these
options must be supported, forward them through make_gemm and include them in
cache_key.
Comment on lines
+344
to
+364
| 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"))) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether cached_single_expert is dereferenced by the wgrad kernel, and how
# other JAX allocations in this cohort choose a device.
rg -nP -C5 'single_expert' --type=py -g '!**/test/**'
rg -nP -C3 'jnp\.(empty|zeros|asarray)\s*\(' --type=py -g '!**/test/**'Repository: NVIDIA/cudnn-frontend
Length of output: 159
🏁 Script executed:
#!/bin/bash
set -e
file=$(fd -t f '_bf16_api\.py$' . | head -n 1)
printf '%s\n' "FILE=$file"
wc -l "$file"
sed -n '320,385p' "$file"
sed -n '475,530p' "$file"
sed -n '545,595p' "$file"
printf '%s\n' '--- related symbols and JAX allocations ---'
rg -n -C4 'single_expert|_generate_wgrad_ptrs|jnp\.(empty|zeros|asarray|array)|get_device|a_desc\.device' "$file"
printf '%s\n' '--- repository-wide matching Python files ---'
rg -n -C3 'single_expert|_generate_wgrad_ptrs|jnp\.(empty|zeros|asarray|array)' python/cudnn --glob '*.py' || trueRepository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -e
file=python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py
printf '%s\n' '--- target implementation ---'
sed -n '70,130p' "$file"
sed -n '330,370p' "$file"
sed -n '395,465p' "$file"
sed -n '490,520p' "$file"
sed -n '550,610p' "$file"
printf '%s\n' '--- device model and allocator ---'
sed -n '1,245p' python/cudnn/tensor_adapter.py
printf '%s\n' '--- callers of the target API ---'
rg -n -C5 'GroupedGemmWgrad|wgrad_ptrs|sample_wgrad_expert|_generate_wgrad_ptrs' python/cudnn --glob '*.py' | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in {
"_allocate_single_expert_placeholder",
"_generate_wgrad_ptrs",
}:
print(f"{node.name}: lines {node.lineno}-{node.end_lineno}")
for call in ast.walk(node):
if isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute):
receiver = ast.unparse(call.func.value)
name = call.func.attr
if receiver in {"jnp", "jax"} and name in {"empty", "asarray", "block_until_ready"}:
keywords = [kw.arg for kw in call.keywords]
print(f" {receiver}.{name} line {call.lineno}, keywords={keywords}")
PY
printf '%s\n' '--- descriptor device normalization and validation ---'
sed -n '245,315p' python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py
rg -n -C8 'def _make_tensor_desc|class TensorDesc|Device\(' python/cudnn --glob '*.py'Repository: NVIDIA/cudnn-frontend
Length of output: 29104
Allocate JAX buffers on the input device. Pass the input device to both jnp.empty and jnp.asarray. Otherwise multi-GPU JAX inputs on a non-default device create buffers on device 0; wgrad_ptrs then fails the device check, while cached_single_expert has no equivalent check.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py` around lines 344 - 364,
Update _allocate_single_expert_placeholder to pass desc.device to the JAX
jnp.empty call, ensuring the placeholder is allocated on the input device. Also
update the related jnp.asarray path for cached_single_expert to pass the same
device, preserving correct placement for multi-GPU inputs.
Comment on lines
+25
to
+33
| def make_inputs(m, n, k, l, sf_vec_size, rng): | ||
| rest_k = ceil_div(ceil_div(k, sf_vec_size), 4) | ||
| a_np = rng.standard_normal((m, k, l), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn) | ||
| b_np = rng.standard_normal((n, k, l), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn) | ||
| # e8m0 value 1.0 == byte 127; physical C-contiguous atom shape (L, MN', K', 32, 4, 4) | ||
| sfa_np = np.full((l, ceil_div(m, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu) | ||
| sfb_np = np.full((l, ceil_div(n, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu) | ||
| prob_np = rng.random((m, 1, l), dtype=np.float32) | ||
| return a_np, b_np, sfa_np, sfb_np, prob_np |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous l variable.
Ruff reports E741 for l. Rename it to batch_size in the helper and each test. This prevents a lint failure.
Also applies to: 45-48, 86-89, 130-133
🧰 Tools
🪛 Ruff (0.16.1)
[error] 25-25: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py` around lines 25 - 33,
Rename the ambiguous l parameter and all related references to batch_size in
make_inputs and each affected test case, including shape construction, helper
calls, and loops. Preserve the existing batch-dimension behavior while
eliminating Ruff E741 violations.
Source: Linters/SAST tools
| 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() |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Prefix unused tuple-unpacked values with underscores.
Ruff reports RUF059 at both error-test setup sites.
test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py#L163-L163: Rename unuseda_npandnorm_const_npto_a_npand_norm_const_np.test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py#L177-L177: Rename unusedb_npandnorm_const_npto_b_npand_norm_const_np.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 163-163: Unpacked variable a_np is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
[warning] 163-163: Unpacked variable norm_const_np is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
📍 Affects 2 files
test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py#L163-L163(this comment)test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py#L177-L177
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py` at
line 163, Prefix the unused tuple-unpacked values with underscores at both
error-test setup sites: in
test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py lines
163-163, rename a_np and norm_const_np to _a_np and _norm_const_np; in
test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py lines 177-177,
rename b_np and norm_const_np to _b_np and _norm_const_np.
Source: Linters/SAST tools
| from fe_api.gemm.test_gemm_amax_jax import skip_unless_sm100 | ||
|
|
||
|
|
||
| def _make_jax_inputs(m=256, n=128, k=128, l=2): |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename l to a descriptive expert-count variable.
Ruff reports E741 for both declarations. Rename l and its local uses to num_experts or an equivalent descriptive name.
Also applies to: 91-91
🧰 Tools
🪛 Ruff (0.16.1)
[error] 23-23: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py` at line 23,
Rename the _make_jax_inputs parameter l to a descriptive expert-count name such
as num_experts, and update all local uses in the function and the additional
affected declaration accordingly, preserving the existing behavior.
Source: Linters/SAST tools
Comment on lines
+35
to
+87
| @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, | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Add the required SM100 capability gate to each CuTeDSL test.
These tests can run and pass on unsupported systems because JAX rejection occurs before kernel support validation. Import and call skip_unless_sm100() at the start of every test.
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L35-L87: import and callskip_unless_sm100()in all three tests.test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L36-L89: import and callskip_unless_sm100()in all three tests.test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L36-L84: import and callskip_unless_sm100()in all three tests.
As per coding guidelines, “Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks.”
📍 Affects 3 files
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L35-L87(this comment)test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L36-L89test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L36-L84
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py` around lines
35 - 87, Gate all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py (lines 35-87),
all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py (lines 36-89),
and all three tests in
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py (lines 36-84) by
importing and calling skip_unless_sm100() at the start of each test.
Source: Coding guidelines
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 <noreply@anthropic.com>
…cy 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 <noreply@anthropic.com>
Anerudhan
force-pushed
the
cutedsl-torch-group
branch
from
August 11, 2026 05:43
6542a01 to
d31d137
Compare
3 tasks
Anerudhan
added a commit
that referenced
this pull request
Aug 11, 2026
…cross the GEMM CuTeDSL APIs (stacked on #534) (#553) * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * Add cudnn.jax.call on CuTeDSL's native JAX bridge; jit entry points across the GEMM CuTeDSL APIs Replace the jax-tvm-ffi backend with cutlass.jax.cutlass_call wrapped as cudnn.jax.call, and add jax.jit-compatible XLA custom-call entry points for every JAX-reachable GEMM API: the four dense fusions (amax, swiglu incl. quantized, srelu, dsrelu), proj_rope_mxfp8 (both input paths), and the discrete-mode grouped family (unfused, glu, dglu, dsrelu, wgrad, discrete-grouped swiglu/dswiglu). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*.Affected area
Build, packaging, or installation
Summary
Removes
torchandtorch-c-dlpack-extfrom the[cutedsl]optional extra. With the CuTeDSL APIs type-erased and torch-lazy (#529/#530), torch is now opt-in exactly like jax, via the PEP 735 dependency groups:AGENTS.md (build/test instructions) and the FE-OSS overview installation docs are updated to match.
Why
The
cutedslextra described the CuTeDSL runtime's needs plus one framework's. After #529/#530 the package core is framework-neutral — a JAX user installing[cutedsl]should not download torch, and the two frameworks should be symmetric opt-ins.Related issues
Stacked on #530 (which stacks on #529).
API and compatibility impact
pip install nvidia-cudnn-frontend[cutedsl]no longer pulls torch. Users of the torch-only OSS APIs that live behind this extra (SDPA, BSA/DSA/NSA, and the torch-only grouped configurations) must install torch separately:pip install --group torchfrom a checkout, orpip install torch torch-c-dlpack-extfor the published wheel (PEP 735 groups are not part of wheel metadata — the same caveat as thejaxgroup).Testing
tomllibparse validated:cutedsl = [nvidia-cutlass-dsl[cu13]>=4.5.0, cuda-python, apache-tvm-ffi>=0.1.11], groupstorch = [torch, torch-c-dlpack-ext],jax = [jax>=0.4.35, jax-tvm-ffi>=0.1.3].🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Tests