Extend JAX support to the grouped/discrete-grouped GEMM APIs and proj_rope_mxfp8 (stacked on #529) - #530
Extend JAX support to the grouped/discrete-grouped GEMM APIs and proj_rope_mxfp8 (stacked on #529)#530Anerudhan wants to merge 1 commit into
Conversation
Anerudhan
added
cat-feature
📝 WalkthroughWalkthroughThis change adds framework-neutral tensor handling across CuTeDSL GEMM APIs. Selected grouped and discrete GEMM paths now support eager JAX execution. Unsupported layouts and modes raise explicit errors. Projection-RoPE MXFP8 also supports JAX inputs. ChangesJAX framework support
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant JAX
participant API
participant TensorAdapter
participant CuTeDSLKernel
JAX->>API: invoke GEMM wrapper
API->>TensorAdapter: detect framework and normalize metadata
TensorAdapter->>API: return dtype, device, stream, and pointer data
API->>CuTeDSLKernel: compile or execute with framework-owned buffers
CuTeDSLKernel->>JAX: write output arrays
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 17
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/discrete_grouped/dswiglu/api.py (1)
383-387: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winUpdate the torch-specific dtype name in the error text.
The API now accepts JAX arrays. The message still names
torch.float8_e8m0fnu. A JAX caller cannot act on that name.✏️ Proposed fix
- "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=torch.float8_e8m0fnu", + "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=float8_e8m0fnu",🤖 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 383 - 387, Update the error message in the _value_error_if call for kernel_generate_sfd to use the framework-neutral or JAX-appropriate dtype name instead of torch.float8_e8m0fnu, while preserving the existing validation condition and required argument guidance.
🧹 Nitpick comments (21)
test/python/fe_api/gemm/test_gemm_amax_jax.py (2)
84-89: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe JAX test suite has no shared helper module. Common helpers now live inside a test module or are copied between test modules. Both problems have one root cause: there is no non-test helper module for the JAX GEMM tests.
test/python/fe_api/gemm/test_gemm_amax_jax.py#L84-L89: movedevice_syncandskip_unless_sm100into a shared helper module (for examplefe_api/gemm/jax_test_utils.py) or a conftest fixture, so other test modules stop importing a test module.test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py#L38-L40: delete the local_packed_jax_ptrsand import it from the shared helper module.test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py#L40-L42: delete the duplicate_packed_jax_ptrsand import it from the shared helper module.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/fe_api/gemm/test_gemm_amax_jax.py` around lines 84 - 89, Create a shared JAX GEMM test helper module containing device_sync, skip_unless_sm100, and _packed_jax_ptrs; update test/python/fe_api/gemm/test_gemm_amax_jax.py lines 84-89 to move device_sync and skip_unless_sm100 there, and replace local/test-module usage with imports. In test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py lines 38-40 and test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py lines 40-42, remove each local _packed_jax_ptrs definition and import the shared helper instead.
139-159: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConfirm the xfail marker cannot mask a real regression.
strict=Falselets this test pass silently if the packed-FP4 container path starts working, and it also hides any new failure type other thanTypeError. The stated plan is to remove the marker when the container path is fixed. Add a tracking issue reference in thereasonstring, or usestrict=Trueonce the failure mode is stable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/python/fe_api/gemm/test_gemm_amax_jax.py` around lines 139 - 159, Update the xfail configuration on test_gemm_amax_jax_wrapper_fp4_uint8 so it cannot silently pass when the expected failure is resolved or mask unrelated failures: add the relevant tracking issue reference to reason and use strict=True if the TypeError failure mode is stable, while preserving removal of the marker once the container path is fixed.docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md (1)
73-80: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd the JAX imports to the snippet.
The example uses
jax.jitandjnp.float32but only importsgemm_swiglu_jax_sm100. A reader who copies the block gets aNameError.📝 Proposed snippet fix
+import jax +import jax.numpy as jnp from cudnn import gemm_swiglu_jax_sm100 `@jax.jit` def swiglu_mlp(a, b): ab12, c = gemm_swiglu_jax_sm100(a, b, alpha=1.0, ab12_dtype=jnp.float32, c_dtype=jnp.bfloat16) return c🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md` around lines 73 - 80, Update the swiglu_mlp example snippet to import both jax and jax.numpy as jnp before using jax.jit and jnp.float32, while retaining the existing gemm_swiglu_jax_sm100 import.python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (3)
1493-1494: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the conditional
import torchnext to its use sites.
import torchbinds the name in the function scope, and every later use (lines 1531-1533, 1549, 1568-1570, 1582, 1587) sits under a matchingframework == "torch"guard, so the current code works. The binding is 36 lines from its first use and depends on the guard structure staying in sync. A future branch that usestorchwithout the guard raisesNameErrorat runtime rather than at import.The other APIs in this PR import inside the branch that uses the symbol. For example
python/cudnn/gemm/cutedsl/dense/srelu/api.pylines 494-495. Apply the same pattern here and to theimport jax.numpy as jnpat lines 1519-1520.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py` around lines 1493 - 1494, The conditional imports in the affected API function are separated from their use sites. Move import torch into the framework == "torch" branches immediately before the guarded torch usages, and move import jax.numpy as jnp into the corresponding JAX branch near its first use, preserving the existing branch behavior.
20-21: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
from __future__ import annotationswas added without aTYPE_CHECKINGtorch import. These three modules dropped the module-levelimport torchand added the future import so the remainingtorch.Tensorandtorch.dtypeannotations stay lazy. The nametorchis now unbound at module scope, so Ruff reports F821 and any consumer that resolves annotations at runtime (typing.get_type_hints, Sphinxautodoc_typehints) fails. Add a guarded import in each module.
python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py#L20-L21: addif TYPE_CHECKING: import torchafter the future import; clears F821 at lines 54, 129, 1408, 1409.python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py#L10-L17: add the same guarded import; clears F821 at lines 78, 888, 889.python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py#L15-L22: add the same guarded import; clears F821 at lines 90, 906, 907, 908.🤖 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 20 - 21, Add a TYPE_CHECKING-guarded torch import after the future import in python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (lines 20-21), python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py (lines 10-17), and python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py (lines 15-22). This binds torch for annotation analysis while keeping the runtime import lazy and resolves the remaining torch.Tensor and torch.dtype references in each module.Source: Linters/SAST tools
1607-1629: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCanonicalize strides in
tensor_signatureand reuse the physical-form predicate.Two points in this block:
tensor_signaturereturnsget_strides(tensor)withoutcanonicalize_unit_dim_strides. The API canonicalizes strides when it builds descriptors, so a torch tensor and a JAX array with the same logical layout but different extent-1 strides produce different cache keys and compile two identical kernels. The sibling wrappers canonicalize here. Seepython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.pyline 1061.
_sf_is_physicalduplicates the static method_sf_desc_is_physicaldefined at lines 285-295 of this file. Only the shape accessor differs.♻️ Proposed fix
def tensor_signature(tensor: Optional[torch.Tensor]) -> Tuple[Optional[Tuple[int, ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: if tensor is None: return None, None, None - return get_shape(tensor), get_strides(tensor), _convert_to_cutlass_data_type(tensor.dtype) + tensor_shape = get_shape(tensor) + return tensor_shape, canonicalize_unit_dim_strides(tensor_shape, get_strides(tensor)), _convert_to_cutlass_data_type(tensor.dtype)For point 2, drop the local
_sf_is_physicaland callGroupedGemmDsreluSm100._sf_desc_is_physicalafter widening it to accept any object with a.shape, or extract a shared module-level helper.Note:
canonicalize_unit_dim_stridesis not currently imported in this module. Add it to thecudnn.tensor_adapterimport block at lines 40-51.🤖 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 1607 - 1629, Update tensor_signature to canonicalize get_strides(tensor) with canonicalize_unit_dim_strides, adding that helper to the cudnn.tensor_adapter imports and preserving the existing signature shape and dtype behavior. Remove the duplicate local _sf_is_physical and reuse GroupedGemmDsreluSm100._sf_desc_is_physical, widening or extracting the predicate as needed so it accepts the tensor shape representation used here.python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py (1)
1082-1086: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe SFA cache signature indexes the permuted layout unconditionally. Both wrappers call
dynamic_m_tensor_signature(sfa_tensor, (get_shape(sfa_tensor)[4], 1), dynamic_stride_dims=(0, 1, 5)). Index 4 isK'for the permuted atom view but the atom constant4for the physical form these APIs now accept from JAX. The key stays correct only becauseK'is derivable from the A-shape entry already present in the key; the relationship is implicit and breaks if that entry changes.python/cudnn/gemm/cutedsl/grouped/dsrelu/api.pylines 1631-1641 already implement a layout-branchingdynamic_m_sf_signature.
python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py#L1082-L1086: replace the call with a layout-aware signature helper.python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py#L1122-L1126: replace the call with the same 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/discrete_grouped/dswiglu/api.py` around lines 1082 - 1086, Replace the unconditional dynamic_m_tensor_signature calls at python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py:1082-1086 and python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py:1122-1126 with the same layout-aware dynamic_m_sf_signature helper pattern used by grouped/dsrelu/api.py:1631-1641. Ensure the helper selects the correct SFA shape and stride dimensions for permuted and physical layouts while preserving the existing None handling.python/cudnn/gemm/cutedsl/dense/srelu/api.py (1)
485-551: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffConsider extracting the shared dense wrapper scaffolding.
gemm_srelu_wrapper_sm100andgemm_dsrelu_wrapper_sm100now carry near-identical blocks: framework detection, torch/JAX output allocation, JAX layout and batch validation, physical SFD allocation,block_until_readycollection, and the four cache-signature helpers (stride_order,tensor_signature,dynamic_compact_signature,dynamic_tensor_signature,dynamic_m_tensor_signature).The two copies have already diverged in small ways. A shared private module under
gemm/cutedsl/dense/would keep future JAX fixes applied in one place.Also applies to: 556-580
🤖 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/api.py` around lines 485 - 551, Extract the duplicated framework setup and cache-signature logic from gemm_srelu_wrapper_sm100 and gemm_dsrelu_wrapper_sm100 into a shared private module under gemm/cutedsl/dense/. Have both wrappers reuse it for framework detection, Torch/JAX allocation and validation, SFD/amax handling, JAX synchronization, and the existing stride/tensor/dynamic signature helpers while preserving current behavior and framework-specific outputs.python/cudnn/gemm/cutedsl/grouped/glu/api.py (1)
1007-1034: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConfirm the JAX allocation ignores
stridesafely for every caller.
_allocate_outputdrops thestrideargument on the JAX branch and returns a C-contiguous buffer. For(valid_m, n_full, 1)the requested stride(n_full, 1, valid_m * n_full)differs from the C-contiguous stride only in the extent-1 batch dimension, whichcanonicalize_unit_dim_stridesnormalizes. If a future shape has more than one non-trivial difference, the descriptor and the buffer diverge silently. Consider asserting that the requested stride matches the C-contiguous stride after unit-dim canonicalization.🤖 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 1007 - 1034, Update the JAX branch of _allocate_output to compute the C-contiguous strides for shape and compare them with the requested stride after canonicalizing extent-1 dimensions; assert they match before allocating, while preserving the existing allocation behavior for valid callers.python/cudnn/datatypes.py (1)
261-267: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMatch the module root exactly in the fallback branch.
str.startswith(("jax", "jaxlib"))matches any module whose name begins with those characters, for examplejaxtypingorjaxton. An unrelated object then reports as a JAX array, anddetect_frameworkroutes it into the JAX code paths. Compare the first module component instead.♻️ Proposed refactor
jax = sys.modules.get("jax") if jax is not None and isinstance(input_tensor, getattr(jax, "Array", ())): return True - return type(input_tensor).__module__.startswith(("jax", "jaxlib")) + return type(input_tensor).__module__.partition(".")[0] in ("jax", "jaxlib")🤖 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/datatypes.py` around lines 261 - 267, Update the fallback in _is_jax_array to split type(input_tensor).__module__ into its first dotted component and compare that component exactly against the supported JAX roots, "jax" and "jaxlib", preventing similarly prefixed modules from matching.python/cudnn/gemm/cutedsl/grouped/dglu/api.py (1)
62-73: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
_block_scaled_dtype_pairsis duplicated verbatim in three modules. Each copy returns the same canonical cutlass dtype vocabulary consumed byselect_grouped_gemm_backend. A future dtype addition must be applied in three places, and a partial update produces silently divergent backend selection between the GLU, dGLU, and wgrad APIs. Move one definition intopython/cudnn/gemm/cutedsl/grouped/backend_utils.py, next toselect_grouped_gemm_backend, and import it in the three call sites.
python/cudnn/gemm/cutedsl/grouped/dglu/api.py#L62-L73: delete the local definition and import the shared helper frombackend_utils.python/cudnn/gemm/cutedsl/grouped/glu/api.py#L64-L75: delete the local definition and import the shared helper frombackend_utils.python/cudnn/gemm/cutedsl/grouped/wgrad/api.py#L36-L45: delete the local definition and import the shared helper frombackend_utils.🤖 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/api.py` around lines 62 - 73, The _block_scaled_dtype_pairs helper is duplicated across the grouped GLU APIs; centralize it to keep backend dtype selection consistent. Add the single definition beside select_grouped_gemm_backend in python/cudnn/gemm/cutedsl/grouped/backend_utils.py, then delete the local definitions and import the shared helper in python/cudnn/gemm/cutedsl/grouped/dglu/api.py#L62-L73, python/cudnn/gemm/cutedsl/grouped/glu/api.py#L64-L75, and python/cudnn/gemm/cutedsl/grouped/wgrad/api.py#L36-L45.python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py (1)
38-40: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSilently ignored parameters:
sf_vec_size,vector_f32, andab12_stages.These three parameters are accepted but never used on the supported path.
make_gemm()does not forwardvector_f32orab12_stagestoGemmSwigluSm100, and they are absent fromcache_key.GemmSwigluSm100._compile_kernelpasses both only toSm100BlockScaledPersistentDenseGemmKernel, which this entry point rejects at line 61. A caller that setsab12_stages=8therefore gets the default behavior with no error.Reject non-default values alongside the existing quantized-input check, so the ignored settings are visible to callers.
Also applies to: 82-92
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py` around lines 38 - 40, The public API must reject non-default sf_vec_size, vector_f32, and ab12_stages values because they are unsupported and ignored. Add validation in the entry-point argument checks near the existing quantized-input validation, raising the established error for any non-default setting while preserving accepted defaults and avoiding changes to make_gemm or cache-key behavior.python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py (1)
59-62: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the repeated torch-only guard into one helper.
The same three-line guard appears in
GemmProjRopeMxfp8Bf16InSm100.__init__,GemmProjRopeMxfp8Mxfp8InSm100.__init__, andgemm_proj_rope_mxfp8_wrapper_sm100, with only the API name changing. A single module-level helper keeps the message format consistent if the JAX support status changes later.♻️ Proposed helper
def _require_torch_tensor(tensor, api_name: str) -> None: from cudnn.tensor_adapter import is_torch_tensor if tensor is not None and not is_torch_tensor(tensor): raise ValueError(f"{api_name} currently supports torch tensors only; JAX support is not yet implemented for this API")Also applies to: 257-260, 544-549
🤖 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 59 - 62, Introduce a module-level _require_torch_tensor helper that accepts the tensor and API name, performs the existing optional-tensor torch check, and raises the consistently formatted ValueError. Replace the duplicated guards in GemmProjRopeMxfp8Bf16InSm100.__init__, GemmProjRopeMxfp8Mxfp8InSm100.__init__, and gemm_proj_rope_mxfp8_wrapper_sm100 with calls to this helper.python/cudnn/gemm/cutedsl/_jax_ffi.py (1)
63-72: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the bare
assertoncheck_support()with an explicit check.Python removes
assertstatements when it runs with-O. In that modegemm.check_support()never runs, so_compile_kernelcompiles a configuration that was never validated, andself._is_supportedstaysFalse._ensure_support_checked()inside_compile_kernelwould then callcheck_support()through anotherassert, which is also removed.♻️ Proposed change
gemm = make_gemm() - assert gemm.check_support() + if not gemm.check_support(): + raise ValueError(f"Unsupported configuration for target prefix {target_prefix!r}") compiled = gemm._compile_kernel(use_tvm_ffi_env_stream=True)🤖 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/_jax_ffi.py` around lines 63 - 72, Replace the bare assert around gemm.check_support() in the registry-miss path with an explicit support validation that always executes, including optimized Python runs. Ensure unsupported configurations are rejected before calling gemm._compile_kernel, and preserve the existing compilation and registry flow for supported configurations.python/cudnn/gemm/cutedsl/dense/amax/jax_api.py (1)
57-60: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRuff E741 on the batch-dimension variable
lin three changed sites. All three sites unpack the batch dimension into a variable namedl, which Ruff reports as an error (E741, ambiguous variable name). The shared root cause is the single-letter name; rename it consistently tobatchacross the new JAX paths, or add a targetednoqaif the domain naming must stay.
python/cudnn/gemm/cutedsl/dense/amax/jax_api.py#L57-L60: renamelinm, _, l = a_tensor.shapeand in thel != 1check.python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py#L56-L59: renamelinm, _, l = a_tensor.shapeand in thel != 1check; also update the later uses at lines 86-87 and 106-110.python/cudnn/gemm/cutedsl/dense/swiglu/api.py#L605-L606: renamelin bothget_shapeunpackings, and prefix the unused reboundkon line 606 with an underscore to clear RUF059.🤖 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 57 - 60, Rename the ambiguous batch-dimension variable l to batch consistently in python/cudnn/gemm/cutedsl/dense/amax/jax_api.py lines 57-60 and python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py lines 56-59, including the later swiglu uses at lines 86-87 and 106-110; in python/cudnn/gemm/cutedsl/dense/swiglu/api.py lines 605-606, rename both l unpackings and prefix the unused rebound k with an underscore to satisfy Ruff.Source: Linters/SAST tools
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py (2)
43-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the endianness contract of the packed pointer form.
_pointer_valuesdecodes the packed array withnp.asarray(ptrs).view(np.int64), which uses the host byte order. The producers pack little-endian bytes (see_generate_wgrad_ptrsinpython/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.pylines 510-515, and the JAX tests). The two agree on every supported CUDA host, so this is correct today. State the little-endian assumption in the_pointer_valuesdocstring so a future reader does not introduce an explicit byte order on one side only.🤖 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 the _pointer_values docstring to explicitly state that packed uint8 JAX pointers are decoded as little-endian 64-bit values, matching the producers’ packing contract. Leave the existing np.asarray(ptrs).view(np.int64) implementation unchanged.
238-265: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftTriplicated BF16 grouped-GEMM helper block.
_output_dtypes,_validate_data_alignment,_validate_pointer_array_alignment,_record_pointer_stream,_copy_values_to_host,_is_validation_cached, and_remember_validationare byte-identical in the three BF16 APIs. This change applied the same framework-adapter edits three times, and thegluanddglucopies already import_pointer_valuesand_validate_pointer_tensorfromunfused/_bf16_api.py, so a shared home exists. Move the block into a shared mixin or module and import it.
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py#L238-L265: promote these helpers into a shared_GroupedGemmBf16HelpersMixin(or a sibling private module) next to_pointer_valuesand_validate_pointer_tensor.python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py#L196-L223: delete the local copies and inherit or import the shared helpers.python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py#L200-L227: delete the local copies and inherit or import the shared helpers.🤖 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 238 - 265, Deduplicate the shared BF16 grouped-GEMM helpers by moving _output_dtypes, _validate_data_alignment, _validate_pointer_array_alignment, _record_pointer_stream, _copy_values_to_host, _is_validation_cached, and _remember_validation from python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py:238-265 into a shared mixin or private module near _pointer_values and _validate_pointer_tensor; update glu/_bf16_api.py:196-223 and dglu/_bf16_api.py:200-227 to inherit or import that shared implementation and remove their local copies.python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py (1)
111-119: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the redundant
sample_a is not Nonecondition.
sample_ais a required positional parameter. When a caller still passesNone,detect_framework(None)returns"unknown"and the guard is skipped, so construction continues withself.a_desc = Noneand fails later with an unclearAttributeError. The sibling APIsGroupedGemmSwigluSm100andGroupedGemmSreluSm100use the unconditional form.♻️ Proposed simplification
framework = detect_framework(sample_a) - if sample_a is not None and framework != "torch": - if framework == "jax": - raise ValueError( - "GroupedGemmDswigluSm100 only supports dense weight mode, whose expert-outermost strided " - "B layout (n, k, l) is not expressible as JAX arrays (row-major only); " - "use torch tensors for this backward API" - ) - raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmDswigluSm100; pass torch tensors") + if framework == "jax": + raise ValueError( + "GroupedGemmDswigluSm100 only supports dense weight mode, whose expert-outermost strided " + "B layout (n, k, l) is not expressible as JAX arrays (row-major only); " + "use torch tensors for this backward API" + ) + if framework != "torch": + raise ValueError(f"Unsupported tensor framework '{framework}' for GroupedGemmDswigluSm100; pass torch tensors")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py` around lines 111 - 119, Update the framework validation around detect_framework(sample_a) to remove the redundant sample_a is not None condition and reject every non-"torch" framework, including "unknown" returned for None. Preserve the existing JAX-specific message and the generic unsupported-framework error for all other values.python/cudnn/gemm/cutedsl/grouped/unfused/api.py (1)
304-321: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the
framework_dtypeimport to the module header.
framework_dtypelives incudnn.tensor_adapter, which this module already imports at lines 17-26. The symbol pulls in no optional dependency, so the function-local import gives no lazy-loading benefit and hides the dependency from readers.♻️ Proposed change
from cudnn.tensor_adapter import ( canonicalize_unit_dim_strides, cuda_is_available, detect_framework, + framework_dtype, get_compute_capability, get_data_ptr, get_device, get_shape, get_strides, )def _allocate_output(dtype): - from cudnn.tensor_adapter import framework_dtype - if framework == "torch":🤖 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 304 - 321, Move the framework_dtype import from _allocate_output to the module-level imports, reusing the existing cudnn.tensor_adapter import section. Remove the function-local import while preserving both Torch and JAX dtype conversion behavior.python/cudnn/tensor_adapter.py (1)
126-133: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDo not let a CPU-only torch build mask an available CUDA device.
cuda_is_availablereturnstorch.cuda.is_available()whenever thetorchmodule is loaded. A JAX-only caller in a process that imports a CPU-only torch build then getsCUDA is not available, even though the CUDA driver and the JAX GPU backend work. The callers in this cohort (for examplecheck_supportinpython/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.pylines 384-385) turn that into a hardRuntimeError. Fall back to the CUDA runtime when torch reports no CUDA.♻️ Proposed fallback
def cuda_is_available() -> bool: torch = sys.modules.get("torch") - if torch is not None: - return torch.cuda.is_available() + if torch is not None and torch.cuda.is_available(): + return True from cuda.bindings import runtime as cudart err, count = cudart.cudaGetDeviceCount() return err == cudart.cudaError_t.cudaSuccess and count > 0🤖 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 126 - 133, Update cuda_is_available so it only returns immediately when torch.cuda.is_available() is true; when torch is loaded but reports no CUDA, fall through to the existing cudart.cudaGetDeviceCount() check instead of returning false. Preserve the current torch fast path and runtime-based detection for environments without torch.python/cudnn/gemm/cutedsl/grouped/quant/api.py (1)
42-46: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTriplicated
_JAX_SF_LAYOUT_ERRORconstant. The same five-line message is defined in three modules. The JAX rejection tests match on the substringnot expressible as JAX arrays, so drift between copies silently weakens that contract. Define the constant once in a shared module such aspython/cudnn/gemm/cutedsl/grouped/moe_utils.py.
python/cudnn/gemm/cutedsl/grouped/quant/api.py#L42-L46: replace the literal with an import of the shared constant.python/cudnn/gemm/cutedsl/grouped/srelu/api.py#L41-L45: replace the literal with an import of the shared constant.python/cudnn/gemm/cutedsl/grouped/swiglu/api.py#L33-L37: replace the literal with an import of the shared constant.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py` around lines 42 - 46, Define _JAX_SF_LAYOUT_ERROR once in python/cudnn/gemm/cutedsl/grouped/moe_utils.py, preserving the existing message and required substring. In python/cudnn/gemm/cutedsl/grouped/quant/api.py:42-46, python/cudnn/gemm/cutedsl/grouped/srelu/api.py:41-45, and python/cudnn/gemm/cutedsl/grouped/swiglu/api.py:33-37, remove the duplicated literals and import the shared constant for each API.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/fe-oss-apis/gemm_fusions/gemm_amax.md`:
- Around line 253-257: Update the JAX-specific constraint bullet in the
gemm_amax documentation to scope “eager use only” exclusively to the eager entry
points, while preserving the documented jax.jit compatibility of
gemm_amax_jax_sm100.
In `@docs/fe-oss-apis/overview.md`:
- Line 8: Update the API support description in the overview to separate Wgrad
from the discrete pointer-array weight APIs. Document Wgrad as accepting plain
BF16 JAX A and B arrays with both dense output and discrete output-pointer
support, and remove it from the statements claiming JAX requires discrete
weights or rejects dense mode.
In `@python/cudnn/api_base.py`:
- Line 690: Update the _get_innermost_stride_dim method annotation to use Any
instead of torch.Tensor, resolving Ruff’s undefined-name error while preserving
the module’s lazy torch dependency and avoiding an eager import.
In `@python/cudnn/gemm/cutedsl/dense/amax/api.py`:
- Line 156: Remove the unnecessary f-string prefix from the unsupported dtype
and scale-vector-size error message, preserving the escaped-brace text
unchanged.
- Around line 57-64: The GemmAmaxSm100 initialization currently creates tensor
descriptors before enabling packed-FP4 interpretation for Uint8 inputs. Set
_interpret_uint8_as_fp4x2 before the _make_tensor_desc calls when the inputs are
Uint8, so descriptor construction matches the Uint8 acceptance already defined
by check_support().
In `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py`:
- Around line 103-112: Raise the declared minimum JAX dependency version to
0.4.36 so the jax.ffi.ffi_call usage in the surrounding API remains compatible
with input_output_aliases.
In `@python/cudnn/gemm/cutedsl/dense/dsrelu/api.py`:
- Around line 537-540: Collapse the JAX major-mode checks in d_major validation
within python/cudnn/gemm/cutedsl/dense/dsrelu/api.py#L537-L540 into one
condition rejecting every value other than "n" with the row-major explanation;
apply the same change to c_major validation in
python/cudnn/gemm/cutedsl/dense/srelu/api.py#L520-L523, so both JAX-specific
messages identify only "n" as supported.
In `@python/cudnn/gemm/cutedsl/dense/srelu/api.py`:
- Around line 553-554: Update the dynamic compile handling in
python/cudnn/gemm/cutedsl/dense/srelu/api.py lines 553-554 and
python/cudnn/gemm/cutedsl/dense/dsrelu/api.py lines 570-571 so physical SF
descriptors are either rejected when use_dynamic_m or use_full_dynamic is
enabled, or handled via a _sf_desc_is_physical branch when constructing
sfa_cute_fake and sfd_cute_fake. Apply the corresponding construction changes in
srelu/api.py lines 256-270 and 315-353, and dsrelu/api.py lines 263-277 and
327-365; preserve the existing permuted-layout path.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/api.py`:
- Around line 167-170: Update the error messages in the dtype validation checks
near lines 169, 183, and 250 to remove unnecessary f-string prefixes and replace
torch-only wording with wording that covers all supported dtype representations:
torch, JAX, NumPy, string, and CUTLASS dtypes. Preserve the existing validation
behavior.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py`:
- Around line 310-311: Update the acc_dtype validation messages to use the
CUTLASS-neutral wording “float32” instead of “torch.float32” in _bf16_api.py at
python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py#L310-L311,
python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py#L300-L301,
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py#L339-L340, and
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L261-L262; preserve the
existing cutlass.Float32 validation and reported self.acc_dtype value at all
sites.
In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py`:
- Around line 103-105: Add a TYPE_CHECKING-guarded import of torch to
python/cudnn/gemm/cutedsl/grouped/glu/api.py (lines 103-105, also covering 180,
982, 994, and 1156-1158), dglu/api.py (104-105, also 188, 1022, 1034, and
1206-1207), glu_hadamard/api.py (45, also 77 and 629-631), wgrad/api.py (108,
also 246-247), dglu/_blockscaled_api.py (127), glu/_blockscaled_api.py (121),
and wgrad/_blockscaled_api.py (62), so all torch annotations resolve for Ruff
without importing torch at runtime.
In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py`:
- Around line 1384-1395: Ensure normalized uint8 FP4 output remains on the
low-precision path by adding cutlass.Uint8 to the is_low_precision_output_config
tuple in python/cudnn/gemm/cutedsl/grouped/quant/api.py:1384-1395 and the
matching tuple in python/cudnn/gemm/cutedsl/grouped/srelu/api.py:1292-1303; keep
both paths consistent with check_support and existing _interpret_uint8_as_fp4x2
handling.
In `@python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py`:
- Around line 344-364: Update _allocate_single_expert_placeholder in
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:344-364 to pass the JAX
device selected by jax.devices("gpu")[desc.device.index or 0] to jnp.empty. Also
update _generate_wgrad_ptrs in
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:499-515 to derive the
device from get_device(wgrad_tensor) and pass it to jnp.asarray, ensuring both
JAX allocations use the input tensor’s device.
- Around line 6-7: Add the typing-only torch import guard to each affected
module: python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py lines 6-7,
dglu/_bf16_api.py lines 6-7, glu/_bf16_api.py lines 6-7, unfused/_bf16_api.py
lines 6-7, unfused/api.py lines 13-28, quant/api.py lines 11-32, srelu/api.py
lines 11-31, swiglu/api.py lines 10-11, and dswiglu/api.py lines 10-11. Import
TYPE_CHECKING, then import torch only within its guard so the existing
torch.Tensor and torch.dtype annotations resolve for Ruff and
typing.get_type_hints without restoring eager runtime imports.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py`:
- Line 25: Rename the ambiguous parameter and corresponding usages of l in
make_inputs and the related test code to batch or num_l, preserving the existing
behavior while resolving Ruff E741 at all affected locations.
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`), updating all references in
the function so Ruff E741 is resolved.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py`:
- Line 27: Correct the second ceiling-division calculations for rest_k in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py:27 and
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py:27 to use (k +
127) // 128; update
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py:27 to use (k +
63) // 64, preserving the documented layout for non-aligned k values.
---
Outside diff comments:
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Around line 383-387: Update the error message in the _value_error_if call for
kernel_generate_sfd to use the framework-neutral or JAX-appropriate dtype name
instead of torch.float8_e8m0fnu, while preserving the existing validation
condition and required argument guidance.
---
Nitpick comments:
In `@docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md`:
- Around line 73-80: Update the swiglu_mlp example snippet to import both jax
and jax.numpy as jnp before using jax.jit and jnp.float32, while retaining the
existing gemm_swiglu_jax_sm100 import.
In `@python/cudnn/datatypes.py`:
- Around line 261-267: Update the fallback in _is_jax_array to split
type(input_tensor).__module__ into its first dotted component and compare that
component exactly against the supported JAX roots, "jax" and "jaxlib",
preventing similarly prefixed modules from matching.
In `@python/cudnn/gemm/cutedsl/_jax_ffi.py`:
- Around line 63-72: Replace the bare assert around gemm.check_support() in the
registry-miss path with an explicit support validation that always executes,
including optimized Python runs. Ensure unsupported configurations are rejected
before calling gemm._compile_kernel, and preserve the existing compilation and
registry flow for supported configurations.
In `@python/cudnn/gemm/cutedsl/dense/amax/jax_api.py`:
- Around line 57-60: Rename the ambiguous batch-dimension variable l to batch
consistently in python/cudnn/gemm/cutedsl/dense/amax/jax_api.py lines 57-60 and
python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py lines 56-59, including the
later swiglu uses at lines 86-87 and 106-110; in
python/cudnn/gemm/cutedsl/dense/swiglu/api.py lines 605-606, rename both l
unpackings and prefix the unused rebound k with an underscore to satisfy Ruff.
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 59-62: Introduce a module-level _require_torch_tensor helper that
accepts the tensor and API name, performs the existing optional-tensor torch
check, and raises the consistently formatted ValueError. Replace the duplicated
guards in GemmProjRopeMxfp8Bf16InSm100.__init__,
GemmProjRopeMxfp8Mxfp8InSm100.__init__, and gemm_proj_rope_mxfp8_wrapper_sm100
with calls to this helper.
In `@python/cudnn/gemm/cutedsl/dense/srelu/api.py`:
- Around line 485-551: Extract the duplicated framework setup and
cache-signature logic from gemm_srelu_wrapper_sm100 and
gemm_dsrelu_wrapper_sm100 into a shared private module under
gemm/cutedsl/dense/. Have both wrappers reuse it for framework detection,
Torch/JAX allocation and validation, SFD/amax handling, JAX synchronization, and
the existing stride/tensor/dynamic signature helpers while preserving current
behavior and framework-specific outputs.
In `@python/cudnn/gemm/cutedsl/dense/swiglu/jax_api.py`:
- Around line 38-40: The public API must reject non-default sf_vec_size,
vector_f32, and ab12_stages values because they are unsupported and ignored. Add
validation in the entry-point argument checks near the existing quantized-input
validation, raising the established error for any non-default setting while
preserving accepted defaults and avoiding changes to make_gemm or cache-key
behavior.
In `@python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py`:
- Around line 1082-1086: Replace the unconditional dynamic_m_tensor_signature
calls at python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py:1082-1086 and
python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py:1122-1126 with the same
layout-aware dynamic_m_sf_signature helper pattern used by
grouped/dsrelu/api.py:1631-1641. Ensure the helper selects the correct SFA shape
and stride dimensions for permuted and physical layouts while preserving the
existing None handling.
In `@python/cudnn/gemm/cutedsl/grouped/dglu/api.py`:
- Around line 62-73: The _block_scaled_dtype_pairs helper is duplicated across
the grouped GLU APIs; centralize it to keep backend dtype selection consistent.
Add the single definition beside select_grouped_gemm_backend in
python/cudnn/gemm/cutedsl/grouped/backend_utils.py, then delete the local
definitions and import the shared helper in
python/cudnn/gemm/cutedsl/grouped/dglu/api.py#L62-L73,
python/cudnn/gemm/cutedsl/grouped/glu/api.py#L64-L75, and
python/cudnn/gemm/cutedsl/grouped/wgrad/api.py#L36-L45.
In `@python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py`:
- Around line 1493-1494: The conditional imports in the affected API function
are separated from their use sites. Move import torch into the framework ==
"torch" branches immediately before the guarded torch usages, and move import
jax.numpy as jnp into the corresponding JAX branch near its first use,
preserving the existing branch behavior.
- Around line 20-21: Add a TYPE_CHECKING-guarded torch import after the future
import in python/cudnn/gemm/cutedsl/grouped/dsrelu/api.py (lines 20-21),
python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py (lines 10-17), and
python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py (lines 15-22). This
binds torch for annotation analysis while keeping the runtime import lazy and
resolves the remaining torch.Tensor and torch.dtype references in each module.
- Around line 1607-1629: Update tensor_signature to canonicalize
get_strides(tensor) with canonicalize_unit_dim_strides, adding that helper to
the cudnn.tensor_adapter imports and preserving the existing signature shape and
dtype behavior. Remove the duplicate local _sf_is_physical and reuse
GroupedGemmDsreluSm100._sf_desc_is_physical, widening or extracting the
predicate as needed so it accepts the tensor shape representation used here.
In `@python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py`:
- Around line 111-119: Update the framework validation around
detect_framework(sample_a) to remove the redundant sample_a is not None
condition and reject every non-"torch" framework, including "unknown" returned
for None. Preserve the existing JAX-specific message and the generic
unsupported-framework error for all other values.
In `@python/cudnn/gemm/cutedsl/grouped/glu/api.py`:
- Around line 1007-1034: Update the JAX branch of _allocate_output to compute
the C-contiguous strides for shape and compare them with the requested stride
after canonicalizing extent-1 dimensions; assert they match before allocating,
while preserving the existing allocation behavior for valid callers.
In `@python/cudnn/gemm/cutedsl/grouped/quant/api.py`:
- Around line 42-46: Define _JAX_SF_LAYOUT_ERROR once in
python/cudnn/gemm/cutedsl/grouped/moe_utils.py, preserving the existing message
and required substring. In python/cudnn/gemm/cutedsl/grouped/quant/api.py:42-46,
python/cudnn/gemm/cutedsl/grouped/srelu/api.py:41-45, and
python/cudnn/gemm/cutedsl/grouped/swiglu/api.py:33-37, remove the duplicated
literals and import the shared constant for each API.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py`:
- Around line 43-81: Update the _pointer_values docstring to explicitly state
that packed uint8 JAX pointers are decoded as little-endian 64-bit values,
matching the producers’ packing contract. Leave the existing
np.asarray(ptrs).view(np.int64) implementation unchanged.
- Around line 238-265: Deduplicate the shared BF16 grouped-GEMM helpers by
moving _output_dtypes, _validate_data_alignment,
_validate_pointer_array_alignment, _record_pointer_stream, _copy_values_to_host,
_is_validation_cached, and _remember_validation from
python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py:238-265 into a shared
mixin or private module near _pointer_values and _validate_pointer_tensor;
update glu/_bf16_api.py:196-223 and dglu/_bf16_api.py:200-227 to inherit or
import that shared implementation and remove their local copies.
In `@python/cudnn/gemm/cutedsl/grouped/unfused/api.py`:
- Around line 304-321: Move the framework_dtype import from _allocate_output to
the module-level imports, reusing the existing cudnn.tensor_adapter import
section. Remove the function-local import while preserving both Torch and JAX
dtype conversion behavior.
In `@python/cudnn/tensor_adapter.py`:
- Around line 126-133: Update cuda_is_available so it only returns immediately
when torch.cuda.is_available() is true; when torch is loaded but reports no
CUDA, fall through to the existing cudart.cudaGetDeviceCount() check instead of
returning false. Preserve the current torch fast path and runtime-based
detection for environments without torch.
In `@test/python/fe_api/gemm/test_gemm_amax_jax.py`:
- Around line 84-89: Create a shared JAX GEMM test helper module containing
device_sync, skip_unless_sm100, and _packed_jax_ptrs; update
test/python/fe_api/gemm/test_gemm_amax_jax.py lines 84-89 to move device_sync
and skip_unless_sm100 there, and replace local/test-module usage with imports.
In test/python/fe_api/grouped_gemm/test_grouped_gemm_glu_jax.py lines 38-40 and
test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu_jax.py lines 40-42,
remove each local _packed_jax_ptrs definition and import the shared helper
instead.
- Around line 139-159: Update the xfail configuration on
test_gemm_amax_jax_wrapper_fp4_uint8 so it cannot silently pass when the
expected failure is resolved or mask unrelated failures: add the relevant
tracking issue reference to reason and use strict=True if the TypeError failure
mode is stable, while preserving removal of the marker once the container path
is fixed.
🪄 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: 5a8b5efd-6e9e-44fa-9465-bfb8e1ea8832
📒 Files selected for processing (78)
docs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/discrete_grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/gemm_amax.mddocs/fe-oss-apis/gemm_fusions/gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.mddocs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.mddocs/fe-oss-apis/overview.mdpyproject.tomlpython/cudnn/__init__.pypython/cudnn/api_base.pypython/cudnn/datatypes.pypython/cudnn/gemm/cutedsl/_jax_ffi.pypython/cudnn/gemm/cutedsl/dense/amax/__init__.pypython/cudnn/gemm/cutedsl/dense/amax/api.pypython/cudnn/gemm/cutedsl/dense/amax/jax_api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/api.pypython/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_bf16in.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pypython/cudnn/gemm/cutedsl/dense/srelu/api.pypython/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.pypython/cudnn/gemm/cutedsl/dense/swiglu/__init__.pypython/cudnn/gemm/cutedsl/dense/swiglu/api.pypython/cudnn/gemm/cutedsl/dense/swiglu/jax_api.pypython/cudnn/gemm/cutedsl/discrete_grouped/discrete_kernel_utils.pypython/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/grouped/backend_utils.pypython/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/dglu/api.pypython/cudnn/gemm/cutedsl/grouped/dsrelu/api.pypython/cudnn/gemm/cutedsl/grouped/dswiglu/api.pypython/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/glu/api.pypython/cudnn/gemm/cutedsl/grouped/glu_hadamard/api.pypython/cudnn/gemm/cutedsl/grouped/glu_hadamard/hadamard_utils.pypython/cudnn/gemm/cutedsl/grouped/moe_kernel_helpers.pypython/cudnn/gemm/cutedsl/grouped/quant/api.pypython/cudnn/gemm/cutedsl/grouped/srelu/api.pypython/cudnn/gemm/cutedsl/grouped/swiglu/api.pypython/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/unfused/api.pypython/cudnn/gemm/cutedsl/grouped/utils.pypython/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.pypython/cudnn/gemm/cutedsl/grouped/wgrad/api.pypython/cudnn/tensor_adapter.pytest/python/conftest.pytest/python/fe_api/gemm/test_cutedsl_jax_guards.pytest/python/fe_api/gemm/test_gemm_amax_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 (3)
- python/cudnn/gemm/cutedsl/dense/dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py
- python/cudnn/gemm/cutedsl/dense/srelu/dense_blockscaled_gemm_persistent_srelu_quant.py
- python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.py
| ### JAX-specific constraints | ||
|
|
||
| - `L == 1`; `A`/`B` k-major; `C` n-major only | ||
| - `SFA`/`SFB` in the physical atom shape `(L, MN', K', 32, 4, 4)` (see "Using JAX arrays") | ||
| - Eager use only (no `jax.jit` over these entry points); synchronize before reading outputs |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Scope the "eager use only" constraint to the eager entry points.
Line 257 states "Eager use only (no jax.jit over these entry points)". The section header is generic ("JAX-specific constraints"), so the statement reads as applying to every JAX entry point. Line 112 and line 134 state the opposite for gemm_amax_jax_sm100, which is jax.jit-compatible.
📝 Proposed wording
-- Eager use only (no `jax.jit` over these entry points); synchronize before reading outputs
+- Eager entry points (`gemm_amax_wrapper_sm100`, `GemmAmaxSm100`): no `jax.jit`; synchronize before reading outputs. `gemm_amax_jax_sm100` is `jax.jit`-compatible and needs no manual synchronization.📝 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.
| ### JAX-specific constraints | |
| - `L == 1`; `A`/`B` k-major; `C` n-major only | |
| - `SFA`/`SFB` in the physical atom shape `(L, MN', K', 32, 4, 4)` (see "Using JAX arrays") | |
| - Eager use only (no `jax.jit` over these entry points); synchronize before reading outputs | |
| ### JAX-specific constraints | |
| - `L == 1`; `A`/`B` k-major; `C` n-major only | |
| - `SFA`/`SFB` in the physical atom shape `(L, MN', K', 32, 4, 4)` (see "Using JAX arrays") | |
| - Eager entry points (`gemm_amax_wrapper_sm100`, `GemmAmaxSm100`): no `jax.jit`; synchronize before reading outputs. `gemm_amax_jax_sm100` is `jax.jit`-compatible and needs no manual synchronization. |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/fe-oss-apis/gemm_fusions/gemm_amax.md` around lines 253 - 257, Update
the JAX-specific constraint bullet in the gemm_amax documentation to scope
“eager use only” exclusively to the eager entry points, while preserving the
documented jax.jit compatibility of gemm_amax_jax_sm100.
| The GEMM CuTeDSL APIs are type-erased and torch-lazy: torch is imported only when torch tensors are passed. JAX arrays are additionally accepted wherever the kernel's tensor layouts are expressible as row-major arrays (each API's page has a "JAX support" section with its exact contract): | ||
|
|
||
| - **Dense fusions** (amax, swiglu, srelu, dsrelu): full JAX eager support, plus `jax.jit`-compatible XLA custom-call entry points for amax and swiglu (see `gemm_amax.md` "Using JAX arrays"). | ||
| - **Grouped / discrete-grouped**: JAX eager support in discrete (pointer-array) weight modes — unfused grouped GEMM, glu/dglu (BF16), dsrelu (FP8), wgrad (BF16), and discrete-grouped swiglu/dswiglu (FP8). Dense weight mode, column-major bias layouts, and kernels whose scale factors are MMA-permuted tensor arguments (grouped swiglu/srelu/quant/dswiglu, glu_hadamard, block-scaled glu/dglu/wgrad backends) reject JAX with clear errors. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the Wgrad JAX support description.
Line 8 states that the listed APIs support JAX only in discrete pointer-array weight modes. It also states that dense weight mode rejects JAX. Both statements conflict with the Wgrad contract. Wgrad accepts plain BF16 JAX A and B arrays and supports both dense output and discrete output pointers.
List Wgrad separately from the discrete-weight APIs.
🤖 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` at line 8, Update the API support description
in the overview to separate Wgrad from the discrete pointer-array weight APIs.
Document Wgrad as accepting plain BF16 JAX A and B arrays with both dense output
and discrete output-pointer support, and remove it from the statements claiming
JAX requires discrete weights or rejects dense mode.
| dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, TensorDesc) or _is_framework_tensor(tensor_or_dtype) else tensor_or_dtype | ||
| return _convert_to_cutlass_data_type_or_none(dtype) in {cutlass.Float16, cutlass.BFloat16} | ||
|
|
||
| def _get_innermost_stride_dim(self, tensor: torch.Tensor, name: str = "") -> int: |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm whether api_base.py defers annotation evaluation and list remaining torch.* annotations.
set -euo pipefail
f=$(fd -p 'python/cudnn/api_base.py' | head -n1)
echo "== file: $f =="
sed -n '1,40p' "$f"
echo "== future import present? =="
rg -n '^from __future__ import annotations' "$f" || echo "NOT PRESENT"
echo "== module-level torch import present? =="
rg -nP '^\s*import torch\b|^\s*from torch\b' "$f" || echo "NOT PRESENT"
echo "== annotations still referencing torch =="
rg -nP '\btorch\.(Tensor|dtype|Size|device)\b' "$f"Repository: NVIDIA/cudnn-frontend
Length of output: 3950
🏁 Script executed:
#!/bin/bash
set -euo pipefail
f=python/cudnn/api_base.py
echo "== lazy torch binding and affected methods =="
sed -n '35,125p' "$f"
sed -n '560,725p' "$f"
echo "== all module-scope torch bindings and references =="
rg -n '(^|[^[:alnum:]_])torch([^[:alnum:]_]|$)|def _torch' "$f"
echo "== AST annotation and name-use probe =="
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("python/cudnn/api_base.py")
tree = ast.parse(path.read_text())
print("future_annotations:", any(
isinstance(node, ast.ImportFrom)
and node.module == "__future__"
and any(alias.name == "annotations" for alias in node.names)
for node in tree.body
))
print("module torch bindings:", [
(type(node).__name__, getattr(node, "lineno", None))
for node in tree.body
if isinstance(node, (ast.Import, ast.ImportFrom))
and any(alias.asname == "torch" or (alias.asname is None and alias.name == "torch")
for alias in node.names)
])
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "_get_innermost_stride_dim":
print("annotation AST:", ast.dump(node.args.args[1].annotation, include_attributes=False))
print("function line:", node.lineno)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 15528
🏁 Script executed:
#!/bin/bash
set -u
echo "== lint configuration =="
git ls-files | rg '(^|/)(pyproject\.toml|ruff\.toml|\.ruff\.toml|setup\.cfg)$' | xargs -r -n1 sh -c 'echo "--- $0"; rg -n -C 4 "ruff|F821|target-version|select|ignore" "$0"'
echo "== ruff availability and targeted diagnostics =="
if command -v ruff >/dev/null 2>&1; then
ruff --version
ruff check python/cudnn/api_base.py --output-format concise
else
echo "ruff: unavailable"
fi
echo "== get_strides definition and call sites =="
rg -n -C 12 'def get_strides|get_strides\(' python test 2>/dev/null | head -240
echo "== affected method call sites =="
rg -n -C 5 '_get_innermost_stride_dim' python test 2>/dev/null | head -160
echo "== annotation evaluation probe =="
python3 - <<'PY'
ns = {}
source = "from __future__ import annotations\n\ndef f(x: torch.Tensor) -> int:\n return 1\n"
exec(compile(source, "<probe>", "exec"), ns)
print("defined_without_torch:", "f" in ns)
print("annotation:", ns["f"].__annotations__)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 39086
Replace torch.Tensor with Any at line 690.
The deferred annotation avoids an import-time failure, but Ruff reports F821 because torch is not defined at module scope. Do not add an eager torch import; preserve the lazy dependency.
🧰 Tools
🪛 Ruff (0.16.1)
[error] 690-690: Undefined name torch
(F821)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/api_base.py` at line 690, Update the _get_innermost_stride_dim
method annotation to use Any instead of torch.Tensor, resolving Ruff’s
undefined-name error while preserving the module’s lazy torch dependency and
avoiding an eager import.
Sources: Coding guidelines, Linters/SAST tools
| self.a_desc = self._make_tensor_desc(sample_a, name="sample_a", canonical=True) | ||
| self.b_desc = self._make_tensor_desc(sample_b, name="sample_b", canonical=True) | ||
| self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa", canonical=True) | ||
| self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb", canonical=True) | ||
| self.c_desc = self._make_tensor_desc(sample_c, name="sample_c", canonical=True) | ||
| self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax", canonical=True) | ||
|
|
||
| self.acc_dtype = acc_dtype | ||
| self.acc_dtype = _convert_to_cutlass_data_type(acc_dtype) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Compare where _interpret_uint8_as_fp4x2 is set relative to _make_tensor_desc calls.
set -euo pipefail
rg -l '_interpret_uint8_as_fp4x2' --glob '*.py' python | while IFS= read -r f; do
echo "===== $f ====="
rg -n '_interpret_uint8_as_fp4x2\s*=|_make_tensor_desc\(' "$f"
doneRepository: NVIDIA/cudnn-frontend
Length of output: 21710
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== APIBase descriptor construction ====="
sed -n '940,1010p' python/cudnn/api_base.py
sed -n '400,440p' python/cudnn/api_base.py
echo "===== dense/amax API ====="
ast-grep outline python/cudnn/gemm/cutedsl/dense/amax/api.py
sed -n '1,240p' python/cudnn/gemm/cutedsl/dense/amax/api.py
echo "===== related FP4 descriptor consumers ====="
rg -n -C 4 '_make_fake_cute_tensor_from_desc|interpret_uint8_as_fp4x2|Uint8 ab_dtype|_tensor_shape|_tensor_stride' python/cudnn python/cudnn/gemm/cutedslRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== APIBase flag and descriptor implementation ====="
sed -n '410,435p' python/cudnn/api_base.py
sed -n '950,1005p' python/cudnn/api_base.py
rg -n -C 8 'def _make_fake_cute_tensor_from_desc|interpret_uint8_as_fp4x2' python/cudnn/api_base.py
echo "===== dense/amax constructor and compile path ====="
sed -n '45,85p' python/cudnn/gemm/cutedsl/dense/amax/api.py
sed -n '115,145p' python/cudnn/gemm/cutedsl/dense/amax/api.py
sed -n '325,365p' python/cudnn/gemm/cutedsl/dense/amax/api.py
echo "===== descriptor model and tensor-shape helpers ====="
rg -n -C 8 'class TensorDesc|def _tensor_shape|def _tensor_stride|interpret_uint8_as_fp4x2' python/cudnnRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== amax tests and call sites ====="
rg -n -C 5 'GemmAmaxSm100|gemm_amax|amax_wrapper|sample_a.*uint8|torch\.uint8|uint8' \
test python docs --glob '*.py' --glob '*.md' --glob '*.rst'
echo "===== amax kernel dtype handling ====="
rg -n -C 8 'def gemm_amax|class .*Amax|Float4E2M1FN|interpret_uint8_as_fp4x2|ab_dtype|a_tensor|b_tensor' \
python/cudnn/gemm/cutedsl/dense/amax python/cudnn/gemm/cutedsl --glob '*.py'
echo "===== amax descriptor shape/stride checks ====="
sed -n '700,755p' python/cudnn/api_base.py
sed -n '180,235p' python/cudnn/gemm/cutedsl/dense/amax/api.pyRepository: NVIDIA/cudnn-frontend
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== files with amax in the name ====="
fd -i 'amax' python test docs
echo "===== direct amax API references ====="
rg -l 'GemmAmaxSm100|gemm_amax' python test docs | sort
echo "===== direct amax context ====="
rg -n -C 6 'GemmAmaxSm100|gemm_amax' \
python/cudnn/gemm/cutedsl/dense/amax \
test/python \
docs/fe-oss-apis \
--glob '*.py' --glob '*.md' --glob '*.rst' \
--max-count 80
echo "===== dense amax kernel definitions ====="
fd -i 'amax' python/cudnn/gemm/cutedsl/dense --exec sh -c 'echo "===== $1 ====="; rg -n -C 6 "class |def |Float4E2M1FN|Uint8|ab_dtype|a_tensor|b_tensor" "$1"' sh {}Repository: NVIDIA/cudnn-frontend
Length of output: 50380
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== documented amax FP4 contract ====="
sed -n '120,136p' docs/fe-oss-apis/gemm_fusions/gemm_amax.md
sed -n '220,258p' docs/fe-oss-apis/gemm_fusions/gemm_amax.md
echo "===== focused amax tests ====="
sed -n '1,260p' test/python/fe_api/gemm/test_gemm_amax.py
sed -n '1,260p' test/python/fe_api/gemm/test_gemm_amax_jax.py
echo "===== amax kernel implementation ====="
sed -n '1,180p' python/cudnn/gemm/cutedsl/dense/amax/dense_blockscaled_gemm_persistent_amax.py
rg -n -C 8 'ab_dtype|Float4E2M1FN|Uint8|make_ptr|gemm_amax' \
python/cudnn/gemm/cutedsl/dense/amax/dense_blockscaled_gemm_persistent_amax.pyRepository: NVIDIA/cudnn-frontend
Length of output: 30878
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "===== amax support validation ====="
sed -n '145,180p' python/cudnn/gemm/cutedsl/dense/amax/api.py
sed -n '235,315p' python/cudnn/gemm/cutedsl/dense/amax/api.py
echo "===== amax test parameterization ====="
rg -n -C 8 'with_gemm_amax_params_fp4|ab_dtype|torch\.uint8|float4|fp4' \
test/python/fe_api/gemm/test_gemm_amax_utils.py \
test/python/fe_api/gemm/test_gemm_amax.py
echo "===== source-level ordering verifier ====="
python3 - <<'PY'
import ast
from pathlib import Path
paths = [
Path("python/cudnn/gemm/cutedsl/dense/amax/api.py"),
Path("python/cudnn/gemm/cutedsl/discrete_grouped/dswiglu/api.py"),
Path("python/cudnn/gemm/cutedsl/discrete_grouped/swiglu/api.py"),
]
for path in paths:
tree = ast.parse(path.read_text())
print(f"===== {path} =====")
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef) and node.name == "__init__":
calls = []
assignments = []
for child in ast.walk(node):
if isinstance(child, ast.Call) and isinstance(child.func, ast.Attribute) and child.func.attr == "_make_tensor_desc":
calls.append(child.lineno)
if isinstance(child, ast.Assign):
for target in child.targets:
if isinstance(target, ast.Attribute) and target.attr == "_interpret_uint8_as_fp4x2":
assignments.append((child.lineno, ast.unparse(child.value)))
print("descriptor_calls:", sorted(calls))
print("flag_assignments:", sorted(assignments))
break
PYRepository: NVIDIA/cudnn-frontend
Length of output: 36800
Align Uint8 FP4 handling in GemmAmaxSm100.
check_support() accepts Uint8 as packed FP4, but descriptors are created before _interpret_uint8_as_fp4x2 is enabled. Compilation therefore treats the inputs as native Uint8, and the packed-FP4 path fails. Enable the flag before descriptor creation, or reject Uint8 in check_support() and update the documentation and tests.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/dense/amax/api.py` around lines 57 - 64, The
GemmAmaxSm100 initialization currently creates tensor descriptors before
enabling packed-FP4 interpretation for Uint8 inputs. Set
_interpret_uint8_as_fp4x2 before the _make_tensor_desc calls when the inputs are
Uint8, so descriptor construction matches the Uint8 acceptance already defined
by check_support().
| self._value_error_if( | ||
| ab_dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and self.sf_vec_size == 16, | ||
| ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and self.sf_vec_size == 16, | ||
| f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported", |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove the extraneous f prefix.
The escaped braces leave no placeholder, so Ruff reports F541. The rendered text does not change.
🧹 Proposed fix
- f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported",
+ "Unsupported ab_dtype and sf_vec_size combination: {float8_e5m2, float8_e4m3fn} and 16 is not supported",📝 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.
| f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported", | |
| "Unsupported ab_dtype and sf_vec_size combination: {float8_e5m2, float8_e4m3fn} and 16 is not supported", |
🧰 Tools
🪛 Ruff (0.16.1)
[error] 156-156: f-string without any placeholders
Remove extraneous f prefix
(F541)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm/cutedsl/dense/amax/api.py` at line 156, Remove the
unnecessary f-string prefix from the unsupported dtype and scale-vector-size
error message, preserving the escaped-brace text unchanged.
Source: Linters/SAST tools
| from __future__ import annotations | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Ruff F821 Undefined name torch across the cohort. Every file removed the eager import torch but kept torch.Tensor and torch.dtype in annotations. from __future__ import annotations prevents a runtime failure, but Ruff 0.16.1 reports F821 on each annotation and typing.get_type_hints cannot resolve them. Add if TYPE_CHECKING: import torch in each file.
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L6-L7: add the guarded import; resolves F821 at lines 63, 145, 193, 198, 202, 499.python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py#L6-L7: add the guarded import; resolves F821 at lines 61, 157, 201, 206, 210.python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py#L6-L7: add the guarded import; resolves F821 at lines 59, 153, 197, 202, 206.python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py#L6-L7: add the guarded import; resolves F821 at lines 100, 191, 244, 248.python/cudnn/gemm/cutedsl/grouped/unfused/api.py#L13-L28: add the guarded import; resolves F821 at lines 50, 132, 262, 263, 264.python/cudnn/gemm/cutedsl/grouped/quant/api.py#L11-L32: add the guarded import; resolves F821 at lines 95, 1266, 1267.python/cudnn/gemm/cutedsl/grouped/srelu/api.py#L11-L31: add the guarded import; resolves F821 at lines 48, 96, 1178, 1179, 1180.python/cudnn/gemm/cutedsl/grouped/swiglu/api.py#L10-L11: add the guarded import; resolves F821 at lines 80, 787, 788, 789.python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py#L10-L11: add the guarded import; resolves F821 at lines 75, 704, 705.
📍 Affects 9 files
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L6-L7(this comment)python/cudnn/gemm/cutedsl/grouped/dglu/_bf16_api.py#L6-L7python/cudnn/gemm/cutedsl/grouped/glu/_bf16_api.py#L6-L7python/cudnn/gemm/cutedsl/grouped/unfused/_bf16_api.py#L6-L7python/cudnn/gemm/cutedsl/grouped/unfused/api.py#L13-L28python/cudnn/gemm/cutedsl/grouped/quant/api.py#L11-L32python/cudnn/gemm/cutedsl/grouped/srelu/api.py#L11-L31python/cudnn/gemm/cutedsl/grouped/swiglu/api.py#L10-L11python/cudnn/gemm/cutedsl/grouped/dswiglu/api.py#L10-L11
🤖 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 6 - 7, Add
the typing-only torch import guard to each affected module:
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py lines 6-7,
dglu/_bf16_api.py lines 6-7, glu/_bf16_api.py lines 6-7, unfused/_bf16_api.py
lines 6-7, unfused/api.py lines 13-28, quant/api.py lines 11-32, srelu/api.py
lines 11-31, swiglu/api.py lines 10-11, and dswiglu/api.py lines 10-11. Import
TYPE_CHECKING, then import torch only within its guard so the existing
torch.Tensor and torch.dtype annotations resolve for Ruff and
typing.get_type_hints without restoring eager runtime imports.
Source: Linters/SAST tools
| def _allocate_single_expert_placeholder(self) -> None: | ||
| """Allocate the real (never read) discrete-mode single-expert template tensor.""" | ||
| desc = self.single_expert_wgrad_desc | ||
| if self._framework == "torch": | ||
| import torch | ||
|
|
||
| self._single_expert_placeholder = torch.empty_strided( | ||
| desc.shape, | ||
| desc.stride, | ||
| dtype=framework_dtype(desc.dtype, "torch"), | ||
| device=desc.device, | ||
| ) | ||
| return | ||
| import jax | ||
| import jax.numpy as jnp | ||
|
|
||
| if canonicalize_unit_dim_strides(desc.shape, desc.stride) != canonicalize_unit_dim_strides( | ||
| desc.shape, TensorDesc._compute_contiguous_stride(desc.shape) | ||
| ): | ||
| raise ValueError(f"single expert placeholder layout {desc.stride} is not expressible as a C-contiguous JAX array") | ||
| self._single_expert_placeholder = jax.block_until_ready(jnp.empty(desc.shape, dtype=framework_dtype(desc.dtype, "jax"))) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
JAX allocations ignore the target device. Both new JAX allocation paths in this file create arrays without a device argument, so the buffers land on JAX's default device instead of the device carried by the input tensors. The torch branches bind the device correctly, and allocate_byte_workspace in python/cudnn/tensor_adapter.py already maps a Device descriptor to a jax.Device. In a multi-GPU JAX process the placeholder pointer is baked into the compiled kernel and the generated pointer array trips the device check at line 576.
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L344-L364: passdevice=jax.devices("gpu")[desc.device.index or 0]tojnp.emptyin_allocate_single_expert_placeholder.python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L499-L515: pass the JAX device derived fromget_device(wgrad_tensor)tojnp.asarrayin_generate_wgrad_ptrs.
📍 Affects 1 file
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L344-L364(this comment)python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py#L499-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 344 - 364,
Update _allocate_single_expert_placeholder in
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:344-364 to pass the JAX
device selected by jax.devices("gpu")[desc.device.index or 0] to jnp.empty. Also
update _generate_wgrad_ptrs in
python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py:499-515 to derive the
device from get_device(wgrad_tensor) and pass it to jnp.asarray, ensuring both
JAX allocations use the input tensor’s device.
| from cudnn.api_base import ceil_div | ||
|
|
||
|
|
||
| def make_inputs(m, n, k, l, sf_vec_size, rng): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous variable l.
Ruff reports E741 for l at these lines. Rename it to batch or num_l to clear the lint error.
♻️ Proposed rename
-def make_inputs(m, n, k, l, sf_vec_size, rng):
+def make_inputs(m, n, k, batch, sf_vec_size, rng):
rest_k = ceil_div(ceil_div(k, sf_vec_size), 4)
- a_np = rng.standard_normal((m, k, l), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn)
- b_np = rng.standard_normal((n, k, l), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn)
+ a_np = rng.standard_normal((m, k, batch), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn)
+ b_np = rng.standard_normal((n, k, batch), dtype=np.float32).astype(ml_dtypes.float8_e4m3fn)
# e8m0 value 1.0 == byte 127; physical C-contiguous atom shape (L, MN', K', 32, 4, 4)
- sfa_np = np.full((l, ceil_div(m, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)
- sfb_np = np.full((l, ceil_div(n, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)
- prob_np = rng.random((m, 1, l), dtype=np.float32)
+ sfa_np = np.full((batch, ceil_div(m, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)
+ sfb_np = np.full((batch, ceil_div(n, 128), rest_k, 32, 4, 4), 127, dtype=np.uint8).view(ml_dtypes.float8_e8m0fnu)
+ prob_np = rng.random((m, 1, batch), dtype=np.float32)Also applies to: 45-45, 86-86, 130-130
🧰 Tools
🪛 Ruff (0.16.1)
[error] 25-25: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/gemm/test_gemm_srelu_dsrelu_jax.py` at line 25, Rename the
ambiguous parameter and corresponding usages of l in make_inputs and the related
test code to batch or num_l, preserving the existing behavior while resolving
Ruff E741 at all affected locations.
Source: Linters/SAST tools
| from fe_api.gemm.test_gemm_amax_jax import skip_unless_sm100 | ||
|
|
||
|
|
||
| def _make_jax_inputs(m=256, n=128, k=128, l=2): |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Rename the ambiguous variable l.
Ruff reports E741 at both lines. Rename the parameter and the local to num_experts (or batch) to clear the lint error.
Also applies to: 91-91
🧰 Tools
🪛 Ruff (0.16.1)
[error] 23-23: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_jax.py` at line 23,
Rename the ambiguous `l` parameter and corresponding local variable in
`_make_jax_inputs` to `num_experts` (or `batch`), updating all references in the
function so Ruff E741 is resolved.
Source: Linters/SAST tools
| """Plausible (physical-layout) inputs; rejection fires before shape validation.""" | ||
| rng = np.random.default_rng(20260809) | ||
| a_j = jnp.asarray(rng.integers(0, 255, (m, k, 1), dtype=np.uint8).view(ml_dtypes.float8_e4m3fn)) | ||
| rest_k = -(-(-(-k // 32) // 4)) # ceil_div(ceil_div(k, 32), 4) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Fix the second ceiling division.
The current expressions round the second division down. For example, with k=160, the FP8 helpers produce rest_k == 1, but the documented layout requires 2.
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L27: use(k + 127) // 128.test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L27: use(k + 127) // 128.test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L27: use(k + 63) // 64.
Proposed fix
- rest_k = -(-(-(-k // 32) // 4))
+ rest_k = (k + 127) // 128📝 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.
| rest_k = -(-(-(-k // 32) // 4)) # ceil_div(ceil_div(k, 32), 4) | |
| rest_k = (k + 127) // 128 |
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 27-27: Python does not support the unary prefix decrement operator (--)
(B002)
📍 Affects 3 files
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py#L27-L27(this comment)test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py#L27-L27test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py#L27-L27
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py` at line 27,
Correct the second ceiling-division calculations for rest_k in
test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_jax.py:27 and
test/python/fe_api/grouped_gemm/test_grouped_gemm_srelu_jax.py:27 to use (k +
127) // 128; update
test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_jax.py:27 to use (k +
63) // 64, preserving the documented layout for non-aligned k values.
Source: Linters/SAST tools
88ca0ba to
366d16d
Compare
Anerudhan
changed the title
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py`:
- Around line 582-585: Replace the public input assertions in the wrapper,
including the dtype check around _convert_to_cutlass_data_type and the
validation checks near lines 620 and 641, with explicit ValueError branches.
Preserve each existing validation condition and message semantics so dtype and
MXFP8 scale checks always execute, including when Python runs with optimizations
or a cached kernel is used.
- Around line 242-266: Synchronize every JAX input with jax.block_until_ready
before the raw CUDA launch in both execution paths of proj_rope_mxfp8/api.py
(242-266 and 487-510), including x, w, cos, sin, and all MXFP8 scale tensors;
retain output synchronization before host reads. Update
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md (5-9) to state that the
wrapper owns input readiness, and extend
test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py (37-40) with an
asynchronous-producer regression test.
🪄 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: 57006d04-52cb-44f3-854c-75c63b1cb23c
📒 Files selected for processing (5)
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.mddocs/fe-oss-apis/overview.mdpython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.pypython/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/gemm_proj_rope_mxfp8_mxfp8in.pytest/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/fe-oss-apis/overview.md
| if current_stream is None: | ||
| # torch inputs stay ordered with the caller's current torch stream; | ||
| # other frameworks (e.g. JAX) default to the CUDA legacy default stream. | ||
| current_stream = default_stream(detect_framework(x)) | ||
| self._runtime_error_if( | ||
| self._compiled_kernel is None, | ||
| "GemmProjRopeMxfp8Bf16InSm100 kernel not compiled; call compile() first", | ||
| ) | ||
| cute_tensors = self._to_cute_tensors(x, w, cos, sin, out_fp8_row, out_scales_row, out_fp8_col, out_scales_col) | ||
| self._compiled_kernel(*cute_tensors, current_stream) | ||
| # TVM-FFI entry point: raw DLPack-capable tensors go straight to the compiled | ||
| # kernel (no per-call from_dlpack). torch inputs are detached views so autograd- | ||
| # tracked tensors stay exportable; w in [in, out] layout is a transposed view. | ||
| w_arg = _maybe_detach(w) | ||
| if not self.w_out_in: | ||
| w_arg = w_arg.transpose(0, 1) | ||
| self._compiled_kernel( | ||
| _maybe_detach(x), | ||
| w_arg, | ||
| _maybe_detach(cos), | ||
| _maybe_detach(sin), | ||
| out_fp8_row, | ||
| out_scales_row, | ||
| out_fp8_col, | ||
| out_scales_col, | ||
| current_stream, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
file='python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py'
test='test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py'
doc='docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md'
printf '%s\n' '--- source outline ---'
ast-grep outline "$file" | sed -n '1,220p'
printf '%s\n' '--- BF16 execute region ---'
sed -n '190,285p' "$file"
printf '%s\n' '--- MXFP8 execute region ---'
sed -n '435,530p' "$file"
printf '%s\n' '--- helper definitions and uses ---'
rg -n -C 3 'def (default_stream|detect_framework|_maybe_detach|_mxfp8_as_e8m0)|default_stream\(|detect_framework\(|block_until_ready|device_sync|synchronize|sync' \
"$file" "$test" "$doc" python/cudnn test/python/fe_api/gemm | sed -n '1,320p'
printf '%s\n' '--- JAX test ---'
cat -n "$test" | sed -n '1,180p'
printf '%s\n' '--- documentation ---'
cat -n "$doc" | sed -n '1,100p'Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -eu
api='python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py'
test='test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py'
printf '%s\n' '--- api imports and local helpers ---'
sed -n '1,70p' "$api"
printf '%s\n' '--- wrapper body ---'
sed -n '553,680p' "$api"
printf '%s\n' '--- stream helper definitions and imports ---'
rg -n -C 8 'default_stream|detect_framework|def device_sync|block_until_ready|external.*stream|legacy default' \
python/cudnn test/python/fe_api/gemm "$test" | sed -n '1,260p'
printf '%s\n' '--- exact synchronization and launch ordering ---'
python3 - <<'PY'
from pathlib import Path
import ast
api_path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
test_path = Path("test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py")
api = api_path.read_text()
tree = ast.parse(api)
for cls_name in ("GemmProjRopeMxfp8Bf16InSm100", "GemmProjRopeMxfp8Mxfp8InSm100"):
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == cls_name)
execute = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "execute")
segment = ast.get_source_segment(api, execute)
print(f"{cls_name}:")
print(" block_until_ready:", "block_until_ready" in segment)
print(" synchronize:", "synchronize" in segment)
print(" selected stream:", "default_stream(" in segment)
print(" compiled launch count:", segment.count("self._compiled_kernel("))
print(" launch line:", api[:execute.lineno - 1].count("\n") + 1)
wrapper = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "gemm_proj_rope_mxfp8_wrapper_sm100")
segment = ast.get_source_segment(api, wrapper)
print("wrapper:")
for token in ("jax.block_until_ready", "self.execute", ".execute(", "current_stream", "return"):
print(f" {token}: {segment.count(token)}")
test = test_path.read_text()
test_tree = ast.parse(test)
run_both = next(n for n in test_tree.body if isinstance(n, ast.FunctionDef) and n.name == "run_both")
run_segment = ast.get_source_segment(test, run_both)
print("test run_both:")
for token in ("jax.block_until_ready", "device_sync", "wrapper_sm100"):
print(f" {token}: {run_segment.count(token)}")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 30419
🌐 Web query:
JAX asynchronous dispatch CUDA external library stream interoperability block_until_ready raw device pointer
💡 Result:
JAX utilizes asynchronous dispatch to decouple Python execution from accelerator compute, allowing the host to enqueue operations on GPU/TPU without waiting for completion [1][2]. Asynchronous Dispatch and Synchronization JAX returns future-like jax.Array objects, meaning control returns to the Python interpreter before kernels finish executing [1][3]. To force synchronization and ensure that computation results are available on the host (or that previous operations are finished), you must call jax.Array.block_until_ready() [1][4]. While this is often necessary for benchmarking or host-side inspection, JAX’s internal queueing mechanisms are designed to keep the accelerator busy without requiring manual blocking [1][5]. External CUDA Library Interoperability For high-performance interoperability with external CUDA libraries, JAX provides the Foreign Function Interface (FFI) [6]. 1. Stream Management: JAX’s FFI exposes the underlying execution stream to custom kernels [7]. In FFI handlers, you can accept a ffi::PlatformStream<cudaStream_t> context to ensure your custom kernels are enqueued on the same stream as JAX’s own operations, preserving correct execution order [7]. 2. Raw Device Pointers: The jax.ffi system provides access to buffer data (which encapsulates raw device pointers) [6][8]. Within a registered FFI handler, the ffi::Buffer objects provide access to the typed data (pointer) [7]. 3. Interoperability Protocols: Beyond FFI, JAX supports standard array protocols such as DLPack and cuda_array_interface [9][10]. These are generally the preferred way to exchange raw device memory buffers with libraries like CuPy, PyTorch, or Warp in a zero-copy manner [10][8]. DLPack is typically recommended for modern cross-library workflows due to its robust support for multi-device and stream synchronization [9][10]. In summary, for deep integration, the recommended path is using the XLA FFI to access JAX's internal streams and buffers [7][6]. For simpler data exchange, use DLPack-based protocols to pass raw device pointers across libraries [9][10].
Citations:
- 1: https://docs.jax.dev/en/latest/async_dispatch.html
- 2: Asynchronous Dispatch in JAX jax-ml/jax#9895
- 3: https://github.com/google/jax/blob/master/docs/async_dispatch.rst
- 4: https://docs.jax.dev/en/latest/%5Fautosummary/jax.block_until_ready.html
- 5: Document async dispatch in the README jax-ml/jax#1219
- 6: https://docs.jax.dev/en/latest/ffi.html
- 7: https://github.com/jax-ml/jax/blob/main/examples/ffi/src/jax_ffi_example/cuda_examples.cu
- 8: https://github.com/jax-ml/jax/blob/main/jaxlib/ffi.h
- 9: Add CUDA Array Interface consumer support jax-ml/jax#19233
- 10: https://nvidia.github.io/warp/stable/user_guide/interoperability.html
🏁 Script executed:
#!/bin/bash
set -eu
api='python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py'
test='test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py'
printf '%s\n' '--- api imports and local helpers ---'
sed -n '1,70p' "$api"
printf '%s\n' '--- wrapper body ---'
sed -n '553,680p' "$api"
printf '%s\n' '--- stream helper definitions and imports ---'
rg -n -C 8 'default_stream|detect_framework|def device_sync|block_until_ready|external.*stream|legacy default' \
python/cudnn test/python/fe_api/gemm "$test" | sed -n '1,260p'
printf '%s\n' '--- exact synchronization and launch ordering ---'
python3 - <<'PY'
from pathlib import Path
import ast
api_path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
test_path = Path("test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py")
api = api_path.read_text()
tree = ast.parse(api)
for cls_name in ("GemmProjRopeMxfp8Bf16InSm100", "GemmProjRopeMxfp8Mxfp8InSm100"):
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == cls_name)
execute = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "execute")
segment = ast.get_source_segment(api, execute)
print(f"{cls_name}:")
print(" block_until_ready:", "block_until_ready" in segment)
print(" synchronize:", "synchronize" in segment)
print(" selected stream:", "default_stream(" in segment)
print(" compiled launch count:", segment.count("self._compiled_kernel("))
print(" launch line:", execute.lineno)
wrapper = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "gemm_proj_rope_mxfp8_wrapper_sm100")
segment = ast.get_source_segment(api, wrapper)
print("wrapper:")
for token in ("jax.block_until_ready", "self.execute", ".execute(", "current_stream", "return"):
print(f" {token}: {segment.count(token)}")
test = test_path.read_text()
test_tree = ast.parse(test)
run_both = next(n for n in test_tree.body if isinstance(n, ast.FunctionDef) and n.name == "run_both")
run_segment = ast.get_source_segment(test, run_both)
print("test run_both:")
for token in ("jax.block_until_ready", "device_sync", "wrapper_sm100"):
print(f" {token}: {run_segment.count(token)}")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 30408
🌐 Web query:
site:docs.jax.dev/en/latest FFI CUDA stream external operation JAX execution stream block_until_ready
💡 Result:
In JAX, when using the Foreign Function Interface (FFI) to execute custom CUDA operations, the integration with CUDA streams and synchronization is managed via the XLA FFI handler interface [1][2]. 1. CUDA Stream Access: When registering an FFI target for the CUDA platform, your handler definition can request the active execution stream [1][2]. You do this by binding a context parameter of type ffi::PlatformStream<cudaStream_t> in your handler's binding specification [1][2]. The JAX runtime then provides the correct cudaStream_t to your implementation function at execution time, allowing you to launch your kernels directly on that stream [1][2]. 2. Asynchronous Execution and block_until_ready: JAX operations, including FFI calls, follow an asynchronous dispatch model [3]. When you call an FFI function, JAX does not wait for the operation to complete on the GPU; instead, it returns a JAX array (a future) and resumes Python execution immediately [3]. Because of this, block_until_ready() is essential for accurate timing and debugging [3][4][5]. Calling block_until_ready() on the output array forces the host to synchronize with the device, effectively waiting for all queued operations—including those in the CUDA stream used by your FFI call—to finish before the Python execution continues [3][6][5]. 3. Integration: By launching kernels on the stream provided by JAX, your custom operation remains correctly ordered within the XLA execution graph [5]. Subsequent JAX operations will be queued on that same stream (or respect the necessary dependencies), ensuring correct execution order [3]. The use of block_until_ready() ensures that you can safely inspect the results or measure performance without the results being skewed by the asynchronous nature of the device dispatch [3][5].
Citations:
- 1: https://docs.jax.dev/en/latest/%5Fsources/ffi.md
- 2: https://docs.jax.dev/en/latest/ffi.html
- 3: https://docs.jax.dev/en/latest/async_dispatch.html
- 4: https://docs.jax.dev/en/latest/profiling.html
- 5: https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html
- 6: https://docs.jax.dev/en/latest/jit-compilation.html
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- all call sites ---'
rg -n -C 4 'gemm_proj_rope_mxfp8_wrapper_sm100|GemmProjRopeMxfp8Bf16InSm100|GemmProjRopeMxfp8Mxfp8InSm100' \
python test docs | sed -n '1,260p'
printf '%s\n' '--- JAX interop tests and stream contracts ---'
rg -l 'jax|block_until_ready|device_sync|CUstream\(0\)|default_stream\("jax"\)' test/python/fe_api python/cudnn | sort | \
xargs -r rg -n -C 3 'block_until_ready|device_sync|CUstream\(0\)|default_stream\(|wrapper.*jax|jax.*wrapper' | sed -n '1,360p'
printf '%s\n' '--- source-level input/launch ordering verifier ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
source = path.read_text()
tree = ast.parse(source)
def call_names(node):
result = []
for child in ast.walk(node):
if isinstance(child, ast.Call):
f = child.func
if isinstance(f, ast.Attribute):
result.append(f"{ast.unparse(f.value)}.{f.attr}")
else:
result.append(ast.unparse(f))
return result
for cls_name in ("GemmProjRopeMxfp8Bf16InSm100", "GemmProjRopeMxfp8Mxfp8InSm100"):
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == cls_name)
execute = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "execute")
calls = call_names(execute)
launch = next(n for n in ast.walk(execute)
if isinstance(n, ast.Call) and ast.unparse(n.func) == "self._compiled_kernel")
launch_line = launch.lineno
print(cls_name)
print(" launch line:", launch_line)
print(" readiness calls:", [c for c in calls if "ready" in c.lower() or "synchron" in c.lower()])
print(" input identifiers in launch:", [ast.unparse(a) for a in launch.args[:-1]])
print(" stream argument:", ast.unparse(launch.args[-1]))
wrapper = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "gemm_proj_rope_mxfp8_wrapper_sm100")
for i, node in enumerate(wrapper.body, 1):
text = ast.get_source_segment(source, node) or ""
if "block_until_ready" in text or ".execute(" in text:
print("wrapper statement", i, "line", node.lineno, ":", " ".join(text.split()))
PYRepository: NVIDIA/cudnn-frontend
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- all call sites ---'
rg -n -C 4 'gemm_proj_rope_mxfp8_wrapper_sm100|GemmProjRopeMxfp8Bf16InSm100|GemmProjRopeMxfp8Mxfp8InSm100' \
python test docs | sed -n '1,260p'
printf '%s\n' '--- JAX interop tests and stream contracts ---'
rg -l 'jax|block_until_ready|device_sync|CUstream\(0\)|default_stream\("jax"\)' test/python/fe_api python/cudnn | sort | \
xargs -r rg -n -C 3 'block_until_ready|device_sync|CUstream\(0\)|default_stream\(|wrapper.*jax|jax.*wrapper' | sed -n '1,360p'
printf '%s\n' '--- source-level input/launch ordering verifier ---'
python3 - <<'PY'
from pathlib import Path
import ast
path = Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
source = path.read_text()
tree = ast.parse(source)
def call_names(node):
result = []
for child in ast.walk(node):
if isinstance(child, ast.Call):
f = child.func
if isinstance(f, ast.Attribute):
result.append(f"{ast.unparse(f.value)}.{f.attr}")
else:
result.append(ast.unparse(f))
return result
for cls_name in ("GemmProjRopeMxfp8Bf16InSm100", "GemmProjRopeMxfp8Mxfp8InSm100"):
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == cls_name)
execute = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "execute")
calls = call_names(execute)
launch = next(n for n in ast.walk(execute)
if isinstance(n, ast.Call) and ast.unparse(n.func) == "self._compiled_kernel")
print(cls_name)
print(" launch line:", launch.lineno)
print(" readiness calls:", [c for c in calls if "ready" in c.lower() or "synchron" in c.lower()])
print(" input identifiers in launch:", [ast.unparse(a) for a in launch.args[:-1]])
print(" stream argument:", ast.unparse(launch.args[-1]))
wrapper = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "gemm_proj_rope_mxfp8_wrapper_sm100")
for i, node in enumerate(wrapper.body, 1):
text = ast.get_source_segment(source, node) or ""
if "block_until_ready" in text or ".execute(" in text:
print("wrapper statement", i, "line", node.lineno, ":", " ".join(text.split()))
PYRepository: NVIDIA/cudnn-frontend
Length of output: 50379
Synchronize all JAX inputs before the raw CUDA launch.
default_stream("jax") selects CUDA stream 0, outside XLA stream tracking. Both execute paths launch without waiting for JAX inputs. Call jax.block_until_ready for every JAX input, including cos, sin, and MXFP8 scales. Keep output synchronization before host reads. Update the documentation and test so input readiness is wrapper-owned and add an asynchronous producer regression case.
📍 Affects 3 files
python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py#L242-L266(this comment)python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py#L487-L510docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md#L5-L9test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py#L37-L40
🤖 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 242 -
266, Synchronize every JAX input with jax.block_until_ready before the raw CUDA
launch in both execution paths of proj_rope_mxfp8/api.py (242-266 and 487-510),
including x, w, cos, sin, and all MXFP8 scale tensors; retain output
synchronization before host reads. Update
docs/fe-oss-apis/gemm_fusions/gemm_proj_rope_mxfp8.md (5-9) to state that the
wrapper owns input readiness, and extend
test/python/fe_api/gemm/test_gemm_proj_rope_mxfp8_jax.py (37-40) with an
asynchronous-producer regression test.
| x_cutlass_dtype = _convert_to_cutlass_data_type(x.dtype) | ||
| assert x_cutlass_dtype == _convert_to_cutlass_data_type( | ||
| w.dtype | ||
| ), f"x and w must share a dtype (both bfloat16 or both float8_e4m3fn); got x {x.dtype}, w {w.dtype}" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
file="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py"
printf '%s\n' '--- target file structure ---'
ast-grep outline "$file" --lang python 2>/dev/null || true
printf '%s\n' '--- relevant assertions and methods ---'
rg -n -C 5 'assert |def (check_support|compile|execute|__call__)|_ensure_support_checked|_compiled_kernel' "$file"
printf '%s\n' '--- APIBase definitions and call paths ---'
rg -n -C 8 'class APIBase|def __call__|_ensure_support_checked|_compiled_kernel' python/cudnn -g '*.py' | head -n 500Repository: NVIDIA/cudnn-frontend
Length of output: 50377
🏁 Script executed:
#!/bin/bash
set -eu
file="python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py"
printf '%s\n' '--- wrapper and support-validation sections ---'
sed -n '100,190p' "$file"
sed -n '300,410p' "$file"
sed -n '553,690p' "$file"
printf '%s\n' '--- APIBase call and support state implementation ---'
sed -n '514,552p' python/cudnn/api_base.py
printf '%s\n' '--- all assert statements in the target file ---'
python3 - "$file" <<'PY'
import ast
import pathlib
import sys
path = pathlib.Path(sys.argv[1])
tree = ast.parse(path.read_text(), filename=str(path))
for node in ast.walk(tree):
if isinstance(node, ast.Assert):
print(f"{node.lineno}: {ast.unparse(node)}")
PY
printf '%s\n' '--- cache keys and wrapper call sites ---'
rg -n -C 4 '_bf16in_obj_cache|_mxfp8in_obj_cache|gemm_proj_rope_mxfp8_wrapper_sm100|check_support\\(' \
python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8 python test -g '*.py' 2>/dev/null | head -n 300Repository: NVIDIA/cudnn-frontend
Length of output: 18657
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
import ast
import dis
import pathlib
path = pathlib.Path("python/cudnn/gemm/cutedsl/dense/proj_rope_mxfp8/api.py")
tree = ast.parse(path.read_text(), filename=str(path))
wrapper = next(
node for node in tree.body
if isinstance(node, ast.FunctionDef)
and node.name == "gemm_proj_rope_mxfp8_wrapper_sm100"
)
asserts = [
node for node in ast.walk(wrapper)
if isinstance(node, ast.Assert)
]
print("assert locations:", [node.lineno for node in asserts])
cache_lookups = [
node for node in ast.walk(wrapper)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "get"
]
print("cache lookup lines:", [node.lineno for node in cache_lookups])
optimized = compile(
ast.Module(body=[wrapper], type_ignores=[]),
str(path),
"exec",
optimize=1,
)
code = next(const for const in optimized.co_consts if hasattr(const, "co_name") and const.co_name == wrapper.name)
assert not any(
instruction.opname == "LOAD_ASSERTION_ERROR"
for instruction in dis.get_instructions(code)
), "assert bytecode remains under optimize=1"
print("optimize=1 removes wrapper assert bytecode: yes")
print("public validation asserts precede cache lookup: yes")
PYRepository: NVIDIA/cudnn-frontend
Length of output: 334
Replace public input assert checks with explicit ValueError exceptions.
Python removes assert under -O. The wrapper then skips dtype and MXFP8 scale validation. Cache hits do not call check_support() again, so invalid tensors can reach a cached kernel. Replace the checks at lines 583, 620, and 641 with explicit ValueError branches.
🤖 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 582 -
585, Replace the public input assertions in the wrapper, including the dtype
check around _convert_to_cutlass_data_type and the validation checks near lines
620 and 641, with explicit ValueError branches. Preserve each existing
validation condition and message semantics so dtype and MXFP8 scale checks
always execute, including when Python runs with optimizations or a cached kernel
is used.
Applies the dense-fusion type-erasure + JAX pattern to the grouped family,
with real JAX eager support wherever the kernel's tensor layouts are
expressible as row-major arrays and clear rejections where they are not.
Per-API JAX support (all eager; discrete/pointer-array weight modes):
- grouped_gemm (unfused): BF16 discrete mode
- grouped_gemm_glu / dglu: BF16 backend, discrete mode, swiglu+geglu /
dswiglu+dgeglu incl. generate_dbias and caller-provided dprob
- grouped_gemm_dsrelu: discrete FP8 (scale factors in the physical
C-contiguous atom shape -- the backward kernels provably rebuild SF
layouts from the GEMM shapes and read only base pointers)
- grouped_gemm_wgrad: BF16 backend, dense (experts, m, n) or discrete
pointer outputs
- discrete_grouped_gemm_swiglu / dswiglu: FP8 (SF physical atom shape for
SFA and the SFD outputs)
Rejected with clear "not expressible as JAX arrays" errors:
- grouped swiglu/srelu/quant: their SFA scale factors are MMA-permuted
strided cute tensor arguments in every mode (unlike amax/dsrelu, the
kernel consumes the full layout, which has no row-major equivalent)
- grouped dswiglu (dense-weight-mode only) and glu_hadamard
(block-scaled only); the block-scaled glu/dglu/wgrad backends
- dense-mode b_tensor (expert-outermost strides), column-major bias, and
packed-fp4 inputs everywhere
Mechanics shared across the family (unfused is the template):
- b_ptrs/sfb_ptrs/wgrad_ptrs pointer arrays from JAX: int64 (jax x64 mode)
or packed little-endian uint8 (8 bytes per pointer), since JAX truncates
int64 without x64; framework-neutral validation + host decoding in
unfused._bf16_api (_validate_pointer_tensor/_pointer_values); pointers
come from jax.Array.unsafe_buffer_pointer() and the arrays must stay
alive until kernel completion (record_stream is torch-only; the JAX path
keeps live references instead)
- internal workspaces via tensor_adapter.allocate_byte_workspace: allocated
in the caller's framework allocator (torch.empty / jnp.zeros +
block_until_ready), written through raw pointers, never surfaced as
arrays; compile-time Int64 pointer placeholders are real bytes retyped
via the from_dlpack element_type override (fake tensors have dummy
iterators)
- new tensor_adapter helpers: get_data_ptr (torch data_ptr / jax
unsafe_buffer_pointer), get_version (0 for immutable arrays),
to_host_list, allocate_byte_workspace
- canonical (cutlass) dtype vocabulary and canonical TensorDescs
throughout, incl. live-tensor validation; expected-stride literals with
extent-1 dims wrapped in canonicalize_unit_dim_strides;
select_grouped_gemm_backend accepts torch/jax/numpy/str dtypes
- execute stream defaulting per framework; wrapper output allocation
branches (torch empty_strided byte-identical; jnp.empty n-major
C-contiguous + block_until_ready)
Also:
- discrete_grouped swiglu/dswiglu now set _interpret_uint8_as_fp4x2 before
descriptor creation (the torch uint8-container path previously built
descs with the flag unset and was silently broken)
- test conftest sets XLA_PYTHON_CLIENT_PREALLOCATE=false: the JAX interop
tests share the pytest process with the torch suites, and XLA's default
75%-of-GPU preallocation starved later torch kernel compiles (12
CUDA_ERROR_OUT_OF_MEMORY failures in full-suite runs)
- per-API "JAX support" docs sections + overview matrix; the blanket
torch-only guard test narrows to proj_rope (each grouped family now has
its own JAX test file)
proj_rope_mxfp8 (added after review): migrated both classes to the TVM-FFI
compile path (--enable-tvm-ffi + fake stream) so raw DLPack tensors go
straight to the compiled kernel -- the per-call from_dlpack(x.detach())
conversion loop is gone from the hot path (~10.8 us/launch CPU after, vs a
per-call conversion protocol that cost 2-3 us per tensor across 8-10
tensors before). torch inputs keep cheap detach views for autograd safety;
the uint8 E8M0 scale inputs keep a per-call element-type reinterpret (now
tvm-ffi-enabled). JAX supported on both input paths with w_out_in=True (the
[in, out] weight reaches the kernel through a transposed strided view --
torch-only, clear error); bit-identical torch-vs-JAX tests for the bf16 and
mxfp8 paths. With proj_rope no longer torch-only, the blanket guard test
file is removed (every API now has its own JAX test file).
Tests: per-family JAX tests assert bit-identical outputs between torch and
JAX wrapper runs on identical input bytes (both paths share one compiled
kernel) for every supported config -- unfused, glu (swiglu+geglu), dglu
(d_row/dprob/dbias), dsrelu (d_row/d_col/d_srelu + all three SFD outputs),
wgrad (dense+discrete), discrete swiglu/dswiglu (fp8, byte-exact) -- with
dprob-style atomic accumulators compared at tight tolerance; rejected
configs assert their clear errors.
Verified on SM100: full fe_api/gemm + fe_api/grouped_gemm + unfused suite
run yields 64 failed / 1437 passed / 1027 skipped / 2 xfailed / 2 errors --
the failure list is byte-identical to the known pre-existing
test_gemm_swiglu env-numerics failures, and the 2 collection errors are the
pre-existing upstream test_grouped_gemm_{glu,dglu}.py missing-module
imports. All grouped/discrete modules import with torch absent.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
366d16d to
6dcef0d
Compare
2 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*.Affected area
FE OSS kernels or CuTeDSL (Python API)
Summary
Extends the type-erased torch+JAX pattern from #529's dense fusions to the grouped / discrete-grouped GEMM APIs and
gemm_proj_rope_mxfp8, with real JAX eager support wherever each kernel's tensor layouts are expressible as row-major arrays, and clear"not expressible as JAX arrays"rejections where they are not.Per-API JAX support matrix after this PR (whole
gemm/cutedsltree; ✅* = added by this PR)jax.jitgemm_amax(dense)gemm_swiglu(dense)gemm_srelu/gemm_dsrelu(dense)gemm_proj_rope_mxfp8w_out_in=True--enable-tvm-ffi: per-callfrom_dlpack(x.detach())loop removed from the hot path (~10.8 µs/launch CPU; benefits torch too).w_out_in=False(transposed weight view) is torch-onlygrouped_gemm(unfused)grouped_gemm_glu/dglugrouped_gemm_dsrelugrouped_gemm_wgraddiscrete_grouped_gemm_swiglu/dswiglugrouped_gemm_swiglu/srelu/quantgrouped_gemm_dswiglugrouped_gemm_glu_hadamardThe dividing line, verified per kernel: kernels that rebuild SF layouts from the GEMM shapes and read only base pointers accept the physical C-contiguous atom form from JAX (same precedent as
gemm_amax); kernels that consume the full MMA-permuted SF layout as a tensor argument cannot — making those kernels rebuild their SF layouts is the follow-up that would unlock them.Key mechanics
b_ptrs/sfb_ptrs/wgrad_ptrs): built fromjax.Array.unsafe_buffer_pointer(), passed as int64 (jax x64 mode) or packed little-endian uint8 (8 bytes/pointer — JAX truncates int64 without x64 mode).record_streamis torch-only; the JAX path holds live references (documented lifetime contract).tensor_adapter.allocate_byte_workspaceallocates in the caller's framework allocator; kernels write through raw pointers; buffers never surface as arrays. Compile-time Int64 pointer placeholders are real bytes retyped via thefrom_dlpack(...).element_typeoverride (fake tensors have dummy iterators).tensor_adapterhelpers:get_data_ptr,get_version(0 for immutable arrays),to_host_list,allocate_byte_workspace.select_grouped_gemm_backendaccepts torch/jax/numpy/str dtypes.Why
MoE-style JAX users can now call the grouped kernels directly (discrete pointer-array mode is the MoE-typical mode) without importing torch, and proj_rope users get both JAX support and a ~2.5× lower per-launch CPU cost from the tvm-ffi migration. Where layouts genuinely cannot be expressed, failing fast at the entry point with an explanatory error beats failing deep inside pointer/workspace machinery.
Related issues
Stacked on #529.
API and compatibility impact
gemm_proj_rope_mxfp8now compiles with--enable-tvm-ffiand passes raw tensors at execute (torch inputs keep cheap detach views for autograd safety) — same numerics (18/18 tests unchanged), ~2.5× lower launch overhead.discrete_groupedswiglu/dswiglu now set_interpret_uint8_as_fp4x2before descriptor creation — the torch uint8-container path previously built descriptors with the flag unset and was silently broken.test/python/conftest.pysetsXLA_PYTHON_CLIENT_PREALLOCATE=false— XLA's default 75%-of-GPU preallocation starved later torch kernel compiles when the JAX tests share the pytest process (12CUDA_ERROR_OUT_OF_MEMORYfailures in full-suite runs).Testing
On a B200-class SM100 (CC 10.0), Python 3.12, torch 2.13+cu130, jax 0.11, nvidia-cutlass-dsl 4.6:
test_gemm_swiglu.pyenv-numerics failure list (present at the base commit); the 2 collection errors are the pre-existing upstreamtest_grouped_gemm_{glu,dglu}.pyimports of missing test-util modules. Zero regressions.sys.modules.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Tests