Add JAX support to block_sparse_attention (SM100/SM110 blk128 paths, stacked on #553) - #555
Closed
Anerudhan wants to merge 1 commit into
Closed
Add JAX support to block_sparse_attention (SM100/SM110 blk128 paths, stacked on #553)#555Anerudhan wants to merge 1 commit into
Anerudhan wants to merge 1 commit into
Conversation
Anerudhan
added
cat-feature
Contributor
📝 WalkthroughWalkthroughThis pull request separates CuTeDSL dependencies, adds framework-neutral Torch and JAX tensor handling, introduces JAX custom-call APIs for GEMM and attention paths, updates exports and documentation, and adds eager, jitted, parity, and rejection tests. ChangesJAX framework integration
Estimated code review effort: 5 (Critical) | ~120 minutes 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 13
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py (1)
152-215: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winDeclare
torchfor static analysis without importing it at runtime.from __future__ import annotationsprevents an import-timeNameError, but Ruff still reports 34F821errors. Add aTYPE_CHECKING-guarded import.🤖 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/_bf16_api.py` around lines 152 - 215, Add a TYPE_CHECKING-guarded torch import in the module containing the static methods, preserving the existing runtime import behavior while allowing annotations such as torch.Tensor to be resolved by static analysis. Do not introduce an unconditional torch import.Sources: Coding guidelines, Linters/SAST tools
🟡 Minor comments (23)
AGENTS.md-47-49 (1)
47-49: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument the minimum
pipversion for--group.
pip install --grouprequires pip 25.1 or newer. Add a bootstrap command or document a direct-install fallback for these dependency groups.🤖 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 `@AGENTS.md` around lines 47 - 49, Update the dependency installation instructions near the pip --group commands to document that pip 25.1 or newer is required, and provide either a bootstrap command to upgrade pip or a direct-install fallback for the torch and jax groups.Source: MCP tools
docs/fe-oss-apis/overview.md-47-53 (1)
47-53: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDocument the JAX GPU wheel requirement.
The JAX dependency group installs only
jax>=0.5. The JAX APIs execute on CUDA, so document the matchingjax[cuda12]orjax[cuda13]extra, or state that users must install it separately.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/fe-oss-apis/overview.md` around lines 47 - 53, Update the JAX installation guidance in the overview documentation to explicitly require the appropriate GPU-enabled JAX package, such as jax[cuda12] or jax[cuda13], for CUDA execution. Cover both the dependency-group and published-wheel instructions, or clearly state that users must install the matching GPU extra separately.Source: MCP tools
pyproject.toml-91-96 (1)
91-96: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAlign the JAX dependency group with the Python floor.
No JAX release satisfying
jax>=0.5supports Python 3.9. Add apython_version >= '3.10'marker and document that JAX support requires Python 3.10+.Possible dependency-group fix
jax = [ - "jax>=0.5", + "jax>=0.5; python_version >= '3.10'", ]🤖 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 `@pyproject.toml` around lines 91 - 96, Update the jax dependency group in pyproject.toml to apply jax>=0.5 only when python_version >= '3.10', and revise the surrounding comments to state that JAX support requires Python 3.10 or newer.Source: MCP tools
python/cudnn/gemm/cutedsl/grouped/unfused/api.py-192-193 (1)
192-193: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUpdate the dtype error messages for JAX callers.
The checks now compare CUTLASS dtypes, 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 who passesjnp.float16now receives advice to use a torch dtype. Name the logical dtype instead.♻️ Proposed fix
if _convert_to_cutlass_data_type(a_tensor.dtype) is not cutlass.BFloat16: - raise ValueError(f"a_tensor must have dtype torch.bfloat16, got {a_tensor.dtype}") + raise ValueError(f"a_tensor must have dtype bfloat16, got {a_tensor.dtype}") @@ if _convert_to_cutlass_data_type(prob_tensor.dtype) is not cutlass.Float32: - raise ValueError(f"prob_tensor must have dtype torch.float32, got {prob_tensor.dtype}") + raise ValueError(f"prob_tensor must have dtype float32, got {prob_tensor.dtype}")Also applies to: 218-219, 225-226, 243-244
🤖 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 - 193, Update the dtype validation messages in the grouped GEMM API checks, including the checks for a_tensor, prob_tensor, b_tensor, and b_dtype, to name the logical dtypes bfloat16 and float32 without the torch. prefix. Keep the existing CUTLASS dtype comparisons and validation behavior unchanged.python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py-536-542 (1)
536-542: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
_check_sm100now validates the current device, not the tensor's device.
get_compute_capability()reads the current CUDA device. The descriptors passed to_check_same_cuda_devicecarry their own device index. On a host with mixed GPU architectures, the check can pass while the tensors live on a non-SM100 device. The error message also no longer reports which device was inspected.Pass the descriptor device index into the capability lookup, or state in a comment that single-architecture hosts are the supported 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/proj_rope_mxfp8/api.py` around lines 536 - 542, Update _check_sm100 to obtain compute capability for the CUDA device associated with the validated tensor descriptors, passing their device index into get_compute_capability rather than inspecting only the current device. Preserve the SM100+ validation and include the inspected device identifier in the failure message; alternatively, explicitly document that only single-architecture hosts are supported.python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py-109-112 (1)
109-112: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRename the ambiguous variable
l.Ruff reports E741 on both lines. Rename
ltobatchand update the uses at Lines 111, 139, 140, 168, 169, 209, 239, 240, 268, and 269.🔧 Proposed fix
- 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: 207-210
🤖 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 - 112, Rename the ambiguous `l` variable to `batch` in the relevant JAX API function, and update every dependent reference at the validation and later listed locations, including lines 111, 139, 140, 168, 169, 209, 239, 240, 268, and 269, while preserving the existing batch-dimension checks and behavior.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py-39-45 (1)
39-45: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRaise the JAX minimum to 0.5.1 and require
ml_dtypes>=0.5.0. JAX 0.5.0 hasjax.Array.view, but does not registerfloat8_e8m0fnu; JAX 0.5.1 adds the registration while still allowingml_dtypes0.4.x.🤖 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 39 - 45, Update the project dependency requirements to set JAX minimum version 0.5.1 and require ml_dtypes>=0.5.0, ensuring the _as_e8m0_array float8_e8m0fnu view path has the required registrations and dtype support.python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py-153-155 (1)
153-155: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winResolve undefined
torchannotations.Ruff reports 137
F821errors because both modules use module-leveltorchannotations without a module-level binding. Add aTYPE_CHECKING-only Torch import, or replace the annotations with framework-neutral types. Keep runtime Torch imports lazy.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py` around lines 153 - 155, Resolve the undefined module-level torch annotations by adding a TYPE_CHECKING-only Torch import, while keeping runtime Torch imports lazy. Apply this to _bf16_api.py at lines 153-155 and 197-206, and api.py at lines 982 and 994; update the annotations in these locations without changing runtime behavior.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py-113-113 (1)
113-113: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winDo not gate validation behind
assert.Python removes
assertstatements when the interpreter runs with-OorPYTHONOPTIMIZE. In that casegemm.check_support()never executes, so all shape, stride, dtype, and architecture validation for the JAX path is skipped and an unsupported configuration reaches kernel construction. Call the method unconditionally.🐛 Proposed fix
- assert gemm.check_support() + if not gemm.check_support(): + raise RuntimeError("Unsupported GemmSwigluSm100 configuration for the JAX entry point")🤖 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() in the JAX API path with an unconditional method call so validation always executes, including optimized Python runs. Preserve the existing validation failure behavior while ensuring unsupported shapes, strides, dtypes, and architectures are rejected before kernel construction.test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py-264-264 (1)
264-264: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winPrefix the unused unpacked variables.
Ruff reports RUF059 for
a_npandnorm_const_np. Neither value is used in this test.🧹 Proposed fix
- a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, norm_const_np = make_problem() + _a_np, b_np, c_np, sfa_u8, sfb_u8, offsets_np, alpha_np, beta_np, prob_np, _norm_const_np = make_problem()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py` at line 264, Update the unpacking of make_problem() in the test to prefix the unused a_np and norm_const_np variables with underscores, while preserving all other returned values and test behavior.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py-302-303 (1)
302-303: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUse
not infor the membership test.Ruff reports E713 on this line. The lint failure blocks CI when Ruff runs on changed files.
🧹 Proposed fix
- if not (d_dtype in _fp8_dtypes) and discrete_col_sfd: + if d_dtype not in _fp8_dtypes and discrete_col_sfd:🤖 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 membership condition in the discrete SFD handling block to use `d_dtype not in _fp8_dtypes` instead of negating an `in` expression, preserving the existing behavior of disabling `discrete_col_sfd` when the dtype is unsupported.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py-261-262 (1)
261-262: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winTwo dtype validation errors still name torch on paths that now accept JAX arrays. Both checks compare against canonical
cutlasstypes, but both messages print atorch.*dtype name. A JAX caller receives a message that references a framework it is not using.
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L261-L262: change "acc_dtype must be torch.float32" to "acc_dtype must be float32".python/cudnn/gemm/cutedsl/grouped/dglu/api.py#L999-L1000: change "dprob_tensor must have dtype torch.float32" to "dprob_tensor must have dtype float32".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py` around lines 261 - 262, Update the dtype validation messages in _bf16_api.py lines 261-262 and dglu/api.py lines 999-1000 to use framework-neutral “float32” wording instead of “torch.float32”; leave the existing canonical cutlass type checks unchanged.python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py-69-74 (1)
69-74: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAlign the rejection message with the shared phrase.
The other three JAX rejection sites use the exact phrase "not expressible as JAX arrays" (
wgrad/api.py_BLOCK_SCALED_JAX_ERROR,dglu/api.py_JAX_BLOCK_SCALED_ERROR, and thedglu/_blockscaled_api.pyguard). This message inserts "row-major" between "as" and "JAX", so it does not contain that substring. The tests assert on the shared phrase. A caller who constructs this class directly gets an error that does not match the documented pattern.🐛 Proposed fix to align the wording
raise ValueError( "The block-scaled wgrad backend supports torch tensors only: its B operand " "(and fp4-packed A/B operands) require K-major (token-innermost) layouts that " - "are not expressible as row-major JAX arrays" + "are not expressible as JAX arrays (no row-major equivalent)" )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py` around lines 69 - 74, Update the ValueError message in the sample_a validation guard to include the exact shared phrase “not expressible as JAX arrays,” removing the inserted “row-major” wording while preserving the existing explanation of the layout limitation.python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py-384-388 (1)
384-388: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winRestore the one-byte minimum for the workspace allocation.
The previous torch code allocated
max(kernel.get_workspace_bytes(), 1)bytes. This call passes the raw value. The wgrad JAX entry point still clamps withmax(kernel.get_workspace_bytes(), 1), which indicates zero is reachable. A zero-byte allocation produces a buffer whose data pointer may be null, and line 388 wraps it withfrom_dlpack.🐛 Proposed fix to clamp the workspace size
- self._workspace = allocate_byte_workspace(self._framework, kernel.get_workspace_bytes(), self.a_desc.device) + self._workspace = allocate_byte_workspace(self._framework, max(kernel.get_workspace_bytes(), 1), self.a_desc.device)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py` around lines 384 - 388, Update the workspace allocation in the initializer around _workspace to pass at least one byte, using max(kernel.get_workspace_bytes(), 1) before calling allocate_byte_workspace. Keep the existing alignment validation and from_dlpack wrapping unchanged.python/cudnn/jax/call.py-14-17 (1)
14-17: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the documented install hint to
cudnn.jaxdependency errors.cudnn.jaxis not in_LAZY_OPTIONAL_IMPORTS, so its errors bypass the root wrapper and do not mentionpip install nvidia-cudnn-frontend[cutedsl].🤖 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 14 - 17, Update the ImportError raised by the cutlass.jax availability guard in cudnn.jax to include the documented installation hint for the cudnn.jax dependency, specifically `pip install nvidia-cudnn-frontend[cutedsl]`, while preserving the existing JAX version and upgrade guidance.Source: Coding guidelines
python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py-6-6 (1)
6-6: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winSort
__all__to satisfy RUF022.Ruff reports
__all__is not sorted. The fix is mechanical and does not change runtime behavior.🔧 Proposed fix
-__all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100", "grouped_gemm_jax_sm100"] +__all__ = ["GroupedGemmSm100", "grouped_gemm_jax_sm100", "grouped_gemm_wrapper_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 `@python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py` at line 6, Sort the entries in __all__ alphabetically to satisfy RUF022, preserving the same three exported symbols and their runtime behavior.Source: Linters/SAST tools
test/python/fe_api/bsa/test_BSA_attention_jax.py-198-198 (1)
198-198: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winPrefix the unused unpacked variables with an underscore.
Ruff RUF059 reports
block_sizeandblock_sizesas unused. The test hardcodessparse_block_size=64and never reads either name.🔧 Proposed fix
- block_size, q, k, v, do, q2k, block_sparse_num, block_sizes = _make_problem(seed=3) + _block_size, q, k, v, do, q2k, block_sparse_num, _block_sizes = _make_problem(seed=3)🤖 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/bsa/test_BSA_attention_jax.py` at line 198, Update the `_make_problem(seed=3)` unpacking in the affected test to prefix the unused `block_size` and `block_sizes` variables with underscores, while preserving all other unpacked values and the hardcoded `sparse_block_size=64` behavior.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py-121-122 (1)
121-122: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDrop the block-scaled suffix from the output-dtype error.
_JAX_BLOCK_SCALED_ERRORexplains that the block-scaled backend is unavailable for JAX. That reason applies to line 118, where a non-bfloat16a_tensorimplies the block-scaled backend. It does not apply here. An unsupportedc_dtypeord_dtypesuch as int32 is rejected by the BF16 backend itself, so the appended text points the user at the wrong cause.♻️ Proposed fix
if c_dtype not in _output_dtypes or d_dtype not in _output_dtypes: - raise ValueError(f"c_dtype/d_dtype must be BF16, FP16, or FP32, got {c_dtype}/{d_dtype}; " + _JAX_BLOCK_SCALED_ERROR) + raise ValueError(f"c_dtype/d_dtype must be BF16, FP16, or FP32, got {c_dtype}/{d_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/glu/jax_api.py` around lines 121 - 122, Remove the appended _JAX_BLOCK_SCALED_ERROR text from the ValueError raised by the c_dtype/d_dtype validation in the grouped GLU JAX API, leaving only the supported-dtype message and received values. Keep the block-scaled suffix on the separate a_tensor validation where it applies.test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py-98-101 (1)
98-101: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winReplace
cutlass.jax.is_available()with an import-based check.
cutlass.jaxdoes not defineis_available(), so the current call raisesAttributeError.pytest.importorskip("cutlass.jax")handles a missing module but not this invalid attribute. Apply the correction to the other JAX tests andpython/cudnn/jax/call.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 `@test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py` around lines 98 - 101, Replace the invalid cutlass.jax.is_available() checks in the grouped GEMM JAX test and the other JAX tests, plus python/cudnn/jax/call.py, with import-based availability checks using pytest.importorskip("cutlass.jax"). Preserve the existing skip behavior and message where applicable, and remove reliance on the nonexistent is_available symbol.python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py-310-311 (1)
310-311: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the stale error text.
The check compares
self.acc_dtypeagainstcutlass.Float32, but the message namestorch.float32. The same mismatch exists at Line 340 ofpython/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py.📝 Proposed fix
if self.acc_dtype is not cutlass.Float32: - raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + raise ValueError(f"acc_dtype must be float32, got {self.acc_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/dglu/_bf16_api.py` around lines 310 - 311, Update the ValueError messages in the acc_dtype validation checks of the grouped dglu and unfused bf16 APIs to name cutlass.Float32 instead of torch.float32, matching the condition being validated.python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py-339-340 (1)
339-340: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the stale error text.
The check compares against
cutlass.Float32, but the message namestorch.float32.📝 Proposed fix
if self.acc_dtype is not cutlass.Float32: - raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + raise ValueError(f"acc_dtype must be float32, got {self.acc_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/_bf16_api.py` around lines 339 - 340, Update the ValueError message in the acc_dtype validation to name cutlass.Float32, matching the type used by the check; leave the validation condition unchanged.test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py-107-112 (1)
107-112: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUse tolerance-based comparisons for atomically accumulated outputs.
dprob_tensorusesatomic_add_float32, anddbias_tensoris also documented as atomically accumulated. Replace exact comparisons for both outputs with dtype-appropriateassert_allclosechecks. Keepd_row_tensorbyte-exact.🤖 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 107 - 112, Update the comparison loop in the grouped dGLU test so d_row_tensor continues using exact assert_array_equal, while dprob_tensor and dbias_tensor use dtype-appropriate np.testing.assert_allclose tolerances for atomic accumulation results. Preserve the existing JAX-to-Torch conversions and error context.python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py-43-81 (1)
43-81: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winReject host-backed JAX pointer arrays.
_validate_pointer_tensoronly enforces CUDA for torch tensors. JAX CPU arrays pass, and discrete SWiGLU, dSWiGLU, and dSReLU paths pass their host addresses to the kernel. Require CUDA device storage for every pointer array in this helper.🤖 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/_bf16_api.py` around lines 43 - 81, Update _validate_pointer_tensor to require CUDA device storage for all pointer-array inputs, including JAX arrays, rather than checking is_cuda only for torch tensors. Use the repository’s existing framework-neutral device check, and reject host-backed JAX arrays before returning the validated pointer count.
🤖 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_proj_rope_mxfp8.md`:
- Around line 7-9: Correct the layout description in the JAX wrapper paragraph:
state that w_out_in=True uses w shaped (NUM_HEADS·HEAD_DIM, Q_LORA), while the
unsupported [in, out] layout corresponds to w_out_in=False. Also add this
explicit weight layout to the gemm_proj_rope_mxfp8_jax_sm100(...) custom-call
paragraph, since it has no w_out_in parameter, and keep the descriptions
consistent with the definitions elsewhere in the document.
In `@docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md`:
- Around line 7-9: Update the JAX API documentation to distinguish the SFA
tensor contract from the SFB pointer contract: document SFA as the C-contiguous
physical shape (1, ceil_div(m, 128), rest_k, 32, 4, 4), and state that
per-expert SFB storage is supplied through the int64 or packed-uint8 sfb_ptrs
pointer array rather than as an SFB tensor. Keep the remaining FP8 layout and
output descriptions unchanged.
In `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py`:
- Around line 93-101: In the kernel setup around gemm.check_support and
gemm._kernel, call gemm.check_support() directly instead of relying on an assert
so validation remains active under optimization. After computing mac, reject
non-positive values before storing the (kernel, mac) entry in _kernel_cache or
allowing it to reach call().
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py`:
- Around line 79-96: Expand the cache keys used by the BF16 and MXFP8 paths to
include every input shape and dtype validated by
GemmProjRopeMxfp8Bf16InSm100.check_support(), including cos, sin, x_scale, and
w_scale, plus the relevant dtypes. Ensure calls with differing validated inputs
create separate entries so check_support() runs rather than reusing an
incompatible cached entry.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Around line 947-1095: Update the JAX branch of
discrete_grouped_gemm_dswiglu_wrapper_sm100 to route kernel execution and
donated outputs through the existing cudnn.jax.call FFI path, preserving launch
ordering and preventing downstream races. If this wrapper cannot support that
path, reject framework == "jax" before allocating JAX outputs instead of using
jax.block_until_ready around jnp.empty.
- Around line 831-834: Update the non-Torch pointer retention in the execute
path around self._live_ptrs so each generated (b_ptrs, sfb_ptrs) generation
remains referenced until its corresponding _compiled_kernel launch completes.
Use a completion event on current_stream to release prior generations only after
completion, or synchronize current_stream before replacing the retained
references; do not overwrite the sole reference while an earlier launch may
still read it.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py`:
- Around line 846-849: Update the lifetime handling in the execute path around
self._live_ptrs so JAX pointer arrays remain referenced until their asynchronous
launches complete, rather than replacing the prior tuple on each call. Use
stream-ordered retention if available; otherwise synchronize before allowing the
next execute, while preserving the existing torch record_stream behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 210-216: The module-level torch removal leaves runtime-evaluated
torch.Tensor annotations that break torch-free imports. In
python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py lines 157, 201, 206, and
210-216, and python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py lines 191,
244, and 248-254, replace the b_ptrs annotations with Any or consistently enable
postponed annotation evaluation via from __future__ import annotations; preserve
the local torch imports and ensure import cudnn works without torch, cutlass, or
cuda-python.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py`:
- Around line 54-55: Add from __future__ import annotations as the first import
in python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (lines 54-55) and
python/cudnn/gemm/cutedsl/grouped/unfused/api.py (lines 13-28). This must cover
annotations on _reinterpret_raw_grouped_fp4_tensor, _tensor_signature,
_validate_output, and GroupedGemmSm100 methods so importing cudnn does not
resolve torch.Tensor annotations without torch installed.
In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py`:
- Around line 1016-1034: Update the JAX branch of _allocate_output to allocate
zero-initialized buffers instead of using jnp.empty, matching the existing JAX
custom-call behavior while leaving the Torch allocation unchanged. Add a
regression case with padded_offsets[-1] less than valid_m and verify that
unwritten trailing rows in the returned outputs remain zero.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py`:
- Around line 357-364: Update both JAX allocations in _bf16_api.py: make the
single-expert placeholder created in the initialization flow use desc.device via
jnp.empty, and make the generated pointer array near lines 510-515 use the wgrad
tensor’s device via jnp.asarray. Apply the changes at lines 357-364 and 510-515
so both buffers remain on the inputs’ devices and pass the existing device
validation.
- Around line 202-207: Update _record_pointer_stream so immutable-framework
pointer arrays are retained per in-flight execution/stream rather than replacing
the single _live_wgrad_ptrs reference. Account for current_stream when storing
and releasing arrays, ensuring each JAX pointer array remains referenced until
its associated kernel has completed.
In `@python/cudnn/tensor_adapter.py`:
- Around line 242-264: Update detect_framework and get_device to recognize
_JaxPermutedDLPackView and transparently delegate to its underlying _array,
ensuring both return the same framework and canonical device as the wrapped JAX
array. Because the class is currently defined after these functions, move
_JaxPermutedDLPackView above them or use a deferred type lookup; then remove
callers’ need to inspect the private _array attribute.
---
Outside diff comments:
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 152-215: Add a TYPE_CHECKING-guarded torch import in the module
containing the static methods, preserving the existing runtime import behavior
while allowing annotations such as torch.Tensor to be resolved by static
analysis. Do not introduce an unconditional torch import.
---
Minor comments:
In `@AGENTS.md`:
- Around line 47-49: Update the dependency installation instructions near the
pip --group commands to document that pip 25.1 or newer is required, and provide
either a bootstrap command to upgrade pip or a direct-install fallback for the
torch and jax groups.
In `@docs/fe-oss-apis/overview.md`:
- Around line 47-53: Update the JAX installation guidance in the overview
documentation to explicitly require the appropriate GPU-enabled JAX package,
such as jax[cuda12] or jax[cuda13], for CUDA execution. Cover both the
dependency-group and published-wheel instructions, or clearly state that users
must install the matching GPU extra separately.
In `@pyproject.toml`:
- Around line 91-96: Update the jax dependency group in pyproject.toml to apply
jax>=0.5 only when python_version >= '3.10', and revise the surrounding comments
to state that JAX support requires Python 3.10 or newer.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 536-542: Update _check_sm100 to obtain compute capability for the
CUDA device associated with the validated tensor descriptors, passing their
device index into get_compute_capability rather than inspecting only the current
device. Preserve the SM100+ validation and include the inspected device
identifier in the failure message; alternatively, explicitly document that only
single-architecture hosts are supported.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/jax_api.py`:
- Around line 39-45: Update the project dependency requirements to set JAX
minimum version 0.5.1 and require ml_dtypes>=0.5.0, ensuring the _as_e8m0_array
float8_e8m0fnu view path has the required registrations and dtype support.
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Around line 109-112: Rename the ambiguous `l` variable to `batch` in the
relevant JAX API function, and update every dependent reference at the
validation and later listed locations, including lines 111, 139, 140, 168, 169,
209, 239, 240, 268, and 269, while preserving the existing batch-dimension
checks and behavior.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py`:
- Line 113: Replace the assert around gemm.check_support() in the JAX API path
with an unconditional method call so validation always executes, including
optimized Python runs. Preserve the existing validation failure behavior while
ensuring unsupported shapes, strides, dtypes, and architectures are rejected
before kernel construction.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.py`:
- Around line 302-303: Update the membership condition in the discrete SFD
handling block to use `d_dtype not in _fp8_dtypes` instead of negating an `in`
expression, preserving the existing behavior of disabling `discrete_col_sfd`
when the dtype is unsupported.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 310-311: Update the ValueError messages in the acc_dtype
validation checks of the grouped dglu and unfused bf16 APIs to name
cutlass.Float32 instead of torch.float32, matching the condition being
validated.
In `@python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py`:
- Around line 153-155: Resolve the undefined module-level torch annotations by
adding a TYPE_CHECKING-only Torch import, while keeping runtime Torch imports
lazy. Apply this to _bf16_api.py at lines 153-155 and 197-206, and api.py at
lines 982 and 994; update the annotations in these locations without changing
runtime behavior.
In `@python/cudnn/gemm/cutedsl/grouped/glu/jax_api.py`:
- Around line 121-122: Remove the appended _JAX_BLOCK_SCALED_ERROR text from the
ValueError raised by the c_dtype/d_dtype validation in the grouped GLU JAX API,
leaving only the supported-dtype message and received values. Keep the
block-scaled suffix on the separate a_tensor validation where it applies.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/__init__.py`:
- Line 6: Sort the entries in __all__ alphabetically to satisfy RUF022,
preserving the same three exported symbols and their runtime behavior.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py`:
- Around line 339-340: Update the ValueError message in the acc_dtype validation
to name cutlass.Float32, matching the type used by the check; leave the
validation condition unchanged.
- Around line 43-81: Update _validate_pointer_tensor to require CUDA device
storage for all pointer-array inputs, including JAX arrays, rather than checking
is_cuda only for torch tensors. Use the repository’s existing framework-neutral
device check, and reject host-backed JAX arrays before returning the validated
pointer count.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/api.py`:
- Around line 192-193: Update the dtype validation messages in the grouped GEMM
API checks, including the checks for a_tensor, prob_tensor, b_tensor, and
b_dtype, to name the logical dtypes bfloat16 and float32 without the torch.
prefix. Keep the existing CUTLASS dtype comparisons and validation behavior
unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py`:
- Around line 261-262: Update the dtype validation messages in _bf16_api.py
lines 261-262 and dglu/api.py lines 999-1000 to use framework-neutral “float32”
wording instead of “torch.float32”; leave the existing canonical cutlass type
checks unchanged.
- Around line 384-388: Update the workspace allocation in the initializer around
_workspace to pass at least one byte, using max(kernel.get_workspace_bytes(), 1)
before calling allocate_byte_workspace. Keep the existing alignment validation
and from_dlpack wrapping unchanged.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py`:
- Around line 69-74: Update the ValueError message in the sample_a validation
guard to include the exact shared phrase “not expressible as JAX arrays,”
removing the inserted “row-major” wording while preserving the existing
explanation of the layout limitation.
In `@python/cudnn/jax/call.py`:
- Around line 14-17: Update the ImportError raised by the cutlass.jax
availability guard in cudnn.jax to include the documented installation hint for
the cudnn.jax dependency, specifically `pip install
nvidia-cudnn-frontend[cutedsl]`, while preserving the existing JAX version and
upgrade guidance.
In `@test/python/fe_api/bsa/test_BSA_attention_jax.py`:
- Line 198: Update the `_make_problem(seed=3)` unpacking in the affected test to
prefix the unused `block_size` and `block_sizes` variables with underscores,
while preserving all other unpacked values and the hardcoded
`sparse_block_size=64` behavior.
In `@test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu_jax.py`:
- Line 264: Update the unpacking of make_problem() in the test to prefix the
unused a_np and norm_const_np variables with underscores, while preserving all
other returned values and test behavior.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py`:
- Around line 107-112: Update the comparison loop in the grouped dGLU test so
d_row_tensor continues using exact assert_array_equal, while dprob_tensor and
dbias_tensor use dtype-appropriate np.testing.assert_allclose tolerances for
atomic accumulation results. Preserve the existing JAX-to-Torch conversions and
error context.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_jax.py`:
- Around line 98-101: Replace the invalid cutlass.jax.is_available() checks in
the grouped GEMM JAX test and the other JAX tests, plus
python/cudnn/jax/call.py, with import-based availability checks using
pytest.importorskip("cutlass.jax"). Preserve the existing skip behavior and
message where applicable, and remove reliance on the nonexistent is_available
symbol.
---
Nitpick comments:
In `@python/cudnn/block_sparse_attention/_interface.py`:
- Around line 117-122: Move the contextlib import from inside _nvtx_range to
module scope, then keep the function’s non-Torch path returning
contextlib.nullcontext() with unchanged behavior.
In `@python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py`:
- Line 2206: Update the call site in the backward attention implementation to
pass q directly to detect_framework without inspecting its private _array
attribute. Modify tensor_adapter.detect_framework to unwrap
_JaxPermutedDLPackView internally, and update any other call sites that
conditionally access _array to use the same direct call pattern.
In `@python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py`:
- Around line 150-153: Centralize the tensor detachment helper in
cudnn.tensor_adapter using the is_torch_tensor semantics, detaching torch
tensors and passing other framework objects through unchanged. Remove the local
_maybe_detach definitions and import the shared helper in
python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py (lines 150-153)
and python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py (lines
251-253); add it to the existing import block in
python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py (lines 53-55).
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 592-617: Replace the duplicated torch/JAX buffer allocation branch
in the wrapper with cudnn.tensor_adapter.allocate_tensor for all four output
buffers, preserving each existing shape and dtype. Keep the w_out_in=False JAX
rejection in place, and add allocate_tensor to the existing tensor_adapter
imports.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py`:
- Around line 1113-1133: Update the cache-key construction around
dynamic_m_tensor_signature for sfa_tensor to use a layout-aware scale-factor
signature, following the _sf_is_physical and dynamic_m_sf_signature pattern from
the grouped dsrelu API. Preserve the existing M-independent behavior while
distinguishing permuted atom and physical contiguous layouts, and use the
resulting signature in place of the current hard-coded dimension assumptions.
- Around line 141-143: Remove the redundant self._interpret_uint8_as_fp4x2 =
True assignment from the later descriptor-building flow, while retaining the
earlier assignment before descriptor construction so uint8 sample tensors
continue using fp4 logical shapes.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 586-590: Replace the duplicated device and 8-byte alignment checks
after _validate_pointer_tensor in the relevant class method with its existing
_validate_pointer_array_alignment helper, matching the unfused sibling’s call
pattern while preserving the current validation order and error behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py`:
- Around line 52-101: Consolidate _grouped_dglu_bf16_adapter and
_grouped_dglu_bf16_dbias_adapter into one adapter accepting dbias and forwarding
it as dbias_tensor, after verifying the bridge does not require distinct traced
arities. Update callers to pass the appropriate value, using None for the
no-dbias path, while preserving all other kernel arguments and behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py`:
- Around line 49-52: Remove the local _prob_spec definition in the dsrelu JAX
API and import _prob_spec from the existing grouped.unfused.jax_api module
alongside _pointer_count. Preserve all current callers and use the shared
implementation for the (m, 1, 1) TensorSpec layout.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py`:
- Around line 111-123: Update the framework validation in the initializer around
detect_framework and self._framework so it runs unconditionally, including when
sample_a is None. Reject the resulting "unknown" framework with the existing
clear unsupported-framework error, while preserving the JAX-specific message and
torch initialization behavior.
In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py`:
- Around line 42-47: Move the duplicated _JAX_SF_LAYOUT_ERROR diagnostic into a
shared internal helper within the grouped GEMM package, keeping it private and
unexported through cudnn. In python/cudnn/gemm/cutedsl/grouped/quant/api.py
lines 42-47, python/cudnn/gemm/cutedsl/grouped/srelu/api.py lines 41-45, and
python/cudnn/gemm/cutedsl/grouped/swiglu/api.py lines 33-37, remove the local
definitions and import the shared constant for each API’s existing rejection
path.
- Around line 985-995: Initialize self._compile_b_ptrs and
self._compile_sfb_ptrs to None in the relevant __init__ method near
self._workspace = None, matching the initialization pattern in
GroupedGemmDsreluSm100.__init__ and DiscreteGroupedGemmSwigluSm100.__init__;
leave _compile_discrete as the writer that replaces them for discrete mode.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py`:
- Around line 43-46: Move the shared _pointer_count and _prob_spec helpers from
grouped/unfused/jax_api.py into a small internal framework-neutral module under
gemm/cutedsl/, then update discrete_grouped/swiglu/jax_api.py and
grouped/unfused callers to import them from that module. Keep both helpers
private and do not expose the new module or symbols through cudnn.
- Around line 195-200: Extract the duplicated max_active_clusters and
workspace-entry logic into an unexported shared helper, such as
_kernel_entry(kernel, cluster_shape_mn), preserving the environment variable,
workspace floor, and error text. Replace the blocks in
python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py:195-200,
grouped/glu/jax_api.py:194-199, grouped/dsrelu/jax_api.py:362-367,
grouped/wgrad/jax_api.py:184-189, and discrete_grouped/swiglu/jax_api.py:313-318
with calls to that helper; define it in the internal shared GEMM helpers module
and do not export it from cudnn.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/api.py`:
- Around line 36-45: Move the shared _block_scaled_dtype_pairs definition from
the grouped API modules into backend_utils.py, then update the grouped wgrad,
dglu, and glu callers to import or reuse that centralized symbol when invoking
select_grouped_gemm_backend. Remove the duplicate local definitions while
preserving the existing dtype pairs and behavior.
In `@python/cudnn/jax/__init__.py`:
- Around line 15-33: Update the __all__ list in the JAX package initializer to
exactly alphabetical order: TensorSpec, call, gemm_operand_spec, neg_inf_init,
row_major_desc, sf_atom_spec, zeros_init. Leave the imports and top-level
package behavior unchanged.
In `@python/cudnn/jax/call.py`:
- Around line 94-100: Update the docstring for initialized_outputs to explicitly
state that its keys are flattened output-leaf indices matching the order from
jax.tree.leaves(output_shape_dtype, ...), not structural tree indices.
- Around line 20-44: Update row_major_desc to use
TensorDesc._compute_contiguous_stride instead of manually calculating row-major
strides, while preserving the existing shape and stride-order handling. Extend
its docstring to explicitly state that the metadata descriptor assumes CUDA
device 0 and may mismatch real tensors on other devices during check_support.
In `@test/python/fe_api/bsa/test_BSA_attention_jax.py`:
- Line 176: Update the assertion in the relevant attention test to verify that
got is an actual JAX array using the public/intended JAX array type, rather than
checking absence of the private "_array" attribute. Preserve the existing
key-specific failure message and ensure the assertion rejects internal view
wrappers directly.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py`:
- Line 149: Rename the ambiguous variable l in the test’s GEMM dimension setup
to batch or num_batches, and update every reference to it within the test so
behavior remains unchanged and Ruff E741 is resolved.
- Around line 178-183: Repeat the backward invocation through jitted_b in the
test, using the same inputs after the first d_jit/dprob_jit result is
synchronized, then synchronize and validate the second outputs against the eager
reference. Preserve the existing exact d_tensor comparison and tight dprob
tolerance checks for the repeated call to exercise donated accumulator reuse.
In `@test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_swiglu_jax.py`:
- Around line 182-207: Update the expected-byte construction and result
comparison in the swiglu test to use the existing as_bytes helper instead of
direct np.asarray(...).view(np.uint8) conversions. Apply it to c_tensor,
d_tensor, and d_col_tensor while preserving the current eager-versus-JIT
assertions.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dsrelu_jax.py`:
- Line 284: Update the unpacking assignment from make_problem in the grouped
GEMM test to rename the unused b_np and norm_const_np variables to _b_np and
_norm_const_np, preserving the existing value order and all other bindings.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py`:
- Line 23: Rename the ambiguous l parameter and corresponding local variable in
_make_jax_inputs to num_experts (or batch), and update all references and
keyword call sites to use the new name while preserving the existing behavior.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_hadamard_jax.py`:
- Line 17: Remove the unused module-level torch import skip guard from
test_grouped_gemm_glu_hadamard_jax.py, while leaving skip_unless_sm100() and
both API-entry-point tests unchanged.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py`:
- Around line 27-28: Replace the duplicated nested-negation ceiling divisions
with the existing ceil_div helper. In
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py lines 27-28,
import ceil_div and use it for rest_k and the m-based dimension; apply the same
import and replacements in
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py lines 27-28.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py`:
- Line 27: Update the rest_k calculation in the grouped GEMM test to use (k +
15) // 64, preserving the existing nested ceiling-division result for every
positive integer k, including k=1.
🪄 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: 7d115f43-d81c-44f3-9261-29ef144e0f2f
📒 Files selected for processing (101)
AGENTS.mddocs/fe-oss-apis/bsa.mddocs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/gemm_amax.mddocs/fe-oss-apis/gemm_fusions/gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.mddocs/fe-oss-apis/gemm_fusions/gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.mddocs/fe-oss-apis/overview.mdpyproject.tomlpython/cudnn/__init__.pypython/cudnn/block_sparse_attention/_interface.pypython/cudnn/block_sparse_attention/api.pypython/cudnn/block_sparse_attention/csrc/bwd/bsa_bwd_prepost.pypython/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.pypython/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.pypython/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.pypython/cudnn/gemm/cutedsl/_jax_ffi.pypython/cudnn/gemm/cutedsl/dense/amax/__init__.pypython/cudnn/gemm/cutedsl/dense/amax/api.pypython/cudnn/gemm/cutedsl/dense/amax/jax_api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/__init__.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/__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/proj_rope_mxfp8/jax_api.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/api.pypython/cudnn/gemm/cutedsl/dense/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/__init__.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/__init__.pypython/cudnn/gemm/cutedsl/grouped/backend_utils.pypython/cudnn/gemm/cutedsl/grouped/dglu/__init__.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/dglu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/__init__.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.pypython/cudnn/gemm/cutedsl/grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/grouped/glu/__init__.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/jax_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/__init__.pypython/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/unfused/api.pypython/cudnn/gemm/cutedsl/grouped/unfused/jax_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/__init__.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/gemm/cutedsl/grouped/wgrad/jax_api.pypython/cudnn/jax/__init__.pypython/cudnn/jax/call.pypython/cudnn/tensor_adapter.pytest/python/conftest.pytest/python/fe_api/bsa/test_BSA_attention_jax.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)
- python/cudnn/gemm/cutedsl/_jax_ffi.py
- test/python/fe_api/gemm/test_cutedsl_jax_guards.py
Comment on lines
+7
to
+9
| Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` (the `[in, out]` weight layout reaches the kernel through a transposed strided view, which has no row-major JAX equivalent and raises a clear error). The E8M0 scale inputs stay `uint8` as with torch. Outputs are allocated as C-contiguous `jnp` arrays. The wrapper is eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. | ||
|
|
||
| For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"): same contract as the wrapper with `w_out_in=True`, dispatching on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs. |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Correct the w_out_in layout description.
Line 7 associates w_out_in=True with [in, out], but Line [146] defines w_out_in=True as [out, in]. Lines [33]-[34] use the same [out, in] interpretation. This contradiction can make JAX callers pass the rejected layout or misinterpret the custom-call input contract. State that JAX supports w_out_in=True with w shaped (NUM_HEADS·HEAD_DIM, Q_LORA), and that the unsupported [in, out] path is w_out_in=False. Add the explicit layout to the custom-call paragraph because gemm_proj_rope_mxfp8_jax_sm100(...) has no w_out_in parameter.
Proposed documentation fix
-Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` (the `[in, out]` weight layout reaches the kernel through a transposed strided view, which has no row-major JAX equivalent and raises a clear error).
+Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` and `w` stored in the `[out, in]` layout. The unsupported `w_out_in=False` (`[in, out]`) path requires a transposed strided view, which has no row-major JAX equivalent and raises a clear error.
-For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"): same contract as the wrapper with `w_out_in=True`, dispatching on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs.
+For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"). The entry point expects `w` in the `[out, in]` layout with shape `(NUM_HEADS·HEAD_DIM, Q_LORA)`, equivalent to the wrapper with `w_out_in=True`. It dispatches on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` (the `[in, out]` weight layout reaches the kernel through a transposed strided view, which has no row-major JAX equivalent and raises a clear error). The E8M0 scale inputs stay `uint8` as with torch. Outputs are allocated as C-contiguous `jnp` arrays. The wrapper is eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. | |
| For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"): same contract as the wrapper with `w_out_in=True`, dispatching on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs. | |
| Supports **JAX arrays** on both input paths (BF16 and MXFP8) with `w_out_in=True` and `w` stored in the `[out, in]` layout. The unsupported `w_out_in=False` (`[in, out]`) path requires a transposed strided view, which has no row-major JAX equivalent and raises a clear error. The E8M0 scale inputs stay `uint8` as with torch. Outputs are allocated as C-contiguous `jnp` arrays. The wrapper is eager only, on the CUDA legacy default stream: `block_until_ready` inputs, synchronize before reading outputs. | |
| For jitted JAX programs use the `jax.jit`-compatible XLA custom-call entry point `gemm_proj_rope_mxfp8_jax_sm100(x, w, cos, sin, x_scale=None, w_scale=None)` (built on `cudnn.jax.call`; see `gemm_amax.md` "Using JAX arrays"). The entry point expects `w` in the `[out, in]` layout with shape `(NUM_HEADS·HEAD_DIM, Q_LORA)`, equivalent to the wrapper with `w_out_in=True`. It dispatches on `x.dtype` (bfloat16 → BF16 GEMM; float8_e4m3fn plus E8M0 scales → MXFP8 GEMM), returning `(out_fp8_row, out_scales_row, out_fp8_col, out_scales_col)` as fresh XLA-managed arrays — no manual synchronization needed, composes with `jax.jit` and CUDA graphs. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md` around lines 7 - 9,
Correct the layout description in the JAX wrapper paragraph: state that
w_out_in=True uses w shaped (NUM_HEADS·HEAD_DIM, Q_LORA), while the unsupported
[in, out] layout corresponds to w_out_in=False. Also add this explicit weight
layout to the gemm_proj_rope_mxfp8_jax_sm100(...) custom-call paragraph, since
it has no w_out_in parameter, and keep the descriptions consistent with the
definitions elsewhere in the document.
Comment on lines
+93
to
+101
| assert gemm.check_support() | ||
| kernel = gemm._kernel( | ||
| sf_vec_size=sf_vec_size, | ||
| mma_tiler_mn=mma_tiler_mn, | ||
| cluster_shape_mn=cluster_shape_mn, | ||
| ) | ||
| mac = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]) - gemm.num_cluster_overlap_margin | ||
| entry = (kernel, mac) | ||
| _kernel_cache[cache_key] = entry |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 3 'assert gemm\.check_support\(\)|max_active_clusters|mac =' \
python/cudnn/gemm/cutedsl/dense/amax/jax_api.py \
python/cudnn/gemm/cutedsl/dense/amax/api.pyRepository: NVIDIA/cudnn-frontend
Length of output: 3906
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- jax_api.py outline ---'
ast-grep outline python/cudnn/gemm/cutedsl/dense/amax/jax_api.py
printf '%s\n' '--- api.py outline ---'
ast-grep outline python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- relevant jax bridge ---'
sed -n '1,180p' python/cudnn/gemm/cutedsl/dense/amax/jax_api.py
printf '%s\n' '--- support and compile implementation ---'
sed -n '250,370p' python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- related call sites and tests ---'
rg -n -C 4 'GemmAmaxSm100|check_support\(\)|num_cluster_overlap_margin|get_max_active_clusters|CUDNNFE_CLUSTER_OVERLAP_MARGIN' \
python testRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- jax_api.py ---'
sed -n '1,135p' python/cudnn/gemm/cutedsl/dense/amax/jax_api.py
printf '%s\n' '--- amax api methods ---'
rg -n -A 45 -B 12 'def check_support|def _compile_kernel|def compile|def _kernel|num_cluster_overlap_margin|_value_error_if' \
python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- JAX custom-call arguments ---'
sed -n '100,180p' python/cudnn/gemm/cutedsl/dense/amax/jax_api.py
printf '%s\n' '--- focused references ---'
rg -n -C 3 'dense\.amax|gemm_amax|CUDNNFE_CLUSTER_OVERLAP_MARGIN|assert gemm\.check_support' \
python/cudnn/gemm/cutedsl/dense/amax test/python/gemm 2>/dev/null || true
printf '%s\n' '--- deterministic Python behavior probe ---'
python3 - <<'PY'
import ast
path = "python/cudnn/gemm/cutedsl/dense/amax/jax_api.py"
tree = ast.parse(open(path, encoding="utf-8").read())
for node in ast.walk(tree):
if isinstance(node, ast.Assert):
print("assert:", ast.unparse(node))
PY
python3 -O - <<'PY'
def check_support():
print("check_support called")
return False
assert check_support()
print("execution continues with -O")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 32944
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- kernel assignment and definitions ---'
rg -n -C 5 'Sm100BlockScaledPersistentDenseGemmKernel|def __call__|max_active_clusters|mac' \
python/cudnn/gemm/cutedsl/dense/amax \
python/cudnn/gemm/cutedsl/dense 2>/dev/null | head -n 240
printf '%s\n' '--- JAX call implementation references ---'
rg -n -C 5 'def call|initialized_outputs|constexpr|kernel=.*mac|output_spec' \
python/cudnn/jax.py python/cudnn/jax 2>/dev/null | head -n 260
printf '%s\n' '--- focused cache-key and environment references ---'
rg -n -C 3 '_kernel_cache|cache_key|CUDNNFE_CLUSTER_OVERLAP_MARGIN|num_cluster_overlap_margin' \
python/cudnn/gemm/cutedsl/dense/amax/jax_api.py \
python/cudnn/gemm/cutedsl/dense/amax/api.pyRepository: NVIDIA/cudnn-frontend
Length of output: 28373
Call validation directly and reject non-positive mac. assert gemm.check_support() is skipped under python -O; call gemm.check_support() directly. The JAX path calls gemm._kernel(...) and does not run _compile_kernel(), so validate mac <= 0 before caching or passing it to 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/amax/jax_api.py` around lines 93 - 101, In
the kernel setup around gemm.check_support and gemm._kernel, call
gemm.check_support() directly instead of relying on an assert so validation
remains active under optimization. After computing mac, reject non-positive
values before storing the (kernel, mac) entry in _kernel_cache or allowing it to
reach call().
Comment on lines
+79
to
+96
| 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 |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Validation is skipped on a cache hit for shapes not covered by the cache key.
Both cache keys use only x.shape and w.shape. check_support() runs only when the entry is missing. check_support() also validates cos/sin shapes, x_scale/w_scale shapes, and the even-num_heads requirement. After the first successful call for a given (x.shape, w.shape), a later call with the same x/w shapes but a wrong cos, sin, or scale shape passes straight to the custom call. The kernel then reads outside the supplied buffers.
Include the remaining input shapes and dtypes in the cache key.
🐛 Proposed fix for the BF16 path
- cache_key = (tuple(x.shape), tuple(w.shape))
+ cache_key = (tuple(x.shape), tuple(w.shape), tuple(cos.shape), tuple(sin.shape), x.dtype, w.dtype)Apply the same change on the MXFP8 path, adding x_scale.shape and w_scale.shape.
Also applies to: 113-132
🤖 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, Expand the cache keys used by the BF16 and MXFP8 paths to include every
input shape and dtype validated by GemmProjRopeMxfp8Bf16InSm100.check_support(),
including cos, sin, x_scale, and w_scale, plus the relevant dtypes. Ensure calls
with differing validated inputs create separate entries so check_support() runs
rather than reusing an incompatible cached entry.
Comment on lines
+831
to
+834
| if not is_torch_tensor(b_ptrs): | ||
| # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the | ||
| # arrays referenced until the next execute so their buffers outlive the launch. | ||
| self._live_ptrs = (b_ptrs, sfb_ptrs) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect _live_ptrs usage and the stream chosen for non-torch frameworks.
rg -nP -C6 '_live_ptrs' --type=py
rg -nP -C10 'def default_stream' --type=pyRepository: NVIDIA/cudnn-frontend
Length of output: 159
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate file symbols ---'
rg -n -C8 'live_ptrs|default_stream|is_torch_tensor|execute\(' python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py || true
printf '%s\n' '--- stream helpers and usages ---'
rg -n -C8 'def default_stream|default_stream\("jax"|default_stream\(.*jax|record_stream|current_stream' python test || true
printf '%s\n' '--- related files ---'
git ls-files | rg '(^|/)(api\.py|.*stream.*|.*jax.*|.*dswiglu.*)$' | head -100Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -eu
file=python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py
printf '%s\n' '--- reviewed API references ---'
rg -n -C12 'live_ptrs|is_torch_tensor|default_stream|detect_framework|_compiled_kernel\(' "$file"
printf '%s\n' '--- helper definitions ---'
rg -l 'def default_stream|def detect_framework|def is_torch_tensor' python/cudnn --glob '*.py' |
while IFS= read -r f; do
echo "FILE $f"
rg -n -C12 'def default_stream|def detect_framework|def is_torch_tensor' "$f"
done
printf '%s\n' '--- reviewed execute range ---'
sed -n '760,875p' "$file"Repository: NVIDIA/cudnn-frontend
Length of output: 16388
🏁 Script executed:
#!/bin/bash
set -eu
api=python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py
printf '%s\n' '--- wrapper allocation and execute call ---'
sed -n '945,1160p' "$api"
printf '%s\n' '--- all reviewed API call sites ---'
rg -n -C8 'discrete_grouped_gemm_dswiglu_wrapper_sm100|DiscreteGroupedGemmDswigluSm100|\.execute\(' python/cudnn test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py |
rg 'discrete_grouped|Dswiglu|b_ptrs|sfb_ptrs|current_stream|execute|^python/' | head -300
printf '%s\n' '--- JAX and pointer-array handling in the repository ---'
rg -n -C8 'jax\.device_put|device_put\(|b_ptrs|sfb_ptrs|pointer.*array|array.*pointer|default_stream\("jax"|default_stream\(framework\)' python/cudnn test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py |
head -500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -eu
api=python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py
test=test/python/fe_api/grouped_gemm/test_discrete_grouped_gemm_dswiglu.py
printf '%s\n' '--- wrapper execute arguments ---'
sed -n '1115,1195p' "$api"
printf '%s\n' '--- pointer-array input construction ---'
rg -n -C12 'b_ptrs_tensor|sfb_ptrs_tensor|stack\(.*data_ptr|data_ptr.*stack|b_ptrs\s*=|sfb_ptrs\s*=' "$test" python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu python/cudnn --glob '*.py' |
head -400
printf '%s\n' '--- direct API stream arguments ---'
sed -n '245,305p' "$test"
sed -n '340,390p' "$test"
printf '%s\n' '--- all _live_ptrs assignments and reads, limited to production code ---'
rg -n -C5 '_live_ptrs' python/cudnn --glob '*.py'Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🌐 Web query:
JAX Python array buffer deallocation stream ordering CUDA external custom call legacy default stream buffer lifetime
💡 Result:
In JAX, managing buffer lifetimes and stream ordering within external custom calls (or the modern Foreign Function Interface, FFI) requires adhering to XLA's static memory management model [1]. Buffer Lifetime and Management XLA uses a destination-passing style where it performs static buffer assignment at compile time based on the live ranges of all values [1]. Custom calls—whether legacy or FFI-based—do not allocate memory for their results; they must write into destination buffers provided by the XLA runtime [1]. Because XLA manages these lifetimes, a custom call must never attempt to free, reallocate, or own the device buffers passed to it as arguments or result destinations [1]. Attempting to perform manual memory management (e.g., calling cudaFree or cudaMalloc inside a handler) is unsafe and can cause corruption of XLA-managed memory, such as closure-captured constants or other active buffers [2]. Stream Ordering and Synchronization Custom call handlers are executed on the host (CPU), but they are responsible for enqueueing asynchronous work onto the device (GPU) [1]. When working with CUDA, the handler is provided with a cudaStream_t (often accessed via ffi::PlatformStream<CUstream> in FFI) [1][3]. 1. Asynchronous Execution: All kernel launches or device-side operations performed within the handler must be submitted to the provided stream to ensure they are correctly ordered with respect to other XLA-dispatched operations [1][4]. 2. Race Conditions: If a handler performs host-side operations that depend on device memory (e.g., cudaMemcpy from device to host), it must ensure proper synchronization (e.g., using cudaStreamSynchronize or managing dependencies correctly) [2]. Failing to synchronize results in race conditions because the FFI handler execution on the CPU may complete before the previously enqueued GPU kernels have finished writing to the buffers [4][2]. 3. Legacy vs. FFI: While the legacy API provided raw pointers and opaque descriptors, the modern XLA FFI (recommended for new development) provides typed interfaces and safer access to context, such as the platform stream, which reduces the potential for synchronization errors [3][5]. If using legacy custom calls, you are responsible for maintaining the same stream ordering principles [6][5]. To avoid segmentation faults or memory corruption, do not access device memory directly from the host-side FFI handler without proper device-to-host transfer/synchronization, and always rely on the provided stream to serialize device work [1][7].
Citations:
- 1: https://openxla.org/xla/custom_call
- 2: FFI handlers can corrupt closure-captured constant buffers in jax.lax.scan jax-ml/jax#36580
- 3: https://docs.jax.dev/en/latest/ffi.html
- 4: https://github.com/jax-ml/jax/blob/main/examples/ffi/src/jax_ffi_example/cuda_examples.cu
- 5: https://docs.jax.dev/en/latest/%5Fsources/ffi.md
- 6: https://github.com/dfm/extending-jax/tree/main?tab=readme-ov-file
- 7: FFI calls on GPU lead to segmentation fault jax-ml/jax#31593
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py")
tree = ast.parse(path.read_text())
execute = next(
node for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "execute"
)
events = []
for node in ast.walk(execute):
if isinstance(node, ast.Assign):
text = ast.get_source_segment(path.read_text(), node) or ""
if "_live_ptrs" in text:
events.append((node.lineno, "live_ptrs assignment"))
if isinstance(node, ast.Call):
func = node.func
if isinstance(func, ast.Attribute) and func.attr == "_compiled_kernel":
events.append((node.lineno, "compiled kernel launch"))
print("execute events:", sorted(events))
assignment = next(line for line, kind in events if kind == "live_ptrs assignment")
launch = next(line for line, kind in events if kind == "compiled kernel launch")
print("release-before-launch:", assignment < launch)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 268
Keep each non-Torch pointer-array generation alive until its launch completes.
self._live_ptrs is overwritten before _compiled_kernel(...) is enqueued. A temporary JAX b_ptrs or sfb_ptrs can therefore lose its last reference while the previous launch still reads it. Retain generations until a completion event on current_stream, or synchronize that stream before releasing them.
🤖 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 831 -
834, Update the non-Torch pointer retention in the execute path around
self._live_ptrs so each generated (b_ptrs, sfb_ptrs) generation remains
referenced until its corresponding _compiled_kernel launch completes. Use a
completion event on current_stream to release prior generations only after
completion, or synchronize current_stream before replacing the retained
references; do not overwrite the sole reference while an earlier launch may
still read it.
Comment on lines
54
to
+55
| def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: | ||
| import torch | ||
|
|
||
| if tensor.dtype == torch.uint8: | ||
| if _convert_to_cutlass_data_type_or_none(tensor.dtype) is cutlass.Uint8: |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Module-level torch annotations remain after the eager torch import was removed. Both modules dropped import torch from their headers but still annotate functions with torch.Tensor. Python evaluates those annotations at definition time, so importing either module raises NameError: name 'torch' is not defined on a torch-free install. The single root cause is the missing from __future__ import annotations declaration; adding it makes all annotations lazy strings and resolves both sites.
python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py#L54-L55: addfrom __future__ import annotationsas the first import so thetorch.Tensorannotations on_reinterpret_raw_grouped_fp4_tensorno longer resolve at import time.python/cudnn/gemm/cutedsl/grouped/unfused/api.py#L13-L28: addfrom __future__ import annotationsto the import block so_tensor_signature,_validate_output, and theGroupedGemmSm100method signatures no longer resolvetorchat import time.
As per coding guidelines, "import cudnn must work without torch, cutlass, or cuda-python installed".
🧰 Tools
🪛 Ruff (0.16.1)
[error] 54-54: Undefined name torch
(F821)
[error] 54-54: Undefined name torch
(F821)
📍 Affects 2 files
python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py#L54-L55(this comment)python/cudnn/gemm/cutedsl/grouped/unfused/api.py#L13-L28
🤖 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 54 - 55, Add
from __future__ import annotations as the first import in
python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (lines 54-55) and
python/cudnn/gemm/cutedsl/grouped/unfused/api.py (lines 13-28). This must cover
annotations on _reinterpret_raw_grouped_fp4_tensor, _tensor_signature,
_validate_output, and GroupedGemmSm100 methods so importing cudnn does not
resolve torch.Tensor annotations without torch installed.
Sources: Coding guidelines, Linters/SAST tools
Comment on lines
+1016
to
+1034
| def _allocate_output(shape, stride, dtype): | ||
| if framework == "torch": | ||
| import torch | ||
|
|
||
| return torch.empty_strided( | ||
| shape, | ||
| stride, | ||
| dtype=framework_dtype(dtype, "torch"), | ||
| device=call.a_tensor.device, | ||
| ) | ||
| import jax | ||
| import jax.numpy as jnp | ||
|
|
||
| # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. | ||
| # The kernel writes into this buffer on the launch stream; materialize it first. | ||
| return jax.block_until_ready(jnp.empty(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device)) | ||
|
|
||
| c_tensor = _allocate_output((valid_m, n_full, 1), (n_full, 1, valid_m * n_full), call.c_dtype) | ||
| d_tensor = _allocate_output((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), call.d_dtype) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Zero-initialize JAX outputs before the eager kernel launch.
The offset contract permits padded_offsets[-1] < valid_m. The kernel can then leave trailing rows unwritten. Line 1031 allocates those JAX buffers with jnp.empty, so the wrapper can return undefined values. Allocate zeroed JAX buffers, as the JAX custom-call path already does.
Proposed fix
- return jax.block_until_ready(jnp.empty(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device))
+ return jax.block_until_ready(jnp.zeros(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device))Add a regression case where padded_offsets[-1] < valid_m and verify that the remaining rows are zero.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _allocate_output(shape, stride, dtype): | |
| if framework == "torch": | |
| import torch | |
| return torch.empty_strided( | |
| shape, | |
| stride, | |
| dtype=framework_dtype(dtype, "torch"), | |
| device=call.a_tensor.device, | |
| ) | |
| import jax | |
| import jax.numpy as jnp | |
| # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. | |
| # The kernel writes into this buffer on the launch stream; materialize it first. | |
| return jax.block_until_ready(jnp.empty(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device)) | |
| c_tensor = _allocate_output((valid_m, n_full, 1), (n_full, 1, valid_m * n_full), call.c_dtype) | |
| d_tensor = _allocate_output((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), call.d_dtype) | |
| def _allocate_output(shape, stride, dtype): | |
| if framework == "torch": | |
| import torch | |
| return torch.empty_strided( | |
| shape, | |
| stride, | |
| dtype=framework_dtype(dtype, "torch"), | |
| device=call.a_tensor.device, | |
| ) | |
| import jax | |
| import jax.numpy as jnp | |
| # n-major C-contiguous; the extent-1 batch dim's stride is unobservable. | |
| # The kernel writes into this buffer on the launch stream; materialize it first. | |
| return jax.block_until_ready(jnp.zeros(shape, dtype=framework_dtype(dtype, "jax"), device=call.a_tensor.device)) | |
| c_tensor = _allocate_output((valid_m, n_full, 1), (n_full, 1, valid_m * n_full), call.c_dtype) | |
| d_tensor = _allocate_output((valid_m, n_out, 1), (n_out, 1, valid_m * n_out), call.d_dtype) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py` around lines 1016 - 1034,
Update the JAX branch of _allocate_output to allocate zero-initialized buffers
instead of using jnp.empty, matching the existing JAX custom-call behavior while
leaving the Torch allocation unchanged. Add a regression case with
padded_offsets[-1] less than valid_m and verify that unwritten trailing rows in
the returned outputs remain zero.
Comment on lines
+202
to
+207
| def _record_pointer_stream(self, pointers: torch.Tensor, current_stream: cuda.CUstream) -> None: | ||
| if not is_torch_tensor(pointers): | ||
| # No record_stream equivalent for immutable frameworks (e.g. JAX): keep the | ||
| # array referenced until the next execute so its buffer outlives the launch. | ||
| self._live_wgrad_ptrs = pointers | ||
| return |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Find execute() call sites that pass an explicit current_stream for wgrad.
set -euo pipefail
rg -nP -C4 '_live_wgrad_ptrs|_record_pointer_stream' python/cudnn --type=py
rg -nP -C3 'current_stream\s*=' python/cudnn/gemm/cutedsl/grouped/wgrad --type=pyRepository: NVIDIA/cudnn-frontend
Length of output: 14028
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- wgrad BF16 implementation ---'
sed -n '90,225p' python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py
sed -n '500,595p' python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py
printf '%s\n' '--- wgrad public API and object construction ---'
sed -n '1,245p' python/cudnn/gemm/cutedsl/grouped/wgrad/api.py
sed -n '330,395p' python/cudnn/gemm/cutedsl/grouped/wgrad/api.py
rg -n -C4 'MoEGroupedGemmWgradBF16API|wgrad.*API|compile\(|execute\(' python/cudnn/gemm/cutedsl/grouped/wgrad test --type=py 2>/dev/null | head -n 500
printf '%s\n' '--- stream and retention-related tests or documentation ---'
rg -n -i -C4 'wgrad_ptrs|current_stream|multiple streams|multi.?stream|JAX|live_wgrad' test python/cudnn/gemm/cutedsl/grouped/wgrad --type=py 2>/dev/null | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- wrapper cache-key and execution path ---'
sed -n '235,390p' python/cudnn/gemm/cutedsl/grouped/wgrad/api.py
rg -n -C5 'grouped_gemm_wgrad_wrapper_sm100\(' python test --type=py 2>/dev/null | head -n 300
printf '%s\n' '--- stream-sensitive cache fields and pointer retention ---'
rg -n -C3 '_cache_of_GroupedGemmWgradSm100Objects|backend_cache_key|current_stream|_live_wgrad_ptrs' \
python/cudnn/gemm/cutedsl/grouped/wgrad/api.py \
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py
printf '%s\n' '--- AST facts: wrapper parameters, cache-key arguments, and execute arguments ---'
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("python/cudnn/gemm/cutedsl/grouped/wgrad/api.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "grouped_gemm_wgrad_wrapper_sm100":
args = [a.arg for a in node.args.args + node.args.kwonlyargs]
print("wrapper parameters:", args)
for child in ast.walk(node):
if isinstance(child, ast.Call):
text = ast.get_source_segment(path.read_text(), child)
if text and ("backend_cache_key" in text or ".execute(" in text):
print("relevant call:", " ".join(text.split()))
PYRepository: NVIDIA/cudnn-frontend
Length of output: 38460
Retain each in-flight JAX pointer array.
The cache key does not include current_stream, so one API object can execute on multiple streams. _live_wgrad_ptrs stores only the latest pointer array. A subsequent execution can release the previous array while its kernel still reads it.
🧰 Tools
🪛 Ruff (0.16.1)
[error] 202-202: Undefined name torch
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py` around lines 202 - 207,
Update _record_pointer_stream so immutable-framework pointer arrays are retained
per in-flight execution/stream rather than replacing the single _live_wgrad_ptrs
reference. Account for current_stream when storing and releasing arrays,
ensuring each JAX pointer array remains referenced until its associated kernel
has completed.
Comment on lines
+357
to
+364
| import jax | ||
| import jax.numpy as jnp | ||
|
|
||
| if canonicalize_unit_dim_strides(desc.shape, desc.stride) != canonicalize_unit_dim_strides( | ||
| desc.shape, TensorDesc._compute_contiguous_stride(desc.shape) | ||
| ): | ||
| raise ValueError(f"single expert placeholder layout {desc.stride} is not expressible as a C-contiguous JAX array") | ||
| self._single_expert_placeholder = jax.block_until_ready(jnp.empty(desc.shape, dtype=framework_dtype(desc.dtype, "jax"))) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
JAX allocations in _bf16_api.py omit device= and fall back to the default JAX device. Every other JAX allocation added in this change pins the device (wgrad/api.py lines 308 and 344, dglu/api.py line 1078). These two do not, so on a host where the inputs are not on the default JAX device the buffers land on the wrong device.
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L357-L364: pass the descriptor's device tojnp.emptyfor the single-expert placeholder; its pointer is captured intocached_single_expertand passed to every kernel invocation.python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L510-L515: pass the wgrad tensor's device tojnp.asarrayfor the generated pointer array; otherwise the device check at line 576 rejects the caller's own valid inputs.
📍 Affects 1 file
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L357-L364(this comment)python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L510-L515
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py` around lines 357 - 364,
Update both JAX allocations in _bf16_api.py: make the single-expert placeholder
created in the initialization flow use desc.device via jnp.empty, and make the
generated pointer array near lines 510-515 use the wgrad tensor’s device via
jnp.asarray. Apply the changes at lines 357-364 and 510-515 so both buffers
remain on the inputs’ devices and pass the existing device validation.
Comment on lines
+242
to
+264
| class _JaxPermutedDLPackView: | ||
| """Zero-copy permuted (transposed) view of a JAX array for DLPack consumers. | ||
|
|
||
| JAX cannot express strided views, but the CuTeDSL/TVM-FFI kernel ABI only | ||
| consumes DLPack metadata. ``tvm_ffi.from_dlpack`` materializes explicit | ||
| strides even for compact arrays, so this wrapper re-exports the underlying | ||
| array through a fresh tvm-ffi capsule whose shape/strides arrays are | ||
| permuted in place — the data pointer is untouched. Each ``__dlpack__`` call | ||
| builds a fresh capsule, so a view can be consumed multiple times. | ||
| """ | ||
|
|
||
| def __init__(self, array: Any, perm: Tuple[int, ...]): | ||
| base_shape = get_shape(array) | ||
| if sorted(perm) != list(range(len(base_shape))): | ||
| raise ValueError(f"perm {perm} is not a permutation of {len(base_shape)} dims") | ||
| self._array = array | ||
| self._perm = tuple(int(p) for p in perm) | ||
| self.shape = tuple(base_shape[p] for p in self._perm) | ||
| base_strides = get_strides(array) | ||
| self.strides = tuple(base_strides[p] for p in self._perm) | ||
| self.ndim = len(self.shape) | ||
| self.dtype = array.dtype | ||
| self.device = array.device |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Make _JaxPermutedDLPackView visible to detect_framework and get_device.
detect_framework returns "unknown" for a _JaxPermutedDLPackView because the view is not a torch, JAX, or NumPy object. get_device also falls through to the generic branch and returns the raw JAX device object instead of the canonical Device("cuda", id) that get_device returns for the underlying array. The two calls therefore disagree for the same buffer.
The consequence is already visible in the codebase: python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py works around it with detect_framework(q) if not hasattr(q, "_array") else detect_framework(q._array), which reaches into a private attribute of another module.
Unwrap the view inside the adapter so callers do not need to know about _array.
♻️ Proposed refactor
def detect_framework(tensor: Any) -> str:
"""Return "torch", "jax", "numpy", or "unknown" for the given tensor."""
+ if isinstance(tensor, _JaxPermutedDLPackView):
+ return detect_framework(tensor._array)
if is_torch_tensor(tensor):
return "torch" def get_device(tensor: Any) -> Any:
"""Device of the tensor: torch tensors keep their native torch.device; others map to Device."""
+ if isinstance(tensor, _JaxPermutedDLPackView):
+ return get_device(tensor._array)
if is_torch_tensor(tensor):
return tensor.deviceNote that _JaxPermutedDLPackView is defined below these functions, so the class must move above them or the checks must use a deferred lookup.
Also applies to: 336-349
🤖 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 242 - 264, Update
detect_framework and get_device to recognize _JaxPermutedDLPackView and
transparently delegate to its underlying _array, ensuring both return the same
framework and canonical device as the wrapped JAX array. Because the class is
currently defined after these functions, move _JaxPermutedDLPackView above them
or use a deferred type lookup; then remove callers’ need to inspect the private
_array attribute.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py (1)
83-180: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider extracting the shared cache-and-build logic.
gemm_srelu_jax_sm100andgemm_dsrelu_jax_sm100duplicate the validation, cache lookup, kernel build, andmaccomputation. A single private helper that accepts the GEMM class, the sample descriptors, and the cache dict would remove the duplication. This is optional and can be deferred.Also applies to: 183-279
🤖 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 83 - 180, Optionally extract the duplicated validation, cache lookup, kernel construction, and mac computation shared by gemm_srelu_jax_sm100 and gemm_dsrelu_jax_sm100 into one private helper. Have the helper accept the GEMM class, sample descriptors, and cache dictionary, while preserving each wrapper’s existing validation and call-specific behavior.
🤖 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 150-151: Replace the bare assert checks before gemm._kernel in
both call sites with explicit validation that raises ValueError when
gemm.check_support() is false. Include a clear message describing the
unsupported GEMM configuration, while preserving kernel compilation for
supported configurations.
- Around line 109-110: Rename the ambiguous variable l to batch in both affected
functions, including the shape unpacking around m, _, l and n, _, _, and every
downstream use in those function bodies; preserve the existing behavior and
tuple shapes while eliminating Ruff E741.
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, following the module’s existing frontend API
patterns. Route the wrapper to the implementation while preserving
grouped_gemm_glu_jax_sm100 as the compatible public entry point, and ensure the
API class is used by its lazy export.
---
Nitpick comments:
In `@python/cudnn/gemm/cutedsl/dense/srelu/jax_api.py`:
- Around line 83-180: Optionally extract the duplicated validation, cache
lookup, kernel construction, and mac computation shared by gemm_srelu_jax_sm100
and gemm_dsrelu_jax_sm100 into one private helper. Have the helper accept the
GEMM class, sample descriptors, and cache dictionary, while preserving each
wrapper’s existing validation and call-specific behavior.
🪄 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: dde7863b-4336-40b8-9c3d-fd90c28afeaa
📒 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)
- python/cudnn/gemm/cutedsl/grouped/dglu/jax_api.py
- python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py
- python/cudnn/gemm/cutedsl/grouped/dsrelu/jax_api.py
- test/python/fe_api/gemm/test_gemm_amax_jax.py
Comment on lines
+109
to
+110
| m, _, l = a_tensor.shape | ||
| n, _, _ = b_tensor.shape |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous loop-free variable l.
Ruff reports E741 at Line 109 and Line 206. Use batch to keep the lint clean and the intent explicit.
♻️ Proposed rename (apply to both functions and their downstream uses)
- 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")Update the remaining l uses in the same function bodies, for example (m, n, l) → (m, n, batch) and (m, 1, l) → (m, 1, batch).
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 both affected functions, including
the shape unpacking around m, _, l and n, _, _, and every downstream use in
those function bodies; preserve the existing behavior and tuple shapes while
eliminating Ruff E741.
Source: Linters/SAST tools
Comment on lines
+150
to
+151
| assert gemm.check_support() | ||
| kernel = gemm._kernel( |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Replace assert gemm.check_support() with an explicit error.
Python removes assert statements when the interpreter runs with -O. The support check is then skipped, and an unsupported configuration reaches kernel compilation and launch. The bare assert also gives the caller no reason for the failure.
Raise a ValueError instead, consistent with the other validation in these functions.
🛡️ Proposed fix for both call sites
- assert gemm.check_support()
+ if not gemm.check_support():
+ raise ValueError(
+ "Unsupported GemmSreluSm100 configuration: "
+ f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, "
+ f"sf_vec_size={sf_vec_size}, d_dtype={d_dtype}"
+ )- assert gemm.check_support()
+ if not gemm.check_support():
+ raise ValueError(
+ "Unsupported GemmDsreluSm100 configuration: "
+ f"mma_tiler_mn={mma_tiler_mn}, cluster_shape_mn={cluster_shape_mn}, "
+ f"sf_vec_size={sf_vec_size}, d_dtype={d_dtype}"
+ )Also applies to: 249-250
🤖 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 150 - 151,
Replace the bare assert checks before gemm._kernel in both call sites with
explicit validation that raises ValueError when gemm.check_support() is false.
Include a clear message describing the unsupported GEMM configuration, while
preserving kernel compilation for supported configurations.
Comment on lines
+76
to
+93
| 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]]: |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Add the required APIBase subclass and wrapper.
grouped_gemm_glu_jax_sm100 is a new public frontend API. It is lazily exported through python/cudnn/__init__.py, but this module provides neither an APIBase subclass nor a wrapper.
Add the API class and wrapper. Keep grouped_gemm_glu_jax_sm100 as the compatible public entry point if required.
As per coding guidelines: “Every new frontend-only Python API must subclass APIBase, provide a wrapper, and 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,
following the module’s existing frontend API patterns. Route the wrapper to the
implementation while preserving grouped_gemm_glu_jax_sm100 as the compatible
public entry point, and ensure the API class is used by its lazy export.
Source: Coding guidelines
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@test/python/fe_api/gemm/test_gemm_amax.py`:
- Around line 322-324: Update the valid layout cases around
gemm_amax_wrapper_sm100 to retain both returned c_tensor and amax_tensor pairs
instead of discarding them. Synchronize torch.cuda.current_stream() after the
calls, then invoke check_ref_gemm_amax for each result pair so both the physical
and permuted scale-factor layouts receive the existing dtype-specific
validation.
🪄 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: 9fc69796-8218-49ef-92e3-eb3f81fe5a6f
📒 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 (5)
- python/cudnn/gemm/cutedsl/dense/amax/jax_api.py
- python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py
- python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py
- python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py
- python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py
Comment on lines
+322
to
+324
| # Valid: the physical form and its (3, 4, 1, 5, 2, 0)-permuted atom view | ||
| gemm_amax_wrapper_sm100(a, b, sfa, sfb, sf_vec_size=sf_vec_size) | ||
| gemm_amax_wrapper_sm100(a, b, sfa.permute(3, 4, 1, 5, 2, 0), sfb.permute(3, 4, 1, 5, 2, 0), sf_vec_size=sf_vec_size) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="test/python/fe_api/gemm/test_gemm_amax.py"
printf '%s\n' '--- file outline ---'
if command -v ast-grep >/dev/null 2>&1; then
ast-grep outline "$file" || true
fi
printf '%s\n' '--- target test region ---'
sed -n '250,360p' "$file"
printf '%s\n' '--- helper and reference symbols ---'
rg -n -C 5 'gemm_amax_wrapper_sm100|check_ref_gemm_amax|current_stream|stream|assert_close|cuda\.synchronize|synchronize' "$file" test/python
printf '%s\n' '--- test guidance files ---'
for f in test/AGENTS.md test/python/AGENTS.md; do
if [ -f "$f" ]; then
printf '\n--- %s ---\n' "$f"
cat "$f"
fi
doneRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- target file: imports and tests ---'
sed -n '1,340p' test/python/fe_api/gemm/test_gemm_amax.py
printf '%s\n' '--- utility file candidates ---'
fd -i 'test_gemm_amax_utils.py' .
for f in $(fd -i 'test_gemm_amax_utils.py' .); do
printf '\n--- %s outline ---\n' "$f"
if command -v ast-grep >/dev/null 2>&1; then ast-grep outline "$f" || true; fi
printf '\n--- %s relevant definitions ---\n' "$f"
rg -n -C 12 'def (check_ref_gemm_amax|gemm_amax_init|allocate_input_tensors|allocate_output_tensors)|reference_gemm_amax|assert_close|synchronize' "$f"
done
printf '%s\n' '--- wrapper definitions and declarations ---'
rg -n -C 8 'def gemm_amax_wrapper_sm100|gemm_amax_wrapper_sm100|gemm_amax_wrapper' --glob '!test/python/**' --glob '!*.so' --glob '!*.pyc' . | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 32563
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- wrapper implementation ---'
sed -n '451,575p' python/cudnn/gemm/cutedsl/dense/amax/api.py
printf '%s\n' '--- execute implementation symbols ---'
rg -n -C 10 'def execute|current_stream|cudaStream|CUstream|stream is None|legacy default|default stream' python/cudnn/gemm/cutedsl/dense/amax
printf '%s\n' '--- wrapper documentation return and stream sections ---'
sed -n '190,225p' docs/fe-oss-apis/gemm_fusions/gemm_amax.md
sed -n '64,115p' docs/fe-oss-apis/gemm_fusions/gemm_amax.md
printf '%s\n' '--- utility reference body ---'
sed -n '178,230p' test/python/fe_api/gemm/test_gemm_amax_utils.pyRepository: NVIDIA/cudnn-frontend
Length of output: 26133
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- default_stream implementation ---'
rg -n -C 12 'def default_stream|default_stream\(' python test/python --glob '*.py' | head -n 400
printf '%s\n' '--- focused Torch wrapper validation patterns ---'
rg -n -C 8 'gemm_amax_wrapper_sm100|check_ref_gemm_amax' test/python/fe_api/gemm --glob '*.py'
printf '%s\n' '--- target test configuration and dtype shape details ---'
sed -n '300,340p' test/python/fe_api/gemm/test_gemm_amax.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50377
Validate both accepted scale-factor layouts. The valid calls discard c_tensor and amax_tensor, so they test only API acceptance. Retain both result pairs, synchronize torch.cuda.current_stream(), and call check_ref_gemm_amax for each pair to apply the existing dtype-specific reference checks and catch asynchronous errors.
🤖 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, Update the
valid layout cases around gemm_amax_wrapper_sm100 to retain both returned
c_tensor and amax_tensor pairs instead of discarding them. Synchronize
torch.cuda.current_stream() after the calls, then invoke check_ref_gemm_amax for
each result pair so both the physical and permuted scale-factor layouts receive
the existing dtype-specific validation.
Sources: Coding guidelines, MCP tools
Type-erase the BSA public API and the SM100/SM110 blk128 forward/backward paths: torch tensors and JAX arrays are both accepted, torch is imported only for torch tensors, and the SM90/SM120/blk64 backends reject JAX with clear errors. The backward's internal transposed views travel as zero-copy permuted DLPack wrappers (new cudnn.tensor_adapter.permuted_view; JAX has no strided views, but tvm-ffi materializes explicit strides so a capsule's shape/strides can be permuted in place over the same data pointer). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Collaborator
Author
|
Closing for now. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*.Affected area
FE OSS kernels or CuTeDSL (Python API)
Summary
Extends the type-erased torch+JAX tensor contract from the GEMM CuTeDSL APIs (#529/#530/#553) to
python/cudnn/block_sparse_attention:cudnn.BSA.block_sparse_attention_forward/backwardnow accept JAX arrays alongside torch tensors, with torch imported only when torch tensors are passed.Per-path status
bhsdandbshd)bhsdandbshd, incl. pre-allocated dq/dk/dv)ValueError(arch-specific torch code paths)kv_splits,use_clc)ValueErrorThe key new mechanism: zero-copy permuted views for JAX
The blk128 kernels consume transposed strided views at every kernel boundary (the backward's logical layout is
bshd, reached frombhsdbuffers viatranspose(1, 2)views). JAX has no strided views — but the TVM-FFI kernel ABI only consumes DLPack metadata, andtvm_ffi.from_dlpackmaterializes explicit strides even for compact arrays. The newcudnn.tensor_adapter.permuted_view(tensor, perm)exploits this:permuteview (unchanged semantics);__dlpack__call builds a fresh capsule, so views are multi-consumable and composable).Result: the JAX paths run the same compiled kernels with the same stride patterns as torch, with zero extra copies. Both layouts work for forward and backward.
Implementation notes
api.py+_interface.py+ the blk128 csrc host functions (bsa_bwd_sm100.py,bucketed_k2q_csr.py,bsa_bwd_prepost.py,cute_dsl_utils.py) are torch-lazy (try: import torch+from __future__ import annotations); the whole BSA import chain — including the optional blk128 backward, whosetry/except ImportErrorguard previously masked the torch dependency — now imports with torch and jax absent.cudnn.tensor_adapter.allocate_tensor(torchempty/zerospreservingempty_likestride semantics on the torch path; materializedjnpbuffers for JAX); validation via cutlass dtypes and adapter metadata helpers; JAX launches on the CUDA legacy default stream (the backward's pre/main/post kernel chain uses the tvm-ffi environment stream, which resolves to the default stream without torch).layout="bshd"JAX backward, dq/dk/dv are pre-allocated in the caller's layout and the kernel writes through permuted views, so callers always receive real JAX arrays.Why
BSA is used from JAX-based training stacks; the CuTeDSL data path was already framework-neutral (tvm-ffi/DLPack) and only host-side metadata, allocation, and stream handling were torch-bound — the same situation the GEMM stack was in before #529.
permuted_viewcloses the one genuinely new gap (strided views), and is reusable for future JAX enablement of the other attention APIs.Related issues
Stacked on #553 (→ #534 → #530).
API and compatibility impact
ValueErrornaming the supported path.cudnn.tensor_adapter.permuted_viewandcudnn.tensor_adapter.allocate_tensor.jax.jitentry point viacudnn.jax.callis a natural follow-up.Testing
On a B200-class SM100 (CC 10.0), Python 3.12, torch 2.13, jax 0.11, nvidia-cutlass-dsl 4.6:
test_BSA_attention_jax.py: forward (bhsd+bshd) bit-identical to torch on identical input bytes; backward (bhsd+bshd) with dq bit-identical and dk/dv at 1-ulp-scale bf16 tolerance (their cross-Q-block accumulation is scheduling-order nondeterministic run-to-run within either framework — verified jax-vs-jax); error-path coverage (blk64 rejection fwd+bwd, non-torch/JAX frameworks).tensor_adapteradditions (test_gemm_amax_jax.py,test_grouped_gemm_jax.py: 8 passed, 2 xfailed).cudnn.BSA.*resolves with torch and jax both absent (sys.modulesnulled).🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
jax.jit-compatible entry points for supported SM100/SM110 workflows.Documentation
Bug Fixes