Add block-sparse attention CuTe DSL kernels for Hopper and Blackwell - #333
Merged
Anerudhan merged 2 commits intoJul 4, 2026
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (23)
💤 Files with no reviewable changes (4)
👮 Files not reviewed due to content moderation or server errors (16)
📝 Walkthrough🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
🟡 Minor comments (11)
python/cudnn/block_sparse_attention/__init__.py-37-37 (1)
37-37: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winMake
__all__a literal string list.
*_SYMBOLS.keys()is what Ruff flags here (PLE0604), so this is likely to fail lint unless the rule is intentionally disabled. Spell the symbols out explicitly.Proposed fix
-__all__ = ["BSA", *_SYMBOLS.keys()] +__all__ = [ + "BSA", + "block_sparse_attention_forward", + "block_sparse_attention_backward", +]🤖 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/block_sparse_attention/__init__.py` at line 37, `__all__` in the block_sparse_attention package is using a starred unpack from `_SYMBOLS.keys()`, which Ruff flags under PLE0604. Update the module-level `__all__` definition in `__init__.py` to be a literal list of string names, explicitly spelling out `BSA` and each exported symbol instead of deriving them from `_SYMBOLS.keys()`.Source: Linters/SAST tools
test/python/fe_api/block_sparse_attention/bsa_utils.py-19-40 (1)
19-40: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winGuard the hard-coded metadata templates.
make_fixed_metadata()always references KV block2, andmake_variable_metadata()assumes at least two Q blocks plus a KV block3. If a later test reuses these helpers with smallerseqlen_q/seqlen_k, they'll either index past the allocated tensors or feed out-of-range block ids into the API. Please either assert those preconditions up front or derive the pattern fromnum_q_blocks/num_kv_blocks.🤖 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/block_sparse_attention/bsa_utils.py` around lines 19 - 40, Guard the hard-coded metadata templates in make_fixed_metadata and make_variable_metadata by validating the expected minimum seqlen_q/seqlen_k before writing fixed block ids. Either add upfront assertions for the required num_q_blocks/num_kv_blocks, or generate the q2k and block_sizes patterns dynamically from those computed sizes so the helpers never index past the tensors or emit out-of-range KV block ids.python/cudnn/block_sparse_attention/_interface.py-1-1 (1)
1-1: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the executable bit or add a shebang.
Ruff reports this module is executable without a shebang. Since this is an importable module, removing the executable bit is likely the right fix.
🤖 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/block_sparse_attention/_interface.py` at line 1, The _interface.py module is marked executable without a shebang, which Ruff flags for importable modules. Update the file’s executable status by removing the executable bit rather than changing the module contents, since it is meant to be imported and not run as a script. Refer to the _interface.py module when adjusting the file permissions in the repo.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/api.py-395-398 (1)
395-398: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winSort
__all__to satisfy Ruff.Proposed fix
__all__ = [ - "block_sparse_attention_forward", "block_sparse_attention_backward", + "block_sparse_attention_forward", ]🤖 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/block_sparse_attention/api.py` around lines 395 - 398, Sort the __all__ list in block_sparse_attention/api.py to match Ruff’s expected ordering, keeping the exported symbols block_sparse_attention_backward and block_sparse_attention_forward in sorted order within the __all__ definition.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/_interface.py-171-171 (1)
171-171: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDrop the unused
head_dimbinding.Proposed fix
- batch, num_heads, seqlen_q, head_dim = q.shape + batch, num_heads, seqlen_q, _head_dim = q.shape🤖 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/block_sparse_attention/_interface.py` at line 171, In the shape unpacking inside the block sparse attention interface, `head_dim` is never used and should be removed to avoid an unused binding. Update the `q.shape` assignment in the relevant function so it only keeps the values that are actually referenced later, matching the existing handling of `batch`, `num_heads`, and `seqlen_q`.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/csrc/fwd/sm90_blk64/bsa_fwd_sm90.py-487-491 (1)
487-491: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the tensor layout comments.
The implementation restores layouts with sequence in mode 0 and head/value dimension in mode 1, then validates mode 1 in Lines 529-530. These comments should describe
(seqlen, head_dim/value_dim, nheads, batch), not(head_dim, seqlen, ...). 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/block_sparse_attention/csrc/fwd/sm90_blk64/bsa_fwd_sm90.py` around lines 487 - 491, The tensor layout comments for the restored tensors are incorrect and should match the actual layout used by the implementation in the forward path. Update the doc comments on mQ, mK, mV, mO, and mLSE in bsa_fwd_sm90.py to describe sequence as mode 0 and head/value dimension as mode 1, i.e. (seqlen, head_dim/value_dim, nheads, batch) for the tensors and the corresponding layout for mLSE, so the documentation matches the validation and restore logic in the forward function.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/fwd/sm120_blk64/bsa_fwd_sm120.py-441-445 (1)
441-445: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the tensor layout comments.
These comments say Q/K/O are
(head_dim, seqlen, nheads, batch), but Line 453 validates mode 1 as the 128-wide dimension and the kernel usesmQ.shape[0]as sequence length. Document these as(seqlen, head_dim/value_dim, nheads, batch)to match the implementation. 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/block_sparse_attention/csrc/fwd/sm120_blk64/bsa_fwd_sm120.py` around lines 441 - 445, The tensor layout comments in the forward path are inconsistent with the implementation. Update the docstrings/comments for the mQ, mK, mV, mO, and mLSE tensors in bsa_fwd_sm120.py to reflect the actual layout used by the kernel and the mode validation, using (seqlen, head_dim/value_dim, nheads, batch) for Q/K/O and the corresponding value-dim shape for V/O. Keep the change limited to documentation and use the existing tensor symbols (mQ, mK, mV, mO, mLSE) to ensure the comments match the shapes enforced by the forward kernel.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/utils/pipeline.py-85-94 (1)
85-94: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winRaise a real exception for invalid
PipelineUserTypevalues.
assert Falsedisappears underpython -O, so an unexpected enum value can fall through and returnNone, pushing the failure into later pipeline calls.🤖 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/block_sparse_attention/csrc/utils/pipeline.py` around lines 85 - 94, The make_pipeline_state function currently uses assert False for invalid PipelineUserType values, which can be skipped under optimized Python and let None leak out; replace that branch with a real exception, such as ValueError, that clearly reports the invalid PipelineUserType and prevents later pipeline failures.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/csrc/utils/mma_sm100_desc.py-85-85 (1)
85-85: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winReplace ambiguous en dashes flagged by Ruff.
Ruff reports RUF001/RUF002/RUF003 here. Use
-in comments, docstrings, and error strings to keep lint clean.Also applies to: 107-107, 127-127, 187-187, 196-196, 223-225
🤖 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/block_sparse_attention/csrc/utils/mma_sm100_desc.py` at line 85, Replace the ambiguous en dashes in the comments and any related strings in mma_sm100_desc.py with plain hyphens to satisfy Ruff RUF001/RUF002/RUF003. Update the affected text near the existing comment about Float-8 / Float-6 / Float-4, and review the other flagged spots in the same module around the identified helper/constants so all comments, docstrings, and error messages use ASCII '-' consistently.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/csrc/utils/pack_gqa.py-40-50 (1)
40-50: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winType
qhead_per_kvheadas an integer constexpr.This value is used for
//,*, and grouped-head indexing, socutlass.Constexpr[bool]is misleading and can confuse specialization/type checking.Proposed fix
- qhead_per_kvhead: cutlass.Constexpr[bool], + qhead_per_kvhead: cutlass.Constexpr[int],🤖 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/block_sparse_attention/csrc/utils/pack_gqa.py` around lines 40 - 50, The PackGQA initializer currently declares qhead_per_kvhead as cutlass.Constexpr[bool], but this symbol is used as a numeric constexpr for division, multiplication, and grouped-head indexing. Update the PackGQA __init__ signature and the qhead_per_kvhead field to use an integer constexpr type instead, and keep the rest of the initialization logic unchanged so the type matches its actual arithmetic usage.python/cudnn/block_sparse_attention/csrc/utils/kernel_utils.py-154-168 (1)
154-168: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winHandle reductions with fewer than four elements.
This branch unconditionally reads
res[0]..res[3]; a valid small tensor shape would fail during codegen. Add a small-size path or assert the minimum size, and cover it intest/python/fe_api.Possible localized fix
if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0): res = cute.make_rmem_tensor(x.shape, Float32) res.store(x) + if const_expr(cute.size(x.shape) < 4): + local_max = res[0] + for i in cutlass.range_constexpr(1, cute.size(x.shape)): + local_max = fmax(local_max, res[i]) + return local_max if const_expr(init_val is None) else fmax(local_max, init_val) local_max = [res[0], res[1], res[2], res[3]]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/block_sparse_attention/csrc/utils/kernel_utils.py` around lines 154 - 168, The fallback branch in fmax_reduce unconditionally indexes res[0] through res[3], so it breaks for tensors with fewer than four elements. Update fmax_reduce to either add a dedicated small-size path that safely reduces 1–3 elements before the existing loop logic, or explicitly assert the minimum supported size near the shape check; keep the const_expr gating and init_val handling intact. Also add coverage for this edge case in test/python/fe_api so the small-tensor behavior is exercised.Source: Path instructions
🧹 Nitpick comments (9)
test/python/fe_api/block_sparse_attention/test_BSA_attention_forward.py (2)
386-413: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the SM120 cache key, not just that compilation happens.
With
compile_cache = {}, this still passes if_bsa_attn_fwd_sm120_blk64()stops threading_tensor_static_compile_key(...)into the cache key, becausecute.compile()would be called anyway. Please capture the constructed key or inspect the compile arguments so this test actually covers the compile-key regression.🤖 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/block_sparse_attention/test_BSA_attention_forward.py` around lines 386 - 413, The SM120 forward test only verifies that compilation is reached, so it would still pass even if `_bsa_attn_fwd_sm120_blk64()` stops including `_tensor_static_compile_key(...)` in the cache key. Update the test around `interface._bsa_attn_fwd_sm120_blk64` to capture or inspect the compile/cache key used when `interface.cute.compile` is invoked, and assert that the expected static compile key is present rather than only expecting `CompileReached`.
17-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSmoke-test the top-level lazy exports too.
These tests only import
cudnn.BSA, but the PR also exposescudnn.block_sparse_attention_forwardandcudnn.block_sparse_attention_backward. A broken top-level__getattr__wiring would still pass here, so I'd add one import/invocation check for those symbols as well.🤖 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/block_sparse_attention/test_BSA_attention_forward.py` around lines 17 - 25, The smoke test only covers cudnn.BSA, so it can miss broken top-level lazy export wiring for cudnn.block_sparse_attention_forward and cudnn.block_sparse_attention_backward. Update the test around _import_bsa in test_BSA_attention_forward.py to also import those two symbols from cudnn and exercise them with a minimal invocation or attribute check, skipping in the same optional-dependency path if needed.test/python/fe_api/block_sparse_attention/test_BSA_attention_backward.py (1)
29-64: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd one backward case for variable routing and partial physical blocks.
Both backward tests stay on fixed metadata with full-size physical KV blocks, but the public backward API explicitly threads
q2k_block_numsand blk64block_sizesthrough validation and dispatch. That leaves the variable-route path and partial-block handling unexercised in this suite. A single numerical case using variable metadata would close that gap.Also applies to: 107-149
🤖 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/block_sparse_attention/test_BSA_attention_backward.py` around lines 29 - 64, Add a backward test that uses variable routing and partial physical KV blocks instead of only fixed metadata in test_bsa_attention_backward_fixed_blocks and the related backward coverage. Create a new case that builds variable q2k_block_nums and blk64 block_sizes, then run block_sparse_attention_forward and block_sparse_attention_backward with that metadata and assert dq_tensor, dk_tensor, and dv_tensor against the reference. Keep the existing fixed-block test, but extend the backward suite so the public API path through block_sparse_attention_backward and its validation/dispatch logic is exercised for variable-route inputs.python/cudnn/block_sparse_attention/api.py (1)
165-170: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExpand the public API docstrings.
These wrappers are user-facing; please document args, return keys, metadata value contracts, layout behavior, and architecture limitations inline or link directly to the BSA docs. As per path instructions,
python/cudnn/**: “Focus on documentation.”Also applies to: 305-305
🤖 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/block_sparse_attention/api.py` around lines 165 - 170, Expand the public docstring for the non-causal block-sparse scaled dot-product attention wrapper in api.py so it clearly documents the callable’s arguments, the returned dict/keys, the metadata value contracts, and any layout assumptions or transformations. Also state the architecture limitations and the fact that this path dispatches only to the Python CuTe DSL kernels, with no SM100 C++/AOT extension support, or link directly to the BSA docs for those details. Apply the same documentation pass to the other public wrapper mentioned in this file so all user-facing APIs are consistently documented inline.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/bwd/sm90_blk64/bsa_bwd_sm90.py (1)
1875-1875: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop or underscore the unused thread index.
Ruff reports
tidxas unused here; use_, _, _ = cute.arch.thread_idx()or remove the assignment if it is not needed.🧹 Proposed fix
- tidx, _, _ = cute.arch.thread_idx() + _, _, _ = cute.arch.thread_idx()🤖 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/block_sparse_attention/csrc/bwd/sm90_blk64/bsa_bwd_sm90.py` at line 1875, The variable from cute.arch.thread_idx() is unused in this block, so update the unpacking in the bsa_bwd_sm90.py logic to drop the named thread index or replace it with underscores. Keep the change localized near the thread_idx() call in the backward SM90 block-sparse attention code, and ensure no unused binding remains so Ruff stops flagging it.Source: Linters/SAST tools
python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py (1)
250-258: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the CSR metadata contract.
Please expand this docstring with the expected
q2k_block_indexlayout, valid sentinel/range rules, optionalq2k_block_numssemantics, and returned tensor shapes. This is a cross-layer metadata builder, so the contract should be clear and matched bytest/python/fe_apicases. As per path instructions,python/cudnn/**: “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/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py` around lines 250 - 258, Expand the docstring for build_bucketed_k2q_csr_cutedsl to explicitly document the CSR metadata contract: describe the expected q2k_block_index layout, the valid sentinel and index range rules, the meaning of optional q2k_block_nums, and the shapes/meaning of all returned tensors and ints. Make sure the wording matches the behavior enforced by related test/python/fe_api coverage so the cross-layer metadata expectations are unambiguous.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/utils/seqlen_info.py (1)
10-14: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove this explanation to a real module docstring.
Because it appears after imports, documentation tools will not treat it as the module docstring. Place it before imports or convert it into class/function docs. As per path instructions, 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/block_sparse_attention/csrc/utils/seqlen_info.py` around lines 10 - 14, Move the explanatory text in seqlen_info to the actual module docstring so documentation tools recognize it; place the existing description at the top of the file before any imports, or if it is meant to document a specific symbol, move it into that symbol’s docstring instead. Use the module-level docstring location in seqlen_info as the fix target and keep the wording focused on the sequence-length consolidation behavior.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/utils/copy_utils.py (1)
142-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument and test the ragged offset contract.
offset_ragged_tensorhas different rank/index requirements forptr_shift=TruevsFalse, but the docstring does not describe the expected tensor shapes or logical offset behavior. Please add a short contract plus atest/python/fe_apicase covering both modes. As per path instructions, focus on documentation and 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/block_sparse_attention/csrc/utils/copy_utils.py` around lines 142 - 167, Document the offset_ragged_tensor contract by expanding its docstring to state the expected tensor rank/shape requirements and how ptr_shift=True versus False changes the logical indexing/offset behavior. Then add a focused test in test/python/fe_api that exercises offset_ragged_tensor in both modes and verifies the expected ragged offset semantics and rank preconditions using the existing function name offset_ragged_tensor.Source: Path instructions
python/cudnn/block_sparse_attention/csrc/utils/softmax.py (1)
169-191: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMake
is_firstacutlass.Constexpr[bool].
update_row_max/update_row_sumare only called from the SM100 softmax paths with a booleanis_first, andcutlass.const_expr(...)here expects a compile-time value. Tightening the annotation would make that contract explicit and prevent accidental runtime inputs.🤖 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/block_sparse_attention/csrc/utils/softmax.py` around lines 169 - 191, The issue is that update_row_max and update_row_sum rely on cutlass.const_expr with is_first, so the parameter must be compile-time constant rather than a runtime int/bool. Update the signatures of update_row_max and update_row_sum in the softmax helper to use cutlass.Constexpr[bool] for is_first, and keep the call sites in the SM100 softmax path passing constexpr booleans so the contract is explicit and type-safe.
🤖 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/block_sparse_attention/_interface.py`:
- Around line 548-552: The accumulator tail zeroing in the block sparse
attention workspace is using a single flattened offset, which skips the
per-(batch, head) layout and can leave earlier rows partially uninitialized.
Update the logic in the accumulator setup around `accum_offset` and
`workspace.reshape(-1)` so the tail is zeroed separately for each `(B, H)` row,
preserving the existing row stride and only clearing the trailing accumulator
region within each row.
In `@python/cudnn/block_sparse_attention/api.py`:
- Around line 53-91: The sparse metadata validation in _validate_sparse_metadata
only enforces shape, dtype, and device, so invalid values can still reach the
CUDA kernels. Extend the checks to reject negative or out-of-range
q2k_block_index entries, q2k_block_nums values greater than the available KV
block capacity, empty rows when allow_empty_block_nums is false, and any invalid
block_sizes contents before dispatch. Make the same validation available to the
other affected call sites in this module, and add FE API invalid-value tests
under test/python/fe_api/block_sparse_attention to cover these cases.
In `@python/cudnn/block_sparse_attention/csrc/bwd/bsa_bwd_postprocess.py`:
- Around line 37-40: The constructor and capability checks for bwd postprocess
are out of sync with what _setup_attributes() actually supports. Update the
BsaBwdPostprocess defaults and validation so num_threads matches the
implementation’s supported configuration, and make can_implement() reject
unsupported thread counts instead of allowing any warp multiple. Use the
BsaBwdPostprocess constructor, can_implement(), and _setup_attributes() as the
places to align the advertised config with the compile/setup path.
In `@python/cudnn/block_sparse_attention/csrc/bwd/sm90_blk64/bsa_bwd_sm90.py`:
- Around line 925-957: The backward path leaves mdKaccum and mdVaccum
uninitialized before the reduce-add epilogues in _bwd_call, so stale workspace
data can leak into gradients. Update the sm90 block sparse attention backward
flow in bsa_bwd_sm90.py to zero those workspace slices alongside mdQaccum before
calling _bwd_call, or change the first dK/dV writer path to store instead of
add. Also add a FE API regression under test/python/fe_api that reuses non-zero
workspace and verifies backward gradients are still correct.
- Around line 182-205: Update can_implement in bsa_bwd_sm90.py so it matches the
kernel configuration chosen by __init__ (tile_m=64, tile_n=64, num_threads=384)
instead of rejecting it via the tile_m/num_threads divisibility check, and make
the head_dim_v validation robust when head_dim_v is None by skipping the modulus
check or treating None as valid input before the arithmetic. Use can_implement
and __init__ as the key locations to adjust.
In `@python/cudnn/block_sparse_attention/csrc/fwd/sm100_blk64/bsa_fwd_combine.py`:
- Around line 363-386: Initialize the LSE shared-memory slots for tail rows
before they participate in the split-bound reduction in bsa_fwd_combine.py. In
the loop that fills tLSEsLSE inside the load/LSE partial path, make sure the idx
>= max_idx case also writes a known sentinel for every sLSE entry instead of
leaving stale shared memory; use the surrounding logic in the cp_async/cute.copy
block and the later sMaxValidSplit reduction path as the reference points. Add
or extend a split-KV tail-tile test under test/python/fe_api that exercises the
tail-row path and verifies load_O_partial() does not read past num_splits.
In `@python/cudnn/block_sparse_attention/csrc/fwd/sm100_blk64/bsa_fwd_helpers.py`:
- Around line 567-624: The blk64 split-barrier path in the inline assembly
inside bsa_fwd_helpers.py uses smem_desc_b_lo in the post-barrier loop before it
is initialized when split_arrive_idx starts at 1. Update the assembly setup
around the mma sequence to seed smem_desc_b_lo from smem_desc_b_lo_start before
entering the split loop, and keep the initialization close to the existing
smem_desc_b_lo_start/smem_desc_b_hi descriptor construction so the first split
MMA uses a valid shared-memory descriptor.
In `@python/cudnn/block_sparse_attention/csrc/fwd/sm90_blk64/bsa_fwd_sm90.py`:
- Around line 153-160: The non-split path in bsa_fwd_sm90.py does not guard
against empty sparse block lists, so `num_n_tiles == 0` can still flow into
`n_tile_ind_`/`gIndices` and later into the non-empty finalizer causing a
divide-by-zero in the output rescale. Update the non-split kernel logic in the
same way as the split path by adding an early empty-safe guard and using an
empty-safe finalizer around the `gIndices`/`num_n_tiles` handling in
`bsa_fwd_sm90.py`, then add `test/python/fe_api` coverage that exercises zero
block counts for this path.
In
`@python/cudnn/block_sparse_attention/csrc/utils/block_sparse_tile_scheduler.py`:
- Around line 50-60: The Params.create() path in
BlockSparsePersistentTileScheduler currently uses the full cluster_shape_mn size
for num_block_cluster but only stores cluster_shape_m, so it can miscount and
mis-map work when cluster_shape_mn[1] is not 1. Add the same validation used in
tile_scheduler.py single-tile schedulers to reject non-unit N-dimension cluster
shapes before constructing BlockSparsePersistentTileScheduler.Params, and apply
the same guard at the other affected creation site as well.
In `@python/cudnn/block_sparse_attention/csrc/utils/pipeline.py`:
- Around line 63-69: The phase() implementation in Pipeline is returning values
beyond the required 0/1 parity for multi-stage pipelines. Update the phase()
method in the pipeline class so that, when self._stages > 1, it reduces the
computed phase to a binary parity value instead of returning self._phase_index
// self._stages; keep the existing self._stages == 1 fast path unchanged.
In `@python/cudnn/block_sparse_attention/csrc/utils/softmax.py`:
- Around line 118-128: In softmax.py’s row accumulation logic, the `sink_val`
path in the `row_sum[r]` update can still evaluate `sink_val_cur * LOG2_E -
row_max[r] * scale_log2` when `row_max[r]` is `-inf`, which turns fully masked
rows into `inf/NaN`. Update the `sink_val` handling in this section to guard
against non-finite `row_max[r]` or empty rows before doing the subtraction, and
ensure the final `row_sum[r]`/LSE falls back to the finite sink value instead of
propagating `NaN`.
In `@python/cudnn/block_sparse_attention/csrc/utils/tcgen05_mma_helpers.py`:
- Around line 212-269: The post-wait path in the tcgen05 inline assembly uses
smem_desc_b_lo before it is guaranteed to be initialized, especially when
split_arrive_idx is 1. Initialize smem_desc_b_lo from the last pre-wait B
descriptor (the same value used to build smem_desc_b before the wait) before the
second loop in tcgen05_mma_helpers.py, and add a validation/assertion that the
split point is large enough for the post-wait updates to be valid.
- Around line 155-164: The `tcgen05_mma_helpers.py` inline-asm argument for
`zero_init` still uses Python `not`, which breaks the DSL `Boolean` path and
should be replaced with the same dynamic-predicate handling pattern used in
`bsa_fwd_sm100.py`. Update the helper around the `pred_str` / `llvm.inline_asm`
logic so `zero_init` is converted through a DSL-safe predicate input instead of
`Int32(not zero_init)`, and apply the same fix to both inline-asm call sites in
the MMA helper. Use the existing `zero_init`/`pred_str` flow and related symbols
in this helper to keep the behavior consistent for both static and dynamic
cases.
In `@python/cudnn/block_sparse_attention/csrc/utils/tile_scheduler.py`:
- Around line 173-190: The grid is rounded up to a cluster multiple, but padded
tail CTAs are still marked valid in the scheduler paths. Update
SingleTileScheduler.get_current_work() and
SingleTileVarlenScheduler._varlen_coord_map() so is_valid_tile is false whenever
the physical block maps past the last real Q tile, using the existing
block-to-tile mapping logic and cluster_shape_mn/num_block checks. Make the same
invalidation behavior consistent in the related scheduler helpers that expose
WorkTileInfo, so downstream kernels never see a valid tile_idx for padded
cluster blocks.
---
Minor comments:
In `@python/cudnn/block_sparse_attention/__init__.py`:
- Line 37: `__all__` in the block_sparse_attention package is using a starred
unpack from `_SYMBOLS.keys()`, which Ruff flags under PLE0604. Update the
module-level `__all__` definition in `__init__.py` to be a literal list of
string names, explicitly spelling out `BSA` and each exported symbol instead of
deriving them from `_SYMBOLS.keys()`.
In `@python/cudnn/block_sparse_attention/_interface.py`:
- Line 1: The _interface.py module is marked executable without a shebang, which
Ruff flags for importable modules. Update the file’s executable status by
removing the executable bit rather than changing the module contents, since it
is meant to be imported and not run as a script. Refer to the _interface.py
module when adjusting the file permissions in the repo.
- Line 171: In the shape unpacking inside the block sparse attention interface,
`head_dim` is never used and should be removed to avoid an unused binding.
Update the `q.shape` assignment in the relevant function so it only keeps the
values that are actually referenced later, matching the existing handling of
`batch`, `num_heads`, and `seqlen_q`.
In `@python/cudnn/block_sparse_attention/api.py`:
- Around line 395-398: Sort the __all__ list in block_sparse_attention/api.py to
match Ruff’s expected ordering, keeping the exported symbols
block_sparse_attention_backward and block_sparse_attention_forward in sorted
order within the __all__ definition.
In `@python/cudnn/block_sparse_attention/csrc/fwd/sm120_blk64/bsa_fwd_sm120.py`:
- Around line 441-445: The tensor layout comments in the forward path are
inconsistent with the implementation. Update the docstrings/comments for the mQ,
mK, mV, mO, and mLSE tensors in bsa_fwd_sm120.py to reflect the actual layout
used by the kernel and the mode validation, using (seqlen, head_dim/value_dim,
nheads, batch) for Q/K/O and the corresponding value-dim shape for V/O. Keep the
change limited to documentation and use the existing tensor symbols (mQ, mK, mV,
mO, mLSE) to ensure the comments match the shapes enforced by the forward
kernel.
In `@python/cudnn/block_sparse_attention/csrc/fwd/sm90_blk64/bsa_fwd_sm90.py`:
- Around line 487-491: The tensor layout comments for the restored tensors are
incorrect and should match the actual layout used by the implementation in the
forward path. Update the doc comments on mQ, mK, mV, mO, and mLSE in
bsa_fwd_sm90.py to describe sequence as mode 0 and head/value dimension as mode
1, i.e. (seqlen, head_dim/value_dim, nheads, batch) for the tensors and the
corresponding layout for mLSE, so the documentation matches the validation and
restore logic in the forward function.
In `@python/cudnn/block_sparse_attention/csrc/utils/kernel_utils.py`:
- Around line 154-168: The fallback branch in fmax_reduce unconditionally
indexes res[0] through res[3], so it breaks for tensors with fewer than four
elements. Update fmax_reduce to either add a dedicated small-size path that
safely reduces 1–3 elements before the existing loop logic, or explicitly assert
the minimum supported size near the shape check; keep the const_expr gating and
init_val handling intact. Also add coverage for this edge case in
test/python/fe_api so the small-tensor behavior is exercised.
In `@python/cudnn/block_sparse_attention/csrc/utils/mma_sm100_desc.py`:
- Line 85: Replace the ambiguous en dashes in the comments and any related
strings in mma_sm100_desc.py with plain hyphens to satisfy Ruff
RUF001/RUF002/RUF003. Update the affected text near the existing comment about
Float-8 / Float-6 / Float-4, and review the other flagged spots in the same
module around the identified helper/constants so all comments, docstrings, and
error messages use ASCII '-' consistently.
In `@python/cudnn/block_sparse_attention/csrc/utils/pack_gqa.py`:
- Around line 40-50: The PackGQA initializer currently declares qhead_per_kvhead
as cutlass.Constexpr[bool], but this symbol is used as a numeric constexpr for
division, multiplication, and grouped-head indexing. Update the PackGQA __init__
signature and the qhead_per_kvhead field to use an integer constexpr type
instead, and keep the rest of the initialization logic unchanged so the type
matches its actual arithmetic usage.
In `@python/cudnn/block_sparse_attention/csrc/utils/pipeline.py`:
- Around line 85-94: The make_pipeline_state function currently uses assert
False for invalid PipelineUserType values, which can be skipped under optimized
Python and let None leak out; replace that branch with a real exception, such as
ValueError, that clearly reports the invalid PipelineUserType and prevents later
pipeline failures.
In `@test/python/fe_api/block_sparse_attention/bsa_utils.py`:
- Around line 19-40: Guard the hard-coded metadata templates in
make_fixed_metadata and make_variable_metadata by validating the expected
minimum seqlen_q/seqlen_k before writing fixed block ids. Either add upfront
assertions for the required num_q_blocks/num_kv_blocks, or generate the q2k and
block_sizes patterns dynamically from those computed sizes so the helpers never
index past the tensors or emit out-of-range KV block ids.
---
Nitpick comments:
In `@python/cudnn/block_sparse_attention/api.py`:
- Around line 165-170: Expand the public docstring for the non-causal
block-sparse scaled dot-product attention wrapper in api.py so it clearly
documents the callable’s arguments, the returned dict/keys, the metadata value
contracts, and any layout assumptions or transformations. Also state the
architecture limitations and the fact that this path dispatches only to the
Python CuTe DSL kernels, with no SM100 C++/AOT extension support, or link
directly to the BSA docs for those details. Apply the same documentation pass to
the other public wrapper mentioned in this file so all user-facing APIs are
consistently documented inline.
In `@python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py`:
- Around line 250-258: Expand the docstring for build_bucketed_k2q_csr_cutedsl
to explicitly document the CSR metadata contract: describe the expected
q2k_block_index layout, the valid sentinel and index range rules, the meaning of
optional q2k_block_nums, and the shapes/meaning of all returned tensors and
ints. Make sure the wording matches the behavior enforced by related
test/python/fe_api coverage so the cross-layer metadata expectations are
unambiguous.
In `@python/cudnn/block_sparse_attention/csrc/bwd/sm90_blk64/bsa_bwd_sm90.py`:
- Line 1875: The variable from cute.arch.thread_idx() is unused in this block,
so update the unpacking in the bsa_bwd_sm90.py logic to drop the named thread
index or replace it with underscores. Keep the change localized near the
thread_idx() call in the backward SM90 block-sparse attention code, and ensure
no unused binding remains so Ruff stops flagging it.
In `@python/cudnn/block_sparse_attention/csrc/utils/copy_utils.py`:
- Around line 142-167: Document the offset_ragged_tensor contract by expanding
its docstring to state the expected tensor rank/shape requirements and how
ptr_shift=True versus False changes the logical indexing/offset behavior. Then
add a focused test in test/python/fe_api that exercises offset_ragged_tensor in
both modes and verifies the expected ragged offset semantics and rank
preconditions using the existing function name offset_ragged_tensor.
In `@python/cudnn/block_sparse_attention/csrc/utils/seqlen_info.py`:
- Around line 10-14: Move the explanatory text in seqlen_info to the actual
module docstring so documentation tools recognize it; place the existing
description at the top of the file before any imports, or if it is meant to
document a specific symbol, move it into that symbol’s docstring instead. Use
the module-level docstring location in seqlen_info as the fix target and keep
the wording focused on the sequence-length consolidation behavior.
In `@python/cudnn/block_sparse_attention/csrc/utils/softmax.py`:
- Around line 169-191: The issue is that update_row_max and update_row_sum rely
on cutlass.const_expr with is_first, so the parameter must be compile-time
constant rather than a runtime int/bool. Update the signatures of update_row_max
and update_row_sum in the softmax helper to use cutlass.Constexpr[bool] for
is_first, and keep the call sites in the SM100 softmax path passing constexpr
booleans so the contract is explicit and type-safe.
In `@test/python/fe_api/block_sparse_attention/test_BSA_attention_backward.py`:
- Around line 29-64: Add a backward test that uses variable routing and partial
physical KV blocks instead of only fixed metadata in
test_bsa_attention_backward_fixed_blocks and the related backward coverage.
Create a new case that builds variable q2k_block_nums and blk64 block_sizes,
then run block_sparse_attention_forward and block_sparse_attention_backward with
that metadata and assert dq_tensor, dk_tensor, and dv_tensor against the
reference. Keep the existing fixed-block test, but extend the backward suite so
the public API path through block_sparse_attention_backward and its
validation/dispatch logic is exercised for variable-route inputs.
In `@test/python/fe_api/block_sparse_attention/test_BSA_attention_forward.py`:
- Around line 386-413: The SM120 forward test only verifies that compilation is
reached, so it would still pass even if `_bsa_attn_fwd_sm120_blk64()` stops
including `_tensor_static_compile_key(...)` in the cache key. Update the test
around `interface._bsa_attn_fwd_sm120_blk64` to capture or inspect the
compile/cache key used when `interface.cute.compile` is invoked, and assert that
the expected static compile key is present rather than only expecting
`CompileReached`.
- Around line 17-25: The smoke test only covers cudnn.BSA, so it can miss broken
top-level lazy export wiring for cudnn.block_sparse_attention_forward and
cudnn.block_sparse_attention_backward. Update the test around _import_bsa in
test_BSA_attention_forward.py to also import those two symbols from cudnn and
exercise them with a minimal invocation or attribute check, skipping in the same
optional-dependency path if needed.
🪄 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: 39161ab5-47a7-448c-ab09-045892787b56
📒 Files selected for processing (50)
README.mddocs/fe-oss-apis/bsa.mddocs/fe-oss-apis/overview.mddocs/operations/Attention.mdpython/cudnn/README.mdpython/cudnn/__init__.pypython/cudnn/block_sparse_attention/__init__.pypython/cudnn/block_sparse_attention/_interface.pypython/cudnn/block_sparse_attention/api.pypython/cudnn/block_sparse_attention/csrc/__init__.pypython/cudnn/block_sparse_attention/csrc/bwd/__init__.pypython/cudnn/block_sparse_attention/csrc/bwd/bsa_bwd_postprocess.pypython/cudnn/block_sparse_attention/csrc/bwd/bsa_bwd_prepost.pypython/cudnn/block_sparse_attention/csrc/bwd/bsa_bwd_preprocess.pypython/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.pypython/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/__init__.pypython/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.pypython/cudnn/block_sparse_attention/csrc/bwd/sm100_blk64/__init__.pypython/cudnn/block_sparse_attention/csrc/bwd/sm100_blk64/bsa_bwd_sm100.pypython/cudnn/block_sparse_attention/csrc/bwd/sm90_blk64/bsa_bwd_sm90.pypython/cudnn/block_sparse_attention/csrc/fwd/__init__.pypython/cudnn/block_sparse_attention/csrc/fwd/sm100_blk128/__init__.pypython/cudnn/block_sparse_attention/csrc/fwd/sm100_blk128/bsa_fwd_sm100.pypython/cudnn/block_sparse_attention/csrc/fwd/sm100_blk64/bsa_fwd_combine.pypython/cudnn/block_sparse_attention/csrc/fwd/sm100_blk64/bsa_fwd_helpers.pypython/cudnn/block_sparse_attention/csrc/fwd/sm100_blk64/bsa_fwd_sm100.pypython/cudnn/block_sparse_attention/csrc/fwd/sm120_blk64/bsa_fwd_sm120.pypython/cudnn/block_sparse_attention/csrc/fwd/sm90_blk64/bsa_fwd_sm90.pypython/cudnn/block_sparse_attention/csrc/utils/__init__.pypython/cudnn/block_sparse_attention/csrc/utils/batched_static_scheduler.pypython/cudnn/block_sparse_attention/csrc/utils/block_info.pypython/cudnn/block_sparse_attention/csrc/utils/block_sparse_tile_scheduler.pypython/cudnn/block_sparse_attention/csrc/utils/copy_utils.pypython/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.pypython/cudnn/block_sparse_attention/csrc/utils/kernel_utils.pypython/cudnn/block_sparse_attention/csrc/utils/layout_utils.pypython/cudnn/block_sparse_attention/csrc/utils/mma_sm100_desc.pypython/cudnn/block_sparse_attention/csrc/utils/named_barrier.pypython/cudnn/block_sparse_attention/csrc/utils/pack_gqa.pypython/cudnn/block_sparse_attention/csrc/utils/pipeline.pypython/cudnn/block_sparse_attention/csrc/utils/seqlen_info.pypython/cudnn/block_sparse_attention/csrc/utils/sm90_utils.pypython/cudnn/block_sparse_attention/csrc/utils/softmax.pypython/cudnn/block_sparse_attention/csrc/utils/tcgen05_mma_helpers.pypython/cudnn/block_sparse_attention/csrc/utils/tile_scheduler.pytest/python/fe_api/block_sparse_attention/__init__.pytest/python/fe_api/block_sparse_attention/bsa_reference.pytest/python/fe_api/block_sparse_attention/bsa_utils.pytest/python/fe_api/block_sparse_attention/test_BSA_attention_backward.pytest/python/fe_api/block_sparse_attention/test_BSA_attention_forward.py
Anerudhan
added
cat-feature
Collaborator
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-333-78ed7d0 |
Collaborator
|
Thanks @jiayus-nvidia for the PR. Will run the CI and get back to you |
jhjpark
approved these changes
Jul 2, 2026
Collaborator
|
@cudnn-ci-bot run |
Anerudhan
approved these changes
Jul 2, 2026
Collaborator
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-333-004c40c |
Merged
Anerudhan
added a commit
that referenced
this pull request
Jul 7, 2026
* test/python: cap peak GPU memory via PYTORCH_CUDA_ALLOC_CONF (#247) Long pytest-xdist runs (e.g. test_mhas_v2 ~2.5k SDPA configs in one worker) hit a much higher GPU memory high-water mark than any single test needs, because the caching allocator retains freed blocks across configs. Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, garbage_collection_threshold:0.6 before torch is imported reduces the peak to roughly the maximum any single test needs, with no change in wall time or test outcome. Use os.environ.setdefault so user-provided values still win, and place it above the transformer_engine import so the env var is visible by the time torch initializes its CUDA allocator. * Fix DSA link in README.md Updated the link for DSA in the README to point to the correct directory. * Remove stale H200 benchmark artifacts (#252) These artifacts were superseded by the newer SDPA benchmark result layout and were already removed from the internal GitLab develop branch. * Change profile_pass from 'fwd' to 'both' * Bump the develop to 1.25.0 * Fix varpack-template lifecycle bugs + add defensive checks Two pre-existing bugs in the VariantPackTemplate, plus one defensive guard: 1. Graph copy -> dangling host pointers. template_ptrs stores raw addresses into cached_pass_by_value storage owned by the source Graph. Default copy propagated prepared=true while the addresses still pointed at the source. Fix: VarpackPrepStateBox copy ctor/assign now always start with prepared=false so the copy re-preps on first use against its own storage. 2. Re-deserialize on the same Graph -> stale template. deserialize(handle,...) rebinds cached_pass_by_value but the existing prepared=true causes the eager prep to short-circuit, leaving the slot layout from the prior deserialize. Fix: reset prepared=false and clear varpack_template before the eager prep call. 3. Null device_ptrs in raw-ptr create_variant_pack overloads. Reject nullptr + non-empty uids instead of forwarding to the cuDNN backend. Adds explicit null-plan guards across detail::execute overloads, returning GRAPH_EXECUTION_FAILED with "No plan found to execute!" instead of dereferencing plan via plan->getTag(). Ports https://gitlab-master.nvidia.com/cudnn/cudnn_frontend/-/merge_requests/2117 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Clear deserialize-owned containers on re-deserialize Addresses review feedback on PR #248: the prior fix reset prepared=false and varpack_template but left deserialized_tensor_properties, deserialized_pass_by_value, deserialized_workspace_modifications, and tensors_to_dump populated from any earlier deserialize(handle, old_data). On re-deserialize, prepare_variant_pack_template() could then ingest the stale entries alongside the new ones. Clear all four containers immediately after json::from_ubjson, before any of the deserialize logic that repopulates them. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Add row-scale support to grouped GEMM quant Signed-off-by: Ziang Li <ziangli@umich.edu> * Tighten row-scale grouped GEMM quant tests Signed-off-by: Ziang Li <ziangli@umich.edu> * feat(python): add get_engine_and_knobs_at_index for structured plan pinning (#259) * feat(python): add get_engine_and_knobs_at_index for structured plan pinning get_plan_name_at_index returns a formatted "engN_kT=V" tag built from the engine global index and knob choices. Callers that want to persist a tuned plan and replay it later are forced to either store the bare plan index (which drifts when the policy=ALL plan list is re-enumerated across cudnn-frontend / backend versions) or parse the tag string. Expose the structured data directly: get_engine_and_knobs_at_index returns (engine_id, {KnobType_t: value}), reading the same backend attributes get_engine_tag stringifies. The result feeds straight into create_execution_plan(engine_id, knobs) to rebuild the exact same kernel on a fresh graph without a heuristics query. - detail::get_engine_id_and_knobs (cudnn_frontend_utils.h): structured reader - Execution_plan_list::get_engine_and_knobs_at_index (plans.h) - Graph::get_engine_and_knobs_at_index (graph_interface.h) - PyGraph binding (pygraph.h/.cpp) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * address review: bounds-check index, add cpp unit test, trim comments - get_engine_and_knobs_at_index: reject out-of-range index (mirrors check_support_at_index) instead of indexing engine_configs OOB. - add test/cpp/get_engine_and_knobs.cpp: enumerate a matmul graph's plans, read (engine_id, knobs) for each, and confirm re-pinning via create_execution_plan reproduces the same plan (matching name); also checks out-of-range indices error. - trim the new doc comments to match neighboring style. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * knobs: add SWAP_AB / INPUT_TMA_ENABLE / OUTPUT_TMA_ENABLE to KnobType_t KnobType_t (and the to/from backend converters) stopped at WARP_SPEC_CFG (42), so engines using SWAP_AB (43, cuDNN 9.18), INPUT_TMA_ENABLE (44) or OUTPUT_TMA_ENABLE (45, cuDNN 9.22) had those knobs mapped to NOT_SET by convert_from_backend_knob_type. Feeding NOT_SET back into create_execution_plan then failed convert_to_backend_knob_type with INVALID_VALUE -- so a plan enumerated with one of these knobs (e.g. via get_engine_and_knobs_at_index) could not be pinned. Add the three knob types to the enum, both converters (version-gated to match the backend @SInCE), and the pybind knob_type enum. The cpp test now compares the structured identity (engine id + knob map) instead of the plan-name tag, since the tag serializes knobs in engine-config order, which differs between the heuristic config and the pinned one even though the kernel is identical. create_execution_plan is now asserted to succeed for every enumerated plan; building it stays best-effort (can fail for unrelated environment reasons such as a ptxas older than the engine's target). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * make get_engine_tag deterministic: sort knob choices by type The plan-name tag was built by iterating CUDNN_ATTR_ENGINECFG_KNOB_CHOICES in stored order, which differs between the heuristics path and create_execution_plan (set_knob_choices iterates a std::unordered_map). So the same engine + knob values could serialize to differently-ordered tags (e.g. eng11_k2=29_k27=0...k43=0 vs eng11_k43=0_k38=0...k2=29) -- the kernel is identical but the string isn't a stable id. Sort the knob choices by type before formatting so the tag is a deterministic function of the engine config regardless of how it was built. This is off the execution hot path (tag is used for logging / plan identity), so no perf impact; the actual knob choices passed to the backend are unchanged. The cpp test now also asserts the pinned plan's tag matches the original's. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Update SDPA Benchmarking Artifacts (#265) * update sdpa benchmark artifacts * update acknowledgement * Adding coderabbit review guide (initial template) * fix: allow overriding libcudart selection via CUDNN_FRONTEND_CUDART_LIB_NAME When dynamic loading is enabled, load_cudart_so() searches for the supported libcudart major versions and aborts with "Multiple libcudart libraries found" when more than one is visible on the library search path. This happens in containerized environments such as GKE, where the TCPXO NCCL plugin mounts a different libcudart major version from the host than the one shipped in the container. Check the CUDNN_FRONTEND_CUDART_LIB_NAME environment variable first; when set to a library name or path, dlopen exactly that library and skip the automatic multi-version detection. Behavior is unchanged when the variable is unset. Fixes #267 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Clean up guardword-flagged comments (xmma path, gitlab URL, P4 label, Perfsim, HACK/Ugly, STS/CGA SASS terms) (#273) Comment-only cleanups, no behaviour change. Replaces guardword-flagged phrasing with neutral equivalents in 7 files: - attention_utils.h:67 — drop internal `xmma/fast_math.h:118-125` path reference; keep the rationale ("matches cuDNN backend's find_divisor_v2 fast-math helper"). - test_sdpa_bwd.py:8 — drop `gitlab-master.nvidia.com` job URL from the module docstring; the rationale (2-CTA + Blackwell TMEM + xdist) is fully self-explanatory above it. - dense_score_recompute_sm90.py — "Perfsim" → "Profiling"; "Weights/LSE LDG" → "Weights/LSE load-from-global" (x2). - indexer_backward_sm90.py — `# P4:` block-pass label → `# Pass 4:` (x2); rephrase 5 "STS" SASS-instruction references in comments to "shared-mem store(s)" / "write to shared mem". - indexer_backward_sm100.py — same STS → shared-mem-store rephrasing in 1 docstring. - dsa_bwd_sm90.py:386 — `# HACK:` → `# Note:` (same meaning). - dsa_bwd_sm90.py:1554 — `STS(dS)` → "storing dS to shared mem". - dsa_bwd_sm100.py:941 — `# Ugly,` → `# Awkward,`. - dense_gemm_persistent_swiglu.py:1049 — "single CGA" → "single cluster". Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * remove_9.99_version_tag * add_protection_flags * fix(windows): consolidate getenv access and fix C4996/C4005 on MSVC The Windows wheel build (deploy:build_bdist_wheels_3.10) failed because the std::getenv call added to load_cudart_so() in cudnn_frontend_shim.h triggers MSVC warning C4996 ('getenv' is unsafe), which is treated as an error under /WX. Root cause and fixes: - Move get_environment() to cudnn_frontend_shim.h (the lowest-level header, included by utils.h before Logging.h) so a single definition is shared by all layers without inverting include dependencies. It wraps std::getenv with a properly scoped #pragma warning(push)/disable(4996)/pop, guarded by _WIN32. - Route all getenv call sites through get_environment(): shim.h, graph_properties.h, scaled_dot_product_flash_attention.h, and sm100_rms_norm_silu_engine.h. These were previously only spared from C4996 by an unscoped pragma leak in Logging.h, and would have started failing once that leak was fixed. - Remove the duplicate get_environment() from cudnn_frontend_Logging.h, which had three issues: an unscoped 'warning(disable:4996)' that leaked to the rest of the TU, a no-op '#define _CRT_SECURE_NO_WARNINGS' (placed after the CRT headers), and a 'WIN32' guard that should be '_WIN32'. Dropping the macro also resolves the C4005 '_CRT_SECURE_NO_WARNINGS macro redefinition' warning for downstream projects. Fixes #139 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(shim): warn instead of throwing when multiple libcudart libraries are found Loading cudart no longer aborts when both libcudart.so.12 and libcudart.so.13 are present in the library search path. Instead, load_cudart_so() emits a warning on stderr and falls back to the first library found. Users can still select a specific library explicitly via CUDNN_FRONTEND_CUDART_LIB_NAME. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Unblock SDPA tests and promote FP8 ragged backward to L0 (#275) * Promote L1 Python tests to L0 * Restore L1 markers except FP8 ragged backward * Add per-expert reduction (group_offset) for MoE grouped GEMM Adds optional group_offset support to the reduction node so cuDNN FE can express per-expert reductions for MoE grouped GEMM workloads. - New Group_offset graph_properties tensor input and Reduction_attributes::set_group_offset setter - INode::reduction and PyGraph::reduction signatures take an optional group_offset tensor - Operation_v8 builder wires CUDNN_ATTR_OPERATION_REDUCTION_GROUP_OFFSET_DESC with runtime version checks (cuDNN >= 9.24.0) - Python binding (pygraph) exposes the optional group_offset argument Mirrors gitlab-master cudnn/cudnn_frontend MR !2111 by @yanqinz. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * Fix the 9.99 bound * Skip flexible-graph SDPA bwd sample on SM120 and above (#284) The fp16 backward-with-flexible-graphs sample guards against SM 120 (consumer Blackwell) where this path is not supported. The guard used an exact == 120 check, which missed SM 121 (GB10 / DGX Spark) and any later consumer Blackwell arch, causing the sample to run and fail there. Change the check to >= 120 so the sample is skipped on SM 120 and above, and update the SKIP message to match. Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * 1 * Add pre-commit hooks (#286) * Fix clang format issues * Fix clang-format * Add pre-commit hooks and fix pre-commit * Fix the black issues * Skip TensorIR MemBound / compile-time-const samples on consumer Blackwell (SM12x) (#285) * Skip TensorIR MemBound / compile-time-const samples on consumer Blackwell (SM12x) The TensorIR MemBound engine (cudnnTensorIrMemBoundEngine) only supports SM100-SM109 (data center Blackwell): its arch gate is [SM_100, SM_110) and the DKG cubins it emits are the sm_100f family-portable target, which the CUDA driver will not load on sm_120. The membound and compile-time-constant samples guarded their device check with check_device_arch_newer_than("blackwell") / is_blackwell_arch(), both of which are true for SM120 consumer Blackwell. So on an RTX 50-series (sm_120) GPU these samples fall through to create_execution_plans() and FAIL with "No valid engine configs returned from heuristics" (no engine serves the graph; the kernelgen runtime-fusion fallback only targets SM70/SM80/SM90). Narrow the guard to is_blackwell_computing_arch() (100 <= cc < 110) so the samples skip cleanly on SM120 and above, matching the backend engine's actual support range. This mirrors PR #283, which skipped the flexible-graph SDPA backward sample on SM120+. Affected test cases (verified on RTX 5080 / sm_120, cuDNN 9.30 -> now SKIP): membound/transpose.cpp "Membound transpose permutes dims" membound/reshape.cpp "Membound reshape ... LOGICAL mode" membound/slice.cpp "Membound slice window with step" membound/concat.cpp "Membound concatenate on channel axis" membound/membound_fusion.cpp "Fusion reshape then ReLU" / "Fusion transpose then add bias tensor" membound/boolean_fusion.cpp "Boolean CMP_GT and LOGICAL_AND fusion" misc/compile_time_constant_example.cpp "Compile-time constant scalar multiply and add" Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Skip boolean_cmp_logic Python notebook on consumer Blackwell (SM12x) Python counterpart of the C++ membound/boolean sample fix. The CMP_GT + LOGICAL_AND boolean fusion runs on the TensorIR mem-bound engine, which only supports SM100-SM109 (data center Blackwell). On SM120 consumer Blackwell the notebook's create_execution_plans([A, FALLBACK]) silently falls back to an engine that produces WRONG results (verified on RTX 5080 / sm_120: 109/512 mismatches -> assertion failure). Gate the cuDNN cells on is_supported_arch so the notebook skips cleanly on SM120 instead of producing wrong results, and fix the prerequisite markdown (SM100+ "or later" -> SM100-SM109). The arch check computes the full compute capability (major*10 + minor) and tests 100 <= cc < 110 to mirror the C++ is_blackwell_computing_arch() helper exactly. This notebook is not part of ci/run_python_samples.sh, so it does not affect CI; the fix is for correctness/consistency with the C++ sample. Committed with --no-verify: the local black-jupyter pre-commit hook reflows the whole .ipynb to indent=1 (repo notebooks are indent=2) and collapses unrelated aligned dicts; CI does not enforce notebook formatting. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Support cu_seqlens in unified SDPA (#266) * use static signature for sfd_col_d_srelu_tensor (#281) Signed-off-by: Jieming Zhang <jiemingz@nvidia.com> * DSA: fix CuTe DSL guards and add SM90 indexer forward (#263) * DSA: fix CuTe DSL guards and add SM90 indexer forward * DSA: allow indexer top-k on SM90 * DSA: trim CuTe DSL compile-cache keys + unify indexer_forward paths Compile-cache keys across the deepseek_sparse_attention kernels included runtime-only values (batch/seqlen/seqlen_k, sm_scale, tensor shapes/strides, num_head, num_threads), forcing spurious recompiles under varlen / changing batch even though one compiled kernel serves them all. Drop those fields and keep only params that change generated code. The two dense_indexer_backward kernels originally baked seqlen into codegen, so to drop it safely they were reworked to take seqlen at runtime: - sm90: the dense K-load looped via range_constexpr(num_topk_blocks = seqlen_k // block_I); it now loops at runtime over num_k_blocks, like the compute warpgroup already did. - sm100: ScoreGradDense baked max_seqlen_q into its launch grid and max_seqlen_q/k into the causal-mask bound via __init__ ints; they are now runtime Int32 args (matching the GEMM kernel), which also fixes a latent bug where a kernel compiled for one max_seqlen_k could be silently reused for another. Collapse the redundant two-layer compile cache (dict-of-closures + per-closure lazy holder) in the indexer_backward factories to the single forward-style dict (key -> compiled kernel), matching indexer_forward. indexer_forward: route the SM100 BSHD path through the same indexer_fwd wrapper as THD instead of the separate IndexerForward APIBase class, which compiled against concrete fake-tensor shapes (recompiling per shape/stride). indexer_fwd marks layouts dynamic and compiles once per config; on B300 the two produce bit-identical output with <2% kernel-time difference at realistic shapes. indexer_fwd gains an optional current_stream arg (also fixing the THD path, which previously dropped the caller's stream). The public IndexerForward class/export is retained. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * DSA: address indexer stream and cache review * DSA: format CuTe DSL indexer files * DSA: key SM100 sparse bwd by num heads --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: mingyangw <mingyangw@nvidia.com> * Fix formatting issues from #263 (#294) * Support static linking of libcudnn (#182) * Support static linking of libcudnn * Fix variable handling * Don't use static zlib for PIC * Rename CUDNN_STATIC_LINK * Make version variables compatible for pytorch * Apply suggestion from @coderabbitai[bot] Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * Apply review suggestions --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * make dgeglu config values compile time constants instead of runtime values (#293) * bench: add autoregressive video DiT SDPA config + GB200/GB300 results (#277) (#295) * bench: add autoregressive video DiT SDPA config + GB200/GB300 results Adds a new benchmark config for the autoregressive (world-model / next-frame) video DiT shape: short query (one new frame, s_q ∈ {985, 1024, 2048, 4096, 8192}) attending a long cached KV history (s_kv=62208) with h=9, d=128 and no operator-level mask. This is a class of workload that prior DiT configs (LTX-2, Wan 2.2) don't cover, because those run bidirectional self-attention with s_q == s_kv. Captured on lyris GB200 and GB300 (cuDNN 9.23.0, FAv4 from the CuTe-DSL build). FAv4 FP8/MXFP8 bars are absent because that build's forward asserts on non-fp16/bf16 inputs; the runner now skips FAv4 cases for both FP8 and MXFP8 (previously only MXFP8) to keep the CSVs free of traceback noise. * bench: add B300 peak comparison for autoregressive DiT (cuDNN split-K vs FAv4 best num_splits) Adds a "peak vs peak" view that complements the existing default-vs-default chart: cuDNN 9.30.0 with prefill split-K enabled on bf16/fp8/mxfp8, paired against FAv4 BF16 swept over num_splits ∈ {1, 2, 4, 8, 16, 32} with the best per-seqlen result annotated on the bar (ks=). For the autoregressive video DiT shape (B=1, h=9, d=128, s_q ∈ {985..8192}, s_kv=62208) on B300 SXM6: s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks) 985 1701 2429 2274 1424 (ks=4) 1024 1767 2526 2367 1485 (ks=4) 2048 1880 2713 2547 1597 (ks=2) 4096 1997 2947 2655 1995 (ks=1) 8192 1998 2974 2681 1980 (ks=1) (TFLOPS, fwd only) cuDNN BF16+split-K beats FAv4-best-num_splits at every seqlen (+19% at the short-Q end, tied at large s_q where neither needs splitting). FP8/MXFP8 dominate by +30-50% over FAv4 BF16 thanks to the higher mma throughput. Changes: * benchmark_single_sdpa.py: --fa4_num_splits flag plumbed end-to-end so callers can force FAv4 into a specific split count (default unchanged: let FAv4 pick automatically). * bench_ar_dit_peak.py: standalone driver that runs the cartesian {seqlens} x {cudnn dtypes} sweep plus the FAv4 num_splits sweep and emits a CSV with one row per (backend, dtype, seqlen) — with the winning num_splits recorded for the FAv4 rows. * results/auto_regressive_dit/b300/: CSV + chart. * README: B300 peak section. * bench: GB200 + GB300 peak comparison for autoregressive DiT (replace B300 preview) Drops the earlier B300 preview chart in favour of the matching peak charts on the production GB200 and GB300 superchip variants (same SM_103 silicon in the GB300 case, fewer SMs / lower clock on GB200). Charts are the same peak-vs-peak view: cuDNN 9.30.0 with prefill split-K enabled on bf16/fp8/mxfp8, paired against FAv4 BF16 swept over num_splits and keeping the best per-seqlen result. GB300 (TFLOPS, fwd only): s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks) 985 1752 2519 2359 1451 (ks=4) 1024 1813 2619 2447 1515 (ks=4) 2048 1923 2768 2598 1613 (ks=2) 4096 2050 2978 2687 2055 (ks=1) 8192 2085 3002 2707 2071 (ks=1) GB200 (TFLOPS, fwd only): s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks) 985 1380 1796 1717 1332 (ks=4) 1024 1429 1870 1785 1389 (ks=4) 2048 1573 1996 1915 1513 (ks=2) 4096 1697 2066 1971 1746 (ks=1) 8192 1762 2080 1988 1802 (ks=1) On GB300 cuDNN BF16+split-K beats FAv4-best-num_splits at every seqlen (+21% at the short-Q end, tied at large s_q where neither needs splitting). On GB200 the short-Q advantage is +4-5% and FAv4 narrowly edges cuDNN BF16 at the large s_q end (-2-3%). FP8/MXFP8 dominate by +30-50% over FAv4 BF16 on both GPUs. * bench: consolidate autoregressive DiT charts to a single canonical view per GPU Drops the cuDNN 9.23 default-vs-default chart pair — those numbers are stale relative to what ships next, and keeping two charts per GPU with two different cuDNN versions is more confusing than informative. The remaining chart on each GPU is the cuDNN 9.30.0 + prefill split-K view paired against FAv4 BF16 with the best num_splits per seqlen, captured on the production GB200 and GB300 superchips. CSV is named auto_regressive_dit_no_mask.csv so the chart and its source data follow the standard <config>_<mask>.{png,csv} convention used by other benchmarks in this suite. * bench: relabel autoregressive DiT charts to cuDNN 9.24.0 (split-K release version) The split-K prefill feature exercised by these charts is cherry-picked onto release/9.24.0 and ships in that release, so the chart labels and the cudnn_backend_version column in the CSVs should reflect that version rather than the dev-branch version they happened to be measured on. --------- Co-authored-by: Vedaanta Agarwalla <142048820+vedaanta@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * - Update the Black version. (#296) - Fix the formatting issues in grouped_gemm_dglu/api.py * Add ragged offset multiplier support (#290) Add frontend support for the per-tensor ragged offset multiplier (CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER), letting ragged offsets be stored in coarser units and scaled back to element offsets by the engine. - Add ragged_offset_multiplier field, getters/setters, and validation to Tensor_attributes; emit the backend attribute (gated on cuDNN >= 9.24.0). - Expose ragged_offset_multiplier through the Python tensor() bindings (appended last to preserve positional backward compatibility). - Serialize/deserialize the multiplier and the ragged offset reference. - Reject a non-default multiplier on the composite SDPA path (unified forward only). - Add C++ and Python (test_mhas_v2) coverage, including a cu_ragged_mult configuration exercising cu_seqlens together with the multiplier. * Fix unused ragged offset version error variable (#299) `NV_CUDNN_FE_DYNAMIC_CHECK_BACKEND_DESCRIPTOR` expands to nothing when `NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` is not defined. So, the variable `ragged_offset_multiplier_cudnn_ver_error` may be unused. * Add the results. Initial script and README.md (#303) * Add acknowledgements for cuteDSL Kernels (#305) * Align DSA indexer kernels and fix dense score-grad clipping (#297) * Fix SM100 dense score grad clip mask * Align DSA indexer kernels with indexer implementation * The reduce_dKV validity guard compared the topk column position (#298) (global_row_idx) against max_seqlen_kv. A column position >= total_S_kv is not invalid -- with a non-compact topk_idxs layout (-1 sentinels, width > total_S_kv) valid indices can sit at any column. Entries past column total_S_kv were silently treated as -1 and their dKV contributions dropped, while dQ (whose load path correctly judges validity by the index value) stayed correct. With a [window | compressed] layout this zeroes the entire original-KV region of dkv bit-exactly. Drop the position-vs-seqlen comparison; the < topk bound plus the topk_idx >= 0 sentinel check in the store helpers already match the load-side and FlashMLA-forward semantics. Remove the now-unused max_seqlen_kv parameter from reduce_dKV. Also fix the test reference _make_topk_mask: without topk_length it clamped -1 sentinels to index 0, spuriously marking KV row 0 as attended, which corrupted out/lse/gradient references for non-compact inputs. Verified on B200: topk width 1024 > S_kv 256 now gives cos_sim(dkv) 0.9996 (was 0.498); wide non-compact layouts pass FP32 autograd checks; fe_api/dsa pytest suite passes (16 tests). Co-Authored-By: Claude Fable 5 noreply@anthropic.com * Update SDPA Benchmarking Artifacts - 9.24.0.27 (#306) * Add docs folder (#308) * Add docs folder Copy the docs folder (operations, fe-oss-apis, and guides) from the internal cudnn_frontend develop branch. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Apply black formatting to python folder Run black (line-length 160) over python/; collapse multi-line ternaries in the deepseek_sparse_attention indexer kernels. Formatting only. * Apply black formatting to dsa_reference.py Collapse two multi-line calls that fit within 160 chars. Formatting only. * Support SReLU in grouped GEMM hadamard fusion (#315) Signed-off-by: Siddhartha Raman <sraman@nvidia.com> * Add byte boolean frontend data type (#302) Add DataType_t::BYTE_BOOLEAN and map it to CUDNN_DATA_BYTE_BOOLEAN for cuDNN 9.30+. Update the boolean membound sample to use byte-backed boolean tensor storage on 9.30+ backends while keeping logical compute precision as BOOLEAN. * fix(sdpa_benchmark): use sampled SM clock + per-arch MMA throughput for SOL% (#314) The MMA SOL% reported by benchmark_single_sdpa.py relied on nvmlDeviceGetMaxClockInfo for the peak-throughput denominator. On some Blackwell datacenter SKUs that value is unreliable: it can read below the boost clock the kernel actually runs at (producing > 100% SOL) or above the sustained clock under power/thermal caps (understating SOL when clocks are locked). Replace it with: * a background pynvml sampler that records the SM clock during the benchmark window, taking max(sampled) as the operating clock; and * a per-data_type FLOPs/clock/SM table (BF16/FP16 dense = 8192, FP8/MXFP8 dense = 16384 on Blackwell DC). Validated on a GB200 node (152 SMs, sm_100, 2062 MHz nvml max): * free clock: baseline 37.5%, patched 37.3% (agree when nvml is correct) * locked 1200: baseline 28.0%, patched 48.3% * locked 900: baseline 21.5%, patched 49.1% Patched SOL is clock-invariant by construction. Limited to Blackwell datacenter for now; other archs report TFLOPS without a SOL suffix rather than fall back to a wrong constant. * Migrate "cute.core.ThrMma" and "cute.make_fragment" (#321) * cute.core.ThrMma is deprecated * cute.make_fragment is deprecated * Fix sort order in block_scale_quantize.h (#319) If I compile and run the `samples/cpp/norm/norm_block_scale.cpp` sample with clang in debug mode I get this error: ``` strict_weak_ordering_check.h:50: libc++ Hardening assertion !__comp(*(__first + __a), *(__first + __b)) failed: Your comparator is not a valid strict-weak ordering ``` The comparator indeed violates strict weak ordering. I.e. it in this case it will report that index 0 is smaller than index 1 and also that index 1 is smaller than index 0: ``` X_stride = {10, 10} X_dim = {1, 1} ``` The fix makes the comparator a strict weak order. * Fix SM100 sparse score recompute compact top-k codegen (#317) * Fix SM100 sparse score recompute compact top-k codegen Summary This fixes the SM100 sparse attention score-recompute kernel when topk_length is provided for compact top-k layouts. The change removes the runtime topk_length branch around the TMEM copy in both attention epilogues: - n_block_size >= 128 / Ld32x32bOp - n_block_size < 128 / Ld16x64bOp The dynamic guard is still kept for score accumulation and output, so blocks past topk_length continue to contribute zero. Why this is needed Downstream DSA sparse indexer loss calls sparse_attn_score_recompute_wrapper(..., topk_length=...) for packed THD / CP workloads. With cuDNN Frontend 1.25.0 and CUTLASS DSL 4.5.0 on SM100, the compact path currently fails during DSL compilation with an ICE like: failed to legalize unresolved materialization from !cute_nvgpu.atom.tmem_load ... to !cute.tiled_copy The failure happens at the TMEM copy construction inside the runtime should_copy_tmem branch. Always materializing the TMEM copy avoids the compiler legalization issue while preserving the existing topk_length masking semantics for the values that are actually accumulated and written. This is needed so the cuDNN DSA sparse indexer-loss path can stay fully on the cuDNN Frontend implementation instead of requiring a framework-side fallback. Signed-off-by: Hollow Man <hollowman@opensuse.org> * fix test cases Now has_topk_length is added to the shared DSA_SCORE_RECOMPUTE_PARAM_MARKS, which is used by both sparse and dense score-recompute tests. Dense test functions do not accept has_topk_length, so pytest collection failed. Signed-off-by: Hollow Man <hollowman@opensuse.org> --------- Signed-off-by: Hollow Man <hollowman@opensuse.org> * grouped gemm dglu dbias reduction dsl 4.5 regression: switch to constexpr loop (#322) * Fix MXFP8 testing sync issue (#325) * fix (#326) * Add enforce_precompiled deserialize option (#323) * fix: IMA on indexer_topk_wrapper (#312) * Add run_warmup opt-out and reuse-parsed-json overload to Graph::deser… (#329) * Add run_warmup opt-out and reuse-parsed-json overload to Graph::deserialize * docstring, clang, warmup level fixes * DSA: add q causal offsets and SM100F support (#316) * DSA: fix ratio length assertions * DSA: support q causal offsets * Add Rubin sm100f support for DSA CuTe DSL kernels * docs: clarify DSA q causal offsets * DSA: skip masked dense K blocks * Update DSA stream handling and SM100 score kernels * Fix SM100 dense indexer backward synchronization Wait for the final dQ MMA before reading TMEM, synchronize q0 TMA store completion before reusing shared memory for q1, and include the pending DSA formatting updates. --------- Co-authored-by: cjerry <cjerry@nvidia.com> * Fix documentation check failures (#332) * Add unified-engine FP8 and MXFP8 forward SDPA support (#301) Wire per-tensor FP8 and block-scaled MXFP8 (E8M0) forward attention through the unified SDPA runtime fusion engine: - scaled_dot_product_flash_attention.h: enable FP8/MXFP8 descale, scale, and amax attributes on the unified path. - sdpa_support_surface.h: gate unified FP8/MXFP8 support and drop constraints no longer required by the unified engine. - python bindings (pygraph.h, sdpa.cpp): expose the new descale/scale/amax inputs and outputs. - tests: extend fp8.py, mxfp8.py, and test_mhas_v2.py to cover the unified-engine path. * rename SMxxx to Blackwell (#334) * Fix grid dim overflow in DSA backward convert kernel on SM100 (#331) The convert kernel grid was configured as [1, convert_grid_x, 1], placing the seq-block dimension on grid.y. CUDA caps grid.y/z at 65535, so large mKV.shape[0] / block_seq values trigger `invalid configuration argument`. grid.x supports up to 2^31-1, so move convert_grid_x to grid.x and update the corresponding block_idx() unpacking in the kernel accordingly. No behavior change for in-range sizes. * Bypass OSS d=256 path on cuDNN 9.23+ (#335) * Bypass cuteDSL d=256 path on cuDNN 9.23+ cuDNN 9.23.0 added native d=256 SDPA fprop and bprop support in the graph backend, so the OSS (cuteDSL) kernels at `cudnn.experimental.ops.sdpa` are no longer required when the linked backend is recent enough. Add `_cudnn_supports_native_d256()` gated on `cudnn.backend_version() >= 92300` and require it to be `False` before routing fprop/bprop through the SM100 OSS wrappers. The pre-existing SM100+ device check is kept so older cuDNN versions still light up the OSS path on Blackwell. The `test_d256_uses_oss_forward_path` test now skips on cuDNN 9.23+ since the OSS bypass is intentional, and a new `test_d256_uses_graph_path_on_cudnn_9_23_plus` asserts that fprop/bprop populate the cuDNN graph cache (proving the OSS path is bypassed). Also: `_skip_if_unsupported_d256` and `test_d256_uses_oss_forward_path` used `import cudnn.sdpa` inside the function body, which made `cudnn` a local variable and shadowed the module-level import as soon as any earlier line referenced `cudnn` (e.g. the new `cudnn.backend_version()` check). Switch to `importlib.import_module("cudnn.sdpa")` to avoid the binding. * Address review: rename to cudnn_backend, harden routing test - Rename `_CUDNN_NATIVE_D256_VERSION` → `_CUDNN_BACKEND_D256_VERSION` and `_cudnn_supports_native_d256()` → `_cudnn_backend_supports_d256()` per @Anerudhan's request that we say "cuDNN backend" instead of "cuDNN native". Update the surrounding log messages and skip strings to match. - Strengthen the cuDNN-backend routing test: replace `sdpa_fwd_d256` and `sdpa_bwd_d256` on the module with a sentinel that fails the test if the OSS path is ever entered. The cache-population assertions stay as corroborating signals, but the sentinel is what guarantees we did not enter the cuteDSL kernels. Rename the test to `test_d256_uses_cudnn_backend_on_cudnn_9_23_plus`. * Fix d=256 tests on Ampere * Tidy SDPA imports and formatting --------- Co-authored-by: Vedaanta Agarwalla <vagarwalla@nvidia.com> * Update the cudnn version to 1.26.0 (#337) * Update conv get-plan sample heuristic config count (#278) * Use BYTE_BOOLEAN for cuDNN 9.25+ (#339) * Use BYTE_BOOLEAN for cuDNN 9.25+ * Lower unified SDPA FP8 gate to cuDNN 9.25 * Add block-sparse attention CuTe DSL kernels for Hopper and Blackwell (#333) * Add block sparse attention CuTe DSL kernels * Refactor block sparse attention kernels * Add optional caller-provided output tensor to grouped_gemm_quant_wrapper_sm100 (#338) Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * optimize dsa bwd sm100 kernel (#318) * optimize dsa bwd sm100 kernel * add dsa bwd benchmark * Test/sample improvements + block-scale & SDPA fixes (9.18–9.24 fuzzer mining) (#330) * test: fuzzer coverage from 9.18-9.24 fixed-bug mining Derived from a triage of the 134 fixed front-end bugs in cuDNN 9.18-9.24. - matmul fuzzer: run-to-run determinism assert (reuses the previously-discarded output hash; re-executes the same built plan into a re-poisoned output+workspace and asserts bit-identical). Deselects NONDETERMINISTIC plans so legitimate atomic split-K cannot false-fail. Env: MATMUL_DET_RERUNS / MATMUL_NUM_TESTS / MATMUL_FUZZ_SEED. - SDPA: add the S_Q>S_KV regime — RandomSequenceLength structurally capped s_q<=s_kv, so it was never exercised (NVBug 5829882). Clamped to s_q_max; wired into 9 suites. Env: MHAS_NUM_TESTS / MHAS_SEED_OFFSET. - MoE grouped-matmul: per-expert numeric oracle (fwd+bwd; was execute-only) plus a randomized variant covering empty experts / offset boundaries. - matmul: opt-in degenerate/GEMV shapes (MATMUL_FUZZ_DEGENERATE=1) — M=1/N=1/tiny-K were structurally unreachable. Gated off by default: it surfaced a real FORT-native matmul IMA on K=1+int8 (filed separately) that crashes the process. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(low-precision-matmul): use canonical block-reduced nvfp4 descale shape The fp4 matmul test passed a full-size descale (1,M,K)=(1,128,64) instead of the canonical F8_128x4 block-reduced (1,M,ceil(K/block) rounded to 4)=(1,128,4) (and B symmetrically). It only "passed" because scales were all 1.0 (identity) and the test does no numeric comparison -- a malformed descale that the backend silently accepted (OOB/NaN with real scales). create_matmul_dequantize_graph also derived M/N/K from the descale shape, conflating it with the data shape. Derive dims from the data tensors and build descales at the canonical block-reduced shape/stride (block dim contiguous), matching the C++ sample and BlockScaleQuantizeOperation. Now passes the new dequant shape guard. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * sample(sdpa-mxfp8): align fwd SF_V to d-contiguous (stride[3]==1) convention The fwd mxfp8 sample was the lone outlier declaring SF_V s_scale-contiguous (stride[2]==1); SF_Q/SF_K, the bwd sample, and test_mhas_v2 all use d-contiguous (stride[3]==1). The kernel reads block-scale factors via the F8_128x4 swizzle, so the declared inner stride is not load-bearing (verified: flipping it with fixed data is bit-identical) -- consistency/clarity fix, behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(matmul-fuzzer): add MATMUL_FUZZ_UNALIGNED for FORT-native widening-cast corner Opt-in: emit non-mult-of-4 K/N so the bits_per_access<32 LDG+STS smem-staging path is reachable, where a widening-cast (int8/fp8->fp16/fp32) operand over-runs the staging buffer (silent wrong-result on unaligned K, IMA on unaligned N). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test: remove broken L2 mxfp8 SDPA test (home-grown swizzle reference) create_scale_factor_tensor_for_sdpa builds the F8_128x4 scale swizzle by hand inconsistently with the kernel, feeding mis-ordered scales -> fails numerically across cuDNN versions (incl. official 9.23.1.3). MXFP8 SDPA fwd+bwd is already covered correctly by test_mhas_v2 (TE-quantized, numeric-validated) + the C++ samples. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: correct mislabeled IS_VIRTUAL tensor descriptor error message The IS_VIRTUAL SetAttribute failure reused the BYTE_ALIGNMENT error string. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Fix formatting issues by various commits before 1.26.0 (#341) --------- Signed-off-by: Ziang Li <ziangli@umich.edu> Signed-off-by: Jieming Zhang <jiemingz@nvidia.com> Signed-off-by: Siddhartha Raman <sraman@nvidia.com> Signed-off-by: Hollow Man <hollowman@opensuse.org> Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Co-authored-by: Vedaanta Agarwalla <142048820+vedaanta@users.noreply.github.com> Co-authored-by: Hwanseo Choi <hwanseoc@nvidia.com> Co-authored-by: Vincent <vinnietombari@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Ziang Li <ziangli@umich.edu> Co-authored-by: Yang Xu <38851819+YangXu1990uiuc@users.noreply.github.com> Co-authored-by: Yang Xu <yanxu@nvidia.com> Co-authored-by: Brandon Zhang <31413216+brandonfzhang@users.noreply.github.com> Co-authored-by: Jane (Jiancheng) Liu <liujane@nvidia.com> Co-authored-by: Yanqin Zhai <yanqinz@nvidia.com> Co-authored-by: Emil Gilliam <egilliam@nvidia.com> Co-authored-by: Jimmy Zhang <133159885+jiemingz@users.noreply.github.com> Co-authored-by: jiayus-nvidia <jiayus@nvidia.com> Co-authored-by: mingyangw <mingyangw@nvidia.com> Co-authored-by: Takeshi Watanabe <take-cheeze@users.noreply.github.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Mingyang Wang <35635157+saltyminty@users.noreply.github.com> Co-authored-by: Shraiysh <svaishay@nvidia.com> Co-authored-by: Jie Fang <jief@nvidia.com> Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com> Co-authored-by: yeliu-oss <yeliu@nvidia.com> Co-authored-by: Vincent <34876120+Vinnie6167@users.noreply.github.com> Co-authored-by: Dimitar (Mitko) Asenov <dimitar.asenov@gmail.com> Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 <j88437182@hotmail.com> Co-authored-by: Josh Park <89948656+jhjpark@users.noreply.github.com> Co-authored-by: yanzhuo607 <yanzhuoc@nvidia.com> Co-authored-by: Haisha Zhao <33570593+Hyaloid@users.noreply.github.com> Co-authored-by: Vince (Junghoon) Han <vincejhan98@gmail.com> Co-authored-by: cjerry <cjerry@nvidia.com> Co-authored-by: Zhenyu <370166800@qq.com> Co-authored-by: Vedaanta Agarwalla <vagarwalla@nvidia.com> Co-authored-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
@jiayus-nvidia Is this block-sparse path intended to accelerate DSA at the token level, by using an indexer to select fine-grained sparse attention regions? Is the overall design similar in spirit to ideas such as HISA: Efficient Hierarchical Indexing for Fine-Grained Sparse Attention, where a higher-level indexing mechanism is used to reduce the cost of fine-grained sparse attention? Are there any plans to integrate this implementation into NVIDIA/Megatron-LM, for the DSA path? Thanks! |
Contributor
Author
|
@xiaoxi-wangfj VSA operates at block/cube granularity: it selects KV blocks for each query block and performs attention over the selected blocks. HISA, in contrast, uses blocks for coarse filtering and then performs an additional token-level refinement step within the selected candidate blocks. Our current implementation does not include this block-to-token selection stage. We currently do not have plans to integrate this implementation into NVIDIA/Megatron-LM. If you are looking for kernels supporting the original DSA path, you are welcome to use the implementation under deepseek_sparse_attention. |
Anerudhan
added a commit
that referenced
this pull request
Aug 6, 2026
* test/python: cap peak GPU memory via PYTORCH_CUDA_ALLOC_CONF (#247)
Long pytest-xdist runs (e.g. test_mhas_v2 ~2.5k SDPA configs in one
worker) hit a much higher GPU memory high-water mark than any single
test needs, because the caching allocator retains freed blocks across
configs.
Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,
garbage_collection_threshold:0.6 before torch is imported reduces the
peak to roughly the maximum any single test needs, with no change in
wall time or test outcome.
Use os.environ.setdefault so user-provided values still win, and
place it above the transformer_engine import so the env var is
visible by the time torch initializes its CUDA allocator.
* Fix DSA link in README.md
Updated the link for DSA in the README to point to the correct directory.
* Remove stale H200 benchmark artifacts (#252)
These artifacts were superseded by the newer SDPA benchmark result layout and were already removed from the internal GitLab develop branch.
* Change profile_pass from 'fwd' to 'both'
* Bump the develop to 1.25.0
* Fix varpack-template lifecycle bugs + add defensive checks
Two pre-existing bugs in the VariantPackTemplate, plus one defensive guard:
1. Graph copy -> dangling host pointers. template_ptrs stores raw addresses
into cached_pass_by_value storage owned by the source Graph. Default copy
propagated prepared=true while the addresses still pointed at the source.
Fix: VarpackPrepStateBox copy ctor/assign now always start with
prepared=false so the copy re-preps on first use against its own storage.
2. Re-deserialize on the same Graph -> stale template. deserialize(handle,...)
rebinds cached_pass_by_value but the existing prepared=true causes the
eager prep to short-circuit, leaving the slot layout from the prior
deserialize. Fix: reset prepared=false and clear varpack_template before
the eager prep call.
3. Null device_ptrs in raw-ptr create_variant_pack overloads. Reject nullptr
+ non-empty uids instead of forwarding to the cuDNN backend.
Adds explicit null-plan guards across detail::execute overloads, returning
GRAPH_EXECUTION_FAILED with "No plan found to execute!" instead of
dereferencing plan via plan->getTag().
Ports https://gitlab-master.nvidia.com/cudnn/cudnn_frontend/-/merge_requests/2117
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Clear deserialize-owned containers on re-deserialize
Addresses review feedback on PR #248: the prior fix reset prepared=false
and varpack_template but left deserialized_tensor_properties,
deserialized_pass_by_value, deserialized_workspace_modifications, and
tensors_to_dump populated from any earlier deserialize(handle, old_data).
On re-deserialize, prepare_variant_pack_template() could then ingest the
stale entries alongside the new ones.
Clear all four containers immediately after json::from_ubjson, before any
of the deserialize logic that repopulates them.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Add row-scale support to grouped GEMM quant
Signed-off-by: Ziang Li <ziangli@umich.edu>
* Tighten row-scale grouped GEMM quant tests
Signed-off-by: Ziang Li <ziangli@umich.edu>
* feat(python): add get_engine_and_knobs_at_index for structured plan pinning (#259)
* feat(python): add get_engine_and_knobs_at_index for structured plan pinning
get_plan_name_at_index returns a formatted "engN_kT=V" tag built from the
engine global index and knob choices. Callers that want to persist a tuned
plan and replay it later are forced to either store the bare plan index
(which drifts when the policy=ALL plan list is re-enumerated across
cudnn-frontend / backend versions) or parse the tag string.
Expose the structured data directly: get_engine_and_knobs_at_index returns
(engine_id, {KnobType_t: value}), reading the same backend attributes
get_engine_tag stringifies. The result feeds straight into
create_execution_plan(engine_id, knobs) to rebuild the exact same kernel on a
fresh graph without a heuristics query.
- detail::get_engine_id_and_knobs (cudnn_frontend_utils.h): structured reader
- Execution_plan_list::get_engine_and_knobs_at_index (plans.h)
- Graph::get_engine_and_knobs_at_index (graph_interface.h)
- PyGraph binding (pygraph.h/.cpp)
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* address review: bounds-check index, add cpp unit test, trim comments
- get_engine_and_knobs_at_index: reject out-of-range index (mirrors
check_support_at_index) instead of indexing engine_configs OOB.
- add test/cpp/get_engine_and_knobs.cpp: enumerate a matmul graph's plans,
read (engine_id, knobs) for each, and confirm re-pinning via
create_execution_plan reproduces the same plan (matching name); also checks
out-of-range indices error.
- trim the new doc comments to match neighboring style.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* knobs: add SWAP_AB / INPUT_TMA_ENABLE / OUTPUT_TMA_ENABLE to KnobType_t
KnobType_t (and the to/from backend converters) stopped at WARP_SPEC_CFG (42),
so engines using SWAP_AB (43, cuDNN 9.18), INPUT_TMA_ENABLE (44) or
OUTPUT_TMA_ENABLE (45, cuDNN 9.22) had those knobs mapped to NOT_SET by
convert_from_backend_knob_type. Feeding NOT_SET back into create_execution_plan
then failed convert_to_backend_knob_type with INVALID_VALUE -- so a plan
enumerated with one of these knobs (e.g. via get_engine_and_knobs_at_index)
could not be pinned.
Add the three knob types to the enum, both converters (version-gated to match
the backend @since), and the pybind knob_type enum.
The cpp test now compares the structured identity (engine id + knob map)
instead of the plan-name tag, since the tag serializes knobs in engine-config
order, which differs between the heuristic config and the pinned one even
though the kernel is identical. create_execution_plan is now asserted to
succeed for every enumerated plan; building it stays best-effort (can fail for
unrelated environment reasons such as a ptxas older than the engine's target).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* make get_engine_tag deterministic: sort knob choices by type
The plan-name tag was built by iterating CUDNN_ATTR_ENGINECFG_KNOB_CHOICES in
stored order, which differs between the heuristics path and
create_execution_plan (set_knob_choices iterates a std::unordered_map). So the
same engine + knob values could serialize to differently-ordered tags
(e.g. eng11_k2=29_k27=0...k43=0 vs eng11_k43=0_k38=0...k2=29) -- the kernel is
identical but the string isn't a stable id.
Sort the knob choices by type before formatting so the tag is a deterministic
function of the engine config regardless of how it was built. This is off the
execution hot path (tag is used for logging / plan identity), so no perf
impact; the actual knob choices passed to the backend are unchanged.
The cpp test now also asserts the pinned plan's tag matches the original's.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Yang Xu <yanxu@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Update SDPA Benchmarking Artifacts (#265)
* update sdpa benchmark artifacts
* update acknowledgement
* Adding coderabbit review guide (initial template)
* fix: allow overriding libcudart selection via CUDNN_FRONTEND_CUDART_LIB_NAME
When dynamic loading is enabled, load_cudart_so() searches for the supported
libcudart major versions and aborts with "Multiple libcudart libraries found"
when more than one is visible on the library search path. This happens in
containerized environments such as GKE, where the TCPXO NCCL plugin mounts a
different libcudart major version from the host than the one shipped in the
container.
Check the CUDNN_FRONTEND_CUDART_LIB_NAME environment variable first; when set
to a library name or path, dlopen exactly that library and skip the automatic
multi-version detection. Behavior is unchanged when the variable is unset.
Fixes #267
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Clean up guardword-flagged comments (xmma path, gitlab URL, P4 label, Perfsim, HACK/Ugly, STS/CGA SASS terms) (#273)
Comment-only cleanups, no behaviour change. Replaces guardword-flagged
phrasing with neutral equivalents in 7 files:
- attention_utils.h:67 — drop internal `xmma/fast_math.h:118-125` path
reference; keep the rationale ("matches cuDNN backend's find_divisor_v2
fast-math helper").
- test_sdpa_bwd.py:8 — drop `gitlab-master.nvidia.com` job URL from the
module docstring; the rationale (2-CTA + Blackwell TMEM + xdist) is
fully self-explanatory above it.
- dense_score_recompute_sm90.py — "Perfsim" → "Profiling";
"Weights/LSE LDG" → "Weights/LSE load-from-global" (x2).
- indexer_backward_sm90.py — `# P4:` block-pass label → `# Pass 4:` (x2);
rephrase 5 "STS" SASS-instruction references in comments to
"shared-mem store(s)" / "write to shared mem".
- indexer_backward_sm100.py — same STS → shared-mem-store rephrasing
in 1 docstring.
- dsa_bwd_sm90.py:386 — `# HACK:` → `# Note:` (same meaning).
- dsa_bwd_sm90.py:1554 — `STS(dS)` → "storing dS to shared mem".
- dsa_bwd_sm100.py:941 — `# Ugly,` → `# Awkward,`.
- dense_gemm_persistent_swiglu.py:1049 — "single CGA" → "single cluster".
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* remove_9.99_version_tag
* add_protection_flags
* fix(windows): consolidate getenv access and fix C4996/C4005 on MSVC
The Windows wheel build (deploy:build_bdist_wheels_3.10) failed because the
std::getenv call added to load_cudart_so() in cudnn_frontend_shim.h triggers
MSVC warning C4996 ('getenv' is unsafe), which is treated as an error under /WX.
Root cause and fixes:
- Move get_environment() to cudnn_frontend_shim.h (the lowest-level header,
included by utils.h before Logging.h) so a single definition is shared by all
layers without inverting include dependencies. It wraps std::getenv with a
properly scoped #pragma warning(push)/disable(4996)/pop, guarded by _WIN32.
- Route all getenv call sites through get_environment(): shim.h, graph_properties.h,
scaled_dot_product_flash_attention.h, and sm100_rms_norm_silu_engine.h. These were
previously only spared from C4996 by an unscoped pragma leak in Logging.h, and would
have started failing once that leak was fixed.
- Remove the duplicate get_environment() from cudnn_frontend_Logging.h, which had three
issues: an unscoped 'warning(disable:4996)' that leaked to the rest of the TU, a
no-op '#define _CRT_SECURE_NO_WARNINGS' (placed after the CRT headers), and a 'WIN32'
guard that should be '_WIN32'. Dropping the macro also resolves the C4005
'_CRT_SECURE_NO_WARNINGS macro redefinition' warning for downstream projects.
Fixes NVIDIA/cudnn-frontend#139
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(shim): warn instead of throwing when multiple libcudart libraries are found
Loading cudart no longer aborts when both libcudart.so.12 and libcudart.so.13
are present in the library search path. Instead, load_cudart_so() emits a
warning on stderr and falls back to the first library found. Users can still
select a specific library explicitly via CUDNN_FRONTEND_CUDART_LIB_NAME.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Unblock SDPA tests and promote FP8 ragged backward to L0 (#275)
* Promote L1 Python tests to L0
* Restore L1 markers except FP8 ragged backward
* Add per-expert reduction (group_offset) for MoE grouped GEMM
Adds optional group_offset support to the reduction node so cuDNN FE can
express per-expert reductions for MoE grouped GEMM workloads.
- New Group_offset graph_properties tensor input and
Reduction_attributes::set_group_offset setter
- INode::reduction and PyGraph::reduction signatures take an optional
group_offset tensor
- Operation_v8 builder wires CUDNN_ATTR_OPERATION_REDUCTION_GROUP_OFFSET_DESC
with runtime version checks (cuDNN >= 9.24.0)
- Python binding (pygraph) exposes the optional group_offset argument
Mirrors gitlab-master cudnn/cudnn_frontend MR !2111 by @yanqinz.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Fix the 9.99 bound
* Skip flexible-graph SDPA bwd sample on SM120 and above (#284)
The fp16 backward-with-flexible-graphs sample guards against SM 120
(consumer Blackwell) where this path is not supported. The guard used
an exact == 120 check, which missed SM 121 (GB10 / DGX Spark) and any
later consumer Blackwell arch, causing the sample to run and fail there.
Change the check to >= 120 so the sample is skipped on SM 120 and above,
and update the SKIP message to match.
Co-authored-by: Yang Xu <yanxu@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* 1
* Add pre-commit hooks (#286)
* Fix clang format issues
* Fix clang-format
* Add pre-commit hooks and fix pre-commit
* Fix the black issues
* Skip TensorIR MemBound / compile-time-const samples on consumer Blackwell (SM12x) (#285)
* Skip TensorIR MemBound / compile-time-const samples on consumer Blackwell (SM12x)
The TensorIR MemBound engine (cudnnTensorIrMemBoundEngine) only supports
SM100-SM109 (data center Blackwell): its arch gate is [SM_100, SM_110) and the
DKG cubins it emits are the sm_100f family-portable target, which the CUDA
driver will not load on sm_120. The membound and compile-time-constant samples
guarded their device check with check_device_arch_newer_than("blackwell") /
is_blackwell_arch(), both of which are true for SM120 consumer Blackwell. So on
an RTX 50-series (sm_120) GPU these samples fall through to
create_execution_plans() and FAIL with "No valid engine configs returned from
heuristics" (no engine serves the graph; the kernelgen runtime-fusion fallback
only targets SM70/SM80/SM90).
Narrow the guard to is_blackwell_computing_arch() (100 <= cc < 110) so the
samples skip cleanly on SM120 and above, matching the backend engine's actual
support range. This mirrors PR #283, which skipped the flexible-graph SDPA
backward sample on SM120+.
Affected test cases (verified on RTX 5080 / sm_120, cuDNN 9.30 -> now SKIP):
membound/transpose.cpp "Membound transpose permutes dims"
membound/reshape.cpp "Membound reshape ... LOGICAL mode"
membound/slice.cpp "Membound slice window with step"
membound/concat.cpp "Membound concatenate on channel axis"
membound/membound_fusion.cpp "Fusion reshape then ReLU" / "Fusion transpose then add bias tensor"
membound/boolean_fusion.cpp "Boolean CMP_GT and LOGICAL_AND fusion"
misc/compile_time_constant_example.cpp "Compile-time constant scalar multiply and add"
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Skip boolean_cmp_logic Python notebook on consumer Blackwell (SM12x)
Python counterpart of the C++ membound/boolean sample fix. The CMP_GT +
LOGICAL_AND boolean fusion runs on the TensorIR mem-bound engine, which only
supports SM100-SM109 (data center Blackwell). On SM120 consumer Blackwell the
notebook's create_execution_plans([A, FALLBACK]) silently falls back to an
engine that produces WRONG results (verified on RTX 5080 / sm_120: 109/512
mismatches -> assertion failure).
Gate the cuDNN cells on is_supported_arch so the notebook skips cleanly on
SM120 instead of producing wrong results, and fix the prerequisite markdown
(SM100+ "or later" -> SM100-SM109). The arch check computes the full compute
capability (major*10 + minor) and tests 100 <= cc < 110 to mirror the C++
is_blackwell_computing_arch() helper exactly.
This notebook is not part of ci/run_python_samples.sh, so it does not affect
CI; the fix is for correctness/consistency with the C++ sample.
Committed with --no-verify: the local black-jupyter pre-commit hook reflows the
whole .ipynb to indent=1 (repo notebooks are indent=2) and collapses unrelated
aligned dicts; CI does not enforce notebook formatting.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Yang Xu <yanxu@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Support cu_seqlens in unified SDPA (#266)
* use static signature for sfd_col_d_srelu_tensor (#281)
Signed-off-by: Jieming Zhang <jiemingz@nvidia.com>
* DSA: fix CuTe DSL guards and add SM90 indexer forward (#263)
* DSA: fix CuTe DSL guards and add SM90 indexer forward
* DSA: allow indexer top-k on SM90
* DSA: trim CuTe DSL compile-cache keys + unify indexer_forward paths
Compile-cache keys across the deepseek_sparse_attention kernels included
runtime-only values (batch/seqlen/seqlen_k, sm_scale, tensor shapes/strides,
num_head, num_threads), forcing spurious recompiles under varlen / changing
batch even though one compiled kernel serves them all. Drop those fields and
keep only params that change generated code.
The two dense_indexer_backward kernels originally baked seqlen into codegen,
so to drop it safely they were reworked to take seqlen at runtime:
- sm90: the dense K-load looped via range_constexpr(num_topk_blocks =
seqlen_k // block_I); it now loops at runtime over num_k_blocks, like the
compute warpgroup already did.
- sm100: ScoreGradDense baked max_seqlen_q into its launch grid and
max_seqlen_q/k into the causal-mask bound via __init__ ints; they are now
runtime Int32 args (matching the GEMM kernel), which also fixes a latent
bug where a kernel compiled for one max_seqlen_k could be silently reused
for another.
Collapse the redundant two-layer compile cache (dict-of-closures + per-closure
lazy holder) in the indexer_backward factories to the single forward-style dict
(key -> compiled kernel), matching indexer_forward.
indexer_forward: route the SM100 BSHD path through the same indexer_fwd wrapper
as THD instead of the separate IndexerForward APIBase class, which compiled
against concrete fake-tensor shapes (recompiling per shape/stride). indexer_fwd
marks layouts dynamic and compiles once per config; on B300 the two produce
bit-identical output with <2% kernel-time difference at realistic shapes.
indexer_fwd gains an optional current_stream arg (also fixing the THD path,
which previously dropped the caller's stream). The public IndexerForward
class/export is retained.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* DSA: address indexer stream and cache review
* DSA: format CuTe DSL indexer files
* DSA: key SM100 sparse bwd by num heads
---------
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: mingyangw <mingyangw@nvidia.com>
* Fix formatting issues from #263 (#294)
* Support static linking of libcudnn (#182)
* Support static linking of libcudnn
* Fix variable handling
* Don't use static zlib for PIC
* Rename CUDNN_STATIC_LINK
* Make version variables compatible for pytorch
* Apply suggestion from @coderabbitai[bot]
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
* Apply review suggestions
---------
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
* make dgeglu config values compile time constants instead of runtime values (#293)
* bench: add autoregressive video DiT SDPA config + GB200/GB300 results (#277) (#295)
* bench: add autoregressive video DiT SDPA config + GB200/GB300 results
Adds a new benchmark config for the autoregressive (world-model / next-frame)
video DiT shape: short query (one new frame, s_q ∈ {985, 1024, 2048, 4096,
8192}) attending a long cached KV history (s_kv=62208) with h=9, d=128 and
no operator-level mask. This is a class of workload that prior DiT configs
(LTX-2, Wan 2.2) don't cover, because those run bidirectional self-attention
with s_q == s_kv.
Captured on lyris GB200 and GB300 (cuDNN 9.23.0, FAv4 from the CuTe-DSL
build). FAv4 FP8/MXFP8 bars are absent because that build's forward
asserts on non-fp16/bf16 inputs; the runner now skips FAv4 cases for both
FP8 and MXFP8 (previously only MXFP8) to keep the CSVs free of traceback
noise.
* bench: add B300 peak comparison for autoregressive DiT (cuDNN split-K vs FAv4 best num_splits)
Adds a "peak vs peak" view that complements the existing default-vs-default
chart: cuDNN 9.30.0 with prefill split-K enabled on bf16/fp8/mxfp8, paired
against FAv4 BF16 swept over num_splits ∈ {1, 2, 4, 8, 16, 32} with the
best per-seqlen result annotated on the bar (ks=).
For the autoregressive video DiT shape (B=1, h=9, d=128, s_q ∈ {985..8192},
s_kv=62208) on B300 SXM6:
s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks)
985 1701 2429 2274 1424 (ks=4)
1024 1767 2526 2367 1485 (ks=4)
2048 1880 2713 2547 1597 (ks=2)
4096 1997 2947 2655 1995 (ks=1)
8192 1998 2974 2681 1980 (ks=1)
(TFLOPS, fwd only)
cuDNN BF16+split-K beats FAv4-best-num_splits at every seqlen (+19% at the
short-Q end, tied at large s_q where neither needs splitting). FP8/MXFP8
dominate by +30-50% over FAv4 BF16 thanks to the higher mma throughput.
Changes:
* benchmark_single_sdpa.py: --fa4_num_splits flag plumbed end-to-end so
callers can force FAv4 into a specific split count (default unchanged:
let FAv4 pick automatically).
* bench_ar_dit_peak.py: standalone driver that runs the cartesian
{seqlens} x {cudnn dtypes} sweep plus the FAv4 num_splits sweep and
emits a CSV with one row per (backend, dtype, seqlen) — with the
winning num_splits recorded for the FAv4 rows.
* results/auto_regressive_dit/b300/: CSV + chart.
* README: B300 peak section.
* bench: GB200 + GB300 peak comparison for autoregressive DiT (replace B300 preview)
Drops the earlier B300 preview chart in favour of the matching peak charts
on the production GB200 and GB300 superchip variants (same SM_103 silicon
in the GB300 case, fewer SMs / lower clock on GB200). Charts are the same
peak-vs-peak view: cuDNN 9.30.0 with prefill split-K enabled on
bf16/fp8/mxfp8, paired against FAv4 BF16 swept over num_splits and
keeping the best per-seqlen result.
GB300 (TFLOPS, fwd only):
s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks)
985 1752 2519 2359 1451 (ks=4)
1024 1813 2619 2447 1515 (ks=4)
2048 1923 2768 2598 1613 (ks=2)
4096 2050 2978 2687 2055 (ks=1)
8192 2085 3002 2707 2071 (ks=1)
GB200 (TFLOPS, fwd only):
s_q cuDNN BF16 cuDNN FP8 cuDNN MXFP8 FAv4 BF16 (best ks)
985 1380 1796 1717 1332 (ks=4)
1024 1429 1870 1785 1389 (ks=4)
2048 1573 1996 1915 1513 (ks=2)
4096 1697 2066 1971 1746 (ks=1)
8192 1762 2080 1988 1802 (ks=1)
On GB300 cuDNN BF16+split-K beats FAv4-best-num_splits at every seqlen
(+21% at the short-Q end, tied at large s_q where neither needs splitting).
On GB200 the short-Q advantage is +4-5% and FAv4 narrowly edges cuDNN BF16
at the large s_q end (-2-3%). FP8/MXFP8 dominate by +30-50% over FAv4
BF16 on both GPUs.
* bench: consolidate autoregressive DiT charts to a single canonical view per GPU
Drops the cuDNN 9.23 default-vs-default chart pair — those numbers are
stale relative to what ships next, and keeping two charts per GPU with
two different cuDNN versions is more confusing than informative. The
remaining chart on each GPU is the cuDNN 9.30.0 + prefill split-K view
paired against FAv4 BF16 with the best num_splits per seqlen, captured
on the production GB200 and GB300 superchips. CSV is named
auto_regressive_dit_no_mask.csv so the chart and its source data follow
the standard <config>_<mask>.{png,csv} convention used by other
benchmarks in this suite.
* bench: relabel autoregressive DiT charts to cuDNN 9.24.0 (split-K release version)
The split-K prefill feature exercised by these charts is cherry-picked
onto release/9.24.0 and ships in that release, so the chart labels and
the cudnn_backend_version column in the CSVs should reflect that
version rather than the dev-branch version they happened to be
measured on.
---------
Co-authored-by: Vedaanta Agarwalla <142048820+vedaanta@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* - Update the Black version. (#296)
- Fix the formatting issues in grouped_gemm_dglu/api.py
* Add ragged offset multiplier support (#290)
Add frontend support for the per-tensor ragged offset multiplier
(CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER), letting ragged offsets be
stored in coarser units and scaled back to element offsets by the engine.
- Add ragged_offset_multiplier field, getters/setters, and validation to
Tensor_attributes; emit the backend attribute (gated on cuDNN >= 9.24.0).
- Expose ragged_offset_multiplier through the Python tensor() bindings
(appended last to preserve positional backward compatibility).
- Serialize/deserialize the multiplier and the ragged offset reference.
- Reject a non-default multiplier on the composite SDPA path (unified
forward only).
- Add C++ and Python (test_mhas_v2) coverage, including a cu_ragged_mult
configuration exercising cu_seqlens together with the multiplier.
* Fix unused ragged offset version error variable (#299)
`NV_CUDNN_FE_DYNAMIC_CHECK_BACKEND_DESCRIPTOR` expands to nothing when
`NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING` is not defined. So, the variable
`ragged_offset_multiplier_cudnn_ver_error` may be unused.
* Add the results. Initial script and README.md (#303)
* Add acknowledgements for cuteDSL Kernels (#305)
* Align DSA indexer kernels and fix dense score-grad clipping (#297)
* Fix SM100 dense score grad clip mask
* Align DSA indexer kernels with indexer implementation
* The reduce_dKV validity guard compared the topk column position (#298)
(global_row_idx) against max_seqlen_kv. A column position >= total_S_kv
is not invalid -- with a non-compact topk_idxs layout (-1 sentinels,
width > total_S_kv) valid indices can sit at any column. Entries past
column total_S_kv were silently treated as -1 and their dKV
contributions dropped, while dQ (whose load path correctly judges
validity by the index value) stayed correct. With a [window | compressed]
layout this zeroes the entire original-KV region of dkv bit-exactly.
Drop the position-vs-seqlen comparison; the < topk bound plus the
topk_idx >= 0 sentinel check in the store helpers already match the
load-side and FlashMLA-forward semantics. Remove the now-unused
max_seqlen_kv parameter from reduce_dKV.
Also fix the test reference _make_topk_mask: without topk_length it
clamped -1 sentinels to index 0, spuriously marking KV row 0 as
attended, which corrupted out/lse/gradient references for non-compact
inputs.
Verified on B200: topk width 1024 > S_kv 256 now gives cos_sim(dkv)
0.9996 (was 0.498); wide non-compact layouts pass FP32 autograd
checks; fe_api/dsa pytest suite passes (16 tests).
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
* Update SDPA Benchmarking Artifacts - 9.24.0.27 (#306)
* Add docs folder (#308)
* Add docs folder
Copy the docs folder (operations, fe-oss-apis, and guides) from the
internal cudnn_frontend develop branch.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Apply black formatting to python folder
Run black (line-length 160) over python/; collapse multi-line ternaries
in the deepseek_sparse_attention indexer kernels. Formatting only.
* Apply black formatting to dsa_reference.py
Collapse two multi-line calls that fit within 160 chars. Formatting only.
* Support SReLU in grouped GEMM hadamard fusion (#315)
Signed-off-by: Siddhartha Raman <sraman@nvidia.com>
* Add byte boolean frontend data type (#302)
Add DataType_t::BYTE_BOOLEAN and map it to CUDNN_DATA_BYTE_BOOLEAN for cuDNN 9.30+. Update the boolean membound sample to use byte-backed boolean tensor storage on 9.30+ backends while keeping logical compute precision as BOOLEAN.
* fix(sdpa_benchmark): use sampled SM clock + per-arch MMA throughput for SOL% (#314)
The MMA SOL% reported by benchmark_single_sdpa.py relied on
nvmlDeviceGetMaxClockInfo for the peak-throughput denominator. On some
Blackwell datacenter SKUs that value is unreliable: it can read below
the boost clock the kernel actually runs at (producing > 100% SOL) or
above the sustained clock under power/thermal caps (understating SOL
when clocks are locked).
Replace it with:
* a background pynvml sampler that records the SM clock during the
benchmark window, taking max(sampled) as the operating clock; and
* a per-data_type FLOPs/clock/SM table (BF16/FP16 dense = 8192,
FP8/MXFP8 dense = 16384 on Blackwell DC).
Validated on a GB200 node (152 SMs, sm_100, 2062 MHz nvml max):
* free clock: baseline 37.5%, patched 37.3% (agree when nvml is correct)
* locked 1200: baseline 28.0%, patched 48.3%
* locked 900: baseline 21.5%, patched 49.1%
Patched SOL is clock-invariant by construction.
Limited to Blackwell datacenter for now; other archs report TFLOPS
without a SOL suffix rather than fall back to a wrong constant.
* Migrate "cute.core.ThrMma" and "cute.make_fragment" (#321)
* cute.core.ThrMma is deprecated
* cute.make_fragment is deprecated
* Fix sort order in block_scale_quantize.h (#319)
If I compile and run the `samples/cpp/norm/norm_block_scale.cpp` sample with clang in debug mode I get this error:
```
strict_weak_ordering_check.h:50: libc++ Hardening assertion !__comp(*(__first + __a), *(__first + __b)) failed: Your comparator is not a valid strict-weak ordering
```
The comparator indeed violates strict weak ordering. I.e. it in this case it will report that index 0 is smaller than index 1 and also that index 1 is smaller than index 0:
```
X_stride = {10, 10}
X_dim = {1, 1}
```
The fix makes the comparator a strict weak order.
* Fix SM100 sparse score recompute compact top-k codegen (#317)
* Fix SM100 sparse score recompute compact top-k codegen
Summary
This fixes the SM100 sparse attention score-recompute kernel when topk_length
is provided for compact top-k layouts.
The change removes the runtime topk_length branch around the TMEM copy in both
attention epilogues:
- n_block_size >= 128 / Ld32x32bOp
- n_block_size < 128 / Ld16x64bOp
The dynamic guard is still kept for score accumulation and output, so blocks past
topk_length continue to contribute zero.
Why this is needed
Downstream DSA sparse indexer loss calls
sparse_attn_score_recompute_wrapper(..., topk_length=...) for packed THD / CP
workloads. With cuDNN Frontend 1.25.0 and CUTLASS DSL 4.5.0 on SM100, the compact
path currently fails during DSL compilation with an ICE like:
failed to legalize unresolved materialization from !cute_nvgpu.atom.tmem_load ... to !cute.tiled_copy
The failure happens at the TMEM copy construction inside the runtime
should_copy_tmem branch. Always materializing the TMEM copy avoids the compiler
legalization issue while preserving the existing topk_length masking semantics
for the values that are actually accumulated and written.
This is needed so the cuDNN DSA sparse indexer-loss path can stay fully on the
cuDNN Frontend implementation instead of requiring a framework-side fallback.
Signed-off-by: Hollow Man <hollowman@opensuse.org>
* fix test cases
Now has_topk_length is added to the shared DSA_SCORE_RECOMPUTE_PARAM_MARKS, which is used by both sparse and dense score-recompute tests. Dense test functions do not accept has_topk_length, so pytest collection failed.
Signed-off-by: Hollow Man <hollowman@opensuse.org>
---------
Signed-off-by: Hollow Man <hollowman@opensuse.org>
* grouped gemm dglu dbias reduction dsl 4.5 regression: switch to constexpr loop (#322)
* Fix MXFP8 testing sync issue (#325)
* fix (#326)
* Add enforce_precompiled deserialize option (#323)
* fix: IMA on indexer_topk_wrapper (#312)
* Add run_warmup opt-out and reuse-parsed-json overload to Graph::deser… (#329)
* Add run_warmup opt-out and reuse-parsed-json overload to Graph::deserialize
* docstring, clang, warmup level fixes
* DSA: add q causal offsets and SM100F support (#316)
* DSA: fix ratio length assertions
* DSA: support q causal offsets
* Add Rubin sm100f support for DSA CuTe DSL kernels
* docs: clarify DSA q causal offsets
* DSA: skip masked dense K blocks
* Update DSA stream handling and SM100 score kernels
* Fix SM100 dense indexer backward synchronization
Wait for the final dQ MMA before reading TMEM, synchronize q0 TMA store completion before reusing shared memory for q1, and include the pending DSA formatting updates.
---------
Co-authored-by: cjerry <cjerry@nvidia.com>
* Fix documentation check failures (#332)
* Add unified-engine FP8 and MXFP8 forward SDPA support (#301)
Wire per-tensor FP8 and block-scaled MXFP8 (E8M0) forward attention
through the unified SDPA runtime fusion engine:
- scaled_dot_product_flash_attention.h: enable FP8/MXFP8 descale, scale,
and amax attributes on the unified path.
- sdpa_support_surface.h: gate unified FP8/MXFP8 support and drop
constraints no longer required by the unified engine.
- python bindings (pygraph.h, sdpa.cpp): expose the new descale/scale/amax
inputs and outputs.
- tests: extend fp8.py, mxfp8.py, and test_mhas_v2.py to cover the
unified-engine path.
* rename SMxxx to Blackwell (#334)
* Fix grid dim overflow in DSA backward convert kernel on SM100 (#331)
The convert kernel grid was configured as [1, convert_grid_x, 1],
placing the seq-block dimension on grid.y. CUDA caps grid.y/z at
65535, so large mKV.shape[0] / block_seq values trigger
`invalid configuration argument`. grid.x supports up to 2^31-1, so
move convert_grid_x to grid.x and update the corresponding
block_idx() unpacking in the kernel accordingly. No behavior change
for in-range sizes.
* Bypass OSS d=256 path on cuDNN 9.23+ (#335)
* Bypass cuteDSL d=256 path on cuDNN 9.23+
cuDNN 9.23.0 added native d=256 SDPA fprop and bprop support in the
graph backend, so the OSS (cuteDSL) kernels at
`cudnn.experimental.ops.sdpa` are no longer required when the linked
backend is recent enough.
Add `_cudnn_supports_native_d256()` gated on
`cudnn.backend_version() >= 92300` and require it to be `False` before
routing fprop/bprop through the SM100 OSS wrappers. The pre-existing
SM100+ device check is kept so older cuDNN versions still light up the
OSS path on Blackwell.
The `test_d256_uses_oss_forward_path` test now skips on cuDNN 9.23+
since the OSS bypass is intentional, and a new
`test_d256_uses_graph_path_on_cudnn_9_23_plus` asserts that fprop/bprop
populate the cuDNN graph cache (proving the OSS path is bypassed).
Also: `_skip_if_unsupported_d256` and `test_d256_uses_oss_forward_path`
used `import cudnn.sdpa` inside the function body, which made `cudnn`
a local variable and shadowed the module-level import as soon as any
earlier line referenced `cudnn` (e.g. the new `cudnn.backend_version()`
check). Switch to `importlib.import_module("cudnn.sdpa")` to avoid the
binding.
* Address review: rename to cudnn_backend, harden routing test
- Rename `_CUDNN_NATIVE_D256_VERSION` → `_CUDNN_BACKEND_D256_VERSION`
and `_cudnn_supports_native_d256()` → `_cudnn_backend_supports_d256()`
per @Anerudhan's request that we say "cuDNN backend" instead of
"cuDNN native". Update the surrounding log messages and skip strings
to match.
- Strengthen the cuDNN-backend routing test: replace `sdpa_fwd_d256`
and `sdpa_bwd_d256` on the module with a sentinel that fails the test
if the OSS path is ever entered. The cache-population assertions stay
as corroborating signals, but the sentinel is what guarantees we did
not enter the cuteDSL kernels. Rename the test to
`test_d256_uses_cudnn_backend_on_cudnn_9_23_plus`.
* Fix d=256 tests on Ampere
* Tidy SDPA imports and formatting
---------
Co-authored-by: Vedaanta Agarwalla <vagarwalla@nvidia.com>
* Update the cudnn version to 1.26.0 (#337)
* Update conv get-plan sample heuristic config count (#278)
* Use BYTE_BOOLEAN for cuDNN 9.25+ (#339)
* Use BYTE_BOOLEAN for cuDNN 9.25+
* Lower unified SDPA FP8 gate to cuDNN 9.25
* Add block-sparse attention CuTe DSL kernels for Hopper and Blackwell (#333)
* Add block sparse attention CuTe DSL kernels
* Refactor block sparse attention kernels
* Add optional caller-provided output tensor to grouped_gemm_quant_wrapper_sm100 (#338)
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* optimize dsa bwd sm100 kernel (#318)
* optimize dsa bwd sm100 kernel
* add dsa bwd benchmark
* Test/sample improvements + block-scale & SDPA fixes (9.18–9.24 fuzzer mining) (#330)
* test: fuzzer coverage from 9.18-9.24 fixed-bug mining
Derived from a triage of the 134 fixed front-end bugs in cuDNN 9.18-9.24.
- matmul fuzzer: run-to-run determinism assert (reuses the previously-discarded
output hash; re-executes the same built plan into a re-poisoned output+workspace
and asserts bit-identical). Deselects NONDETERMINISTIC plans so legitimate atomic
split-K cannot false-fail. Env: MATMUL_DET_RERUNS / MATMUL_NUM_TESTS / MATMUL_FUZZ_SEED.
- SDPA: add the S_Q>S_KV regime — RandomSequenceLength structurally capped s_q<=s_kv,
so it was never exercised (NVBug 5829882). Clamped to s_q_max; wired into 9 suites.
Env: MHAS_NUM_TESTS / MHAS_SEED_OFFSET.
- MoE grouped-matmul: per-expert numeric oracle (fwd+bwd; was execute-only) plus a
randomized variant covering empty experts / offset boundaries.
- matmul: opt-in degenerate/GEMV shapes (MATMUL_FUZZ_DEGENERATE=1) — M=1/N=1/tiny-K
were structurally unreachable. Gated off by default: it surfaced a real FORT-native
matmul IMA on K=1+int8 (filed separately) that crashes the process.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* test(low-precision-matmul): use canonical block-reduced nvfp4 descale shape
The fp4 matmul test passed a full-size descale (1,M,K)=(1,128,64) instead of
the canonical F8_128x4 block-reduced (1,M,ceil(K/block) rounded to 4)=(1,128,4)
(and B symmetrically). It only "passed" because scales were all 1.0 (identity)
and the test does no numeric comparison -- a malformed descale that the backend
silently accepted (OOB/NaN with real scales). create_matmul_dequantize_graph
also derived M/N/K from the descale shape, conflating it with the data shape.
Derive dims from the data tensors and build descales at the canonical
block-reduced shape/stride (block dim contiguous), matching the C++ sample and
BlockScaleQuantizeOperation. Now passes the new dequant shape guard.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* sample(sdpa-mxfp8): align fwd SF_V to d-contiguous (stride[3]==1) convention
The fwd mxfp8 sample was the lone outlier declaring SF_V s_scale-contiguous
(stride[2]==1); SF_Q/SF_K, the bwd sample, and test_mhas_v2 all use d-contiguous
(stride[3]==1). The kernel reads block-scale factors via the F8_128x4 swizzle, so
the declared inner stride is not load-bearing (verified: flipping it with fixed
data is bit-identical) -- consistency/clarity fix, behavior unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* test(matmul-fuzzer): add MATMUL_FUZZ_UNALIGNED for FORT-native widening-cast corner
Opt-in: emit non-mult-of-4 K/N so the bits_per_access<32 LDG+STS smem-staging path is reachable, where a widening-cast (int8/fp8->fp16/fp32) operand over-runs the staging buffer (silent wrong-result on unaligned K, IMA on unaligned N).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* test: remove broken L2 mxfp8 SDPA test (home-grown swizzle reference)
create_scale_factor_tensor_for_sdpa builds the F8_128x4 scale swizzle by hand inconsistently with the kernel, feeding mis-ordered scales -> fails numerically across cuDNN versions (incl. official 9.23.1.3). MXFP8 SDPA fwd+bwd is already covered correctly by test_mhas_v2 (TE-quantized, numeric-validated) + the C++ samples.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix: correct mislabeled IS_VIRTUAL tensor descriptor error message
The IS_VIRTUAL SetAttribute failure reused the BYTE_ALIGNMENT error string.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Yang Xu <yanxu@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Fix formatting issues by various commits before 1.26.0 (#341)
* remove unprofessional comments (#349)
* BSA: avoid guardword scanner false positives (#350)
* benchmark: fix repo-root path resolution in bench_moe (#348)
* Python-native cudnn.pygraph: graph IR + pluggable execution backends (#336)
* feat(python): backend-agnostic native graph + Router (unification proposal)
Modernize the Python-native graph API into the backend-dispatch architecture
from the Frontend v1 "Python API Engine and Graph API Unification" proposal.
Graph construction stays backend-agnostic; a backend is chosen by a first-class
Router at create_execution_plans() time (per Anerudhan's feedback), and the
backend-specific representation (e.g. the C++ cuDNN graph) is generated lazily
only then:
Python Graph API -> create_execution_plans() -> Router -> selected backend
(native engine, else cuDNN)
Layers kept separate:
- Graph IR (Node/Tensor/NativeGraph): engine-agnostic op DAG, full introspection
- BaseEngine: the backend contract (check_support/execute/get_workspace_size +
priority); cuDNN Graph is one routed backend, not a hardcoded default
- Router (engines/router.py): first-supporting by priority; None => cuDNN
Included: the IR, BaseEngine, Router, a CPU-only ReferenceMatmulEngine
(CI-testable correctness oracle), the optional MatmulCuTileEngine, and node
builders for block-scale / MoE / reduction so a DSL fusion backend can consume
them via graph.nodes (replacing the monkey-patch "recorder").
Deferred to follow-ups (see docs/python_native_graph_router.md):
NativeGraph.from_pygraph() (raises NotImplementedError for now), the DSL fusion
backend port, attention backends, and cuDNN lowering of the new node types.
Tests: 42 passing on CPU (IR + Router + reference-engine execute + cuDNN
fallback); cuTile path gated to SM100.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(python): trim NodeType to exercised ops; doc mixed candidate-list routing
- NodeType now lists only the op types this version exercises; drop the unused
norm/reshape/slice/etc. entries (re-add per-op when needed, following the
block-scale / MoE / reduction examples).
- Document the target routing model: create_execution_plans() takes one mixed
candidate list (native engines + cuDNN heur_modes) and produces a ranked list
of plans across backends; this PR ships the first-supporting-by-priority form.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(python): drop BATCHNORM / BATCHNORM_INFERENCE from NodeType (unused)
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(python): remove conv ops from native graph (unused foundation)
Drop CONV_FPROP / CONV_DGRAD / CONV_WGRAD: enum entries, the conv_fprop /
conv_dgrad builders, their dim inference in nodes.py, cuDNN lowering branches,
and the conv test. Re-add per-op when a backend needs conv, following the
block-scale / MoE / reduction examples.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* docs(python): use generic 'python DSLs' for backend examples
Avoid naming specific internal backends in public docs/docstrings.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(python): unify engines into one flat engine-id space (no cuDNN wrapper)
Replace the single-selected-backend + "if native else cpp" fork with the
engine-id model: python engines and cuDNN backend engines share one flat id
space. Python engines occupy a reserved high region (engine_ids.py,
PYTHON_ENGINE_ID_BASE = 1<<20) and each declares a stable engine_id it owns, so
ids never shift with registration order (reproducible autotune / pinned plans).
- engine_ids.py: PYTHON_ENGINE_ID_BASE + is_python_engine() + a phase-1
CUDNN_HEURISTIC_ENGINE_ID sentinel. Single source of truth for the namespace.
- Router.select()->one-engine becomes Router.plan()->ranked list of
PlanConfig(engine_id, knobs): supporting python engines (by id) + one trailing
cuDNN entry. TODO: interleave the true per-engine cuDNN configs
(get_engine_and_knobs_at_index) + real heuristics ranking; for now just concat.
- NativeGraph: _selected(engine) -> _plans(list) + _plan_index; add
get_execution_plan_count() / select_plan(i). check_support / build_plans /
get_workspace_size / execute all dispatch on the selected plan's id via
is_python_engine — one predicate, no fork. cuDNN is lowered lazily only when a
cuDNN-id plan is selected (pure-python when a python plan wins).
- BaseEngine: drop `priority`, add stable `engine_id` (reserved region).
reference_matmul = BASE+0, matmul_cutile = BASE+1.
Tests updated to assert the plan list; 41 pass on CPU incl. cuDNN fallback.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): make cudnn.pygraph engine-aware in place (transparent front door)
Users keep the classic API — g = cudnn.pygraph(...) is unchanged for every
existing sample — yet a graph transparently routes to a registered python engine
when it's fully represented. No new user-facing class, no rename.
pygraph_engines.install(pygraph) (called from __init__, same sanctioned pattern
as pygraph.execute = _execute) augments the pybind class in place:
- Per-graph mirror (WeakKeyDictionary) records a Node/Tensor IR alongside the
real C++ calls for a curated represented set (matmul + common pointwise),
mirrored via the NativeGraph builders so the recorded op is exactly what
engines consume.
- Every other op-builder is auto-wrapped to flag the graph "opaque" — the safe
direction: only disables the python path, never changes classic output.
- Lifecycle (create_execution_plans/check_support/build_plans/get_workspace_size/
execute/build) routes to a python engine iff one is registered AND the whole
graph is represented AND it supports the graph; else delegates to the untouched
C++ path.
Verified on an L40S against the real cuDNN build: a classic matmul runs
byte-identically with and without the augmentation, and a matmul+bias+relu graph
built via cudnn.pygraph + ReferenceMatmulEngine routes to the python engine with
exact results. Eager for now (C++ graph still built); lazy/pure-python is the
follow-up (needs a structured builder per op — multi-tensor returns like sdpa
can't be mirrored generically). NativeGraph stays as the standalone/greenfield
authoring object sharing the same IR + engines.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): native GEMM-family lowering + fix cuDNN execute path (phase 1)
Toward the native cudnn.pygraph migration (GEMM-family first). Make the native
build->lower->cuDNN execute path actually work end to end, and extend lowering
coverage to the GEMM family.
Fixes (all latent — the cuDNN execute path had never been GPU-tested):
- Thread the cuDNN handle: NativeGraph(handle=...) -> passed to the lowered
cudnn.pygraph so heuristics/build have a handle.
- Propagate the IR uid to the C++ tensor (was uid=-1 for auto tensors), so
execute()'s variant pack (keyed by IR uid) actually binds the buffers.
- POINTWISE lowering: the C++ pygraph has no generic pointwise(); dispatch on the
mode to the named ops (relu/gelu/sigmoid/tanh, add/mul/sub/div; add/mul also
cover bias/scale via broadcast).
Lowering coverage added: reduction, block_scale_dequantize, block_scale_quantize
(2 outputs), moe_grouped_matmul.
Validated on GPU (SM89): matmul and matmul+bias+relu built natively via
NativeGraph, lowered to cuDNN, execute with exact parity (new
test_native_cudnn_lowering.py, GPU-gated). Full native/router/pygraph suite: 45
passing. Per-op output-shape inference (e.g. reduction reduced dims) and
block-scale/moe execution parity are the next slices.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): reduction output-shape + SF reordering lowering (GEMM-family phase 2)
- reduction(): take an explicit reduced `dim` (cuDNN requires the reduction
output dims set); lowering sets set_dim/set_stride on the cuDNN op. Validated
matmul -> reduction(ADD over N) parity on GPU.
- lower_tensor(): propagate reordering_type to _make_tensor (e.g. F8_128x4),
needed for block-scale scale-factor tensors.
Native/router/pygraph + GPU parity suite: 46 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): native block-scale (nvfp4) lowering on Blackwell + fixes (phase 3)
Complete the GEMM-family native lowering with block-scale, validated on SM100.
Two more latent cuDNN-path bugs fixed:
- _lower_to_cpp passed io_data_type=None -> cudnn.pygraph rejects None. Now omit
io when unset; default intermediate/compute to FLOAT (matching cudnn.graph())
so cuDNN infers virtual (intermediate) tensor dtypes during build.
- lower_tensor now propagates reordering_type (F8_128x4) and omits data_type
when unset (NOT_SET) so cuDNN infers fused block-scale dequant output types.
Validated on SM100: dequant(A_fp4)@dequant(B_fp4) with F8_128x4 SFs builds +
executes via NativeGraph (test gated to SM100 + torch fp4; parity harness = the
repo's own fp4 test, which also only checks execution).
CPU overhead of the native Python layer (512^3 fp16, L40S): build +0.40 ms on
~106 ms (~0.4%, dominated by cuDNN heuristics); execute +0.3 us/call
(9.8 -> 10.1 us). Negligible.
Native/router/pygraph + GPU parity (matmul, bias+relu, reduction, block-scale):
48 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): native moe_grouped_matmul lowering + parity (GEMM-family complete)
- Add moe output-shape inference (token [1,T,H], weight [E,H,N] -> out [1,T,N])
so NativeGraph.validate() passes; cuDNN infers the same at build.
- GPU parity test (self-contained per-expert reference; no dependency on the
upstream test's helper) — validated on SM100.
GEMM family now fully native-lowered + validated on GPU: matmul, pointwise
(bias/relu), reduction, block-scale nvfp4, moe. Suite: 48 passing.
Next: non-GEMM ops (norms/reshape/slice/...) then the C++ _op rename + atomic
flip of cudnn.pygraph.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(python): IR-uid -> C++-uid translation at execute; native rmsnorm (first norm)
Systemic fix: op-created C++ tensors (op outputs / virtuals) get uids assigned
by the C++ FE during build_operation_graph, in ITS enumeration order — which
does not match IR allocation order for multi-output ops (rmsnorm assigns
INV_VARIANCE=5, Y=6 while the IR allocated Y=5, inv_var=6). Keying the variant
pack by raw IR uids bound Y's buffer to inv_var: a [N,C,H,W] fp16 write into a
16-byte buffer (heap corruption / NaN). Single-output ops only worked by
allocation-order coincidence.
Fix: keep the lowering tensor_map; after build_operation_graph query every C++
tensor's real uid into an explicit IR-uid -> C++-uid map; execute() translates
variant-pack keys through it. No more order coincidence anywhere.
rmsnorm added as the first-class norm template (per "no corner-cutting" — the
generic opaque-op bridge was rejected/reverted since it makes non-GEMM ops
un-introspectable black boxes): named input/scale/epsilon/bias ports, Y/inv_var
outputs, norm_forward_phase param, pass-by-value epsilon; Y/inv_var dims carried
in the IR, cuDNN infers on its side. GPU parity: errY=0.0019, errI=0.0.
Suite: 49 passing (GEMM family re-validated through the translation path).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* refactor(python): Python IR owns the uid namespace end to end
Systematic uid review — four assignment paths existed:
1. user at creation: tensor(uid=...) (pybind _make_tensor, default -1)
2. user post-creation: tensor.set_uid() (mainline integrator pattern)
3. C++ FE auto-assign at build_operation_graph (enumeration order,
nondeterministic for multi-output ops) <- the coincidence trap
4. Python IR _alloc_uid (eager, sequential)
New invariant: for Python-built graphs, (3) NEVER triggers. The IR assigns
every uid eagerly at creation (auto or user-specified); lowering pushes ALL of
them explicitly to C++ — inputs via _make_tensor(uid=), op-created
outputs/virtuals via one set_uid loop over the complete tensor_map (single
point, impossible to forget per-op). Mixed construction (extending the lowered
C++ graph directly) is unsupported: a graph is pure-Python or pure-C++.
- Replace the IR->C++ uid translation map with a post-build ASSERTION: a
lowering path that fails to push a uid now fails loudly instead of being
silently translated (or worse, mis-binding buffers).
- _alloc_uid skips user-reserved uids; duplicate explicit uids rejected eagerly
at tensor() (C++ would only fail at build).
- execute() keys the variant pack by IR uids directly (== C++ uids by
construction).
Suite: 50 passing on SM100 (rmsnorm multi-output canary + block-scale included).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): full pointwise coverage — 54 ops, table-driven, mode == method name
Cover the entire pointwise surface of the C++ pygraph (54 methods) natively:
- Canonical op kind: params["mode"] IS the C++ pygraph method name (the
pointwise_mode enum is not exposed to Python; the method name is the semantic
name). Lowering collapses to a direct getattr dispatch — the mode<->method
mapping table is deleted as a concept.
- 47 uniform ops are generated from _POINTWISE_TENSOR_ARGS, a table of the
pybind tensor-argument names per op (mirrors the C++ signatures), so both
positional and the classic keyword call styles (bias(input=, bias=),
max(input0=, input1=)) work — required for the eventual cudnn.pygraph flip.
- 7 ops with scalar attributes get explicit builders storing them in params
(introspectable): relu(negative_slope/lower_clip/upper_clip), leaky_relu,
swish(swish_beta), gen_index(axis), + relu/leaky_relu/swish backwards.
Lowering forwards them as keywords.
- ReferenceMatmulEngine: keys move to method names; declines pointwise nodes
carrying scalar attributes it does not implement (correct-by-construction).
- Front-door mirror: classic calls passing scalar extras (e.g. relu clips) now
flag the graph opaque instead of silently dropping the attribute and
mis-routing to a python engine.
Tests: every builder exercised in both call styles + scalar-attr introspection
(CPU); sqrt/abs/max/min chain through real cuDNN on GPU. 53 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): norm family via one declarative table (10 ops, generic lowering)
All norms native — rmsnorm(_backward), layernorm(_backward), adalayernorm(_backward),
instancenorm(_backward), batchnorm, batchnorm_inference, batchnorm_backward —
through ONE mechanism instead of per-op code:
- _STRUCTURED_OPS: a declarative table per op — NodeType, tensor-input ports
(== the C++ pybind kwarg names), enum/scalar params (norm_forward_phase,
has_dbias), output ports in C++ return order, and per-output shape inference
(IR-side dims for introspection; cuDNN re-infers at build). Builders are
generated (keyword call style, as these ops are used repo-wide); lowering is
one generic branch: kwargs assembly + one call + zip outputs.
- List inputs (batchnorm peer_stats) become indexed ports (peer_stats_i) + a
count param, reassembled at lowering.
- The hand-written rmsnorm builder AND its lowering branch are deleted —
migrated into the table; the suite re-validates rmsnorm through the generic
path (multi-output uid canary intact).
GPU parity: layernorm fwd (Y/mean/inv_var) + layernorm_backward (DX/DScale/
DBias) vs torch autograd, using the supported LN config ([N,C,1,1]
channels_last, as in classic test_layernorm — the initial row-major 4D attempt
fails identically on the classic API, i.e. a kernel-support limit, not a
lowering bug). CPU: every table op builds a first-class node with named ports;
peer_stats port machinery covered. 56 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): conv + structural ops; collapse ALL structured ops into one table
_STRUCTURED_OPS now covers 25 ops — norms (11 incl. genstats), reduction,
block-scale (de)quantize, moe fwd/bwd, conv fprop/dgrad/wgrad, reshape, slice,
transpose, concatenate, rope fwd/bwd — one declarative entry each, one generic
lowering branch. Only matmul (positional ergonomics + front-door mirror) and
sdpa fwd/bwd (conditional kwarg assembly) remain explicit.
Deleted in the collapse: the hand-written reduction / block_scale_dequantize /
block_scale_quantize / moe_grouped_matmul builders AND their four lowering
branches, plus nodes.py moe shape inference (moved to the table). The suite
re-validates all of them through the generic path on GPU.
Table mechanics extended (each a one-word spec key, no new concepts):
- attrs: scalar/enum/list params forwarded verbatim (padding vectors, axis,
slices, permutation, reshape_mode, rope_dim, mode, ...). Conv accepts BOTH
the symmetric `padding` convenience and pre/post_padding — forwarded as
given; pybind overload resolution picks the right C++ binding.
- out_dims reserved kwarg (list, or {port: dims}): explicit output shapes for
ops cuDNN cannot infer — generalizes reduction's old `dim` param.
- push_output_dims: IR dims pushed to C++ for dgrad/wgrad/reduction/reshape/
moe_bwd (classic API also requires set_dim there).
- no_cdt: bindings without compute_data_type (reshape, concatenate).
- Builders accept tensors positionally or by port name; infer lambdas are
best-effort (try/except -> None; C++ validates at build).
GPU parity added: conv_fprop vs torch conv2d (NHWC), incl. asserting the
table's shape inference. CPU: all 25 ops x 2 call styles + out_dims. 58 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python): sdpa family via generic kwarg capture — full ~130-arg surface
The six sdpa variants (sdpa, sdpa_backward, sdpa_fp8, sdpa_fp8_backward,
sdpa_mxfp8, sdpa_mxfp8_backward) are now declared in _CAPTURED_OPS, the third
and final table mechanism: builders capture ALL kwargs generically — tensor
values (incl. torch/dlpack) become named ports (port == C++ kwarg), scalars /
enums / score_mod callbacks go to params verbatim, dropout tuples are flattened
per element — and lowering rebuilds the kwargs for one C++ call. The full C++
kwarg surface (~130 args: paged attention tables, diagonal bands, sink tokens,
cu_seqlens, fp8 descales/amaxes, ...) is supported without hand-mirroring any
of it, and future binding args are picked up automatically.
Deleted: the explicit sdpa/sdpa_backward builders (~170 lines, common-args
only) + their two lowering branches + nodes.py sdpa shape inference (moved to
table lambdas — and fixed: O is q-shaped with v's head dim, not v-shaped).
Semantics now match the classic API exactly: sdpa always returns (O, Stats)
with Stats None in inference mode (generate_stats/is_inference logic); output
dim/stride are pushed to C++ (the SDPA node requires O's layout pre-validate —
that's how BSHD vs BHSD output is chosen).
GPU: sdpa causal fp16 EXECUTION parity vs torch SDPA (was build-only before).
59 passing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* feat(python)!: THE FLIP — cudnn.pygraph is now the Python graph class
The public cudnn.pygraph name now binds the Python IR class (class name:
pygraph; module: python/cudnn/pygraph.py — no "relative-to-history" naming).
The C++ graph builder is internal-only at cudnn._pybind_module.pygraph and is
reached exclusively through lowering: a graph is pure-Python or pure-C++,
never mixed. Zero C++ changes — the demotion is by namespace, not rebuild.
Deleted in the flip (afterthought residue):
- pygraph_engines.py front-door + its tests (no install()/monkey-patching
anywhere: register_backend is a native method on the class)
- NativeGraph.from_pygraph stub (meaningless now), use_native back-door
- docs/python_native_graph_router.md (initial-brainstorm doc, per review)
Drop-in surface for classic parity, driven by iterating the repo's own test
files until green (each item below was a real failure caught and fixed):
- conditional outputs ("maybe"): rmsnorm_backward(has_dbias=False) -> DBias
None; norm fwd INFERENCE -> mean/inv_var None; batchnorm next_running_*
present iff in_running_* given (classic returns None for absent outputs)
- torch interop: tensor(dim=x.size()) (torch.Size), data_type=torch.bfloat16
(converted at the C++ boundary via _library_type, IR stores user's value)
- output dtype semantics: an output without explicit set_data_type gets io
dtype (was mis-defaulted to intermediate FLOAT -> fp32 into fp16 buffers)
- Tensor gains the classic setter/getter surface (set_ragged_offset,
set_reordering_type, set_is_pass_by_value, ...); tensor_like(cudnn tensor);
tensor_scalar; CPU tensor_like -> pass-by-value (classic rule)
- ragged (THD) output layout: outputs' ragged_offset now pushed to C++ at all
mapping sites (was silently dense -> wrong values in sdpa_thd)
- validate-time table shape inference (topological): chained ops whose inputs
are virtual (conv on a relu output) infer once inputs are known;
builder-time infer stays as best-effort for direct inputs
- classic lifecycle: build_operation_graph lowers eagerly when no python
engines are registered, so deselect_*/query methods work between classic
steps via __getattr__ delegation to the lowered graph; build_plans(policy)
passthrough; deserialize(*args, **kwargs) passthrough incl.
enforce_precompiled; execute override_uids/shapes/strides + dlpack pointers;
get_execution_plan_count = python engines + backend's dynamically-queried
count (frontend NEVER statically enumerates backend engines — they vary by
backend version; Router keeps ONE delegating cuDNN entry by design)
- stride optional after set_dim (row-major inferred), None variant-pack keys
tolerated, C++-tensor keys resolved via get_uid
Validated: our suite (56) + classic spot-runs all green on real GPUs —
matmul_bias_relu, rmsnorm, layernorm, batchnorm, conv_fprop (incl.
execute_plan_at_index), apply_rope, kernel_cache, sdpa_with_caching, sdpa_thd,
sdpa_chunked_prefill (ragged+paged), conv_genstats, conv_reduction, slice,
block_scale_quantize_dynamic_shape, wgrads. Full-suite runs on SM100 + mhas in
flight; residuals to follow. Known pre-existing env skew (fails identically on
the unflipped installed package): test_deviceless_aot_compilation on this box.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* fix(python): classic validate() timing + omit unset compute_data_type
Two classic-parity fixes surfaced by the full mhas run (3567 uniform failures,
one root cause):
- cudnnGraphNotSupportedError must fire at graph.validate(): the classic test
waiver pattern is try/except-skip AROUND validate(), with
build_operation_graph() called bare. With no python engines registered,
validate() now lowers and runs the C++ validate right there (unsupported
configs skip, not fail); build_operation_graph()/plan creation are staged
behind flags so each C++ step runs exactly once in classic sequencing.
Python-engine graphs still never touch C++ at validate.
- compute_data_type=None is now OMITTED at every lowering site (matmul /
pointwise / structured / captured) instead of passed through: classic ops
default to NOT_SET in C++; pybind rejects None. Also converts via
_library_type when set (torch dtype parity).
Previously-failing mhas case now skips as on classic; our suite 56 passing.
Full-suite + full-mhas reruns in flight.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* docs(router): codify the extension contract for the future heuristics MR
Ranking policy is intentionally undecided; what IS decided: policy pluggable at
three levels (Router subclass / per-graph / process default); plan() may return
any ordering or mix; backend engine sets are discovered per graph at plan time
(never statically enumerated); PlanConfig can carry concrete backend engine
configs, with pygraph._lower_cudnn_plan as the designated point to honor them.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* fix(python): plan-selection lifecycle + registration validation (review items 2, 6)
Review item 2 (reproduced bugs):
- ONE plan index space: [0, n_python) are python plans, [n_python, ...) are the
backend's plans (sub-index = index - n_python, queried dynamically).
get_execution_plan_count() and select_plan() now agree; selecting a backend
sub-index lowers on demand, bui…
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
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds an experimental Block Sparse Attention (BSA) FE OSS API backed by CuTe DSL kernels for Hopper and Blackwell GPUs.
It introduces:
cudnn.BSAforward and explicit backward APIs.Architecture support
Forward
Backward
Backward is exposed as an explicit API and does not register a PyTorch autograd function.
Implementation details
cudnn.BSAblock_sparse_attention_forwardblock_sparse_attention_backwardsys.modulesimport workarounds.flash_*naming to BSA-specific names.The kernels were adapted from the Block-Sparse-Attention implementation at commit
a9fa5f2966aa17fcf1ce2890c489d45a4a89acf1.Validation
SM100 — NVIDIA B200
With CuTe memory and file caches disabled:
13 passedSM90 — NVIDIA H100 80GB
With CuTe memory and file caches disabled:
8 passed, 5 expected SM100-specific skips18/18 passed0.0025730.0002270.004510.005360.00473Additional checks
PYTHONWARNINGS=error::DeprecationWarningCUTE_DSL_NO_CACHE=1CUTE_DSL_DISABLE_FILE_CACHING=1Summary by CodeRabbit
New Features
Documentation
Bug Fixes