Add cudnn.jax.call on CuTeDSL's native JAX bridge; jit entry points across the GEMM CuTeDSL APIs (stacked on #534) - #553
Conversation
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
added
cat-feature
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds native CuTeDSL JAX custom-call APIs, framework-neutral Torch/JAX tensor handling, lazy JAX exports, updated dependency groups, expanded JAX documentation, and eager, jitted, parity, and validation tests. ChangesCuTeDSL JAX integration
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant JAX
participant cudnn.jax.call
participant cutlass.jax
participant CuTeDSLKernel
JAX->>cudnn.jax.call: provide tensors and output specifications
cudnn.jax.call->>cutlass.jax: configure descriptors, aliases, and initializers
cutlass.jax->>CuTeDSLKernel: compile and invoke custom call
CuTeDSLKernel-->>JAX: return donated and external output buffers
Possibly related PRs
Suggested reviewers: 🚥 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}
There was a problem hiding this comment.
Actionable comments posted: 14
🧹 Nitpick comments (15)
test/python/fe_api/gemm/test_gemm_amax_jax.py (1)
247-250: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHandle the optional CuTeDSL JAX extension consistently across the dense-GEMM JIT tests. All three tests exercise the same XLA custom-call bridge, but they treat the optional
cutlassdependency differently: two importcutlass.jaxwithout a guard, and one omits the availability check. Adopt one convention:pytest.importorskip("cutlass.jax")followed by anis_available()skip.
test/python/fe_api/gemm/test_gemm_amax_jax.py#L247-L250: replaceimport cutlass.jaxwithcutlass_jax = pytest.importorskip("cutlass.jax")and callcutlass_jax.is_available().test/python/fe_api/gemm/test_gemm_swiglu_jax.py#L109-L112: apply the same replacement.test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py#L143-L147: add theimportorskipandis_available()skip beforeskip_unless_sm100().As per coding guidelines: "Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks".
🤖 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_amax_jax.py` around lines 247 - 250, Standardize optional CuTeDSL JAX gating across the three dense-GEMM JIT tests: in test/python/fe_api/gemm/test_gemm_amax_jax.py lines 247-250 and test/python/fe_api/gemm/test_gemm_swiglu_jax.py lines 109-112, replace the direct import with pytest.importorskip("cutlass.jax"), then use the returned cutlass_jax symbol for is_available(); in test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py lines 143-147, add the same importorskip and availability skip before skip_unless_sm100().Source: Coding guidelines
test/python/fe_api/gemm/test_gemm_swiglu_jax.py (1)
141-150: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRepeat the quantized jitted call to cover donation safety.
The bf16 path runs
jittedtwice at lines 127-131. The docstring at line 108 states that the repetition covers donation safety. The quantized path callsquantonce at line 146. The quantized adapter also uses donated pre-initialized outputs, so a second invocation would cover the same failure mode for this path.♻️ Proposed change
- ab12q, cq = quant(a2, b2, sfa, sfb) - jax.block_until_ready((ab12q, cq)) ab12q_ref = (a2_ref * sfa_expanded) @ (b2_ref * sfb_expanded).T - np.testing.assert_allclose(np.asarray(ab12q).astype(np.float32)[:, :, 0], ab12q_ref, atol=0.5, rtol=0.05) - np.testing.assert_allclose(np.asarray(cq).astype(np.float32)[:, :, 0], swiglu_block_ref(ab12q_ref, n), atol=1.0, rtol=0.05) + for _ in range(2): + ab12q, cq = quant(a2, b2, sfa, sfb) + jax.block_until_ready((ab12q, cq)) + np.testing.assert_allclose(np.asarray(ab12q).astype(np.float32)[:, :, 0], ab12q_ref, atol=0.5, rtol=0.05) + np.testing.assert_allclose(np.asarray(cq).astype(np.float32)[:, :, 0], swiglu_block_ref(ab12q_ref, n), atol=1.0, rtol=0.05)🤖 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_swiglu_jax.py` around lines 141 - 150, Repeat the quantized adapter invocation in the test after the initial quant(a2, b2, sfa, sfb) call, preserving the existing block_until_ready synchronization and assertions so the donated pre-initialized outputs are exercised for donation safety.python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py (2)
27-28: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider bounding the kernel caches.
_srelu_kernel_cacheand_dsrelu_kernel_cacheare unbounded module-level dicts. Each distinct shape and configuration retains a compiled kernel for the lifetime of the process. Workloads with many shape variants grow host memory without a limit. A bounded LRU policy or a documented cache-clear helper would make the growth predictable.🤖 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/srelu/jax_api.py` around lines 27 - 28, Bound the module-level _srelu_kernel_cache and _dsrelu_kernel_cache so compiled kernels do not accumulate indefinitely across shape and configuration variants. Use an LRU or equivalent bounded policy, applying it consistently to both caches while preserving their existing lookup and reuse behavior.
30-32: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winShare the epilogue definitions with the class APIs.
The comment states that these lambdas are identical to the ones in
GemmSreluSm100.compile()andGemmDsreluSm100.compile(). Two copies can drift. A change on the class API side would silently change only the eager numeric path. Move the two epilogues into a shared module and import them in both places.🤖 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/srelu/jax_api.py` around lines 30 - 32, Move the shared SReLU and DReLU epilogue definitions out of the JAX API and class API implementations into a common module, then import and reuse them in GemmSreluSm100.compile(), GemmDsreluSm100.compile(), and the JAX API. Remove the duplicated local lambdas while preserving their current behavior.python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py (1)
13-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBoth package initializers extend
__all__with a separateappendcall. The name can be declared directly in the__all__literal, which is easier to read and to grep.
python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py#L13-L25: add"gemm_dsrelu_jax_sm100"to the__all__literal and remove the__all__.append(...)line.python/cudnn/gemm/cutedsl/dense/srelu/__init__.py#L13-L25: add"gemm_srelu_jax_sm100"to the__all__literal and remove the__all__.append(...)line.🤖 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/dsrelu/__init__.py` around lines 13 - 25, Declare "gemm_dsrelu_jax_sm100" directly in the __all__ literal and remove the separate __all__.append call in python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py (lines 13-25). Apply the same change for "gemm_srelu_jax_sm100" in python/cudnn/gemm/cutedsl/dense/srelu/__init__.py (lines 13-25), leaving each module’s __getattr__ lazy import behavior unchanged.python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py (1)
153-164: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the reversed output order in the quantized branch.
The quantized branch declares
output_shape_dtype=(out_types[1], out_types[0])and unpacksc_tensor, ab12_tensor, while the standard branch uses the opposite order. The reason is the kernel signature:_swiglu_quant_adapterreceivescbeforeab12. A short comment prevents an accidental reordering later.♻️ Proposed comment
else: sf = sf_atom_spec() + # The quantized kernel takes c before ab12, so the bridge outputs are + # declared and unpacked in that order. c_tensor, ab12_tensor = call(🤖 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 153 - 164, Add a concise comment in the quantized branch immediately before the _swiglu_quant_adapter call explaining that outputs are intentionally reversed because the kernel signature receives c before ab12. Preserve the existing output_shape_dtype ordering and c_tensor, ab12_tensor unpacking.python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py (2)
133-135: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicate
_interpret_uint8_as_fp4x2assignment.Line 135 sets
self._interpret_uint8_as_fp4x2 = True. Line 191 sets the same attribute to the same value. The new assignment at line 135 is the meaningful one, because it must precede descriptor creation. Delete the later assignment at line 191 so the ordering requirement stays clear.🤖 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/dswiglu/api.py` around lines 133 - 135, Remove the redundant self._interpret_uint8_as_fp4x2 = True assignment later in the initialization flow, while retaining the assignment before descriptor creation. Keep the earlier assignment as the sole initialization so the required ordering remains explicit.
1073-1095: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd
frameworktocache_key.The cached
DiscreteGroupedGemmDswigluSm100object storesself._frameworkand allocates its workspace and pointer placeholders in that framework.cache_keydoes not contain the framework. Today the torch and JAX calls differ in the SFA signature, so they land in separate entries by accident. That coupling is implicit. An explicit key entry makes the invariant hold regardless of future SF layout changes.♻️ Proposed change
cache_key = ( + framework, get_shape(a_tensor)[1:], stride_order(a_tensor), ab_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/discrete_grouped/dswiglu/api.py` around lines 1073 - 1095, Add the framework identifier to the cache_key assembled for DiscreteGroupedGemmDswigluSm100, using the same framework value used by the cached object's self._framework and workspace allocations. Keep the existing key components unchanged and ensure framework-specific torch and JAX entries remain distinct independently of tensor signature differences.python/cudnn/tensor_adapter.py (1)
180-182: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueState the mutability assumption for the version fallback.
get_versionreturns 0 for every object without_version. Callers treat the value as a mutation counter and skip revalidation when it is unchanged. That is correct for immutable JAX arrays. It is unsafe for any mutable framework added later that also lacks_version, because the validation cache would never invalidate. Record that constraint in the docstring.🤖 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/tensor_adapter.py` around lines 180 - 182, Update the get_version docstring to explicitly state that the 0 fallback is valid only for immutable arrays such as JAX, and must not be used for mutable frameworks lacking _version because cached validation would not invalidate.python/cudnn/jax/call.py (1)
37-44: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the fixed device index.
row_major_descalways reportsDevice("cuda", 0). The docstring explains that tracers expose no device, but it does not state that any device comparison insidecheck_supportbecomes a check against GPU 0. Add that note so callers do not read the descriptor device as authoritative.🤖 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/jax/call.py` around lines 37 - 44, Update the docstring for row_major_desc to state that its fixed Device("cuda", 0) value causes device comparisons in check_support to check against GPU 0, and clarify that this descriptor device is not authoritative.python/cudnn/gemm/cutedsl/grouped/srelu/api.py (1)
942-949: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInitialize the keep-alive pointer buffers in
__init__.
self._compile_b_ptrsandself._compile_sfb_ptrsare first assigned here. In dense weight mode they never exist. Sibling APIs (discrete_grouped/swiglu/api.py,grouped/dsrelu/api.py) set both toNonein__init__. Match that pattern so the keep-alive intent is visible from the constructor.🤖 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/srelu/api.py` around lines 942 - 949, Initialize self._compile_b_ptrs and self._compile_sfb_ptrs to None in the class __init__ method, matching the sibling grouped API patterns. Preserve the existing assignments in the compilation path so dense weight mode and keep-alive pointer handling continue to work.python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (1)
1493-1494: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBind the framework modules where they are used.
import torchat Line 1494 andimport jax.numpy as jnpat Line 1520 create function-local names that stay unbound on the other framework path. Every later use is currently guarded, so the code is correct. Any future unguarded use raisesUnboundLocalErrorinstead of a clearImportError. Move each import into the branch that uses it, or resolve both through a small helper.Also applies to: 1519-1528
🤖 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/dsrelu/api.py` around lines 1493 - 1494, Update the framework-specific import handling in the containing API function: keep torch resolution within the torch branch and jax.numpy as jnp within the JAX branch, or centralize both through a helper that explicitly reports unavailable dependencies. Ensure each framework path binds only its required module and missing dependencies fail with a clear ImportError rather than an unbound local name.python/cudnn/gemm/cutedsl/grouped/swiglu/api.py (1)
33-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_JAX_SF_LAYOUT_ERRORis duplicated verbatim. Both modules define the same scale-factor diagnostic. The two copies will drift when the wording is updated. Move the constant into a shared module such aspython/cudnn/gemm/cutedsl/grouped/backend_utils.pyand import it in both places.
python/cudnn/gemm/cutedsl/grouped/swiglu/api.py#L33-L37: replace the local definition with an import of the shared constant.python/cudnn/gemm/cutedsl/grouped/srelu/api.py#L41-L45: replace the local definition with an import of the shared constant.🤖 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/swiglu/api.py` around lines 33 - 37, Move the shared _JAX_SF_LAYOUT_ERROR constant into python/cudnn/gemm/cutedsl/grouped/backend_utils.py, then remove the local definitions and import it in both python/cudnn/gemm/cutedsl/grouped/swiglu/api.py lines 33-37 and python/cudnn/gemm/cutedsl/grouped/srelu/api.py lines 41-45. Preserve the diagnostic text unchanged.python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py (1)
187-193: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the message match the detected framework.
The guard rejects every non-torch framework, but the message always names JAX. A numpy input receives a message about JAX arrays. Include the detected framework name so the message stays accurate for direct construction of this class.
♻️ Proposed refactor
from cudnn.tensor_adapter import detect_framework - if sample_a is not None and detect_framework(sample_a) != "torch": + framework = detect_framework(sample_a) + if sample_a is not None and framework != "torch": raise ValueError( - "GroupedGemmDgluBlockScaledAPI supports torch tensors only: the block-scaled " - "scale-factor tensors use an MMA-interleaved layout that is not expressible as JAX arrays" + f"GroupedGemmDgluBlockScaledAPI supports torch tensors only (got '{framework}'): the block-scaled " + "scale-factor tensors use an MMA-interleaved layout that is not expressible as JAX arrays" )🤖 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/dglu/_blockscaled_api.py` around lines 187 - 193, Update the framework validation in GroupedGemmDgluBlockScaledAPI to store the result of detect_framework(sample_a) and use that detected framework name in the ValueError message, while preserving the existing torch-only rejection behavior.python/cudnn/gemm/cutedsl/grouped/unfused/api.py (1)
192-244: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUpdate the dtype error messages for the JAX path.
_normalize_callnow accepts JAX arrays, but the messages still name torch dtypes: "a_tensor must have dtype torch.bfloat16", "prob_tensor must have dtype torch.float32", "b_tensor must have dtype torch.bfloat16", and "b_dtype must be torch.bfloat16 for the BF16 backend". A JAX caller passing"bfloat16"receives a torch-specific message. Use a framework-neutral term such as "bfloat16" and "float32".♻️ Proposed message change
- raise ValueError(f"a_tensor must have dtype torch.bfloat16, got {a_tensor.dtype}") + raise ValueError(f"a_tensor must have dtype bfloat16, got {a_tensor.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/grouped/unfused/api.py` around lines 192 - 244, Update the dtype validation messages in the visible validation flow to use framework-neutral names: replace torch.bfloat16 with bfloat16 and torch.float32 with float32 in the a_tensor, prob_tensor, b_tensor, and b_dtype errors. Keep the validation logic and exception behavior unchanged.
🤖 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`:
- Line 112: Update the “JAX-specific constraints” bullet to apply only to the
eager entry points, excluding gemm_amax_jax_sm100. Preserve the eager-only and
synchronization requirements for those eager APIs while documenting that the
jitted entry point composes with jax.jit without manual synchronization.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 615-617: Update the raw kernel launch flow around the existing
jax.block_until_ready call to establish stream ordering for all JAX input
tensors and output consumers. Use the supported TVM-FFI mechanism to pass the
JAX/XLA stream into the launch; if unavailable, synchronize every producer input
(x, w, cos, sin, and scale tensors) before launching and synchronize the kernel
completion before returning outputs, while preserving the existing output
buffers.
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Line 151: Replace the bare support assertions in both GemmSreluSm100 and
GemmDsreluSm100 with explicit checks that call check_support() and raise
ValueError when validation fails, ensuring validation runs even under Python -O.
Apply this at python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py lines 151-151 and
251-251.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py`:
- Line 113: Replace the assert around gemm.check_support() with an explicit
support check that always executes, raising the same appropriate error style
used by the sibling discrete_grouped wrapper when the configuration is
unsupported.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Around line 831-834: Update the initialization near self._live_ptrs in
__init__ to add self._prev_live_ptrs = None, then in execute rotate the existing
self._live_ptrs into self._prev_live_ptrs before storing the current (b_ptrs,
sfb_ptrs). Preserve the tensor path and ensure non-torch buffers remain retained
through at least the next launch.
- Around line 203-218: Update _check_sf_shape so the physical-form branch
requires both the existing shape pattern and C-contiguous strides before setting
is_physical true. Keep non-physical validation on the permuted shape path, and
ensure non-contiguous torch tensors matching the physical shape are rejected
rather than treated as physical.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 310-311: Update the ValueError message in the accumulator dtype
validation of the relevant API class to name the expected framework-neutral
`Float32`/`cutlass.Float32` type instead of `torch.float32`, while preserving
the existing validation and reported value.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py`:
- Around line 248-253: Replace the single-slot self._live_b_ptrs retention in
_record_pointer_stream for
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py:248-253 with stream-aware
or CUDA-event-based retention so each non-torch pointer array remains referenced
until its launch stream completes; apply the same change to
_record_pointer_stream in
python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py:206-211.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py`:
- Around line 357-364: Update the JAX allocations in _single_expert_placeholder
initialization at python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:357-364
to pass a device derived from desc.device, and update the jnp.asarray allocation
at python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:510-515 to pass a device
derived from get_device(wgrad_tensor). Ensure both buffers are created on the
input device rather than JAX’s default device.
In `@python/cudnn/jax/call.py`:
- Around line 108-119: Update wrapper around full_input_spec so initialized
outputs always contribute their extra_specs, even when input_spec is None. When
inits exist, construct the full input specification by padding the original
inputs with None entries before appending extra_specs, preserving the existing
combined-spec behavior when input_spec is provided.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py`:
- Around line 143-147: Add the CuTeDSL JAX availability guard to
test_gemm_srelu_dsrelu_jax_jit_matches_eager, matching the existing checks in
test_gemm_amax_jax_ffi_sm100 and test_gemm_swiglu_jax_ffi_sm100. Skip the test
when cutlass.jax.is_available() is false, before importing or invoking the JAX
custom-call wrappers.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py`:
- Around line 107-112: Update the result comparison loop in the grouped dGLU JAX
test to exclude dprob_tensor from exact equality checks. Compare dprob_tensor
separately with np.testing.assert_allclose using the tolerance pattern
established by the sibling grouped dSReLU and dSWiGLU tests, while retaining
exact equality for d_row_tensor and dbias_tensor.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py`:
- Around line 23-35: Rename the ambiguous l parameter to experts in both
input-builder functions, update all dependent shape, range, and array
expressions, and adjust every call site accordingly while preserving the
existing expert-count behavior.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py`:
- Around line 27-28: Replace the negative nested ceiling-division used for
rest_k with positive ceiling division in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py:27-28 and
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py:27-28, using (k +
127) // 128; update
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py:27-29 to use (k
+ 63) // 64. Keep the existing sfa_j shape construction unchanged.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py`:
- Around line 13-25: Declare "gemm_dsrelu_jax_sm100" directly in the __all__
literal and remove the separate __all__.append call in
python/cudnn/gemm/cutedsl/dense/dsrelu/__init__.py (lines 13-25). Apply the same
change for "gemm_srelu_jax_sm100" in
python/cudnn/gemm/cutedsl/dense/srelu/__init__.py (lines 13-25), leaving each
module’s __getattr__ lazy import behavior unchanged.
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Around line 27-28: Bound the module-level _srelu_kernel_cache and
_dsrelu_kernel_cache so compiled kernels do not accumulate indefinitely across
shape and configuration variants. Use an LRU or equivalent bounded policy,
applying it consistently to both caches while preserving their existing lookup
and reuse behavior.
- Around line 30-32: Move the shared SReLU and DReLU epilogue definitions out of
the JAX API and class API implementations into a common module, then import and
reuse them in GemmSreluSm100.compile(), GemmDsreluSm100.compile(), and the JAX
API. Remove the duplicated local lambdas while preserving their current
behavior.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py`:
- Around line 153-164: Add a concise comment in the quantized branch immediately
before the _swiglu_quant_adapter call explaining that outputs are intentionally
reversed because the kernel signature receives c before ab12. Preserve the
existing output_shape_dtype ordering and c_tensor, ab12_tensor unpacking.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Around line 133-135: Remove the redundant self._interpret_uint8_as_fp4x2 =
True assignment later in the initialization flow, while retaining the assignment
before descriptor creation. Keep the earlier assignment as the sole
initialization so the required ordering remains explicit.
- Around line 1073-1095: Add the framework identifier to the cache_key assembled
for DiscreteGroupedGemmDswigluSm100, using the same framework value used by the
cached object's self._framework and workspace allocations. Keep the existing key
components unchanged and ensure framework-specific torch and JAX entries remain
distinct independently of tensor signature differences.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py`:
- Around line 187-193: Update the framework validation in
GroupedGemmDgluBlockScaledAPI to store the result of detect_framework(sample_a)
and use that detected framework name in the ValueError message, while preserving
the existing torch-only rejection behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py`:
- Around line 1493-1494: Update the framework-specific import handling in the
containing API function: keep torch resolution within the torch branch and
jax.numpy as jnp within the JAX branch, or centralize both through a helper that
explicitly reports unavailable dependencies. Ensure each framework path binds
only its required module and missing dependencies fail with a clear ImportError
rather than an unbound local name.
In `@python/cudnn/gemm/cutedsl/grouped/srelu/api.py`:
- Around line 942-949: Initialize self._compile_b_ptrs and
self._compile_sfb_ptrs to None in the class __init__ method, matching the
sibling grouped API patterns. Preserve the existing assignments in the
compilation path so dense weight mode and keep-alive pointer handling continue
to work.
In `@python/cudnn/gemm/cutedsl/grouped/swiglu/api.py`:
- Around line 33-37: Move the shared _JAX_SF_LAYOUT_ERROR constant into
python/cudnn/gemm/cutedsl/grouped/backend_utils.py, then remove the local
definitions and import it in both
python/cudnn/gemm/cutedsl/grouped/swiglu/api.py lines 33-37 and
python/cudnn/gemm/cutedsl/grouped/srelu/api.py lines 41-45. Preserve the
diagnostic text unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/api.py`:
- Around line 192-244: Update the dtype validation messages in the visible
validation flow to use framework-neutral names: replace torch.bfloat16 with
bfloat16 and torch.float32 with float32 in the a_tensor, prob_tensor, b_tensor,
and b_dtype errors. Keep the validation logic and exception behavior unchanged.
In `@python/cudnn/jax/call.py`:
- Around line 37-44: Update the docstring for row_major_desc to state that its
fixed Device("cuda", 0) value causes device comparisons in check_support to
check against GPU 0, and clarify that this descriptor device is not
authoritative.
In `@python/cudnn/tensor_adapter.py`:
- Around line 180-182: Update the get_version docstring to explicitly state that
the 0 fallback is valid only for immutable arrays such as JAX, and must not be
used for mutable frameworks lacking _version because cached validation would not
invalidate.
In `@test/python/fe_api/gemm/test_gemm_amax_jax.py`:
- Around line 247-250: Standardize optional CuTeDSL JAX gating across the three
dense-GEMM JIT tests: in test/python/fe_api/gemm/test_gemm_amax_jax.py lines
247-250 and test/python/fe_api/gemm/test_gemm_swiglu_jax.py lines 109-112,
replace the direct import with pytest.importorskip("cutlass.jax"), then use the
returned cutlass_jax symbol for is_available(); in
test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py lines 143-147, add the
same importorskip and availability skip before skip_unless_sm100().
In `@test/python/fe_api/gemm/test_gemm_swiglu_jax.py`:
- Around line 141-150: Repeat the quantized adapter invocation in the test after
the initial quant(a2, b2, sfa, sfb) call, preserving the existing
block_until_ready synchronization and assertions so the donated pre-initialized
outputs are exercised for donation safety.
🪄 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: 89e08d16-dfb9-4028-92e1-1c6fb0c670d0
📒 Files selected for processing (74)
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.mddocs/python_operation_package_structure_v1.mdpyproject.tomlpython/cudnn/__init__.pypython/cudnn/gemm/cutedsl/_jax_ffi.pypython/cudnn/gemm/cutedsl/dense/amax/__init__.pypython/cudnn/gemm/cutedsl/dense/amax/jax_api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/__init__.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pypython/cudnn/gemm/cutedsl/dense/srelu/__init__.pypython/cudnn/gemm/cutedsl/dense/srelu/jax_api.pypython/cudnn/gemm/cutedsl/dense/swiglu/__init__.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/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/wgrad/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/api.pypython/cudnn/jax/__init__.pypython/cudnn/jax/call.pypython/cudnn/tensor_adapter.pytest/python/conftest.pytest/python/fe_api/gemm/test_cutedsl_jax_guards.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)
- test/python/fe_api/gemm/test_cutedsl_jax_guards.py
- python/cudnn/gemm/cutedsl/_jax_ffi.py
| # The kernel writes into these buffers on the launch stream; make sure XLA has | ||
| # finished materializing them before the kernel runs. | ||
| jax.block_until_ready((out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
file="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py"
printf '%s\n' '--- file outline ---'
ast-grep outline "$file" --lang python || true
printf '%s\n' '--- imports and relevant symbols ---'
rg -n "jax|block_until_ready|execute|stream|tvm|ffi|out_fp8|def __call__|class " "$file"
printf '%s\n' '--- surrounding implementation ---'
sed -n '1,220p' "$file"
sed -n '560,660p' "$file"
printf '%s\n' '--- related launch and synchronization patterns ---'
rg -n -C 4 "block_until_ready|cudaStream|stream_ptr|tvm_ffi|launch|current_stream" python/cudnn test/python docs | head -n 400Repository: NVIDIA/cudnn-frontend
Length of output: 2032
🌐 Web query:
NVIDIA cudnn-frontend python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py jax.block_until_ready TVM FFI stream
💡 Result:
The path python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py in the NVIDIA cudnn-frontend repository contains the implementation for the GEMM + RoPE + MXFP8 Projection operation, which targets NVIDIA Hopper (SM90) and Blackwell (SM100) GPU architectures [1][2][3]. This operation is part of the cuDNN Frontend Open Source Software (FE-OSS) APIs [4][5]. The cutedsl path indicates that the kernel is implemented using NVIDIA's CuTe DSL, a domain-specific language for defining high-performance GPU kernels [6]. Key technical aspects include: - Optional Dependency: The cutedsl functionality requires the nvidia-cudnn-frontend[cutedsl] optional dependency, which includes apache-tvm-ffi among other packages [4][7]. - API Structure: The api.py file typically provides an APIBase class that handles the kernel lifecycle, including check_support(), compile(), and execute() [2][4]. - Execution: The execute method requires a current_stream parameter (often passed as None to use the default stream), which allows the integration to specify the CUDA stream on which the kernel should run [2][4]. - JAX/TVM Integration: The mention of jax.block_until_ready and TVM FFI in your query relates to the underlying infrastructure used to invoke these kernels from Python environments like JAX or PyTorch. Specifically, apache-tvm-ffi (TVM Foreign Function Interface) acts as the bridge that enables the Python frontend to interact with the compiled kernel binaries and manage asynchronous execution, where block_until_ready() is a standard JAX/NumPy pattern used to synchronize the host with device-side operations [7].
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.html
- 2: https://docs.nvidia.com/deeplearning/cudnn/v1.27.0/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.html
- 3: https://github.com/NVIDIA/cudnn-frontend
- 4: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/fe-oss-apis/overview.html
- 5: https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/overview.html
- 6: https://github.com/NVIDIA/cudnn-frontend/blob/develop/skills/cutedsl-kernel-integration/SKILL.md
- 7: https://pypi.org/project/nvidia-cudnn-frontend/
🌐 Web query:
site:github.com/NVIDIA/cudnn-frontend "jax.block_until_ready((out_fp8_row"
💡 Result:
The code snippet jax.block_until_ready(out_fp8_row) is a standard JAX operation used to ensure that a computation is complete before proceeding [1]. In this context: 1. jax.block_until_ready() is a method called on a JAX array or computation result to force the host to wait until the asynchronous device execution (typically on a GPU or TPU) is finished [1]. Because JAX uses asynchronous execution by default to improve performance, the CPU returns control immediately while the GPU continues computing. block_until_ready() is essential for accurate profiling, debugging, or synchronizing with non-JAX code. 2. out_fp8_row likely refers to a tensor output containing FP8 (8-bit floating point) data, often encountered in high-performance machine learning workflows involving NVIDIA's cuDNN or Transformer Engine [2][3][4]. In advanced deep learning libraries, "FP8 row" typically suggests a quantization or normalization output where data is scaled or stored with per-row granularity [2][4]. This pattern is common when debugging or timing kernels, such as those involving FP8 operations or cuDNN graph executions, where the user needs to ensure that the asynchronous compute task is fully finalized before reading the resulting data [1][2].
Citations:
- 1: Tracking: cuDNN Python SDPA API Integration #309
- 2: https://github.com/NVIDIA/cudnn-frontend/blob/develop/samples/cpp/norm/norm_block_scale.cpp
- 3: https://github.com/NVIDIA/cudnn-frontend/blob/develop/samples/cpp/sdpa/fp8_fwd.cpp
- 4: Add row-scale support to grouped GEMM quant #251
🌐 Web query:
"proj_rope_mxfp8" "tvm_ffi" JAX stream
💡 Result:
The components in your query relate to high-performance deep learning infrastructure for modern Transformer architectures. proj_rope_mxfp8: This refers to a fused kernel operation, specifically the Fused Projection GEMM + RoPE + MXFP8 quantization [1]. This kernel is designed for NVIDIA Blackwell (SM100+) GPUs and is used to optimize attention mechanisms, such as the Megatron MLA-YARN rotary embedding [1]. It projects activations, applies rotary embeddings to the relevant head dimensions, and performs MXFP8 (E4M3, block=32) quantization on the output [1]. The kernel implementation uses CUTLASS/CuTe and is often associated with libraries like cuDNN [1]. tvm_ffi: This is the TVM Foreign Function Interface (FFI) [2]. It provides a standardized ABI and Python package that allows deep learning frameworks (such as JAX, PyTorch, and others) to call into low-level, high-performance kernels (often written in C++, CUDA, or TVM’s TIR) [3][2][4]. It facilitates the integration of custom GPU kernels with framework tensors by handling stream management and memory layout [5][2]. JAX stream: In the context of TVM FFI and JAX, managing the execution stream is critical for performance [5]. JAX provides mechanisms to expose its compute stream to FFI-called kernels, ensuring that the TVM-compiled operations execute correctly on the same stream as the producer framework [3][2]. The tvm_ffi library, including tools like use_raw_stream or use_torch_stream, provides the infrastructure to bridge these environments, allowing developers to register and invoke custom FFI targets within JAX's JIT-compiled pipelines [5][6][3][4]. In summary, these components represent a pipeline for deploying high-performance, quantized Transformer layers (like MLA) on next-generation hardware, using JAX as the frontend framework and TVM FFI as the bridge to execute fused, optimized kernels., top_results=[1][7][3][2][4]
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/v1.27.0/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.html
- 2: https://tvm.apache.org/ffi/guides/python_lang_guide.html
- 3: https://tvm.apache.org/ffi/get_started/quickstart
- 4: https://github.com/nvidia/jax-tvm-ffi
- 5: https://tvm.apache.org/ffi/_modules/tvm_ffi/stream.html
- 6: https://tvm.apache.org/ffi/reference/python/generated/tvm_ffi.use_raw_stream.html
- 7: https://docs.flashinfer.ai/generated/flashinfer.rope.mla_rope_quantize_fp8.html
🌐 Web query:
site:tvm.apache.org/ffi JAX use_raw_stream current CUDA stream TVM FFI
💡 Result:
In the context of TVM-FFI, tvm_ffi.use_raw_stream is a utility used to explicitly set the CUDA stream for a specific device, providing a way to communicate which stream a kernel should execute on [1][2]. This is particularly useful when the automatic stream detection mechanisms—which typically capture streams from frameworks like PyTorch or through DLPack-compliant tensors—are insufficient for your specific workflow [3][2]. For JAX integration, the jax-tvm-ffi library facilitates the registration of custom kernels for use with JAX's ffi_call [4][5]. When using these FFI targets, the stream context ensures that kernels launched via TVM-FFI are scheduled on the appropriate CUDA stream [6][3]. Key technical points include: 1. Usage: tvm_ffi.use_raw_stream(device, stream) creates a context manager where stream is an integer handle (e.g., a cudaStream_t value) for the specified device [1][2]. 2. Stream Capture: TVM-FFI automatically captures the current stream for a device when using frameworks like PyTorch or when converting DLPack tensors [6][3]. If you are working with manual CUDA streams that fall outside of this automatic tracking, use_raw_stream allows you to override or set the stream manually [3][2]. 3. Kernel Execution: Within C++ or CUDA kernel code, the stream is retrieved using the TVMFFIEnvGetStream() API [6][3]. This function reads the current FFI stream context established in the Python environment, allowing the kernel to launch execution on that specific stream [6][4]. 4. Integration Pattern: For JAX, typically you define your kernel, load it as a tvm_ffi.Module, and register it via jax_tvm_ffi.register_ffi_target [4]. When calling these kernels in a JAX program, the system relies on the established FFI stream environment to handle the underlying CUDA stream execution correctly [6][3]. If you are working with higher-level framework objects like PyTorch streams or CUDA graphs, use tvm_ffi.use_torch_stream instead, as it handles the extraction of the underlying stream information automatically [2][7].
Citations:
- 1: https://tvm.apache.org/ffi/reference/python/generated/tvm_ffi.use_raw_stream.html
- 2: https://tvm.apache.org/ffi/_modules/tvm_ffi/stream.html
- 3: https://tvm.apache.org/ffi/concepts/tensor.html
- 4: https://tvm.apache.org/ffi/get_started/quickstart
- 5: https://tvm.apache.org/ffi/get_started/quickstart.html
- 6: https://tvm.apache.org/ffi/guides/kernel_library_guide.html
- 7: https://tvm.apache.org/ffi/reference/python/generated/tvm_ffi.use_torch_stream.html
🌐 Web query:
github NVIDIA cudnn-frontend proj_rope_mxfp8 api.py "block_until_ready" "execute("
💡 Result:
The query refers to components within the NVIDIA cuDNN Frontend library, which provides a Python interface for high-performance open-source kernels [1][2]. proj_rope_mxfp8: This refers to the GEMM + RoPE + MXFP8 Projection kernel, an experimental operator designed for high-performance training, such as the DeepSeek-V3 MLA (Multi-Head Latent Attention) Q up-projection [3][4]. It fuses a matrix multiplication (GEMM), Rotary Position Embedding (RoPE), and dual-direction MXFP8 quantization [3][4]. It is available in two input-precision variants: - GemmProjRopeMxfp8Bf16InSm100: For BF16 inputs [3]. - GemmProjRopeMxfp8Mxfp8InSm100: For pre-quantized MXFP8 inputs [3]. execute(...): In the cuDNN Frontend Python class API, the execute method is used to run a compiled operation on a given set of inputs and outputs [5]. When using this API, developers first define the operation with sample tensors, call compile to build the kernel, and then call execute to perform the actual computation on a specified stream [5]. block_until_ready: While not a core member of the cuDNN Frontend API itself, this is a standard method in asynchronous computing frameworks like JAX [6]. In the context of GPU programming and cuDNN, it is frequently used to force host-side synchronization by blocking execution until an asynchronously dispatched GPU kernel has finished computing its results [7][8][6]. It is commonly used when benchmarking or debugging performance, as it ensures that the wall-clock time measured includes the actual GPU computation rather than just the time taken to enqueue the task [7][9].
Citations:
- 1: https://github.com/NVIDIA/cudnn-frontend
- 2: https://github.com/NVIDIA/cudnn-frontend/blob/develop/README.md
- 3: https://docs.nvidia.com/deeplearning/cudnn/v1.27.0/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.html
- 4: https://github.com/NVIDIA/cudnn-frontend/releases
- 5: https://docs.nvidia.com/deeplearning/cudnn/frontend/v1.23.0/fe-oss-apis/overview.html
- 6: https://d2l.smola.org/chapter_computational-performance/async-computation.html
- 7: Document async dispatch in the README jax-ml/jax#1219
- 8: https://stackoverflow.com/questions/79816992/jax-crashes-with-cudnn-status-internal-error-when-using-joblib-or-multiproc
- 9: Asynchronous Dispatch in JAX jax-ml/jax#9895
🌐 Web query:
"legacy default stream" "proj_rope_mxfp8" cuDNN frontend
💡 Result:
In the context of the NVIDIA cuDNN Frontend (cuDNN FE), specifically regarding APIs like the gemm_proj_rope_mxfp8 kernel, the interaction with CUDA streams is a critical performance and functional consideration [1][2]. Legacy Default Stream The legacy default stream is a specific implicit CUDA stream that synchronizes with all other streams in the same CUDA context (with the exception of non-blocking streams) [3][4]. When an operation is enqueued in the legacy default stream, it waits for all blocking streams to complete, and subsequently, all blocking streams wait for the legacy default stream to finish [4][5]. This synchronization behavior is often the cause of unintended serialization in high-performance GPU applications [3]. cuDNN Frontend Stream Management The NVIDIA cuDNN Frontend provides mechanisms to control the CUDA stream used for kernel execution in operations like the MXFP8 RoPE projection [1][6]. 1. Default Behavior: If no specific stream is provided to the cuDNN FE wrapper or class API (often by passing None), the implementation defaults to using the current CUDA stream [1][2]. 2. Torch Integration: When integrated with PyTorch, it is standard practice to pass torch.cuda.current_stream(device).cuda_stream to these APIs [7]. This ensures the operation is scheduled on the active PyTorch stream, enabling proper synchronization with other PyTorch operations and avoiding the heavy, implicit synchronization overhead associated with the legacy default stream [7]. 3. CUDA Graph Capture: Explicitly setting the current Torch stream is a requirement for successful CUDA-graph capture, as it allows the graph to correctly encapsulate the execution flow without falling back to implicit legacy stream behaviors [1]. In summary, users should avoid relying on the implicit legacy default stream when using cuDNN FE high-performance kernels (such as gemm_proj_rope_mxfp8) to maintain optimal concurrency and support advanced features like CUDA-graph capture [1][3]. Instead, explicitly pass the active CUDA stream to the API's stream or current_stream parameter [1][2].
Citations:
- 1: https://docs.nvidia.com/deeplearning/cudnn/v1.27.0/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.html
- 2: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/fe-oss-apis/overview.html
- 3: https://docs.nvidia.com/cuda/archive/11.4.0/cuda-runtime-api/stream-sync-behavior.html
- 4: https://docs.nvidia.com/cuda/cuda-driver-api/stream-sync-behavior.html
- 5: https://docs.nvidia.com/cuda/archive/12.9.0/cuda-runtime-api/stream-sync-behavior.html
- 6: https://docs.nvidia.com/deeplearning/cudnn/v1.27.0/fe-oss-apis/overview.html
- 7: https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/utilities/adding_torch_custom_ops.html
🌐 Web query:
JAX external CUDA kernel launch stream ordering returned buffers block_until_ready custom call FFI synchronization
💡 Result:
In JAX, when using the Foreign Function Interface (FFI) for external CUDA kernels, stream ordering and synchronization are primarily managed by binding the XLA platform stream to your FFI handler [1][2]. 1. Stream Ordering and Integration: When you define an FFI custom call in JAX, you must use the ffi::PlatformStream<cudaStream_t> type in your handler's signature [3][1][2]. This ensures that XLA provides the current CUDA stream being used for the computation graph. By using this provided cudaStream_t to launch your CUDA kernel (e.g., my_kernel<<<..., stream>>>(...)), you ensure that your kernel is correctly ordered within the XLA execution sequence [3][1]. Because JAX and XLA enqueue operations asynchronously, using the provided stream is essential for maintaining correct dependency ordering without requiring explicit CPU-side synchronization [1][4]. 2. Synchronization and block_until_ready(): JAX operations, including FFI custom calls, are asynchronous by default [5]. The jax.block_until_ready() method is used to wait for the completion of all asynchronous operations, including the kernel launched by your FFI custom call [5][6]. It forces the CPU to wait until the GPU has finished the execution of the command queue up to that point [4][5]. 3. Handling Buffers and Data Dependency: FFI handlers receive device pointers (via ffi::Buffer or ffi::ResultBuffer) that point to memory already allocated and managed by JAX [1][2]. Since the kernel is launched on the provided stream, XLA automatically manages the dependencies between these buffers and other JAX operations [1]. If your custom C++/CUDA code performs its own allocations or interacts with external streams, you must manually synchronize these using CUDA events (cudaEventRecord/cudaStreamWaitEvent) to avoid race conditions [7][4]. For simple kernel launches where the inputs/outputs are provided as arguments, utilizing the passed cudaStream_t is sufficient to ensure correctness within the JAX graph [3][1]. 4. Performance Considerations: Frequent use of block_until_ready() introduces CPU-GPU synchronization points that can significantly degrade performance by stalling the GPU pipeline [4]. In production workloads, rely on JAX's asynchronous execution model and allow XLA to manage stream ordering through its internal graph scheduling [1][5]. If custom kernels are not being integrated into the main XLA command buffer (which can impact performance), ensure you are using modern FFI patterns and consider whether your operations can be JIT-compiled through JAX's native mechanisms [2][8][9].
Citations:
- 1: https://openxla.org/xla/custom_call
- 2: https://docs.jax.dev/en/latest/ffi.html
- 3: https://github.com/jax-ml/jax/blob/main/examples/ffi/src/jax_ffi_example/cuda_examples.cu
- 4: pytorch interop without cuda synchronization jax-ml/jax#3699
- 5: https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html
- 6: https://docs.jax.dev/en/latest/_autosummary/jax.block_until_ready.html
- 7: Multiple invocation of a jax FFI binded function leads to incorrect results jax-ml/jax#28005
- 8: Update docs/cuda_custom_call (P1) jax-ml/jax#24341
- 9: [Pallas] command buffer is not created for __gpu$xla.gpu.triton by default jax-ml/jax#27988
Establish JAX stream ordering around the raw kernel launch.
jax.block_until_ready() only waits for the output buffers before launch. It does not order pending JAX producers for x, w, cos, sin, or scale tensors with the raw TVM-FFI launch, and later JAX consumers have no dependency on the external write.
Pass the JAX/XLA stream to the launch through the supported TVM-FFI handoff. If no handoff is available, synchronize all inputs before launch and synchronize the launch before returning the outputs.
🤖 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 615 -
617, Update the raw kernel launch flow around the existing jax.block_until_ready
call to establish stream ordering for all JAX input tensors and output
consumers. Use the supported TVM-FFI mechanism to pass the JAX/XLA stream into
the launch; if unavailable, synchronize every producer input (x, w, cos, sin,
and scale tensors) before launching and synchronize the kernel completion before
returning outputs, while preserving the existing output buffers.
| sf_vec_size=sf_vec_size, | ||
| vector_f32=vector_f32, | ||
| ) | ||
| assert gemm.check_support() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Both JAX entry points validate support with a bare assert. Python removes assert statements under -O. In that mode check_support() never runs, so no dtype, shape, stride, or architecture validation happens before kernel construction.
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py#L151-L151: replaceassert gemm.check_support()with an explicitif not gemm.check_support(): raise ValueError(...)forGemmSreluSm100.python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py#L251-L251: apply the same replacement forGemmDsreluSm100.
📍 Affects 1 file
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py#L151-L151(this comment)python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py#L251-L251
🤖 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/srelu/jax_api.py` at line 151, Replace the
bare support assertions in both GemmSreluSm100 and GemmDsreluSm100 with explicit
checks that call check_support() and raise ValueError when validation fails,
ensuring validation runs even under Python -O. Apply this at
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py lines 151-151 and 251-251.
| c_jax_dtype = framework_dtype(c_dtype, "jax") | ||
| ab12_buf = jnp.zeros((m, n, l), dtype=ab12_jax_dtype) | ||
| c_buf = jnp.zeros((m, n // 2, l), dtype=c_jax_dtype) | ||
| assert gemm.check_support() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Replace the assert with an explicit error.
Python removes assert statements when the interpreter runs with -O. In that mode check_support() never runs, and an unsupported configuration reaches kernel construction with a lower-level failure. The sibling wrapper in this PR uses an explicit raise; see python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py lines 1146-1147.
🛠️ Proposed fix
- assert gemm.check_support()
+ if not gemm.check_support():
+ raise RuntimeError("Unsupported gemm_swiglu_jax_sm100 configuration")📝 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.
| assert gemm.check_support() | |
| if not gemm.check_support(): | |
| raise RuntimeError("Unsupported gemm_swiglu_jax_sm100 configuration") |
🤖 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` at line 113, Replace the
assert around gemm.check_support() with an explicit support check that always
executes, raising the same appropriate error style used by the sibling
discrete_grouped wrapper when the configuration is unsupported.
| def wrapper(*arrays: Any) -> Any: | ||
| inits = [] | ||
| aliases = dict(input_output_aliases) | ||
| extra_specs = [] | ||
| for offset, (out_index, init_fn) in enumerate(sorted(initialized_outputs.items())): | ||
| inits.append(init_fn(output_leaves[out_index])) | ||
| aliases[len(arrays) + offset] = out_index | ||
| extra_specs.append(output_spec[out_index] if output_spec is not None else None) | ||
|
|
||
| full_input_spec = input_spec | ||
| if inits and input_spec is not None: | ||
| full_input_spec = tuple(input_spec) + tuple(extra_specs) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Apply the initializer specs even when input_spec is None.
extra_specs carries the output_spec entry for each initialized output. Line 118 discards it whenever the caller passes output_spec without input_spec. The initializer buffers are donated and aliased to those outputs, so the same memory is then described by the output spec on one side and by default inference on the other. For a non-trivial spec such as sf_atom_spec(), which remaps dimensions with mode, the two views disagree.
Build the full input spec whenever initializers exist, padding the original inputs with None.
🐛 Proposed fix
full_input_spec = input_spec
- if inits and input_spec is not None:
- full_input_spec = tuple(input_spec) + tuple(extra_specs)
+ if inits:
+ base_spec = tuple(input_spec) if input_spec is not None else (None,) * len(arrays)
+ full_input_spec = base_spec + tuple(extra_specs)📝 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.
| def wrapper(*arrays: Any) -> Any: | |
| inits = [] | |
| aliases = dict(input_output_aliases) | |
| extra_specs = [] | |
| for offset, (out_index, init_fn) in enumerate(sorted(initialized_outputs.items())): | |
| inits.append(init_fn(output_leaves[out_index])) | |
| aliases[len(arrays) + offset] = out_index | |
| extra_specs.append(output_spec[out_index] if output_spec is not None else None) | |
| full_input_spec = input_spec | |
| if inits and input_spec is not None: | |
| full_input_spec = tuple(input_spec) + tuple(extra_specs) | |
| def wrapper(*arrays: Any) -> Any: | |
| inits = [] | |
| aliases = dict(input_output_aliases) | |
| extra_specs = [] | |
| for offset, (out_index, init_fn) in enumerate(sorted(initialized_outputs.items())): | |
| inits.append(init_fn(output_leaves[out_index])) | |
| aliases[len(arrays) + offset] = out_index | |
| extra_specs.append(output_spec[out_index] if output_spec is not None else None) | |
| full_input_spec = input_spec | |
| if inits: | |
| base_spec = tuple(input_spec) if input_spec is not None else (None,) * len(arrays) | |
| full_input_spec = base_spec + tuple(extra_specs) |
🤖 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/jax/call.py` around lines 108 - 119, Update wrapper around
full_input_spec so initialized outputs always contribute their extra_specs, even
when input_spec is None. When inits exist, construct the full input
specification by padding the original inputs with None entries before appending
extra_specs, preserving the existing combined-spec behavior when input_spec is
provided.
| def test_gemm_srelu_dsrelu_jax_jit_matches_eager(): | ||
| """The XLA custom-call entry points must agree bit-for-bit with the eager wrappers.""" | ||
| skip_unless_sm100() | ||
| import cudnn | ||
| from cudnn import gemm_dsrelu_jax_sm100, gemm_srelu_jax_sm100 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the CuTeDSL JAX availability guard used by the other JIT tests.
This test exercises the same XLA custom-call bridge as test_gemm_amax_jax_ffi_sm100 and test_gemm_swiglu_jax_ffi_sm100. Those tests skip when cutlass.jax.is_available() returns False. This test has no such guard, so it fails instead of skipping on an environment without the CuTeDSL JAX extensions.
♻️ Proposed change
`@pytest.mark.L0`
def test_gemm_srelu_dsrelu_jax_jit_matches_eager():
"""The XLA custom-call entry points must agree bit-for-bit with the eager wrappers."""
+ cutlass_jax = pytest.importorskip("cutlass.jax")
+ if not cutlass_jax.is_available():
+ pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)")
skip_unless_sm100()As per coding guidelines: "Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks".
📝 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.
| def test_gemm_srelu_dsrelu_jax_jit_matches_eager(): | |
| """The XLA custom-call entry points must agree bit-for-bit with the eager wrappers.""" | |
| skip_unless_sm100() | |
| import cudnn | |
| from cudnn import gemm_dsrelu_jax_sm100, gemm_srelu_jax_sm100 | |
| def test_gemm_srelu_dsrelu_jax_jit_matches_eager(): | |
| """The XLA custom-call entry points must agree bit-for-bit with the eager wrappers.""" | |
| cutlass_jax = pytest.importorskip("cutlass.jax") | |
| if not cutlass_jax.is_available(): | |
| pytest.skip("CuTeDSL JAX extensions unavailable (jax >= 0.5 required)") | |
| skip_unless_sm100() | |
| import cudnn | |
| from cudnn import gemm_dsrelu_jax_sm100, gemm_srelu_jax_sm100 |
🤖 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 143 -
147, Add the CuTeDSL JAX availability guard to
test_gemm_srelu_dsrelu_jax_jit_matches_eager, matching the existing checks in
test_gemm_amax_jax_ffi_sm100 and test_gemm_swiglu_jax_ffi_sm100. Skip the test
when cutlass.jax.is_available() is false, before importing or invoking the JAX
custom-call wrappers.
Source: Coding guidelines
| def _make_jax_inputs(m=256, n=128, k=128, l=2): | ||
| rng = np.random.default_rng(20260809) | ||
| a_j = jnp.asarray((rng.integers(-4, 5, size=(m, k, 1)) * 0.25).astype(ml_dtypes.float8_e4m3fn)) | ||
| b_j = jnp.asarray((rng.integers(-4, 5, size=(n, k, l)) * 0.25).astype(ml_dtypes.float8_e4m3fn)) | ||
| c_j = jnp.asarray(rng.standard_normal((m, n * 2, 1), dtype=np.float32).astype(ml_dtypes.bfloat16)) | ||
| rk = ((k + 31) // 32 + 3) // 4 | ||
| sfa_j = jnp.asarray((2.0 ** rng.integers(-2, 3, size=(1, (m + 127) // 128, rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu)) | ||
| sfb_j = jnp.asarray((2.0 ** rng.integers(-2, 3, size=(l, (n + 127) // 128, rk, 32, 4, 4))).astype(ml_dtypes.float8_e8m0fnu)) | ||
| offsets_j = jnp.asarray(np.arange(m // l, m + 1, m // l, dtype=np.int32)) | ||
| alpha_j = jnp.asarray(np.ones(l, dtype=np.float32)) | ||
| beta_j = jnp.asarray(np.ones(l, dtype=np.float32)) | ||
| prob_j = jnp.asarray(np.ones((m, 1, 1), dtype=np.float32)) | ||
| return a_j, b_j, c_j, sfa_j, sfb_j, offsets_j, alpha_j, beta_j, prob_j |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous l variable.
Ruff reports E741 at Lines 23 and 91. Rename l to experts in both input builders and their call sites.
Also applies to: 91-104
🧰 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` around
lines 23 - 35, Rename the ambiguous l parameter to experts in both input-builder
functions, update all dependent shape, range, and array expressions, and adjust
every call site accordingly while preserving the existing expert-count behavior.
Source: Linters/SAST tools
| rest_k = -(-(-(-k // 32) // 4)) # ceil_div(ceil_div(k, 32), 4) | ||
| sfa_j = jnp.asarray(rng.integers(0, 127, (32, 4, -(-m // 128), 4, rest_k, 1), dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Compute rest_k with a positive ceiling division.
The nested unary-minus expressions create negative shape dimensions. For k=128, Quant and SReLU produce -1; SwiGLU produces -2. rng.integers fails before these tests call the wrapper or API class.
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L27-L28: replace the expression withrest_k = (k + 127) // 128.test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L27-L28: replace the expression withrest_k = (k + 127) // 128.test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L27-L29: replace the expression withrest_k = (k + 63) // 64.
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 27-27: Python does not support the unary prefix decrement operator (--)
(B002)
📍 Affects 3 files
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L27-L28(this comment)test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L27-L28test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L27-L29
🤖 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
27 - 28, Replace the negative nested ceiling-division used for rest_k with
positive ceiling division in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py:27-28 and
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py:27-28, using (k +
127) // 128; update
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py:27-29 to use (k
+ 63) // 64. Keep the existing sfa_j shape construction unchanged.
Source: Linters/SAST tools
f2165f7 to
64ad2f2
Compare
Anerudhan
changed the title
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py (1)
37-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winImport
_prob_specinstead of redefining it.Line 37 already imports
_pointer_countfrom..unfused.jax_api. The local_prob_specat Lines 49-52 duplicates..unfused.jax_api._prob_specexactly.python/cudnn/gemm/cutedsl/grouped/glu/jax_api.pyandpython/cudnn/gemm/cutedsl/grouped/dglu/jax_api.pyimport that helper. Use the same import here so the layout contract has one definition.♻️ Proposed change
-from ..unfused.jax_api import _pointer_count +from ..unfused.jax_api import _pointer_count, _prob_spec from .moe_blockscaled_grouped_gemm_dsrelu_quant import BlockScaledMoEGroupedGemmQuantBwdKernel, EpilogueType @@ -def _prob_spec() -> TensorSpec: - # (m, 1, 1) with m innermost: explicit ranks because trailing unit dims make - # leading-dim inference ambiguous - return TensorSpec(layout=(0, 1, 2)) - -🤖 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/dsrelu/jax_api.py` around lines 37 - 52, Update the import from ..unfused.jax_api to include _prob_spec, then remove the local _prob_spec definition so this module reuses the shared helper and its single layout contract.python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py (1)
52-101: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider merging the two adapters into one with a
has_dbiasconstexpr.
_grouped_dglu_bf16_adapterand_grouped_dglu_bf16_dbias_adapterdiffer only in thedbiasparameter and thedbias_tensor=argument.python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.pyhandles the analogous optionalamaxoutput with a single adapter andcutlass.const_expr. A single adapter here would remove the duplicated kernel-argument list. Note that the optional tensor must stay a trailing positional so the output ordering still matches.🤖 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/dglu/jax_api.py` around lines 52 - 101, Merge _grouped_dglu_bf16_adapter and _grouped_dglu_bf16_dbias_adapter into one adapter with a has_dbias constexpr, selecting dbias_tensor through cutlass.const_expr while preserving the shared kernel arguments. Keep the optional dbias tensor trailing in the adapter’s positional arguments so output ordering remains unchanged, and pass None when has_dbias is false.python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py (1)
313-318: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider extracting the repeated kernel-cache and
max_active_clustersboilerplate.The same block appears in
python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py(Lines 195-200),python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py(Lines 194-199),python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py(Lines 228-233),python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py(Lines 362-367), andpython/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.py(Lines 184-189). A small internal helper (for example_resolve_kernel(cache, cache_key, factory)) would remove six copies of theCUDNNFE_CLUSTER_OVERLAP_MARGINparsing and themac <= 0check. Keep the helper internal to the family package so it is not exported throughcudnn.As per coding guidelines: "Shared GEMM helpers, including schedulers and metadata utilities, must remain internal to their family package and must not be exported through
cudnn."🤖 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/jax_api.py` around lines 313 - 318, The repeated kernel-cache resolution and max_active_clusters logic should be consolidated into one internal helper shared by the six grouped GEMM JAX APIs. Extract the CUDNNFE_CLUSTER_OVERLAP_MARGIN parsing, cluster calculation, nonpositive validation, workspace normalization, and cache insertion into a helper such as _resolve_kernel(cache, cache_key, factory), then replace each duplicated block with calls to it. Keep the helper internal to the family package and do not export it through cudnn.Source: Coding guidelines
python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py (1)
204-224: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUnconditional C output allocation when
generate_cis False. Both entry points always declare a(m, n)C output and zero-fill it throughinitialized_outputs, then discard the result whengenerate_cis False. This costs one extra device allocation and fill per call. Buildoutput_shape_dtype/output_specconditionally, aspython/cudnn/gemm/cutedsl/grouped/dglu/jax_api.pydoes fordbias, if the kernel acceptsc=None.
python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py#L204-L224: append the C output and its spec only whengenerate_cis True, and passc=Nonein_grouped_glu_bf16_adapterotherwise.python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py#L204-L223: apply the same conditional output construction for_grouped_bf16_adapter.🤖 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/jax_api.py` around lines 204 - 224, In python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py lines 204-224, update the _grouped_glu_bf16_adapter call to construct the C output shape and output_spec only when generate_c is true, pass c=None otherwise, and omit C from initialized_outputs when disabled; preserve the existing D output and return behavior. Apply the same conditional output construction to _grouped_bf16_adapter in python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py lines 204-223.
🤖 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 `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py`:
- Around line 79-96: Update the cache key used by the bf16-in grid path to
include every descriptor-relevant shape and dtype for x, w, cos, sin, and all
output/scale tensors, preventing reuse across incompatible contracts. In the
cache-miss initialization around GemmProjRopeMxfp8Bf16InSm100, replace assert
obj.check_support() with an explicit runtime validation that raises an
appropriate error when unsupported, remaining active under -O.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py`:
- Around line 302-303: Update the dtype membership condition in the
discrete_col_sfd handling block to use the direct `d_dtype not in _fp8_dtypes`
form, preserving the existing behavior and assignment.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py`:
- Line 6: Sort the entries in the module-level __all__ list using isort-style
ordering, updating only the ordering of GroupedGemmSm100,
grouped_gemm_wrapper_sm100, and grouped_gemm_jax_sm100 so Ruff RUF022 passes.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py`:
- Around line 43-85: Move the shared _pointer_count and _prob_spec helpers out
of the grouped/unfused jax_api module into a neutral, unexported internal
CuTeDSL GEMM helper module that can serve both grouped and discrete_grouped
families. Update the imports in grouped/glu/jax_api.py, grouped/dglu/jax_api.py,
grouped/dsrelu/jax_api.py, grouped/wgrad/jax_api.py, and
discrete_grouped/swiglu/jax_api.py to use the new module, without exposing the
helpers through cudnn.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py`:
- Around line 155-181: Update the check function to compare dprob_tensor and
dbias_tensor using the established numerical tolerance rather than raw-byte
assert_array_equal, while retaining exact byte comparison for d_row_tensor.
Apply the same tolerant comparison to dprob_only in the generate_dbias=False
path, matching the sibling grouped GEMM tests’ existing tolerance settings.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.py`:
- Around line 313-318: The repeated kernel-cache resolution and
max_active_clusters logic should be consolidated into one internal helper shared
by the six grouped GEMM JAX APIs. Extract the CUDNNFE_CLUSTER_OVERLAP_MARGIN
parsing, cluster calculation, nonpositive validation, workspace normalization,
and cache insertion into a helper such as _resolve_kernel(cache, cache_key,
factory), then replace each duplicated block with calls to it. Keep the helper
internal to the family package and do not export it through cudnn.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py`:
- Around line 52-101: Merge _grouped_dglu_bf16_adapter and
_grouped_dglu_bf16_dbias_adapter into one adapter with a has_dbias constexpr,
selecting dbias_tensor through cutlass.const_expr while preserving the shared
kernel arguments. Keep the optional dbias tensor trailing in the adapter’s
positional arguments so output ordering remains unchanged, and pass None when
has_dbias is false.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py`:
- Around line 37-52: Update the import from ..unfused.jax_api to include
_prob_spec, then remove the local _prob_spec definition so this module reuses
the shared helper and its single layout contract.
In `@python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py`:
- Around line 204-224: In python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py lines
204-224, update the _grouped_glu_bf16_adapter call to construct the C output
shape and output_spec only when generate_c is true, pass c=None otherwise, and
omit C from initialized_outputs when disabled; preserve the existing D output
and return behavior. Apply the same conditional output construction to
_grouped_bf16_adapter in python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py
lines 204-223.
🪄 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: c6ce6585-a155-41ea-a874-1a187c40f265
📒 Files selected for processing (36)
docs/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_proj_rope_mxfp8.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_glu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__init__.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/__init__.pypython/cudnn/gemm/cutedsl/grouped/dglu/__init__.pypython/cudnn/gemm/cutedsl/grouped/dglu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/glu/__init__.pypython/cudnn/gemm/cutedsl/grouped/glu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/unfused/__init__.pypython/cudnn/gemm/cutedsl/grouped/unfused/jax_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/__init__.pypython/cudnn/gemm/cutedsl/grouped/wgrad/jax_api.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_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_glu_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_jax.pytest/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_jax.py
🚧 Files skipped from review as they are similar to previous changes (6)
- docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md
- docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.md
- docs/fe-oss-apis/gemm_fusions/grouped_gemm.md
- docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.md
- docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md
- docs/fe-oss-apis/overview.md
| cache_key = (tuple(x.shape), tuple(w.shape)) | ||
| entry = _bf16in_grid_cache.get(cache_key) | ||
| if entry is None: | ||
| obj = GemmProjRopeMxfp8Bf16InSm100( | ||
| sample_x=_make_desc(tuple(x.shape), x.dtype, "sample_x"), | ||
| sample_w=_make_desc(tuple(w.shape), w.dtype, "sample_w"), | ||
| sample_cos=_make_desc(tuple(cos.shape), cos.dtype, "sample_cos"), | ||
| sample_sin=_make_desc(tuple(sin.shape), sin.dtype, "sample_sin"), | ||
| sample_out_fp8_row=_make_desc(out_types[0].shape, cutlass.Float8E4M3FN, "sample_out_fp8_row"), | ||
| sample_out_scales_row=_make_desc(out_types[1].shape, cutlass.Uint8, "sample_out_scales_row"), | ||
| sample_out_fp8_col=_make_desc(out_types[2].shape, cutlass.Float8E4M3FN, "sample_out_fp8_col"), | ||
| sample_out_scales_col=_make_desc(out_types[3].shape, cutlass.Uint8, "sample_out_scales_col"), | ||
| w_out_in=True, | ||
| ) | ||
| assert obj.check_support() | ||
| mac = cutlass.utils.HardwareInfo().get_max_active_clusters(1) | ||
| entry = (tokens // TILE_M, num_heads, mac, 8) | ||
| _bf16in_grid_cache[cache_key] = entry |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py"
rg -n -C 6 'cache_key|_bf16in_grid_cache|_mxfp8in_grid_cache|assert obj\.check_support|obj\.check_support' "$target"
rg -n -C 5 'gemm_proj_rope_mxfp8_jax_sm100' \
test/python python/cudnnRepository: NVIDIA/cudnn-frontend
Length of output: 11052
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py"
api="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py"
test="test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py"
printf '%s\n' '--- jax_api.py outline and implementation ---'
ast-grep outline "$target"
sed -n '1,180p' "$target"
printf '%s\n' '--- API support checks ---'
ast-grep outline "$api"
rg -n -C 12 'class GemmProjRopeMxfp8|def check_support|sample_(x|w|cos|sin|x_scale|w_scale)|_grid_params' "$api"
printf '%s\n' '--- JAX tests and invalid-input coverage ---'
sed -n '1,240p' "$test"
rg -n -i -C 4 'invalid|dtype|shape|scale|cos|sin|raises|support' "$test" test/python/fe_api/gemmRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
log="/tmp/coderabbit-shell-logs/shell-output-H99VSA"
printf '%s\n' '--- jax_api.py section from captured output ---'
rg -n -m 1 'jax_api.py:48|jax_api.py-43' "$log" || true
line=$(rg -n -m 1 'jax_api.py:48|jax_api.py-43' "$log" | cut -d: -f1 || true)
if [ -n "${line:-}" ]; then
start=$((line-8)); end=$((line+115))
sed -n "${start},${end}p" "$log"
fi
printf '%s\n' '--- api.py support-check definitions from captured output ---'
rg -n -m 20 'api.py.*(def check_support|class GemmProjRopeMxfp8|sample_(x|w|cos|sin|x_scale|w_scale)|_grid_params)' "$log" || true
printf '%s\n' '--- dedicated test file section from repository ---'
sed -n '1,230p' test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.pyRepository: NVIDIA/cudnn-frontend
Length of output: 8156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
target="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py"
api="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py"
printf '%s\n' '--- jax_api.py ---'
sed -n '1,165p' "$target"
printf '%s\n' '--- api.py declarations and support methods ---'
rg -n '^(class |def | def | def )|check_support|_grid_params|sample_' "$api"
sed -n '1,260p' "$api"Repository: NVIDIA/cudnn-frontend
Length of output: 24433
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
jax_path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py")
api_path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
jax_tree = ast.parse(jax_path.read_text())
api_tree = ast.parse(api_path.read_text())
keys = []
assert_calls = []
for node in ast.walk(jax_tree):
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == "cache_key" for t in node.targets
):
keys.append(ast.unparse(node.value))
if isinstance(node, ast.Assert) and isinstance(node.test, ast.Call):
if isinstance(node.test.func, ast.Attribute) and node.test.func.attr == "check_support":
assert_calls.append(ast.unparse(node.test))
print("cache keys:", keys)
print("check_support assertions:", assert_calls)
for cls in api_tree.body:
if isinstance(cls, ast.ClassDef) and cls.name in {
"GemmProjRopeMxfp8Bf16InSm100",
"GemmProjRopeMxfp8Mxfp8InSm100",
}:
methods = {n.name: n for n in cls.body if isinstance(n, ast.FunctionDef)}
body = methods["check_support"].body
checks = [
ast.unparse(n)
for n in ast.walk(ast.Module(body=body, type_ignores=[]))
if isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr in {"_check_dtype", "_value_error_if"}
]
print(f"{cls.name} contract checks:")
for check in checks:
print(" ", check)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 4370
Validate descriptors before reusing cached launch parameters.
The caches use only x.shape and w.shape, although check_support() validates weight, RoPE, and scale descriptors. Include each contract-relevant shape and dtype in the cache key. Replace assert obj.check_support() with an explicit check that remains active under -O.
🤖 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/jax_api.py` around lines 79 -
96, Update the cache key used by the bf16-in grid path to include every
descriptor-relevant shape and dtype for x, w, cos, sin, and all output/scale
tensors, preventing reuse across incompatible contracts. In the cache-miss
initialization around GemmProjRopeMxfp8Bf16InSm100, replace assert
obj.check_support() with an explicit runtime validation that raises an
appropriate error when unsupported, remaining active under -O.
| if not (d_dtype in _fp8_dtypes) and discrete_col_sfd: | ||
| discrete_col_sfd = False # eager parity: ignored when SFD is not generated |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the membership test flagged by Ruff (E713).
Ruff reports this line as an error. Use not in to keep the lint gate green.
🔧 Proposed fix
- if not (d_dtype in _fp8_dtypes) and discrete_col_sfd:
+ if d_dtype not in _fp8_dtypes and discrete_col_sfd:
discrete_col_sfd = False # eager parity: ignored when SFD is not generated📝 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.
| if not (d_dtype in _fp8_dtypes) and discrete_col_sfd: | |
| discrete_col_sfd = False # eager parity: ignored when SFD is not generated | |
| if d_dtype not in _fp8_dtypes and discrete_col_sfd: | |
| discrete_col_sfd = False # eager parity: ignored when SFD is not generated |
🧰 Tools
🪛 Ruff (0.16.1)
[error] 302-302: Test for membership should be not in
Convert to not in
(E713)
🤖 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/dswiglu/jax_api.py` around lines
302 - 303, Update the dtype membership condition in the discrete_col_sfd
handling block to use the direct `d_dtype not in _fp8_dtypes` form, preserving
the existing behavior and assignment.
Source: Linters/SAST tools
| from .api import GroupedGemmSm100, grouped_gemm_wrapper_sm100 | ||
|
|
||
| __all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100"] | ||
| __all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100", "grouped_gemm_jax_sm100"] |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Sort __all__ to satisfy Ruff RUF022.
Ruff reports that __all__ is not sorted. Apply isort-style ordering so the lint gate passes.
🧹 Proposed fix
-__all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100", "grouped_gemm_jax_sm100"]
+__all__ = ["GroupedGemmSm100", "grouped_gemm_jax_sm100", "grouped_gemm_wrapper_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.
| __all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100", "grouped_gemm_jax_sm100"] | |
| __all__ = ["GroupedGemmSm100", "grouped_gemm_jax_sm100", "grouped_gemm_wrapper_sm100"] |
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 6-6: __all__ is not sorted
Apply an isort-style sorting to __all__
(RUF022)
🤖 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/unfused/__init__.py` at line 6, Sort the
entries in the module-level __all__ list using isort-style ordering, updating
only the ordering of GroupedGemmSm100, grouped_gemm_wrapper_sm100, and
grouped_gemm_jax_sm100 so Ruff RUF022 passes.
Source: Linters/SAST tools
| def _prob_spec() -> TensorSpec: | ||
| # (m, 1, 1) with m innermost: explicit ranks because trailing unit dims make | ||
| # leading-dim inference ambiguous | ||
| return TensorSpec(layout=(0, 1, 2)) | ||
|
|
||
|
|
||
| @cute.jit | ||
| def _grouped_bf16_adapter(stream, a, b_ptrs, padded_offsets, alpha, prob, d, c, workspace, *, kernel, n, k, mac): | ||
| # Discrete-mode b is a raw pointer to the device int64[] of per-expert base | ||
| # addresses; the packed uint8 (or int64) input buffer recasts for free. | ||
| b_arg = cute.recast_ptr(b_ptrs.iterator, dtype=cutlass.Int64) | ||
| kernel( | ||
| a=a, | ||
| b=b_arg, | ||
| n=cutlass.Int32(n), | ||
| k=cutlass.Int32(k), | ||
| b_stride_size=cutlass.Int64(k), # uniform k-major per-expert (n, k) weights | ||
| b_major_mode=OperandMajorMode.K, | ||
| workspace_ptr=workspace.iterator, | ||
| c=c, | ||
| d=d, | ||
| padded_offsets=padded_offsets, | ||
| alpha=alpha, | ||
| bias=None, | ||
| prob=prob, | ||
| max_active_clusters=mac, | ||
| stream=stream, | ||
| ) | ||
|
|
||
|
|
||
| def _pointer_count(b_ptrs: Any, name: str = "b_ptrs") -> int: | ||
| """Tracing-safe pointer-array shape/dtype check; returns the pointer count.""" | ||
| shape = tuple(b_ptrs.shape) | ||
| if len(shape) != 1: | ||
| raise ValueError(f"{name} must be 1-D, got shape={shape}") | ||
| dtype = _convert_to_cutlass_data_type(b_ptrs.dtype) | ||
| if dtype is cutlass.Int64: | ||
| return shape[0] | ||
| if dtype is cutlass.Uint8: | ||
| if shape[0] % 8 != 0: | ||
| raise ValueError(f"{name} packed uint8 length must be a multiple of 8, got {shape[0]}") | ||
| return shape[0] // 8 | ||
| raise ValueError(f"{name} must be int64 (or, without x64 mode, packed uint8), got {b_ptrs.dtype}") |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Move the shared JAX helpers out of the unfused kernel module.
_pointer_count and _prob_spec are now imported by grouped/glu/jax_api.py, grouped/dglu/jax_api.py, grouped/dsrelu/jax_api.py, grouped/wgrad/jax_api.py, and discrete_grouped/swiglu/jax_api.py (Line 37 of that file). The last import crosses the operand-layout family boundary: a discrete_grouped module depends on a grouped/unfused kernel module. Place these helpers in a neutral internal module shared by the CuTeDSL GEMM families, and keep them unexported.
As per coding guidelines: "GEMM fusions must be organized under gemm/cutedsl/ according to operand layout: dense, grouped, or discrete_grouped" and "Shared GEMM helpers, including schedulers and metadata utilities, must remain internal to their family package and must not be exported through cudnn."
🤖 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/unfused/jax_api.py` around lines 43 - 85,
Move the shared _pointer_count and _prob_spec helpers out of the grouped/unfused
jax_api module into a neutral, unexported internal CuTeDSL GEMM helper module
that can serve both grouped and discrete_grouped families. Update the imports in
grouped/glu/jax_api.py, grouped/dglu/jax_api.py, grouped/dsrelu/jax_api.py,
grouped/wgrad/jax_api.py, and discrete_grouped/swiglu/jax_api.py to use the new
module, without exposing the helpers through cudnn.
Source: Coding guidelines
| expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_row_tensor", "dprob_tensor", "dbias_tensor")} | ||
|
|
||
| def check(d_row_tensor, dprob_tensor, dbias_tensor): | ||
| jax.block_until_ready((d_row_tensor, dprob_tensor, dbias_tensor)) | ||
| for got, key in ((d_row_tensor, "d_row_tensor"), (dprob_tensor, "dprob_tensor"), (dbias_tensor, "dbias_tensor")): | ||
| np.testing.assert_array_equal( | ||
| np.asarray(got).view(np.uint8), | ||
| expected[key], | ||
| err_msg=f"grouped dGLU {key}: jit output differs from eager wrapper output on identical input bytes", | ||
| ) | ||
|
|
||
| # Eager custom call | ||
| check(*grouped_gemm_dglu_jax_sm100(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, n_weight, prob_j, generate_dbias=True)) | ||
|
|
||
| # Under jax.jit, twice (compiled-kernel / registration cache). n stays static. | ||
| jitted = jax.jit( | ||
| lambda a, c, offsets, alpha, beta, ptrs, prob: grouped_gemm_dglu_jax_sm100(a, c, offsets, alpha, beta, ptrs, n_weight, prob, generate_dbias=True), | ||
| ) | ||
| check(*jitted(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, prob_j)) | ||
| check(*jitted(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, prob_j)) | ||
|
|
||
| # generate_dbias=False returns (d_row, dprob, None) | ||
| d_row_only, dprob_only, dbias_none = grouped_gemm_dglu_jax_sm100(a_j, c_j, offsets_j, alpha_j, beta_j, b_ptrs_j, n_weight, prob_j) | ||
| jax.block_until_ready((d_row_only, dprob_only)) | ||
| assert dbias_none is None | ||
| np.testing.assert_array_equal(np.asarray(d_row_only).view(np.uint8), expected["d_row_tensor"]) | ||
| np.testing.assert_array_equal(np.asarray(dprob_only).view(np.uint8), expected["dprob_tensor"]) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Compare dprob_tensor (and dbias_tensor) with a tolerance, not bitwise.
check compares every key with np.testing.assert_array_equal on raw bytes. The grouped dGLU kernel accumulates dprob and dbias with floating-point atomic adds. The comment in python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py states this explicitly: the outputs are donated zero-initialized buffers "because the kernel accumulates into them (atomic add)". Atomic-add ordering is not deterministic, so bitwise equality between the eager run and the jit run can fail intermittently.
The sibling tests already use a tolerance for this exact tensor: test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py lines 240-247 and test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py lines 240-246. A previous review raised the same concern on the torch-parity test in this file; the jit test repeats the bitwise pattern.
Line 181 has the same problem for dprob_only.
🛡️ Proposed change
- expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_row_tensor", "dprob_tensor", "dbias_tensor")}
+ expected = {key: np.asarray(result_eager[key]).view(np.uint8) for key in ("d_row_tensor",)}
+ # dprob/dbias accumulate through floating-point atomics; ordering is nondeterministic.
+ expected_atomic = {key: np.asarray(result_eager[key], dtype=np.float32) for key in ("dprob_tensor", "dbias_tensor")}
def check(d_row_tensor, dprob_tensor, dbias_tensor):
jax.block_until_ready((d_row_tensor, dprob_tensor, dbias_tensor))
- for got, key in ((d_row_tensor, "d_row_tensor"), (dprob_tensor, "dprob_tensor"), (dbias_tensor, "dbias_tensor")):
- np.testing.assert_array_equal(
- np.asarray(got).view(np.uint8),
- expected[key],
- err_msg=f"grouped dGLU {key}: jit output differs from eager wrapper output on identical input bytes",
- )
+ np.testing.assert_array_equal(
+ np.asarray(d_row_tensor).view(np.uint8),
+ expected["d_row_tensor"],
+ err_msg="grouped dGLU d_row_tensor: jit output differs from eager wrapper output on identical input bytes",
+ )
+ for got, key in ((dprob_tensor, "dprob_tensor"), (dbias_tensor, "dbias_tensor")):
+ np.testing.assert_allclose(
+ np.asarray(got, dtype=np.float32),
+ expected_atomic[key],
+ rtol=1e-4,
+ atol=1e-4,
+ err_msg=f"grouped dGLU {key}: jit output differs from eager wrapper output beyond atomic-add tolerance",
+ )Apply the same change at line 181:
- np.testing.assert_array_equal(np.asarray(dprob_only).view(np.uint8), expected["dprob_tensor"])
+ np.testing.assert_allclose(np.asarray(dprob_only, dtype=np.float32), expected_atomic["dprob_tensor"], rtol=1e-4, atol=1e-4)🤖 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_dglu_jax.py` around lines
155 - 181, Update the check function to compare dprob_tensor and dbias_tensor
using the established numerical tolerance rather than raw-byte
assert_array_equal, while retaining exact byte comparison for d_row_tensor.
Apply the same tolerant comparison to dprob_only in the generate_dbias=False
path, matching the sibling grouped GEMM tests’ existing tolerance settings.
64ad2f2 to
bd52582
Compare
| _convert_to_cutlass_data_type(b_tensor.dtype), | ||
| _convert_to_cutlass_data_type(sfa_tensor.dtype), | ||
| _convert_to_cutlass_data_type(sfb_tensor.dtype), | ||
| alpha, |
There was a problem hiding this comment.
another occurance of alpha in cache key
There was a problem hiding this comment.
Fixed in 3b5ab5f — alpha dropped from both cache keys in this file (srelu and dsrelu). The cached (kernel, mac) never depended on it: alpha reaches the kernel only as a constexpr kwarg of the custom call (where it participates in the bridge's own compile cache), and was otherwise only forwarded to check_support, which doesn't validate it.
| cache_key = ( | ||
| m, | ||
| n, | ||
| k, |
There was a problem hiding this comment.
python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py:168-170
python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py:165-167
python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py:202-204
python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py:325-327
Claude flags these MNK cache keys as pure waste
the outputs of the compilation does not depend on kernel cache, cache value is created from
kernel = MoEGroupedGemmBf16Kernel(acc_dtype, mma_tiler_mn, cluster_shape_mn,
expert_cnt, generate_c, use_dynamic_sched, ...)
mac = f(cluster_shape_mn)
ws = kernel.get_workspace_bytes()
mma_tiler_mn, cluster_shape_mn does not vary with mnk dims
There was a problem hiding this comment.
correction: the outputs of the compilation does not VARY with kernel cache MNK
There was a problem hiding this comment.
Fixed in 3b5ab5f — m/n/k dropped from all four cache keys (unfused, glu, dglu, dsrelu). You're right that the cached (kernel, mac, workspace_bytes) triple depends only on the constructor/cluster config; can_implement (the only shape-dependent step) already runs before the cache lookup on every call, so shape validation is unaffected.
There was a problem hiding this comment.
cudnn.jax unreachable from a bare import, and the error is a bare AttributeError: jax with no hint that import cudnn.jax is the fix.
# works
import cudnn.jax
# doesn't work
import cudnn
cudnn.jax.TensorSpec (attributeError jax)
There was a problem hiding this comment.
Fixed in 3b5ab5f — cudnn/__init__.py now lazily imports the submodule on attribute access, so import cudnn; cudnn.jax.TensorSpec works without a separate import cudnn.jax (asserted in test_gemm_amax_jax_jit_sm100). jax itself stays unimported until first access, and the fallback AttributeError now names the module (module 'cudnn' has no attribute ...).
hwanseoc
left a comment
hwanseoc
left a comment
There was a problem hiding this comment.
There is a functional bug if a scale-factor tensor is non-contiguous. It does not reproduce today, because every SF tensor in-tree is contiguous.
_sf_desc_is_physical (grouped/dsrelu/api.py:285) decides between the physical and atom layouts by shape alone, and _check_sf_shape (:297) validates shape alone. The physical branches (:1057, :1082, :1103) then hardcode row-major strides for the kernel (:1065, :1087, :1109, :1116). Both of these tensors get the same verdict:
contiguous : stride=(4096,1024,512,16,4,1) contig=True -> "physical"
permuted : stride=(4096, 512,2048,16,4,1) contig=False -> "physical"
JAX cannot reach this, since its arrays are always row-major. A torch caller can, and the kernel then reads at the wrong addresses and returns wrong numbers with no error.
The same pattern appears in discrete_grouped/swiglu/api.py:202 and dswiglu/api.py:203, of which CodeRabbit flagged only dswiglu.
The fix is a _check_tensor_stride call in the physical branch.
^ Coderabbit also flags this as a medium-nice-to have bug, but seems like a high functionality bug to me for Torch. Please change
bd52582 to
3b5ab5f
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py (1)
27-28: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueModule-level kernel caches grow without a bound and are not thread-safe.
_srelu_kernel_cacheand_dsrelu_kernel_cachekeep every compiled kernel for the process lifetime. Each distinct shape or dtype combination adds an entry. Two threads that call the same entry point with a new key both build the kernel, because the get-then-set sequence is not atomic. The duplicate build wastes compile time; the last writer wins.Consider a bounded cache and a lock around the miss path if multi-threaded dispatch is expected.
🤖 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/srelu/jax_api.py` around lines 27 - 28, Replace the unbounded module-level dictionaries _srelu_kernel_cache and _dsrelu_kernel_cache with bounded caches, and protect cache-miss lookup/build/store operations with a shared lock so concurrent calls for the same key do not compile duplicate kernels. Preserve existing cache-key behavior and kernel reuse for both dispatch paths.
🤖 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 `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Around line 109-110: Rename the ambiguous variable l to batch in the relevant
functions, including gemm_srelu_jax_sm100, and update every listed reference so
shape handling and subsequent computations use batch consistently.
In `@python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py`:
- Around line 76-93: Add an APIBase subclass and corresponding wrapper for
grouped_gemm_glu_jax_sm100 in this module before exposing the public frontend
API. Have the wrapper route calls to grouped_gemm_glu_jax_sm100 while preserving
its arguments and return values, following the existing APIBase/wrapper pattern
used by nearby frontend-only APIs; retain the lazy export in cudnn.__init__.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Around line 27-28: Replace the unbounded module-level dictionaries
_srelu_kernel_cache and _dsrelu_kernel_cache with bounded caches, and protect
cache-miss lookup/build/store operations with a shared lock so concurrent calls
for the same key do not compile duplicate kernels. Preserve existing cache-key
behavior and kernel reuse for both dispatch paths.
🪄 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: 41bda07a-8ffe-4dfb-bf6a-8386cdff9ee4
📒 Files selected for processing (7)
python/cudnn/__init__.pypython/cudnn/gemm/cutedsl/dense/srelu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/dglu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/glu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/unfused/jax_api.pytest/python/fe_api/gemm/test_gemm_amax_jax.py
🚧 Files skipped from review as they are similar to previous changes (4)
- test/python/fe_api/gemm/test_gemm_amax_jax.py
- python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py
- python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py
- python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py
| m, _, l = a_tensor.shape | ||
| n, _, _ = b_tensor.shape |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous variable l.
Ruff reports E741 at Line 109 and Line 206. The name l is easy to confuse with 1. Use batch instead, and update the uses at Line 111, Line 138-139, Line 167-168, Line 208, Line 237-238, and Line 266-267.
♻️ Proposed rename for `gemm_srelu_jax_sm100`
- m, _, l = a_tensor.shape
+ m, _, batch = a_tensor.shape
n, _, _ = b_tensor.shape
- if l != 1:
+ if batch != 1:
raise ValueError("JAX inputs must have batch dim L == 1; batch-outermost (L-major) layouts are not expressible as JAX arrays")Also applies to: 206-207
🧰 Tools
🪛 Ruff (0.16.1)
[error] 109-109: 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 `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py` around lines 109 - 110,
Rename the ambiguous variable l to batch in the relevant functions, including
gemm_srelu_jax_sm100, and update every listed reference so shape handling and
subsequent computations use batch consistently.
Source: Linters/SAST tools
| def grouped_gemm_glu_jax_sm100( | ||
| a_tensor: Any, | ||
| padded_offsets: Any, | ||
| alpha_tensor: Any, | ||
| b_ptrs: Any, | ||
| n: int, | ||
| prob_tensor: Any, | ||
| c_dtype: Any = cutlass.BFloat16, | ||
| d_dtype: Any = cutlass.BFloat16, | ||
| acc_dtype: Any = cutlass.Float32, | ||
| mma_tiler_mn: Tuple[int, int] = (256, 256), | ||
| cluster_shape_mn: Optional[Tuple[int, int]] = None, | ||
| vector_f32: bool = False, | ||
| act_func: str = "swiglu", | ||
| linear_offset: Optional[float] = None, | ||
| generate_c: bool = False, | ||
| use_dynamic_sched: bool = False, | ||
| ) -> Tuple[Any, Optional[Any]]: |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Add the required API class and wrapper.
grouped_gemm_glu_jax_sm100 is a new public frontend-only API. This file has no APIBase subclass or wrapper for it. The lazy mapping in python/cudnn/__init__.py only satisfies the export requirement.
Add an APIBase subclass and a wrapper for this JAX API before exposing it.
As per coding guidelines, “Every new frontend-only Python API must include an APIBase subclass and wrapper, and must be lazily exported from python/cudnn/__init__.py.”
🤖 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/jax_api.py` around lines 76 - 93, Add
an APIBase subclass and corresponding wrapper for grouped_gemm_glu_jax_sm100 in
this module before exposing the public frontend API. Have the wrapper route
calls to grouped_gemm_glu_jax_sm100 while preserving its arguments and return
values, following the existing APIBase/wrapper pattern used by nearby
frontend-only APIs; retain the lazy export in cudnn.__init__.
Source: Coding guidelines
|
There seems to be a lot of redundant cache key vectors, |
3b5ab5f to
da4c69b
Compare
…cross 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>
da4c69b to
2ff2239
Compare
|
Re: non-contiguous scale-factor tensors (review #pullrequestreview-4910023689) — fixed in
Dense srelu/dsrelu/swiglu-quantized build their compile-time SF layouts from the descriptor's actual strides rather than hardcoding row-major, so they were left as-is. |
|
Re: cache-key audit (comment) — audited every
Residual dtype components (e.g. |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py (1)
166-166: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueVerify Black formatting at 160 columns.
Line 166 and Line 268 appear to exceed 160 characters. The coding guidelines require Black with a line length of 160 for
python/**/*.py. Extract the hardware-info lookup into a local variable, or let Black split the expression.♻️ Proposed formatting
- mac = cutlass.utils.HardwareInfo().get_max_active_clusters(gemm.cluster_shape_mn[0] * gemm.cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin + cluster_size = gemm.cluster_shape_mn[0] * gemm.cluster_shape_mn[1] + mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size) - gemm.num_cluster_overlap_marginAs per coding guidelines: "Format Python code with Black using a line length of 160."
Also applies to: 268-268
🤖 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/srelu/jax_api.py` at line 166, Format the long expressions in the relevant GEMM setup logic, including the `mac` calculation and the corresponding line around `gemm.num_cluster_overlap_margin`, with Black using a 160-character line length. If needed, extract the `HardwareInfo().get_max_active_clusters(...)` result into a local variable while preserving behavior.Source: Coding guidelines
test/python/fe_api/gemm/test_gemm_amax.py (2)
322-324: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winVerify that accepted layouts execute successfully.
The test discards both outputs and does not wait for the CUDA stream after Lines 323-324. Add
torch.cuda.current_stream().synchronize()after each accepted call, or validate the outputs withcheck_ref_gemm_amax.Based on
python/cudnn/gemm/cutedsl/dense/amax/api.py:451-563, the Torch wrapper callsexecuteon the supplied stream without an explicit synchronization before returning.🤖 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_amax.py` around lines 322 - 324, The accepted-layout cases in the gemm_amax test currently discard results without waiting for asynchronous execution. After each gemm_amax_wrapper_sm100 call, synchronize torch.cuda.current_stream(), or validate the returned output using check_ref_gemm_amax, so both the physical and permuted layouts are confirmed to execute successfully.Source: Coding guidelines
326-335: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd independent
sfb_tensorrejection cases.Both invalid calls modify only
sfa_tensor;sfb_tensorremains valid. A regression insfb_tensorstride validation would pass this test. Create matching invalidsfb_tensorlayouts usingn // 128, then repeat bothpytest.raiseschecks withsfa_tensor=sfa.The upstream
gemm_amax_wrapper_sm100contract includes separate layout inputs forsfa_tensorandsfb_tensor.🤖 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_amax.py` around lines 326 - 335, Extend the layout-validation test around gemm_amax_wrapper_sm100 with independent invalid sfb_tensor cases: construct non-contiguous and atom-view-shaped sfb layouts using n // 128, while passing the valid sfa tensor unchanged. Repeat both pytest.raises(ValueError, match="stride") assertions to verify sfb_tensor validation separately from sfa_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 `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py`:
- Around line 105-106: Update the active-cluster calculation in the JAX kernel
compilation path to reject any zero or negative mac after subtracting
gemm.num_cluster_overlap_margin, matching GemmAmaxSm100._compile_kernel().
Perform this validation before storing (kernel, mac) in _kernel_cache.
- Line 98: Replace the assert around gemm.check_support() with a direct call so
support validation always executes, including under Python optimization, before
gemm._kernel is accessed.
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Line 150: Validate the runtime dtype of prob_tensor against the compiled
cutlass.Float32 contract before constructing the kernel in both affected
functions, including gemm_srelu_jax_sm100. Also validate c_tensor against the
requested c_dtype where sample_c is built, or include both runtime dtypes and
shapes in validation_key so check_support() performs the comparison.
In `@test/python/fe_api/gemm/test_gemm_amax.py`:
- Around line 308-318: The SM100 test must skip when either FP8 dtype is
unavailable before creating tensors. In the test setup around
gemm_amax_wrapper_sm100, check for torch.float8_e5m2 and torch.float8_e8m0fnu
after the existing hardware/import gates and skip with an appropriate message;
retain the existing GemmAmaxSm100.check_support() flow without adding a
cudnn.backend_version() requirement.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Line 166: Format the long expressions in the relevant GEMM setup logic,
including the `mac` calculation and the corresponding line around
`gemm.num_cluster_overlap_margin`, with Black using a 160-character line length.
If needed, extract the `HardwareInfo().get_max_active_clusters(...)` result into
a local variable while preserving behavior.
In `@test/python/fe_api/gemm/test_gemm_amax.py`:
- Around line 322-324: The accepted-layout cases in the gemm_amax test currently
discard results without waiting for asynchronous execution. After each
gemm_amax_wrapper_sm100 call, synchronize torch.cuda.current_stream(), or
validate the returned output using check_ref_gemm_amax, so both the physical and
permuted layouts are confirmed to execute successfully.
- Around line 326-335: Extend the layout-validation test around
gemm_amax_wrapper_sm100 with independent invalid sfb_tensor cases: construct
non-contiguous and atom-view-shaped sfb layouts using n // 128, while passing
the valid sfa tensor unchanged. Repeat both pytest.raises(ValueError,
match="stride") assertions to verify sfb_tensor validation separately from
sfa_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: 1d85e352-b1d5-4c63-abd2-cd58a27ab5e1
📒 Files selected for processing (8)
python/cudnn/gemm/cutedsl/dense/amax/api.pypython/cudnn/gemm/cutedsl/dense/amax/jax_api.pypython/cudnn/gemm/cutedsl/dense/srelu/jax_api.pypython/cudnn/gemm/cutedsl/dense/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/api.pytest/python/fe_api/gemm/test_gemm_amax.py
🚧 Files skipped from review as they are similar to previous changes (3)
- python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py
- python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py
- python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py
| target, | ||
| ( | ||
| jax.ShapeDtypeStruct((m, n, l), c_jax_dtype), | ||
| assert gemm.check_support() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python - <<'PY'
import ast
from pathlib import Path
path = Path("python/cudnn/gemm/cutedsl/dense/amax/jax_api.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, ast.Assert):
print(f"{path}:{node.lineno}: assert expression found")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 235
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '--- jax_api.py ---'
cat -n python/cudnn/gemm/cutedsl/dense/amax/jax_api.py | sed -n '1,150p'
echo '--- dense amax API definitions ---'
cat -n python/cudnn/gemm/cutedsl/dense/amax/api.py | sed -n '1,240p'
echo '--- relevant symbol usage ---'
rg -n -C 4 'check_support|_kernel|_compile_kernel|GemmAmaxSm100' \
python/cudnn/gemm/cutedsl/dense/amax python/cudnn/gemm/cutedsl/denseRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
import dis
from pathlib import Path
jax_path = Path("python/cudnn/gemm/cutedsl/dense/amax/jax_api.py")
api_path = Path("python/cudnn/gemm/cutedsl/dense/amax/api.py")
jax_tree = ast.parse(jax_path.read_text())
api_tree = ast.parse(api_path.read_text())
assert_nodes = [
node for node in ast.walk(jax_tree)
if isinstance(node, ast.Assert) and node.lineno == 98
]
assert len(assert_nodes) == 1
check_support = next(
node for node in ast.walk(api_tree)
if isinstance(node, ast.FunctionDef) and node.name == "check_support"
)
kernel_assignments = [
node.lineno
for node in ast.walk(check_support)
if isinstance(node, ast.Assign)
and any(isinstance(target, ast.Attribute) and target.attr == "_kernel" for target in node.targets)
]
print("jax_assert_line:", assert_nodes[0].lineno)
print("check_support_kernel_assignment_lines:", kernel_assignments)
print("optimized_assert_bytecode:")
dis.dis(compile("assert side_effect()", "<probe>", "exec", optimize=1))
api_source = api_path.read_text()
print("api_base_kernel_initialization_matches:", "_kernel =" in api_source[:api_source.index("class GemmAmaxSm100")])
PYRepository: NVIDIA/cudnn-frontend
Length of output: 376
Do not use assert for required support validation.
Under Python optimization, the assertion is removed. On a cache miss, gemm._kernel is not initialized before access. Call gemm.check_support() directly.
🤖 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` at line 98, Replace the
assert around gemm.check_support() with a direct call so support validation
always executes, including under Python optimization, before gemm._kernel is
accessed.
| mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin | ||
| _kernel_cache[config_key] = (kernel, mac) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject a non-positive active-cluster count.
GemmAmaxSm100._compile_kernel() rejects this value after it applies CUDNNFE_CLUSTER_OVERLAP_MARGIN. This JAX path bypasses that check and can cache a zero or negative mac.
Proposed fix
mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin
+ if mac <= 0:
+ raise ValueError(
+ "max_active_clusters must be > 0 after applying overlap margin; "
+ "reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN"
+ )
_kernel_cache[config_key] = (kernel, mac)📝 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.
| mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin | |
| _kernel_cache[config_key] = (kernel, mac) | |
| mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin | |
| if mac <= 0: | |
| raise ValueError( | |
| "max_active_clusters must be > 0 after applying overlap margin; " | |
| "reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN" | |
| ) | |
| _kernel_cache[config_key] = (kernel, mac) |
🤖 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 105 - 106,
Update the active-cluster calculation in the JAX kernel compilation path to
reject any zero or negative mac after subtracting
gemm.num_cluster_overlap_margin, matching GemmAmaxSm100._compile_kernel().
Perform this validation before storing (kernel, mac) in _kernel_cache.
| sample_d=_make_desc((m, n, l), d_dtype, "sample_d"), | ||
| sample_sfa=_make_desc(tuple(sfa_tensor.shape), sfa_tensor.dtype, "sample_sfa"), | ||
| sample_sfb=_make_desc(tuple(sfb_tensor.shape), sfb_tensor.dtype, "sample_sfb"), | ||
| sample_prob=_make_desc((m, 1, 1), cutlass.Float32, "sample_prob"), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
prob_tensor dtype is never validated against the compiled Float32 contract.
Both functions build sample_prob as a fixed cutlass.Float32 descriptor. validation_key omits the real prob_tensor dtype and shape. _prob_spec() carries only a layout, not a dtype. A caller that passes a non-float32 prob_tensor therefore reaches a kernel compiled for float32 without any check. The same gap applies to c_tensor in gemm_srelu_jax_sm100, where sample_c uses the requested c_dtype and no runtime tensor is compared.
Add an explicit dtype check on prob_tensor before kernel construction, or include its dtype in validation_key so check_support() sees the real value.
Run the following script to confirm that the bridge does not already enforce operand dtypes:
#!/bin/bash
# Description: Check whether cudnn.jax.call / TensorSpec validate input dtypes.
set -euo pipefail
fd -t f 'call.py' python/cudnn/jax --exec cat -n {}
fd -t f '__init__.py' python/cudnn/jax --exec cat -n {}
# Compare with a sibling jax_api that may validate prob dtype explicitly.
rg -n -C3 'prob_tensor' --glob 'python/cudnn/gemm/cutedsl/**/jax_api.py'Also applies to: 252-252
🤖 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/srelu/jax_api.py` at line 150, Validate the
runtime dtype of prob_tensor against the compiled cutlass.Float32 contract
before constructing the kernel in both affected functions, including
gemm_srelu_jax_sm100. Also validate c_tensor against the requested c_dtype where
sample_c is built, or include both runtime dtypes and shapes in validation_key
so check_support() performs the comparison.
| try: | ||
| from cudnn import gemm_amax_wrapper_sm100 | ||
| except ImportError: | ||
| pytest.skip("Environment not supported: cudnn optional dependencies not installed") | ||
| if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10: | ||
| pytest.skip("requires SM100+") | ||
|
|
||
| m, n, k, sf_vec_size = 512, 256, 256, 32 | ||
| a = torch.randn(m, k, 1, device="cuda").to(torch.float8_e5m2) | ||
| b = torch.randn(n, k, 1, device="cuda").to(torch.float8_e5m2) | ||
| sf_dtype = torch.float8_e8m0fnu |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'backend_version\(\)|get_device_capability\(\)|check_support|float8_e8m0fnu|pytest\.skip' \
test/python/fe_api python/cudnnRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test/AGENTS.md ---'
if [ -f test/AGENTS.md ]; then
cat -n test/AGENTS.md
fi
printf '%s\n' '--- target test outline ---'
ast-grep outline test/python/fe_api/gemm/test_gemm_amax.py
printf '%s\n' '--- target test ---'
sed -n '1,380p' test/python/fe_api/gemm/test_gemm_amax.py
printf '%s\n' '--- relevant helper definitions and nearby tests ---'
rg -n -C 6 'def (skip_unless_sm100|.*support.*|.*backend.*)|backend_version\(\)|check_support\(\)|gemm_amax_wrapper_sm100|test_.*amax' test/python/fe_api test/python/conftest.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- wrapper definitions and support gates ---'
rg -n -C 12 'def gemm_amax_wrapper_sm100|gemm_amax_wrapper_sm100|class GemmAmaxSm100|def check_support|backend_version' \
python/cudnn test/python/fe_api/gemm/test_gemm_amax_utils.py
printf '%s\n' '--- related wrapper rejection tests ---'
sed -n '140,230p' test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8.py
sed -n '1,130p' test/python/fe_api/norm/test_rmsnorm_rht_amax.py
printf '%s\n' '--- GEMM Amax utility configuration ---'
sed -n '1,280p' test/python/fe_api/gemm/test_gemm_amax_utils.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- GemmAmax API support validation ---'
sed -n '1,360p' python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- GemmAmax wrapper implementation ---'
sed -n '480,575p' python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- Amax test utility ---'
wc -l test/python/fe_api/gemm/test_gemm_amax_utils.py
sed -n '1,280p' test/python/fe_api/gemm/test_gemm_amax_utils.pyRepository: NVIDIA/cudnn-frontend
Length of output: 28457
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("test/python/fe_api/gemm/test_gemm_amax.py")
api_path = Path("python/cudnn/gemm/cutedsl/dense/amax/api.py")
test_tree = ast.parse(test_path.read_text())
api_tree = ast.parse(api_path.read_text())
test_fn = next(
node for node in ast.walk(test_tree)
if isinstance(node, ast.FunctionDef)
and node.name == "test_gemm_amax_rejects_noncontiguous_scale_factors"
)
calls = [
(node.lineno, ast.unparse(node.func))
for node in ast.walk(test_fn)
if isinstance(node, ast.Call)
]
print("target test calls:")
for item in calls:
print(item)
api_calls = [
(node.lineno, ast.unparse(node.func))
for node in ast.walk(api_tree)
if isinstance(node, ast.Call)
]
print("\nGemmAmax API backend-version calls:")
print([item for item in api_calls if "backend_version" in item[1]])
wrapper = next(
node for node in ast.walk(api_tree)
if isinstance(node, ast.FunctionDef)
and node.name == "gemm_amax_wrapper_sm100"
)
wrapper_calls = [
(node.lineno, ast.unparse(node.func))
for node in ast.walk(wrapper)
if isinstance(node, ast.Call)
]
print("\nwrapper support calls:")
print([item for item in wrapper_calls if "check_support" in item[1]])
PYRepository: NVIDIA/cudnn-frontend
Length of output: 1446
Gate unavailable FP8 dtypes before allocation.
The SM100 check does not cover whether PyTorch exposes torch.float8_e5m2 and torch.float8_e8m0fnu; missing attributes fail during tensor creation. Add a dtype capability skip before allocation. gemm_amax_wrapper_sm100 already calls GemmAmaxSm100.check_support() before compilation, and this API has no cudnn.backend_version() requirement.
🤖 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_amax.py` around lines 308 - 318, The SM100
test must skip when either FP8 dtype is unavailable before creating tensors. In
the test setup around gemm_amax_wrapper_sm100, check for torch.float8_e5m2 and
torch.float8_e8m0fnu after the existing hardware/import gates and skip with an
appropriate message; retain the existing GemmAmaxSm100.check_support() flow
without adding a cudnn.backend_version() requirement.
Source: Coding guidelines
2 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*.Affected area
FE OSS kernels or CuTeDSL (Python API)
Summary
Replaces the jax-tvm-ffi backend of the
jax.jitentry points with CuTeDSL's native JAX integration (cutlass.jax.cutlass_call), wrapped as a newcudnn.jax.call, and extends jit coverage to every JAX-reachable GEMM API inpython/cudnn/gemm/cutedsl.Per-API status
Dense fusions — full JAX eager support (from #529), now all with jit entry points:
jax.jitentry pointgemm_amax_jax_sm100(backend swap, contract unchanged)gemm_swiglu_jax_sm100— now incl. blockscaled MXFP8 quantized (wasNotImplementedError)gemm_srelu_jax_sm100— new (was impossible over jax-tvm-ffi)gemm_dsrelu_jax_sm100— new (dprob = bridge-managed zero-init atomic-add accumulator)w_out_in=True(from #530)gemm_proj_rope_mxfp8_jax_sm100— new, both input paths (BF16 and MXFP8), dispatch onx.dtypeGrouped / discrete-grouped — JAX eager support in discrete (pointer-array) modes (from #530), now all with jit entry points:
jax.jitentry pointgrouped_gemm_jax_sm100— newgrouped_gemm_glu_jax_sm100— new (no bias;b_major="k";linear_offsetis a compile-time constant)grouped_gemm_dglu_jax_sm100— new (dprob/dbias are bridge-managed zero-init accumulators, no caller-zeroed buffers)grouped_gemm_dsrelu_jax_sm100— new (full d/SFD/dprob/dbias output set)grouped_gemm_wgrad_jax_sm100— new (discrete output pointers; per-expert dgrad buffers are caller-owned external memory, so the entry returns a completion token toblock_until_ready)discrete_grouped_gemm_swiglu_jax_sm100— new (k-major weights; c/d/d_col/SFD/amax all XLA-managed)discrete_grouped_gemm_dswiglu_jax_sm100— new (relu/srelu epilogues; d_row/d_col/SFD/amax/dprob/dbias)Not reachable from JAX (eager or jit; unchanged from #530, clear
ValueErrors):TensorSpec.modeon this bridge is the future path.cudnn.jax.callcutlass_callplus the conveniences the cuDNN kernels need:initialized_outputs: pre-initialized donated buffers for accumulator outputs (amax/dprob/dbias) and for internal workspaces the kernels write (per-expert TMA descriptor tables — XLA inputs are immutable, so grouped workspaces are modeled as donated zero-init outputs that are never surfaced). The bridge drops aliased inputs from the kernel argument list, so the@cute.jitadapters keep the kernels' exact destination-passing signatures.TensorSpecpresets:gemm_operand_spec()(explicit stride ranks for(MN, K, 1)operands — trailing-unit-dim buffers defeat the bridge's leading-dim inference, on inputs and pure results alike, hence all outputs are donated) andsf_atom_spec()(presents physical C-contiguous SF buffers in the logical MMA atom view viaTensorSpec.mode, no transpose materialized).Why srelu/dsrelu/quantized-swiglu (and the optional-output grouped kernels) are now possible: their kernels carry optional None-typed parameters that the jax-tvm-ffi ABI could not supply. With per-kernel
@cute.jitadapters, thoseNones are compile-time constants inside the adapter — the bridge never sees them. Kernel instance /max_active_clusters/ scalar knobs travel as constexpr kwargs (participating in the bridge's compile cache; the instances are cached per config on our side).Grouped-specific jit mechanics: per-expert pointer arrays (
b_ptrs,sfb_ptrs,wgrad_ptrs) pass as regular packed-uint8 (or x64 int64) input arrays and are recast toInt64pointers inside the adapters viacute.recast_ptr; rows at/pastpadded_offsets[-1]come back zero-filled (donated zero-init outputs). Two contract caveats, documented per API:padded_offsetsvalues cannot be host-validated under tracing (shapes/dtypes still are), and the weight buffers behind pointer arrays must stay alive and unmoved across every execution of the traced computation.Why
The previous backend needed an extra dependency (jax-tvm-ffi) and hit a hard ABI limit on optional kernel parameters.
cutlass.jaxships insidenvidia-cutlass-dsl(already provided by thecutedslextra), declares CUDA-graph compatibility to XLA, and the constexpr-Noneadapter pattern unlocked every remaining JAX-reachable kernel — jit coverage now matches eager JAX coverage exactly.Related issues
Stacked on #534 (→ #530).
API and compatibility impact
NotImplementedErrorto working; ten new entry points (srelu, dsrelu, proj_rope, unfused, glu, dglu, grouped dsrelu, wgrad, discrete swiglu, discrete dswiglu). Eager paths and torch behavior untouched.cudnn.jax(call,TensorSpecre-export, spec presets, initializers) — lazily imported, requires jax ≥ 0.5.jaxdependency group dropsjax-tvm-ffiand moves the floor tojax>=0.5(cutlass.jaxrequirement).gemm/cutedsl/_jax_ffi.pyis removed.b_major="k"and treatslinear_offsetas a compile-time constant; discrete swiglu/dswiglu jit require k-major weights; wgrad jit offers discrete output mode only.Testing
On a B200-class SM100 (CC 10.0), Python 3.12, jax 0.11, nvidia-cutlass-dsl 4.6:
.view(np.uint8)comparison on identical input bytes — both paths run the same kernel): once eagerly, thenjax.jit'd twice (donation safety + compile/registration cache). Atomic-add accumulators (dprob) compare at the existing tight fp32 tolerance.import cudnnsucceeds and all jit entry points raise a clearImportErrorwith torch and jax absent (sys.modulesnulled); eager symbols unaffected.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
jax.jitexecution.Documentation
Tests