[DRAFT] Jax Support for FE OSS API#351
Draft
mgoldfarb-nvidia wants to merge 35 commits into
Draft
Conversation
…jax-fe-oss-poc # Conflicts: # docs/fe-oss-apis/dsa.md
Use stable module-level CuTe launchers and shared TupleDict pytree results across the JAX wrappers. Consolidate dtype validation into ApiBaseJax and remove obsolete compatibility modules.
Add and refine JAX bindings across GEMM, DSA, NSA, and SDPA paths, including layout-aware descriptors and sliding-window attention. Consolidate array and dtype validation, add optional JAX test selection, and refresh the design and operator documentation.
Contributor
📝 WalkthroughWalkthroughThis PR adds an optional ChangesDocumentation
Core Infrastructure
GEMM and Grouped-GEMM Fusion Ops
DeepSeek Sparse Attention (DSA)
Estimated code review effort: 5 (Critical) | ~150 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/fe-oss-apis/dsa.md (1)
179-184: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winUse a defined head-dim in the backward example.
math.sqrt(D)references an undefined name here, so the snippet is not runnable as written. This section already fixes the head dimension at 512, so use that constant (or derive it from the tensor shape) instead.🛠️ Proposed fix
result = DSA.sparse_attention_backward_wrapper( q, kv, out, dout, lse, attn_sink, topk_idxs, - softmax_scale=1.0 / math.sqrt(D), + softmax_scale=1.0 / math.sqrt(512), topk_length=topk_length, )🤖 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/dsa.md` around lines 179 - 184, The backward example uses an undefined head-dim symbol in DSA.sparse_attention_backward_wrapper by passing math.sqrt(D), so update this snippet to use the already established 512 head dimension or derive it from q/k/v tensor shape instead. Keep the fix localized to the sparse_attention_backward_wrapper call and the softmax_scale argument so the example remains runnable and consistent with the rest of the section.
🧹 Nitpick comments (10)
python/cudnn/_jax/gemm.py (1)
183-197: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__per Ruff RUF022.🧹 Proposed fix
__all__ = [ "GEMM_A_LAYOUTS", "GEMM_B_LAYOUTS", "GEMM_C_LAYOUTS", "as_gemm_tensor_desc", "block_scale_tensor_spec", "gemm_a_tensor_spec", "gemm_b_tensor_spec", "gemm_c_tensor_spec", "probability_tensor_spec", - "require_fp8_block_scales", - "require_gemm_inputs", "require_16_byte_extent", + "require_fp8_block_scales", + "require_gemm_inputs", "require_layout", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/_jax/gemm.py` around lines 183 - 197, The __all__ list in gemm.py is unsorted and should be ordered to satisfy Ruff RUF022. Update the exported names in the __all__ definition near as_gemm_tensor_desc and require_layout so they are in sorted order, keeping all current symbols but rearranging them alphabetically.Source: Linters/SAST tools
python/cudnn/gemm_dsrelu/jax.py (2)
96-349: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd unit tests and expand docstrings for the new JAX API.
Per path instructions for
python/cudnn/**, please confirm test coverage exists (or add it) undertest/python/fe_apifor this newGemmDsreluSm100/gemm_dsrelu_wrapper_sm100JAX API, and consider expanding the class/__init__docstrings with parameter documentation (currently a one-line class docstring only), consistent with the emphasis on documentation for this path.🤖 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_dsrelu/jax.py` around lines 96 - 349, Add test coverage for the new JAX squared-ReLU API by creating or updating tests under test/python/fe_api to exercise both GemmDsreluSm100 and gemm_dsrelu_wrapper_sm100 with valid inputs and basic signature/dispatch behavior. Also expand the class and __init__ docstrings in GemmDsreluSm100 to document the parameters and usage details, since the current one-line class docstring is too sparse for this new API.Source: Path instructions
142-214: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMissing explicit sf_vec_size/sf_dtype combination validation (parity gap vs Torch API).
_check_supportvalidates block-scale shapes viarequire_fp8_block_scalesbut doesn't replicate the Torch API's explicitsf_vec_size in SF_VEC_SIZEScheck or the kernel'sis_valid_dtypes_and_scale_factor_vec_sizecombo rules (e.g.Float8E4M3FN+sf_vec_size=32orFloat8E8M0FNU+sf_vec_size=16are invalid). An invalidsf_vec_size/sf_dtypecombo that happens to produce matching tensor shapes would surface as an opaque CUTLASS-level compile error rather than a clearValueError.Consider calling
kernel.is_valid_dtypes_and_scale_factor_vec_size(...)(or an equivalent explicit check) in_check_supportfor parity with the Torch API'scheck_support.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/gemm_dsrelu/jax.py` around lines 142 - 214, The _check_support method in the JAX squared-ReLU backward path only validates block-scale shapes, but it still misses the explicit sf_vec_size/sf_dtype compatibility rules that the Torch API enforces. Add an early validation in _check_support using Sm100BlockScaledPersistentDenseGemmKernel (for example, its is_valid_dtypes_and_scale_factor_vec_size check or equivalent) so invalid sf_vec_size and scale-factor dtype combinations fail with a clear ValueError before reaching CUTLASS compilation. Keep the new check alongside the existing require_fp8_block_scales and dtype validation logic to preserve parity with the Torch check_support behavior.python/cudnn/gemm_srelu/jax.py (2)
137-217: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSame sf_vec_size/sf_dtype combination validation gap as gemm_dsrelu/jax.py.
_check_supportdoesn't replicate the kernel'sis_valid_dtypes_and_scale_factor_vec_sizecombo validation (sf_vec_size membership + sf_dtype/sf_vec_size compatibility) that the Torch API enforces explicitly. Same recommendation applies here for parity and clearer error messages.🤖 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_srelu/jax.py` around lines 137 - 217, The _check_support method in the JAX GEMM + squared-ReLU path is missing the same sf_vec_size/sf_dtype combination validation that the Torch API enforces. Update _check_support to mirror the kernel’s is_valid_dtypes_and_scale_factor_vec_size checks so it rejects unsupported sf_vec_size values and incompatible sf_dtype/sf_vec_size pairings before continuing. Use the existing JaxTensorDesc, require_fp8_block_scales, and Sm100BlockScaledPersistentDenseGemmKernel flow as the place to add this parity validation and keep the error messages explicit.
94-331: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConfirm test coverage under test/python/fe_api for this new JAX API.
As per path instructions, please confirm dedicated tests exist (or add them) for
GemmSreluSm100/gemm_srelu_wrapper_sm100, and consider adding parameter-level docstrings similar to the underlying kernel class.🤖 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_srelu/jax.py` around lines 94 - 331, Confirm that dedicated coverage exists under test/python/fe_api for the new JAX entrypoints, and add or update tests for GemmSreluSm100 and gemm_srelu_wrapper_sm100 if they are missing. Make sure the tests exercise the wrapper and class behavior through their public APIs, including basic construction and call-path validation, so the new symbols are covered by the FE API test suite. While touching GemmSreluSm100/__init__ and gemm_srelu_wrapper_sm100, add brief parameter-level docstrings or inline doc comments consistent with the underlying kernel class to document the constructor and wrapper arguments.Source: Path instructions
python/cudnn/gemm_validation.py (1)
12-181: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSolid, well-factored shared validation module; consider fuller docstrings and confirm test coverage.
Logic is sound (short-circuit
orinrequire_mma_tilercorrectly avoids IndexError on malformed inputs; length/positivity checks are ordered safely throughout). Given this module is now the single source of truth for GEMM shape/tiler/cluster validation shared across amax/dsrelu/srelu/swiglu, two suggestions:
- Docstrings here are one-liners; the downstream kernel classes use full Sphinx-style
:param:/:return:docs. Aligning this shared module with that convention would help since it's now imported broadly.- No test file is included for this new module. Given it centralizes correctness-critical validation for multiple GEMM fusion ops, dedicated unit tests (e.g., under
test/python/fe_api) covering boundary cases (non-power-of-two clusters,sf_vec_size<=0, mismatched K/L, etc.) would materially reduce regression risk across all its consumers.As per path instructions, this file should "Focus on documentation" and "Focus on whether there are test cases in test/python/fe_api."
🤖 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_validation.py` around lines 12 - 181, The shared validation helpers in the GEMM validation module are correct, but the review asks for clearer documentation and dedicated coverage. Update the one-line docstrings in the validation functions such as require_gemm_shapes, require_block_scale_shapes, require_mma_tiler, require_cluster_shape, and resolve_max_active_clusters to use fuller Sphinx-style parameter/return docs consistent with the downstream kernel classes. Also add unit tests for this module under test/python/fe_api covering key boundary cases like mismatched shapes, non-power-of-two cluster dimensions, invalid sf_vec_size, and other validation failures to protect all consumers.Source: Path instructions
python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py (1)
163-201: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winClarify the interaction between
MMA_TILER_MandKNOWN_HANG_MMA_TILER_M.
MMA_TILER_M = (128, 256)advertises 256 as an allowed tile, butKNOWN_HANG_MMA_TILER_M = (256,)makesrequire_mma_tilerunconditionally reject anymma_tiler_mn[0] == 256withNotImplementedError. Since this is the only public path (Torch/JAX APIs both callrequire_mma_tilerbefore instantiating the kernel),TWO_CTA_MMA_TILER_Mand the entire 2-CTA code path in this kernel are currently unreachable. If this is an intentional temporary block on a known hang, consider a short comment onMMA_TILER_Mitself cross-referencingKNOWN_HANG_MMA_TILER_Mso future maintainers don't misread 256 as usable. As per path instructions, focus on documentation for this path.def check_mma_tiler_reachability(): from cudnn.gemm_amax.dense_blockscaled_gemm_persistent_amax import ( Sm100BlockScaledPersistentDenseGemmKernel as K, ) try: K.require_mma_tiler((256, 128)) print("256 M tile is reachable") except NotImplementedError as e: print("256 M tile blocked:", e)🤖 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_amax/dense_blockscaled_gemm_persistent_amax.py` around lines 163 - 201, The `Sm100BlockScaledPersistentDenseGemmKernel.require_mma_tiler` path currently blocks any M tile of 256 via `KNOWN_HANG_MMA_TILER_M`, while `MMA_TILER_M` still advertises 256 as supported, which makes the 2-CTA path misleading. Update the documentation/comments around `MMA_TILER_M`, `KNOWN_HANG_MMA_TILER_M`, and `require_mma_tiler` to clearly state that 256 is intentionally disabled as a known hang, and cross-reference the restriction so `require_cluster_shape` and the `TWO_CTA_MMA_TILER_M` path are not misread as reachable.Source: Path instructions
python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py (1)
420-544: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider documenting the public constructor parameters.
GroupedGemmDgluSm100exposes ~25 constructor parameters (dtypes, layouts, tiler/cluster shapes, dGLU clamp/offset knobs) with only a one-line class docstring. A parameter-level docstring (mirroring the TorchGroupedGemmDgluSm100.__init__doc) would make the JAX entry point self-documenting.As per path instructions: "
python/cudnn/**: Focus on documentation."🤖 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/grouped_gemm/grouped_gemm_dglu/jax.py` around lines 420 - 544, The public constructor of GroupedGemmDgluSm100 is under-documented, making its many configuration knobs hard to understand. Add a parameter-level docstring to the __init__ method for this JAX wrapper, mirroring the Torch GroupedGemmDgluSm100.__init__ documentation and covering the layout, dtype, tiling, cluster, clamp, offset, and scheduling arguments so the API is self-describing.Source: Path instructions
python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py (1)
289-294: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
get_workspace_bytesnot updated to use the newget_dense_workspace_bytesclassmethod.The new
get_dense_workspace_bytesclassmethod (lines 289-294) centralizes the dynamic-scheduler workspace size, but the existingget_workspace_bytesinstance method (lines 674-679) still hardcodes4 if self.use_dynamic_sched else 0instead of calling it. The siblingBlockScaledMoEGroupedGemmGluHadamardKernelinmoe_blockscaled_grouped_gemm_glu_hadamard.pywas updated to callself.get_dense_workspace_bytes(self.use_dynamic_sched)at the same location — this file appears to have been missed, leaving the two constants (DYNAMIC_SCHED_WORKSPACE_BYTESand the inline4) to potentially drift apart.♻️ Proposed fix
def get_workspace_bytes(self) -> int: """Return descriptor workspace plus optional dynamic scheduler state.""" desc_workspace_bytes = self.get_desc_workspace_bytes() - dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + dynamic_sched_bytes = self.get_dense_workspace_bytes(self.use_dynamic_sched) return desc_workspace_bytes + dynamic_sched_bytesAlso applies to: 674-679
🤖 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/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py` around lines 289 - 294, The workspace calculation in `BlockScaledMoEGroupedGemmGluBiasKernel.get_workspace_bytes` still hardcodes the dense-mode size instead of using the new shared helper. Update that method to call `self.get_dense_workspace_bytes(self.use_dynamic_sched)` so it matches `BlockScaledMoEGroupedGemmGluHadamardKernel` and keeps `DYNAMIC_SCHED_WORKSPACE_BYTES` as the single source of truth.python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py (1)
97-136: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate MMA/cluster validation logic across kernel classes.
require_mma_tiler,require_cluster_shape, andget_dense_workspace_byteshere are byte-for-byte identical to the same three classmethods added tomoe_blockscaled_grouped_gemm_srelu_quant.py(lines 109-148), and the samekernel.require_mma_tiler(...)/kernel.require_cluster_shape(...)pattern is repeated acrossgemm_amax,gemm_dsrelu, and presumably the other grouped-GEMM ops in this cohort. Since each kernel class re-declares its ownMMA_TILER_M/MMA_TILER_N/CLUSTER_TILER_M/etc. constants and wrapper methods, any future change to the shared validation contract (e.g. a new allowed tile shape) requires updating every kernel class individually and risks silent drift.Consider extracting this into a small mixin/base class (e.g. in
gemm_validation.py) that kernel classes compose or inherit from, parameterized by the tiler/cluster constants.🤖 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/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py` around lines 97 - 136, The grouped-GEMM validation methods are duplicated across kernel classes, so extract the shared logic from require_mma_tiler, require_cluster_shape, and get_dense_workspace_bytes into a common mixin/base helper (for example a gemm_validation module). Make the helper parameterized by each kernel’s constants such as MMA_TILER_M, MMA_TILER_N, CLUSTER_TILER_M, TWO_CTA_MMA_TILER_M, MAX_CLUSTER_CTAS, MAX_CLUSTER_DIMENSION, and DYNAMIC_SCHED_WORKSPACE_BYTES, then update grouped_gemm_quant.py and the other ops that currently call kernel.require_mma_tiler(...) and kernel.require_cluster_shape(...) to inherit from or delegate to the shared implementation.
🤖 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/attention/sdpa_bwd_d256.md`:
- Around line 108-110: The JAX example in sdpa_bwd_d256 has an unpacking typo in
the result handling: result["dq"]_tensor is invalid Python and should be
corrected to the proper documented key access used by sdpa_bwd. Update the
example to reference result["dq_tensor"] alongside result["dk_tensor"] and
result["dv_tensor"] so the snippet is copy-paste runnable and consistent with
the return structure shown in the JAX API docs.
In `@docs/fe-oss-apis/cutedsl-jax-design.md`:
- Around line 45-52: The module tree under the operation bindings is currently
presented as a concrete layout, but it should be marked as illustrative or
updated to match the actual operation-specific module names used in the PR. In
the section describing each operation’s bindings, revise the tree near the
api.py/jax.py/kernel.py example so readers understand it is an example, or
replace the placeholder kernel.py entry with the real module names referenced by
the implementation.
- Around line 177-199: The gap list in the design doc is too mixed together and
does not clearly separate the two release blockers from the remaining scope
notes. Update the section around the grouped-GEMM, sliding-window, and DSA
paragraphs to use explicit bullets or a clearly numbered list, with each gap
described separately and the NSA/DSA follow-up items kept as distinct scope
notes. Keep the ownership and count unambiguous by ensuring each blocker maps to
one bullet and the remaining coverage notes are not folded into the blocker
text.
In `@docs/fe-oss-apis/dsa.md`:
- Around line 507-524: The example currently computes softmax_scale from an
undefined D, so update the recompute_target JIT wrapper to derive the scale from
the input tensor shape instead of relying on a free variable. Use the existing
symbols recompute_target and DSA.sparse_attn_score_recompute_wrapper to locate
the snippet, and make the scale expression self-contained by calculating it from
q_attn (or the relevant attention dimension) inside the function before calling
the wrapper.
In `@docs/fe-oss-apis/nsa.md`:
- Around line 771-774: The JAX availability note conflicts with the existing
Sliding Window section, so update the summary in the table description to match
the documented JAX support. In the nsa.md text that mentions Compression
Attention, Top-K Reduction, Selection Attention, and Sliding Window Attention,
remove the blanket “unavailable” claim for Sliding Window or narrow it to only
the unsupported subset, keeping the wording consistent with the JAX Sliding
Window section.
In `@python/cudnn/deepseek_sparse_attention/score_recompute/jax.py`:
- Around line 233-321: Dense score-recompute currently bypasses kernel-config
validation during the `_validate_only` path, so `_check_support()` can succeed
for unsupported configs. Update `_dense_score_recompute` to eagerly call
`resolve_dense_score_kernel_config` before the `_validate_only` early return,
using the same inputs passed later by `_launch_dense_with_q_causal_offsets` /
`_launch_dense_without_q_causal_offsets`. Ensure
`DenseIndexerScoreRecompute._check_support()` and
`DenseAttnScoreRecompute._check_support()` now reject invalid `head_dim` and
`qhead_per_kv_head` combinations up front, matching `_sparse_score_recompute`
behavior.
In `@python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py`:
- Around line 23-136: `max_topk` is not being forwarded into the JAX sparse
backward launch path, so `FlashAttentionDSABackwardSm100` is built with its
default value and KV loading can be wrong. Update `_launch_with_topk_length` and
`_launch_without_topk_length` to accept and pass `max_topk` when constructing
`FlashAttentionDSABackwardSm100`, and thread the computed `topk_shape[1]`
through `_sparse_attention_backward_impl` via `static_args`/`call_cutedsl` so
the kernel cache key includes it. Use the existing `topk_shape` and the
`_launch_*` helpers as the entry points for the fix.
In `@python/cudnn/gemm_amax/api.py`:
- Around line 228-234: The batch dimension from b_tensor is currently unpacked
as _b_batch in api.py’s c_tensor allocation path but never checked against
a_tensor’s batch before using batch to size c_tensor. In the gemm_amax
allocation logic, validate that both tensors share the same batch dimension
before calling torch.empty_strided, or discard the unused value with _ if the
batch is guaranteed to be checked elsewhere; keep the fix near the m, n, batch
unpacking and c_tensor setup so the mismatch is caught before allocation.
In `@python/cudnn/gemm_srelu/api.py`:
- Line 91: Rename the ambiguous loop/dimension variable `l` in
`require_gemm_shapes` usage within `api.py` to `batch` (or another clearer name)
to satisfy Ruff E741 and match the existing shape naming. Update all related
references in the same flow, including the later uses in the GEMM/SReLU API
logic at the spots where `l` is read again, so the renamed variable stays
consistent across the function.
In `@python/cudnn/gemm_swiglu/api.py`:
- Line 123: Rename the ambiguous loop variable in require_gemm_shapes assignment
within gemm_swiglu/api.py from l to batch to satisfy Ruff E741 and keep the
naming consistent with require_gemm_shapes; this change is localized because the
variable is not used afterward, so only update the unpacking in the relevant API
function.
---
Outside diff comments:
In `@docs/fe-oss-apis/dsa.md`:
- Around line 179-184: The backward example uses an undefined head-dim symbol in
DSA.sparse_attention_backward_wrapper by passing math.sqrt(D), so update this
snippet to use the already established 512 head dimension or derive it from
q/k/v tensor shape instead. Keep the fix localized to the
sparse_attention_backward_wrapper call and the softmax_scale argument so the
example remains runnable and consistent with the rest of the section.
---
Nitpick comments:
In `@python/cudnn/_jax/gemm.py`:
- Around line 183-197: The __all__ list in gemm.py is unsorted and should be
ordered to satisfy Ruff RUF022. Update the exported names in the __all__
definition near as_gemm_tensor_desc and require_layout so they are in sorted
order, keeping all current symbols but rearranging them alphabetically.
In `@python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py`:
- Around line 163-201: The
`Sm100BlockScaledPersistentDenseGemmKernel.require_mma_tiler` path currently
blocks any M tile of 256 via `KNOWN_HANG_MMA_TILER_M`, while `MMA_TILER_M` still
advertises 256 as supported, which makes the 2-CTA path misleading. Update the
documentation/comments around `MMA_TILER_M`, `KNOWN_HANG_MMA_TILER_M`, and
`require_mma_tiler` to clearly state that 256 is intentionally disabled as a
known hang, and cross-reference the restriction so `require_cluster_shape` and
the `TWO_CTA_MMA_TILER_M` path are not misread as reachable.
In `@python/cudnn/gemm_dsrelu/jax.py`:
- Around line 96-349: Add test coverage for the new JAX squared-ReLU API by
creating or updating tests under test/python/fe_api to exercise both
GemmDsreluSm100 and gemm_dsrelu_wrapper_sm100 with valid inputs and basic
signature/dispatch behavior. Also expand the class and __init__ docstrings in
GemmDsreluSm100 to document the parameters and usage details, since the current
one-line class docstring is too sparse for this new API.
- Around line 142-214: The _check_support method in the JAX squared-ReLU
backward path only validates block-scale shapes, but it still misses the
explicit sf_vec_size/sf_dtype compatibility rules that the Torch API enforces.
Add an early validation in _check_support using
Sm100BlockScaledPersistentDenseGemmKernel (for example, its
is_valid_dtypes_and_scale_factor_vec_size check or equivalent) so invalid
sf_vec_size and scale-factor dtype combinations fail with a clear ValueError
before reaching CUTLASS compilation. Keep the new check alongside the existing
require_fp8_block_scales and dtype validation logic to preserve parity with the
Torch check_support behavior.
In `@python/cudnn/gemm_srelu/jax.py`:
- Around line 137-217: The _check_support method in the JAX GEMM + squared-ReLU
path is missing the same sf_vec_size/sf_dtype combination validation that the
Torch API enforces. Update _check_support to mirror the kernel’s
is_valid_dtypes_and_scale_factor_vec_size checks so it rejects unsupported
sf_vec_size values and incompatible sf_dtype/sf_vec_size pairings before
continuing. Use the existing JaxTensorDesc, require_fp8_block_scales, and
Sm100BlockScaledPersistentDenseGemmKernel flow as the place to add this parity
validation and keep the error messages explicit.
- Around line 94-331: Confirm that dedicated coverage exists under
test/python/fe_api for the new JAX entrypoints, and add or update tests for
GemmSreluSm100 and gemm_srelu_wrapper_sm100 if they are missing. Make sure the
tests exercise the wrapper and class behavior through their public APIs,
including basic construction and call-path validation, so the new symbols are
covered by the FE API test suite. While touching GemmSreluSm100/__init__ and
gemm_srelu_wrapper_sm100, add brief parameter-level docstrings or inline doc
comments consistent with the underlying kernel class to document the constructor
and wrapper arguments.
In `@python/cudnn/gemm_validation.py`:
- Around line 12-181: The shared validation helpers in the GEMM validation
module are correct, but the review asks for clearer documentation and dedicated
coverage. Update the one-line docstrings in the validation functions such as
require_gemm_shapes, require_block_scale_shapes, require_mma_tiler,
require_cluster_shape, and resolve_max_active_clusters to use fuller
Sphinx-style parameter/return docs consistent with the downstream kernel
classes. Also add unit tests for this module under test/python/fe_api covering
key boundary cases like mismatched shapes, non-power-of-two cluster dimensions,
invalid sf_vec_size, and other validation failures to protect all consumers.
In `@python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py`:
- Around line 420-544: The public constructor of GroupedGemmDgluSm100 is
under-documented, making its many configuration knobs hard to understand. Add a
parameter-level docstring to the __init__ method for this JAX wrapper, mirroring
the Torch GroupedGemmDgluSm100.__init__ documentation and covering the layout,
dtype, tiling, cluster, clamp, offset, and scheduling arguments so the API is
self-describing.
In
`@python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py`:
- Around line 289-294: The workspace calculation in
`BlockScaledMoEGroupedGemmGluBiasKernel.get_workspace_bytes` still hardcodes the
dense-mode size instead of using the new shared helper. Update that method to
call `self.get_dense_workspace_bytes(self.use_dynamic_sched)` so it matches
`BlockScaledMoEGroupedGemmGluHadamardKernel` and keeps
`DYNAMIC_SCHED_WORKSPACE_BYTES` as the single source of truth.
In `@python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py`:
- Around line 97-136: The grouped-GEMM validation methods are duplicated across
kernel classes, so extract the shared logic from require_mma_tiler,
require_cluster_shape, and get_dense_workspace_bytes into a common mixin/base
helper (for example a gemm_validation module). Make the helper parameterized by
each kernel’s constants such as MMA_TILER_M, MMA_TILER_N, CLUSTER_TILER_M,
TWO_CTA_MMA_TILER_M, MAX_CLUSTER_CTAS, MAX_CLUSTER_DIMENSION, and
DYNAMIC_SCHED_WORKSPACE_BYTES, then update grouped_gemm_quant.py and the other
ops that currently call kernel.require_mma_tiler(...) and
kernel.require_cluster_shape(...) to inherit from or delegate to the shared
implementation.
🪄 Autofix (Beta)
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: 442bcbd5-5425-4631-8205-449a46e81e78
📒 Files selected for processing (173)
docs/fe-oss-apis/attention/sdpa_bwd_d256.mddocs/fe-oss-apis/attention/sdpa_fwd_d256.mddocs/fe-oss-apis/cutedsl-jax-design.mddocs/fe-oss-apis/dsa.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_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_quant_unified.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/nsa.mddocs/fe-oss-apis/overview.mddocs/fe-oss-apis/rmsnorm_rht_amax.mdpyproject.tomlpython/cudnn/README.mdpython/cudnn/__init__.pypython/cudnn/_experimental_warnings.pypython/cudnn/_jax/__init__.pypython/cudnn/_jax/api_base.pypython/cudnn/_jax/gemm.pypython/cudnn/_jax/grouped_gemm.pypython/cudnn/_operation_api.pypython/cudnn/api_base.pypython/cudnn/api_base_torch.pypython/cudnn/deepseek_sparse_attention/indexer_backward/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_backward/api.pypython/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_backward/jax.pypython/cudnn/deepseek_sparse_attention/indexer_forward/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/jax.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/api.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/jax.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.pypython/cudnn/deepseek_sparse_attention/score_recompute/__init__.pypython/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/api.pypython/cudnn/deepseek_sparse_attention/score_recompute/config.pypython/cudnn/deepseek_sparse_attention/score_recompute/jax.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.pypython/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.pypython/cudnn/deepseek_sparse_attention/utils/compiler.pypython/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/api.pypython/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/api.pypython/cudnn/gemm_amax/__init__.pypython/cudnn/gemm_amax/api.pypython/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.pypython/cudnn/gemm_amax/jax.pypython/cudnn/gemm_amax/validation.pypython/cudnn/gemm_dsrelu/__init__.pypython/cudnn/gemm_dsrelu/api.pypython/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.pypython/cudnn/gemm_dsrelu/jax.pypython/cudnn/gemm_srelu/__init__.pypython/cudnn/gemm_srelu/api.pypython/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.pypython/cudnn/gemm_srelu/jax.pypython/cudnn/gemm_swiglu/__init__.pypython/cudnn/gemm_swiglu/api.pypython/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.pypython/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.pypython/cudnn/gemm_swiglu/jax.pypython/cudnn/gemm_swiglu/validation.pypython/cudnn/gemm_validation.pypython/cudnn/grouped_gemm/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/api.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.pypython/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_dsrelu/api.pypython/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.pypython/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_dswiglu/api.pypython/cudnn/grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.pypython/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_glu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_glu/api.pypython/cudnn/grouped_gemm/grouped_gemm_glu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.pypython/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/api.pypython/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.pypython/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.pypython/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.pypython/cudnn/grouped_gemm/grouped_gemm_quant/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_quant/api.pypython/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.pypython/cudnn/grouped_gemm/grouped_gemm_quant/jax.pypython/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_srelu/api.pypython/cudnn/grouped_gemm/grouped_gemm_srelu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.pypython/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_swiglu/api.pypython/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.pypython/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/api.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.pypython/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.pypython/cudnn/grouped_gemm/moe_kernel_helpers.pypython/cudnn/grouped_gemm/utils.pypython/cudnn/jax/__init__.pypython/cudnn/native_sparse_attention/__init__.pypython/cudnn/native_sparse_attention/compression/__init__.pypython/cudnn/native_sparse_attention/compression/api.pypython/cudnn/native_sparse_attention/compression/jax.pypython/cudnn/native_sparse_attention/jax_utils.pypython/cudnn/native_sparse_attention/selection/NSA_select_attn_fwd_hmma.pypython/cudnn/native_sparse_attention/selection/__init__.pypython/cudnn/native_sparse_attention/selection/api.pypython/cudnn/native_sparse_attention/selection/jax.pypython/cudnn/native_sparse_attention/sliding_window_attention/__init__.pypython/cudnn/native_sparse_attention/sliding_window_attention/api.pypython/cudnn/native_sparse_attention/sliding_window_attention/jax.pypython/cudnn/native_sparse_attention/top_k/__init__.pypython/cudnn/native_sparse_attention/top_k/api.pypython/cudnn/native_sparse_attention/top_k/jax.pypython/cudnn/rmsnorm_rht_amax/__init__.pypython/cudnn/rmsnorm_rht_amax/api.pypython/cudnn/rmsnorm_rht_amax/config.pypython/cudnn/rmsnorm_rht_amax/jax.pypython/cudnn/sdpa/__init__.pypython/cudnn/sdpa/bwd/__init__.pypython/cudnn/sdpa/bwd/api.pypython/cudnn/sdpa/bwd/jax.pypython/cudnn/sdpa/fwd/__init__.pypython/cudnn/sdpa/fwd/api.pypython/cudnn/sdpa/fwd/jax.pypython/cudnn/sdpa/jax_utils.pytest/python/fe_api/conftest.pytest/python/fe_api/dsa/test_DSA_dense_indexer_backward.pytest/python/fe_api/test_api_base_framework_split.pytest/python/fe_api/test_gemm_amax_validation.pytest/python/fe_api/test_gemm_swiglu_validation.pytest/python/fe_api/test_gemm_validation.pytest/python/fe_api/test_jax_api_surface.pytest/python/fe_api/test_jax_cutedsl_adapter.pytest/python/fe_api/test_jax_dsa.pytest/python/fe_api/test_jax_dsa_contract.pytest/python/fe_api/test_jax_dsa_indexer_backward_contract.pytest/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.pytest/python/fe_api/test_jax_gemm.pytest/python/fe_api/test_jax_gemm_contract.pytest/python/fe_api/test_jax_grouped_gemm_contract.pytest/python/fe_api/test_jax_grouped_gemm_wgrad_contract.pytest/python/fe_api/test_jax_import_contract.pytest/python/fe_api/test_jax_nsa.pytest/python/fe_api/test_jax_nsa_contract.pytest/python/fe_api/test_jax_rmsnorm_rht_amax.pytest/python/fe_api/test_jax_sdpa_contract.pytest/python/fe_api/test_jax_validation.pytest/python/fe_api/test_optional_jax_test_support.pytest/python/fe_api/test_rmsnorm_rht_amax.pytest/python/fe_api/test_rmsnorm_rht_amax_config.py
Comment on lines
+108
to
+110
| result = sdpa_bwd(q, k, v, o, do, lse) | ||
| dq, dk, dv = result["dq"]_tensor, result["dk_tensor"], result["dv_tensor"] | ||
| ``` |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Fix the unpacking typo in the JAX example.
result["dq"]_tensor is invalid Python, so the snippet won't run if copied. It should be result["dq_tensor"] to match the documented JAX return keys.
🛠️ Proposed fix
-dq, dk, dv = result["dq"]_tensor, result["dk_tensor"], result["dv_tensor"]
+dq, dk, dv = result["dq_tensor"], result["dk_tensor"], result["dv_tensor"]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| result = sdpa_bwd(q, k, v, o, do, lse) | |
| dq, dk, dv = result["dq"]_tensor, result["dk_tensor"], result["dv_tensor"] | |
| ``` | |
| result = sdpa_bwd(q, k, v, o, do, lse) | |
| dq, dk, dv = result["dq_tensor"], result["dk_tensor"], result["dv_tensor"] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/fe-oss-apis/attention/sdpa_bwd_d256.md` around lines 108 - 110, The JAX
example in sdpa_bwd_d256 has an unpacking typo in the result handling:
result["dq"]_tensor is invalid Python and should be corrected to the proper
documented key access used by sdpa_bwd. Update the example to reference
result["dq_tensor"] alongside result["dk_tensor"] and result["dv_tensor"] so the
snippet is copy-paste runnable and consistent with the return structure shown in
the JAX API docs.
Comment on lines
+45
to
+52
| Each operation co-locates its bindings: | ||
|
|
||
| ```text | ||
| cudnn/<operation>/ | ||
| api.py # PyTorch API | ||
| jax.py # JAX API | ||
| kernel.py # shared CuTe DSL implementation | ||
| ``` |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Mark this tree as illustrative.
This reads like a concrete filesystem contract, but the PR summary shows operation-specific kernel modules rather than a shared kernel.py. Please either label the tree as an example or replace the placeholder with real module names so readers don't chase a path that doesn't exist.
🤖 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/cutedsl-jax-design.md` around lines 45 - 52, The module tree
under the operation bindings is currently presented as a concrete layout, but it
should be marked as illustrative or updated to match the actual
operation-specific module names used in the PR. In the section describing each
operation’s bindings, revise the tree near the api.py/jax.py/kernel.py example
so readers understand it is an example, or replace the placeholder kernel.py
entry with the real module names referenced by the implementation.
Comment on lines
+177
to
+199
| The two current gaps need separate ABI work, not another generic adapter: | ||
|
|
||
| 1. Two discrete grouped-GEMM wrappers consume device-resident tables of raw | ||
| addresses. JAX must receive the expert buffers as real operands and build | ||
| pointer tables in XLA-owned device workspace. | ||
|
|
||
| The sliding-window binding deliberately covers fixed-shape BHSD inference with | ||
| `right_bound=0`. Training statistics and packed THD/ragged inputs still need a | ||
| native frontend XLA FFI adapter; depending on JAX's private cuDNN APIs would be | ||
| an unstable OSS contract. | ||
|
|
||
| Dense DSA indexer backward now uses a runtime `grad_loss` operand, an | ||
| XLA-owned score-gradient workspace, and a kernel-cleared FP32 accumulation | ||
| buffer. Its first JAX binding deliberately covers fixed-shape SM100 BSHD | ||
| inputs; the broader PyTorch SM90 and packed-THD paths remain unchanged. | ||
|
|
||
| Only RMSNorm, the four dense GEMMs, five DSA entry points, and NSA | ||
| sliding-window inference currently have real GPU execution tests. These tests | ||
| include numerical reference coverage for dense indexer backward across multiple | ||
| query and key tiles, runtime loss scaling, and the registered cuDNN | ||
| sliding-window lowering. The other families have source and CPU contract | ||
| coverage but have not been qualified through real CUTLASS lowering and SM100 | ||
| execution. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Split the gap list into explicit bullets.
The section says there are two gaps, but only one is numbered and the NSA/DSA notes are mixed into prose. Please separate release blockers from scope notes so the count and ownership stay unambiguous.
🤖 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/cutedsl-jax-design.md` around lines 177 - 199, The gap list
in the design doc is too mixed together and does not clearly separate the two
release blockers from the remaining scope notes. Update the section around the
grouped-GEMM, sliding-window, and DSA paragraphs to use explicit bullets or a
clearly numbered list, with each gap described separately and the NSA/DSA
follow-up items kept as distinct scope notes. Keep the ownership and count
unambiguous by ensuring each blocker maps to one bullet and the remaining
coverage notes are not folded into the blocker text.
Comment on lines
+507
to
+524
| ```python | ||
| import math | ||
|
|
||
| import jax | ||
| from cudnn.jax import DSA | ||
|
|
||
| softmax_scale = 1.0 / math.sqrt(D) | ||
|
|
||
| @jax.jit | ||
| def recompute_target(q_attn, k_attn, lse, topk_indices, topk_length): | ||
| return DSA.sparse_attn_score_recompute_wrapper( | ||
| q_attn, | ||
| k_attn, | ||
| lse, | ||
| topk_indices, | ||
| softmax_scale, | ||
| topk_length=topk_length, | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the scale expression self-contained.
This example also uses D without defining it, so readers will hit a NameError if they copy/paste it. Compute the scale from the input shape inside the JIT wrapper instead.
🛠️ Proposed fix
import math
import jax
from cudnn.jax import DSA
-
-softmax_scale = 1.0 / math.sqrt(D)
`@jax.jit`
def recompute_target(q_attn, k_attn, lse, topk_indices, topk_length):
return DSA.sparse_attn_score_recompute_wrapper(
q_attn,
k_attn,
lse,
topk_indices,
- softmax_scale,
+ softmax_scale=1.0 / math.sqrt(q_attn.shape[-1]),
topk_length=topk_length,
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| ```python | |
| import math | |
| import jax | |
| from cudnn.jax import DSA | |
| softmax_scale = 1.0 / math.sqrt(D) | |
| @jax.jit | |
| def recompute_target(q_attn, k_attn, lse, topk_indices, topk_length): | |
| return DSA.sparse_attn_score_recompute_wrapper( | |
| q_attn, | |
| k_attn, | |
| lse, | |
| topk_indices, | |
| softmax_scale, | |
| topk_length=topk_length, | |
| ) |
🤖 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/dsa.md` around lines 507 - 524, The example currently
computes softmax_scale from an undefined D, so update the recompute_target JIT
wrapper to derive the scale from the input tensor shape instead of relying on a
free variable. Use the existing symbols recompute_target and
DSA.sparse_attn_score_recompute_wrapper to locate the snippet, and make the
scale expression self-contained by calculating it from q_attn (or the relevant
attention dimension) inside the function before calling the wrapper.
Comment on lines
+771
to
+774
| The table above describes the PyTorch surface. In JAX, Compression Attention | ||
| and Top-K Reduction support fixed BHSD only, Selection Attention supports | ||
| packed THD self-attention only, and Sliding Window Attention is unavailable. | ||
|
|
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Fix the JAX Sliding Window availability note.
This contradicts the dedicated JAX Sliding Window section above, which documents a JAX binding and cuDNN lowering. Please remove the unavailable claim or narrow it to the unsupported subset.
🔧 Proposed fix
-The table above describes the PyTorch surface. In JAX, Compression Attention
-and Top-K Reduction support fixed BHSD only, Selection Attention supports
-packed THD self-attention only, and Sliding Window Attention is unavailable.
+The table above describes the PyTorch surface. In JAX, Compression Attention
+and Top-K Reduction support fixed BHSD only, Selection Attention supports
+packed THD self-attention only, and Sliding Window Attention is documented
+separately above.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| The table above describes the PyTorch surface. In JAX, Compression Attention | |
| and Top-K Reduction support fixed BHSD only, Selection Attention supports | |
| packed THD self-attention only, and Sliding Window Attention is unavailable. | |
| The table above describes the PyTorch surface. In JAX, Compression Attention | |
| and Top-K Reduction support fixed BHSD only, Selection Attention supports | |
| packed THD self-attention only, and Sliding Window Attention is documented | |
| separately above. |
🤖 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/nsa.md` around lines 771 - 774, The JAX availability note
conflicts with the existing Sliding Window section, so update the summary in the
table description to match the documented JAX support. In the nsa.md text that
mentions Compression Attention, Top-K Reduction, Selection Attention, and
Sliding Window Attention, remove the blanket “unavailable” claim for Sliding
Window or narrow it to only the unsupported subset, keeping the wording
consistent with the JAX Sliding Window section.
Comment on lines
+233
to
+321
| def _dense_score_recompute( | ||
| q: Any, | ||
| k: Any, | ||
| per_head: Any, | ||
| *, | ||
| score_type: str, | ||
| per_head_name: str, | ||
| per_head_dtype: Any, | ||
| scale: float, | ||
| qhead_per_kv_head: Optional[int], | ||
| ratio: int, | ||
| q_causal_offsets: Any | None, | ||
| _validate_only: bool = False, | ||
| ) -> TupleDict: | ||
| q_shape = require_array(q, name="q", rank=4, dtype=jnp.bfloat16) | ||
| k_shape = require_array(k, name="k", rank=4, dtype=jnp.bfloat16) | ||
|
|
||
| batch, seqlen_q, num_query_heads, head_dim = q_shape | ||
| k_batch, seqlen_k, num_kv_heads, k_head_dim = k_shape | ||
| dimensions = { | ||
| "batch": batch, | ||
| "S_q": seqlen_q, | ||
| "S_k": seqlen_k, | ||
| "H_q": num_query_heads, | ||
| "H_kv": num_kv_heads, | ||
| "head dimension": head_dim, | ||
| } | ||
| nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] | ||
| if nonpositive: | ||
| raise ValueError("Dense score-recompute dimensions must be positive, got " + ", ".join(nonpositive)) | ||
| if k_batch != batch: | ||
| raise ValueError(f"q and k batch dimensions must match, got {batch} and {k_batch}") | ||
| if k_head_dim != head_dim: | ||
| raise ValueError("q and k head dimensions must match, got " f"{head_dim} and {k_head_dim}") | ||
| require_array( | ||
| per_head, | ||
| name=per_head_name, | ||
| shape=(batch, seqlen_q, num_query_heads), | ||
| dtype=per_head_dtype, | ||
| ) | ||
| if num_query_heads % num_kv_heads: | ||
| raise ValueError(f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})") | ||
|
|
||
| inferred_qhead_per_kv_head = num_query_heads // num_kv_heads | ||
| if qhead_per_kv_head is None: | ||
| qhead_per_kv_head = inferred_qhead_per_kv_head | ||
| if qhead_per_kv_head != inferred_qhead_per_kv_head: | ||
| raise ValueError("qhead_per_kv_head must equal H_q / H_kv, got " f"{qhead_per_kv_head} and {num_query_heads} / {num_kv_heads}") | ||
| if ratio < 1: | ||
| raise ValueError(f"ratio must be at least 1, got {ratio}") | ||
|
|
||
| inputs = (q, k, per_head) | ||
| if q_causal_offsets is not None: | ||
| require_array( | ||
| q_causal_offsets, | ||
| name="q_causal_offsets", | ||
| shape=(batch,), | ||
| dtype=jnp.int32, | ||
| ) | ||
| inputs += (q_causal_offsets,) | ||
|
|
||
| resolved_scale = float(scale) | ||
| if _validate_only: | ||
| return None | ||
|
|
||
| launcher = _launch_dense_with_q_causal_offsets if q_causal_offsets is not None else _launch_dense_without_q_causal_offsets | ||
| out, denom = call_cutedsl( | ||
| launcher, | ||
| inputs, | ||
| outputs=( | ||
| BufferSpec( | ||
| "out", | ||
| (batch, seqlen_q, seqlen_k), | ||
| jnp.float32, | ||
| fill_value=float("-inf"), | ||
| ), | ||
| BufferSpec("denom", (batch, seqlen_q), jnp.float32), | ||
| ), | ||
| static_args={ | ||
| "score_type": str(score_type), | ||
| "head_dim": int(head_dim), | ||
| "qhead_per_kv_head": int(qhead_per_kv_head), | ||
| "ratio": int(ratio), | ||
| "max_seqlen_q": int(seqlen_q), | ||
| "max_seqlen_k": int(seqlen_k), | ||
| "scale": resolved_scale, | ||
| }, | ||
| ) | ||
| return TupleDict(out=out, denom=denom) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Dense score-recompute skips kernel-config validation in _check_support.
_sparse_score_recompute validates feasibility via resolve_sparse_score_kernel_config before the _validate_only early return (lines 394-402), but _dense_score_recompute never calls resolve_dense_score_kernel_config — that call only happens inside the launcher helpers at actual trace/execute time. resolve_dense_score_kernel_config enforces head_dim % 64 == 0 and qhead_per_kv_head in {32, 64, 128}, neither of which _dense_score_recompute checks itself. As a result, DenseIndexerScoreRecompute._check_support() / DenseAttnScoreRecompute._check_support() can report success for configs that will only fail later during kernel launch, breaking the intended "gate invocation" contract.
🛠️ Proposed fix: validate config eagerly
if ratio < 1:
raise ValueError(f"ratio must be at least 1, got {ratio}")
inputs = (q, k, per_head)
if q_causal_offsets is not None:
require_array(
q_causal_offsets,
name="q_causal_offsets",
shape=(batch,),
dtype=jnp.int32,
)
inputs += (q_causal_offsets,)
resolved_scale = float(scale)
+ # Fail fast in check_support() on infeasible tile/SMEM configs, matching
+ # the eager validation done in `_sparse_score_recompute`.
+ resolve_dense_score_kernel_config(
+ score_type=score_type,
+ head_dim=head_dim,
+ qhead_per_kv_head=qhead_per_kv_head,
+ )
if _validate_only:
return NoneAlso applies to: 624-637, 683-696
🤖 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/deepseek_sparse_attention/score_recompute/jax.py` around lines
233 - 321, Dense score-recompute currently bypasses kernel-config validation
during the `_validate_only` path, so `_check_support()` can succeed for
unsupported configs. Update `_dense_score_recompute` to eagerly call
`resolve_dense_score_kernel_config` before the `_validate_only` early return,
using the same inputs passed later by `_launch_dense_with_q_causal_offsets` /
`_launch_dense_without_q_causal_offsets`. Ensure
`DenseIndexerScoreRecompute._check_support()` and
`DenseAttnScoreRecompute._check_support()` now reject invalid `head_dim` and
`qhead_per_kv_head` combinations up front, matching `_sparse_score_recompute`
behavior.
Comment on lines
+23
to
+136
| def _launch_with_topk_length( | ||
| stream, | ||
| q, | ||
| kv, | ||
| output, | ||
| doutput, | ||
| lse, | ||
| attn_sink, | ||
| topk_idxs, | ||
| topk_length, | ||
| dq, | ||
| dkv, | ||
| d_sink, | ||
| workspace_lse_odo, | ||
| workspace_dkv, | ||
| *, | ||
| total_seqlen_q: int, | ||
| total_seqlen_kv: int, | ||
| num_heads: int, | ||
| head_dim: int, | ||
| block_tile: int, | ||
| softmax_scale: float, | ||
| ): | ||
| from cutlass import Float32, Int32 | ||
|
|
||
| from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 | ||
|
|
||
| kernel = FlashAttentionDSABackwardSm100( | ||
| head_dim=head_dim, | ||
| head_dim_v=head_dim, | ||
| block_tile=block_tile, | ||
| ) | ||
|
|
||
| problem_shape = ( | ||
| Int32(total_seqlen_q), | ||
| Int32(total_seqlen_kv), | ||
| Int32(head_dim), | ||
| (Int32(num_heads), Int32(1)), | ||
| ) | ||
| kernel( | ||
| problem_shape, | ||
| q, | ||
| kv, | ||
| output, | ||
| doutput, | ||
| lse, | ||
| attn_sink, | ||
| topk_idxs, | ||
| topk_length, | ||
| dq, | ||
| dkv, | ||
| d_sink, | ||
| workspace_lse_odo, | ||
| workspace_dkv, | ||
| Float32(softmax_scale), | ||
| stream, | ||
| ) | ||
|
|
||
|
|
||
| def _launch_without_topk_length( | ||
| stream, | ||
| q, | ||
| kv, | ||
| output, | ||
| doutput, | ||
| lse, | ||
| attn_sink, | ||
| topk_idxs, | ||
| dq, | ||
| dkv, | ||
| d_sink, | ||
| workspace_lse_odo, | ||
| workspace_dkv, | ||
| *, | ||
| total_seqlen_q: int, | ||
| total_seqlen_kv: int, | ||
| num_heads: int, | ||
| head_dim: int, | ||
| block_tile: int, | ||
| softmax_scale: float, | ||
| ): | ||
| from cutlass import Float32, Int32 | ||
|
|
||
| from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 | ||
|
|
||
| kernel = FlashAttentionDSABackwardSm100( | ||
| head_dim=head_dim, | ||
| head_dim_v=head_dim, | ||
| block_tile=block_tile, | ||
| ) | ||
| problem_shape = ( | ||
| Int32(total_seqlen_q), | ||
| Int32(total_seqlen_kv), | ||
| Int32(head_dim), | ||
| (Int32(num_heads), Int32(1)), | ||
| ) | ||
| kernel( | ||
| problem_shape, | ||
| q, | ||
| kv, | ||
| output, | ||
| doutput, | ||
| lse, | ||
| attn_sink, | ||
| topk_idxs, | ||
| None, | ||
| dq, | ||
| dkv, | ||
| d_sink, | ||
| workspace_lse_odo, | ||
| workspace_dkv, | ||
| Float32(softmax_scale), | ||
| stream, | ||
| ) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Critical: max_topk is never forwarded to FlashAttentionDSABackwardSm100, silently corrupting KV loading.
Neither _launch_with_topk_length nor _launch_without_topk_length passes max_topk when constructing FlashAttentionDSABackwardSm100(head_dim=..., head_dim_v=..., block_tile=...), and _sparse_attention_backward_impl's static_args dict (lines 290-297) never includes it either — even though topk_shape[1] (the required value) is already computed at line 204-213. The kernel constructor defaults max_topk=0 (see dsa_bwd_sm100.py), and load_KV/_load_kv_rows in that same file guard every real K/V row load with if idx < self.max_topk: — with max_topk=0 this is always false, so topk_idx stays -1 for every row. That means real K/V rows are never loaded (silently zero-filled dKV/dQ gradients), or — in the topk_length-provided + first-tile branch — _copy_kv_row is invoked with topk_idx=-1, indexing mKV[-1, ...], a likely out-of-bounds access.
The Torch counterpart (sparse_attention_backward/_interface_sm100.py) correctly threads max_topk=max_topk (derived from topk_idxs.shape[1]) into the same kernel class, so this is a JAX-binding-only regression. This also means the call_cutedsl compile-cache key is missing a dimension that materially changes kernel behavior.
🐛 Proposed fix: thread `max_topk` through both launchers and static_args
def _launch_with_topk_length(
stream,
q,
kv,
output,
doutput,
lse,
attn_sink,
topk_idxs,
topk_length,
dq,
dkv,
d_sink,
workspace_lse_odo,
workspace_dkv,
*,
total_seqlen_q: int,
total_seqlen_kv: int,
num_heads: int,
head_dim: int,
block_tile: int,
softmax_scale: float,
+ max_topk: int,
):
from cutlass import Float32, Int32
from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100
kernel = FlashAttentionDSABackwardSm100(
head_dim=head_dim,
head_dim_v=head_dim,
block_tile=block_tile,
+ max_topk=max_topk,
) def _launch_without_topk_length(
stream,
q,
kv,
output,
doutput,
lse,
attn_sink,
topk_idxs,
dq,
dkv,
d_sink,
workspace_lse_odo,
workspace_dkv,
*,
total_seqlen_q: int,
total_seqlen_kv: int,
num_heads: int,
head_dim: int,
block_tile: int,
softmax_scale: float,
+ max_topk: int,
):
from cutlass import Float32, Int32
from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100
kernel = FlashAttentionDSABackwardSm100(
head_dim=head_dim,
head_dim_v=head_dim,
block_tile=block_tile,
+ max_topk=max_topk,
) static_args={
"total_seqlen_q": int(total_seqlen_q),
"total_seqlen_kv": int(total_seqlen_kv),
"num_heads": int(num_heads),
"head_dim": int(head_dim),
"block_tile": int(block_tile),
"softmax_scale": float(resolved_scale),
+ "max_topk": int(topk_shape[1]),
},Also applies to: 227-298
🤖 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/deepseek_sparse_attention/sparse_attention_backward/jax.py`
around lines 23 - 136, `max_topk` is not being forwarded into the JAX sparse
backward launch path, so `FlashAttentionDSABackwardSm100` is built with its
default value and KV loading can be wrong. Update `_launch_with_topk_length` and
`_launch_without_topk_length` to accept and pass `max_topk` when constructing
`FlashAttentionDSABackwardSm100`, and thread the computed `topk_shape[1]`
through `_sparse_attention_backward_impl` via `static_args`/`call_cutedsl` so
the kernel cache key includes it. Use the existing `topk_shape` and the
`_launch_*` helpers as the entry points for the fix.
Comment on lines
+228
to
+234
| m, _, batch = a_tensor.shape | ||
| n, _, _b_batch = b_tensor.shape | ||
| c_tensor = None | ||
| if c_major == "m": | ||
| c_tensor = torch.empty_strided((m, n, l), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) | ||
| c_tensor = torch.empty_strided((m, n, batch), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) | ||
| elif c_major == "n": | ||
| c_tensor = torch.empty_strided((m, n, l), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) | ||
| c_tensor = torch.empty_strided((m, n, batch), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Unused/unchecked _b_batch before allocating c_tensor.
b_tensor.shape is unpacked into _b_batch but it's never validated against batch (from a_tensor.shape) before c_tensor is sized using batch. If A/B batch dims disagree, an incorrectly-shaped c_tensor is allocated here and the mismatch is only caught later inside check_support() (via require_gemm_shapes), after the wasted allocation. Consider validating upfront or just discarding the unused value with _.
🔧 Proposed fix
m, _, batch = a_tensor.shape
- n, _, _b_batch = b_tensor.shape
+ n, _, b_batch = b_tensor.shape
+ if b_batch != batch:
+ raise ValueError(f"a_tensor and b_tensor batch dims must match, got {batch} and {b_batch}")
c_tensor = None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| m, _, batch = a_tensor.shape | |
| n, _, _b_batch = b_tensor.shape | |
| c_tensor = None | |
| if c_major == "m": | |
| c_tensor = torch.empty_strided((m, n, l), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) | |
| c_tensor = torch.empty_strided((m, n, batch), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) | |
| elif c_major == "n": | |
| c_tensor = torch.empty_strided((m, n, l), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) | |
| c_tensor = torch.empty_strided((m, n, batch), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) | |
| m, _, batch = a_tensor.shape | |
| n, _, b_batch = b_tensor.shape | |
| if b_batch != batch: | |
| raise ValueError(f"a_tensor and b_tensor batch dims must match, got {batch} and {b_batch}") | |
| c_tensor = None | |
| if c_major == "m": | |
| c_tensor = torch.empty_strided((m, n, batch), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) | |
| elif c_major == "n": | |
| c_tensor = torch.empty_strided((m, n, batch), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm_amax/api.py` around lines 228 - 234, The batch dimension
from b_tensor is currently unpacked as _b_batch in api.py’s c_tensor allocation
path but never checked against a_tensor’s batch before using batch to size
c_tensor. In the gemm_amax allocation logic, validate that both tensors share
the same batch dimension before calling torch.empty_strided, or discard the
unused value with _ if the batch is guaranteed to be checked elsewhere; keep the
fix near the m, n, batch unpacking and c_tensor setup so the mismatch is caught
before allocation.
| d_m, d_n, d_l = self._tensor_shape(self.d_desc, name="sample_d") | ||
|
|
||
| self._value_error_if((b_k, b_l) != (k, l), f"B shape mismatch: expected (*, {k}, {l}), got {(n, b_k, b_l)}") | ||
| m, n, k, l = require_gemm_shapes( |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Ambiguous loop variable l.
Ruff (E741) flags l on this changed line. Since require_gemm_shapes already names this dimension batch, renaming improves clarity and clears the lint. Note l is used again at Line 174/187, so update those references too.
🧰 Tools
🪛 Ruff (0.15.20)
[error] 91-91: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm_srelu/api.py` at line 91, Rename the ambiguous
loop/dimension variable `l` in `require_gemm_shapes` usage within `api.py` to
`batch` (or another clearer name) to satisfy Ruff E741 and match the existing
shape naming. Update all related references in the same flow, including the
later uses in the GEMM/SReLU API logic at the spots where `l` is read again, so
the renamed variable stays consistent across the function.
Source: Linters/SAST tools
| self._value_error_if( | ||
| self.ab12_desc.stride_order != self.c_desc.stride_order, | ||
| f"AB12 and C tensor stride orders must match, got {self.ab12_desc.stride_order} and {self.c_desc.stride_order}", | ||
| m, n, k, l = require_gemm_shapes( |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Ambiguous loop variable l.
Ruff (E741) flags this changed line. Rename to batch (consistent with require_gemm_shapes) to clear the lint; l is unused afterward here so the change is localized.
🧰 Tools
🪛 Ruff (0.15.20)
[error] 123-123: Ambiguous variable name: l
(E741)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/gemm_swiglu/api.py` at line 123, Rename the ambiguous loop
variable in require_gemm_shapes assignment within gemm_swiglu/api.py from l to
batch to satisfy Ruff E741 and keep the naming consistent with
require_gemm_shapes; this change is localized because the variable is not used
afterward, so only update the unpacking in the relevant API function.
Source: Linters/SAST tools
mgoldfarb-nvidia
marked this pull request as draft
July 9, 2026 21:57
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is an experimental branch to show one possible integration of JAX using the FE OSS API.
Summary by CodeRabbit