DSA: Add FP8/MXFP8 and compressed Top-K indexer paths - #370
Open
jiayus-nvidia wants to merge 7 commits into
Open
DSA: Add FP8/MXFP8 and compressed Top-K indexer paths#370jiayus-nvidia wants to merge 7 commits into
jiayus-nvidia wants to merge 7 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (5)
📝 WalkthroughWalkthroughAdds FP8/MXFP8 precision support across DSA indexer forward and dense score recompute, introduces a combined SM100 compressed-logits Top-K path, and updates wrappers, kernels, scale utilities, documentation, and tests. ChangesDSA precision and compressed Top-K forward
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant TopKWrapper
participant IndexerInterface
participant Stage1IndexerScore
participant Stage2TopK
Caller->>TopKWrapper: Call compressed indexer forward
TopKWrapper->>IndexerInterface: Dispatch compressed logits
IndexerInterface->>Stage1IndexerScore: Compute compact candidate logits
Stage1IndexerScore-->>IndexerInterface: Return candidate buffer and optional LSE
IndexerInterface->>Stage2TopK: Select Top-K and optional softmax
Stage2TopK-->>TopKWrapper: Return indices, logits, and softmax
TopKWrapper-->>Caller: Return TupleDict
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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.
🧹 Nitpick comments (5)
python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py (2)
151-224: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo dedicated unit tests for the new blockscaled descriptor-packing logic.
make_blockscaled_instr_desc/blockscaled_mma_op_to_idescimplement intricate bit-packing for the SM100 MXF8F6F4 idesc, but no test file analogous totest_DSA_mxfp8_scale_utils.pyexercises this packing directly (only indirectly through downstream kernel correctness tests, if any). A focused unit test asserting expected bit values for representative(M, N, sf_id, major)combinations would catch regressions independent of full kernel execution.As per path instructions for
python/cudnn/**, which say to check "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/deepseek_sparse_attention/utils/sm100/mma_desc.py` around lines 151 - 224, Add focused unit coverage for the new blockscaled descriptor packing by creating tests under test/python/fe_api for make_blockscaled_instr_desc and blockscaled_mma_op_to_idesc. Exercise representative M/N, sf_id, and major combinations, and assert the returned descriptor has the expected bit fields, so the SM100 MXF8F6F4 packing logic is validated directly instead of only through downstream kernel tests.
82-89: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFallback
getattr(..., None)could silently misclassify ifcutlass_typeis unexpectedlyNone.If none of the type attributes exist in the installed
cutlassmodule andcutlass_typeitself happens to beNone, the firstis Nonecomparison would incorrectly match and returnMXF8F6F4Format.E4M3instead of raising the intendedTypeError. This is unlikely in normal use (real dtype objects are always passed) but the redundant dual-name checks (FloatE4M3FN/Float8E4M3FN) suggest uncertainty about which name is available in the installed cutlass version.🛡️ Safer alternative
- if cutlass_type is getattr(cutlass, "FloatE4M3FN", None): - return MXF8F6F4Format.E4M3 - if cutlass_type is getattr(cutlass, "Float8E4M3FN", None): - return MXF8F6F4Format.E4M3 - if cutlass_type is getattr(cutlass, "FloatE5M2", None): - return MXF8F6F4Format.E5M2 - if cutlass_type is getattr(cutlass, "Float8E5M2", None): - return MXF8F6F4Format.E5M2 + e4m3_names = ("FloatE4M3FN", "Float8E4M3FN") + e5m2_names = ("FloatE5M2", "Float8E5M2") + if any(cutlass_type is getattr(cutlass, name, object()) for name in e4m3_names): + return MXF8F6F4Format.E4M3 + if any(cutlass_type is getattr(cutlass, name, object()) for name in e5m2_names): + return MXF8F6F4Format.E5M2🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py` around lines 82 - 89, The dtype-to-format mapping in the helper that checks `cutlass_type` against `FloatE4M3FN`, `Float8E4M3FN`, `FloatE5M2`, and `Float8E5M2` can misclassify a `None` input because `getattr(..., None)` makes the first comparison match accidentally. Update this conversion logic to avoid treating a missing cutlass symbol or a `None` dtype as a valid match, and keep the intended `TypeError` path in place for unsupported inputs while still handling the alternate cutlass type names.python/cudnn/deepseek_sparse_attention/utils/tensor_conversion.py (1)
15-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring doesn't cover the new parameter annotations.
to_cute_tensornow has full type hints but the docstring still only describes the return in one line;assumed_align,leading_dim,fully_dynamic,enable_tvm_ffi, anddivisibilityaren't documented.As per path instructions for
python/cudnn/**, which say to "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/deepseek_sparse_attention/utils/tensor_conversion.py` around lines 15 - 32, The to_cute_tensor docstring is missing documentation for its newly annotated parameters. Update the docstring in to_cute_tensor to describe assumed_align, leading_dim, fully_dynamic, enable_tvm_ffi, and divisibility, while keeping the return description accurate; use the existing function signature in tensor_conversion.py as the reference point.python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py (1)
52-98: 📐 Maintainability & Code Quality | 🔵 TrivialMissing test coverage for the generic FP8 epilogue and varlen+FP8 on SM90.
The new
test_DSA_indexer_forward_wrapper_fp8_sm90_matches_dequant_referencetest only exercisesqhead_per_kv_head=64with divisible shapes, which routes throughuse_fast_qh64_epilogue/use_unchecked_qh64_full/use_unchecked_qh64_masked(lines 95-97). The generic_epilogue_store_to_gmemFP8 path (qhead_per_kv_head=32, or shapes that don't satisfy the "unchecked" divisibility conditions) and the varlen (THD) + FP8 combination on SM90 appear untested intest/python/fe_api/dsa/.Given the amount of new FP8-specific branching (split-warpgroup producer/consumer,
rKScaleshared-memory scale application, LSE accumulation), consider adding at least one test withqhead_per_kv_head=32and one varlen+FP8 SM90 test.Based on path instructions: "Focus on whether there are test cases in test/python/fe_api" for
python/cudnn/**files.Also applies to: 369-535
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py` around lines 52 - 98, Add test coverage in test/python/fe_api for the SM90 FP8 forward path beyond the current qhead_per_kvhead=64 fast epilogue case: add one case that forces the generic _epilogue_store_to_gmem branch by using qhead_per_kvhead=32 or a shape that does not satisfy use_unchecked_qh64_full/use_unchecked_qh64_masked in IndexerFwdSM90, and add one varlen (THD) + FP8 SM90 test that exercises the is_varlen path with FP8 scales/LSE handling. Keep the existing reference-matching style used by test_DSA_indexer_forward_wrapper_fp8_sm90_matches_dequant_reference so the new cases verify the generic FP8 epilogue and varlen+FP8 behavior.Source: Path instructions
python/cudnn/deepseek_sparse_attention/indexer_forward/api.py (1)
209-215: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the newly added public parameters. The
indexer_forward_wrapperdocstring still only describes{'scores'}andq_causal_offsets, but the signature now exposesprecision,q_scale,k_scale,sf_vec_size,return_lse, andlse_out. It's also worth noting thatreturn_lse/lse_outare only honored on SM90 (SM100 raisesNotImplementedError) so callers aren't surprised. 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/deepseek_sparse_attention/indexer_forward/api.py` around lines 209 - 215, Update the indexer_forward_wrapper docstring to document all newly exposed public parameters: precision, q_scale, k_scale, sf_vec_size, return_lse, and lse_out, along with any expected behavior for the outputs. Also note in the wrapper docs that return_lse and lse_out are only supported on SM90, while SM100 raises NotImplementedError, so callers understand the hardware-specific limitation.Source: Path instructions
🤖 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.
Nitpick comments:
In `@python/cudnn/deepseek_sparse_attention/indexer_forward/api.py`:
- Around line 209-215: Update the indexer_forward_wrapper docstring to document
all newly exposed public parameters: precision, q_scale, k_scale, sf_vec_size,
return_lse, and lse_out, along with any expected behavior for the outputs. Also
note in the wrapper docs that return_lse and lse_out are only supported on SM90,
while SM100 raises NotImplementedError, so callers understand the
hardware-specific limitation.
In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py`:
- Around line 52-98: Add test coverage in test/python/fe_api for the SM90 FP8
forward path beyond the current qhead_per_kvhead=64 fast epilogue case: add one
case that forces the generic _epilogue_store_to_gmem branch by using
qhead_per_kvhead=32 or a shape that does not satisfy
use_unchecked_qh64_full/use_unchecked_qh64_masked in IndexerFwdSM90, and add one
varlen (THD) + FP8 SM90 test that exercises the is_varlen path with FP8
scales/LSE handling. Keep the existing reference-matching style used by
test_DSA_indexer_forward_wrapper_fp8_sm90_matches_dequant_reference so the new
cases verify the generic FP8 epilogue and varlen+FP8 behavior.
In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py`:
- Around line 151-224: Add focused unit coverage for the new blockscaled
descriptor packing by creating tests under test/python/fe_api for
make_blockscaled_instr_desc and blockscaled_mma_op_to_idesc. Exercise
representative M/N, sf_id, and major combinations, and assert the returned
descriptor has the expected bit fields, so the SM100 MXF8F6F4 packing logic is
validated directly instead of only through downstream kernel tests.
- Around line 82-89: The dtype-to-format mapping in the helper that checks
`cutlass_type` against `FloatE4M3FN`, `Float8E4M3FN`, `FloatE5M2`, and
`Float8E5M2` can misclassify a `None` input because `getattr(..., None)` makes
the first comparison match accidentally. Update this conversion logic to avoid
treating a missing cutlass symbol or a `None` dtype as a valid match, and keep
the intended `TypeError` path in place for unsupported inputs while still
handling the alternate cutlass type names.
In `@python/cudnn/deepseek_sparse_attention/utils/tensor_conversion.py`:
- Around line 15-32: The to_cute_tensor docstring is missing documentation for
its newly annotated parameters. Update the docstring in to_cute_tensor to
describe assumed_align, leading_dim, fully_dynamic, enable_tvm_ffi, and
divisibility, while keeping the return description accurate; use the existing
function signature in tensor_conversion.py as the reference point.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7d40af09-f3c0-42ba-80c9-1f1a4f4868cc
📒 Files selected for processing (31)
docs/fe-oss-apis/dsa.mdpython/cudnn/deepseek_sparse_attention/README.mdpython/cudnn/deepseek_sparse_attention/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_backward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_interface.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/api.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/compress_top_k_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/api.pypython/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/utils/runtime.pypython/cudnn/deepseek_sparse_attention/utils/sm100/gemm.pypython/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.pypython/cudnn/deepseek_sparse_attention/utils/sm100/mxfp8_scale_utils.pypython/cudnn/deepseek_sparse_attention/utils/tensor_conversion.pytest/python/fe_api/dsa/dsa_reference.pytest/python/fe_api/dsa/dsa_utils.pytest/python/fe_api/dsa/test_DSA_dense_score_recompute.pytest/python/fe_api/dsa/test_DSA_indexer_forward.pytest/python/fe_api/dsa/test_DSA_indexer_forward_compressed.pytest/python/fe_api/dsa/test_DSA_mxfp8_scale_utils.py
Collaborator
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-370-8c3ea50 |
jiayus-nvidia
force-pushed
the
jiayus/import-indexer-fp8-support
branch
from
July 13, 2026 06:34
6bda0b9 to
d7579a8
Compare
3 tasks
jiayus-nvidia
marked this pull request as draft
July 14, 2026 07:54
Contributor
Author
|
WIP for adding deterministic and e2e test |
jiayus-nvidia
force-pushed
the
jiayus/import-indexer-fp8-support
branch
from
July 24, 2026 02:58
0821bf2 to
061e6cd
Compare
Share the SM100 unified score kernels between indexer forward and dense recompute, add the SM90 FP8 path, and port the MXFP8 scale helpers and coverage. Keep compressed-logits/top-k support out of scope.
Port the SM100 compact-logits path and fold in the latest indexer optimizations: fused Top-K softmax, THD MXFP8, BF16/MXFP8 LSE for BSHD and THD, caller-owned output buffers, MQA validation, and the backward softmax fast path. Remove the superseded decode KV-split and partial-LSE merge path.
Port /code/indexer commit b731ca5 to the cudnn-frontend DSA layout.
jiayus-nvidia
force-pushed
the
jiayus/import-indexer-fp8-support
branch
from
July 29, 2026 09:38
7e07396 to
b03ff43
Compare
jiayus-nvidia
marked this pull request as ready for review
July 29, 2026 09:49
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (10)
python/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.py (2)
996-999: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDrop the unused participants and
neg_log2_e.
K_consumer/KScale_consumerare never used — the MMA warp drives those pipelines throughpipeline_K.consumer_wait(K_mma_state)/pipeline_KScale.consumer_wait(KScale_mma_state)directly (Lines 1400, 1421).neg_log2_eat Line 998 is also dead; the epilogues compute their ownlog2_e. Ruff flags the first two as RUF059.♻️ Proposed cleanup
- neg_log2_e = Float32(-math.log2(math.e)) -Q_producer, Q_consumer = pipeline_Q.make_participants() - K_producer, K_consumer = pipeline_K.make_participants() - KScale_producer, KScale_consumer = pipeline_KScale.make_participants() + K_producer, _K_consumer = pipeline_K.make_participants() + KScale_producer, _KScale_consumer = pipeline_KScale.make_participants()Also applies to: 1146-1148
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.py` around lines 996 - 999, Remove the unused K_consumer and KScale_consumer participants from the relevant setup, and delete the unused neg_log2_e assignment near warp_idx and tidx in the unified score recompute flow. Preserve the direct pipeline_K.consumer_wait and pipeline_KScale.consumer_wait calls using their existing MMA states.Source: Linters/SAST tools
502-527: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
max_seqlen_kis threaded through both epilogue helpers but never read.Both signatures accept
max_seqlen_kand every call site passes it (Lines 1534, 1559, 1628, 1653), yet neither body uses it — masking usesseqlen_kandcol_limit. If it is kept only for signature parity with the BF16 base epilogues, a short comment saying so avoids a reader assuming padded-extent masking happens somewhere here; otherwise drop it.Also applies to: 723-748
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.py` around lines 502 - 527, Remove the unused max_seqlen_k parameter from both dense single-query epilogue helpers and update all call sites accordingly, unless signature parity with the BF16 base epilogues is required. If retaining it, add a concise comment in each helper clarifying that masking uses seqlen_k and col_limit rather than max_seqlen_k.python/cudnn/deepseek_sparse_attention/utils/sm100/mxfp8_scale_utils.py (1)
196-229: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePadded scale rows are left as zero E8M0.
outis zero-initialized and only[q_scale0, q_scale0 + q_len*qhpkv)is written, so padded rows carry E8M00x00(= 2^-127). That is only safe because the kernel masks those packed-M rows out; if any future tile path reads them, the tiny-but-nonzero factor multiplies garbage FP8 data rather than being ignored. Worth a short comment recording that invariant here so the zero-fill is not mistaken for "unused".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mxfp8_scale_utils.py` around lines 196 - 229, The zero-initialized padding in the output of the scale-packing function should be documented as intentionally retained E8M0 zero values. Add a short comment adjacent to the `torch.zeros` initialization or the per-batch write explaining that padded rows remain zero and are safe only because downstream kernels mask them out; do not change the allocation or packing behavior.python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py (2)
979-991: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRuff RUF005: build the tuple with unpacking.
♻️ Proposed fix
- base_compile_key = ( + compile_key = ( + "dense_indexer", q.dtype, head_dim, qhead_per_kv_head, m_block_size, n_block_size, k_block_size, kv_stage, ratio, is_varlen, q_causal_offsets is not None, ) - compile_key = ("dense_indexer",) + base_compile_key🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py` around lines 979 - 991, Update the compile_key construction in the surrounding score-recompute interface to use tuple unpacking when prefixing "dense_indexer" to base_compile_key, preserving the existing key contents and ordering.Source: Linters/SAST tools
1485-1499: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePublic docstrings weren't updated for the new precision/MXFP8 parameters. Both public entry points gained
precision,q_scale,k_scale,cu_seqlens_q_scale_padded,cu_seqlens_k_scale_padded,sf_vec_size, andk_block_size, but their Args blocks still describe only the BF16 contract.
python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py#L1485-L1499: document the new params indense_indexer_score_recomputeand note that Q/K are FP8 (float8_e4m3fn) whenprecision="mxfp8", with packed E8M0 scale shape expectations.python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py#L1567-L1581: same additions fordense_attn_score_recompute, plus the MXFP8outprefill semantics (zeros rather than-inf).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py` around lines 1485 - 1499, Update the public Args docstrings for dense_indexer_score_recompute at python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py:1485-1499 and dense_attn_score_recompute at :1567-1581 to document precision, q_scale, k_scale, cu_seqlens_q_scale_padded, cu_seqlens_k_scale_padded, sf_vec_size, and k_block_size, including packed E8M0 scale shape expectations. State that precision="mxfp8" uses float8_e4m3fn Q/K, and additionally document for dense_attn_score_recompute that MXFP8 out is prefilled with zeros rather than -inf.test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py (1)
725-725: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUnused
(s_q, s_k)unpacking in both per-batch loops (Ruff B007). Both loops derive their bounds fromcu_q_host/cu_k_host, so the unpacked shape values are dead.
test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py#L725-L725: replacefor batch, (s_q, s_k) in enumerate(shapes):withfor batch in range(len(shapes)):.test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py#L831-L831: apply the same change to this loop.🤖 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/dsa/test_DSA_indexer_forward_compressed.py` at line 725, The two per-batch loops in test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py at lines 725-725 and 831-831 unnecessarily unpack unused shape values. Update both loops to iterate over batch indices using the length of shapes, while preserving their existing bodies and bounds derived from cu_q_host and cu_k_host.Source: Linters/SAST tools
python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py (1)
11-17: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to satisfy Ruff RUF022.♻️ Proposed fix
__all__ = [ "IndexerForward", - "indexer_forward_wrapper", - "indexer_forward_top_k_wrapper", "compress_topk_cand_buffer_size", "compress_topk_cand_buffer_size_thd", + "indexer_forward_top_k_wrapper", + "indexer_forward_wrapper", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py` around lines 11 - 17, Sort the exported names in `__all__` alphabetically in the module containing `IndexerForward` and the wrapper functions, preserving all existing entries and their export behavior.Source: Linters/SAST tools
python/cudnn/deepseek_sparse_attention/indexer_top_k/compress_top_k_sm100.py (1)
992-998: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to satisfy Ruff RUF022.♻️ Proposed fix
__all__ = [ "CompressTopkStage2", + "build_compact_buffer", "compress_stage2_topk", "compress_stage2_topk_varlen", - "build_compact_buffer", "per_batch_floats", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/indexer_top_k/compress_top_k_sm100.py` around lines 992 - 998, Sort the exported names in __all__ alphabetically to satisfy Ruff RUF022, while preserving the same symbols and export behavior.Source: Linters/SAST tools
python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100_mxfp8.py (1)
543-547: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value
rLSE_allmaterializes all 128 tile rows but only 64 are used.The consumer loop indexes
rLSE_all[q_token_stage * qhpkv + ...], i.e. a single 64-entry window, yet a 128-element register tensor is allocated and copied. Withnum_regs_epilogue = 96this is a likely spill source. The*_head_halfvariants already offsetsLSEand size the register tensor to the used range; consider the same here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100_mxfp8.py` around lines 543 - 547, Resize rLSE_all in the score recompute path to the 64-entry range actually consumed by the q_token_stage indexing, matching the offset-and-sized approach used by the *_head_half variants. Update the associated sLSE view or copy bounds as needed so only that window is materialized, while preserving the existing consumer indexing behavior.python/cudnn/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.py (1)
1425-1428: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDead branch and stale comment for
sq_b.
sq_bis unconditionallyNoneat Line 1428, so the guard at Line 1511 is always taken and the comment ("capture path skipped the eager per-batch compute above") describes code that no longer exists. Collapse it to a direct computation.♻️ Proposed cleanup
- # Only needed for the optional GPU-side local-to-global conversion. Device - # prefix values are caller-owned and are not copied to the host for - # validation. - sq_b = None -- if sq_b is None: # capture path skipped the eager per-batch compute above - sq_b = (cu_q32[1:] - cu_q32[:-1]).to(torch.int64) + # Derived on device only for the local->global conversion; the caller's + # prefix values are never read back to the host. + sq_b = (cu_q32[1:] - cu_q32[:-1]).to(torch.int64)Also applies to: 1508-1512
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.py` around lines 1425 - 1428, Remove the unconditional sq_b = None assignment and its stale explanatory comment. In the local-to-global conversion logic around the sq_b guard, eliminate the dead conditional and execute the existing computation directly, preserving the behavior from the currently reachable branch.
🤖 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/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.py`:
- Around line 106-121: Update _get_fwd_unified_denom_placeholder so THD calls
with varying one-dimensional shapes reuse a bounded per-device float32 buffer
instead of caching every full shape. Grow the buffer only when the requested
size exceeds its capacity, return a slice matching shape, and preserve the
existing behavior for other shapes and devices.
In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py`:
- Line 134: Resize sLseReduce_struct at line 134 to 8 * self.q_tokens_per_tile
and derive lse_reduce_base at each _write_lse_from_local_accum call from
q_stage_idx * self.q_tokens_per_stage * 8, preserving the existing stage
indexing. At line 1316, update get_tensor’s layout extent to the same
geometry-derived size instead of (32,); both sites must support
qhead_per_kvhead=16 without out-of-bounds writes.
In `@python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py`:
- Around line 1249-1252: Update the MXFP8 initialization condition in the
relevant score recompute interface to include the head_dim=512 and
qhead_per_kv_head=64 dual-epilogue configuration, while preserving the existing
handling for the other MXFP8 cases. Keep denom_out.zero_() in this path, and
remove the redundant out.zero_() call.
In `@test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py`:
- Around line 42-60: Gate
test_compressed_indexer_rejects_unsupported_qhead_group_before_launch to SM100
support and skip it when the optional cudnn[cutedsl] dependency is unavailable,
using the same capability checks and import pattern as the other tests in this
file. Ensure unsupported architectures or backends are skipped before importing
or invoking DSA, while preserving the existing qhead_per_kv_head ValueError
assertion on supported environments.
---
Nitpick comments:
In `@python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py`:
- Around line 11-17: Sort the exported names in `__all__` alphabetically in the
module containing `IndexerForward` and the wrapper functions, preserving all
existing entries and their export behavior.
In
`@python/cudnn/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.py`:
- Around line 1425-1428: Remove the unconditional sq_b = None assignment and its
stale explanatory comment. In the local-to-global conversion logic around the
sq_b guard, eliminate the dead conditional and execute the existing computation
directly, preserving the behavior from the currently reachable branch.
In
`@python/cudnn/deepseek_sparse_attention/indexer_top_k/compress_top_k_sm100.py`:
- Around line 992-998: Sort the exported names in __all__ alphabetically to
satisfy Ruff RUF022, while preserving the same symbols and export behavior.
In `@python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py`:
- Around line 979-991: Update the compile_key construction in the surrounding
score-recompute interface to use tuple unpacking when prefixing "dense_indexer"
to base_compile_key, preserving the existing key contents and ordering.
- Around line 1485-1499: Update the public Args docstrings for
dense_indexer_score_recompute at
python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py:1485-1499
and dense_attn_score_recompute at :1567-1581 to document precision, q_scale,
k_scale, cu_seqlens_q_scale_padded, cu_seqlens_k_scale_padded, sf_vec_size, and
k_block_size, including packed E8M0 scale shape expectations. State that
precision="mxfp8" uses float8_e4m3fn Q/K, and additionally document for
dense_attn_score_recompute that MXFP8 out is prefilled with zeros rather than
-inf.
In
`@python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100_mxfp8.py`:
- Around line 543-547: Resize rLSE_all in the score recompute path to the
64-entry range actually consumed by the q_token_stage indexing, matching the
offset-and-sized approach used by the *_head_half variants. Update the
associated sLSE view or copy bounds as needed so only that window is
materialized, while preserving the existing consumer indexing behavior.
In
`@python/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.py`:
- Around line 996-999: Remove the unused K_consumer and KScale_consumer
participants from the relevant setup, and delete the unused neg_log2_e
assignment near warp_idx and tidx in the unified score recompute flow. Preserve
the direct pipeline_K.consumer_wait and pipeline_KScale.consumer_wait calls
using their existing MMA states.
- Around line 502-527: Remove the unused max_seqlen_k parameter from both dense
single-query epilogue helpers and update all call sites accordingly, unless
signature parity with the BF16 base epilogues is required. If retaining it, add
a concise comment in each helper clarifying that masking uses seqlen_k and
col_limit rather than max_seqlen_k.
In `@python/cudnn/deepseek_sparse_attention/utils/sm100/mxfp8_scale_utils.py`:
- Around line 196-229: The zero-initialized padding in the output of the
scale-packing function should be documented as intentionally retained E8M0 zero
values. Add a short comment adjacent to the `torch.zeros` initialization or the
per-batch write explaining that padded rows remain zero and are safe only
because downstream kernels mask them out; do not change the allocation or
packing behavior.
In `@test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py`:
- Line 725: The two per-batch loops in
test/python/fe_api/dsa/test_DSA_indexer_forward_compressed.py at lines 725-725
and 831-831 unnecessarily unpack unused shape values. Update both loops to
iterate over batch indices using the length of shapes, while preserving their
existing bodies and bounds derived from cu_q_host and cu_k_host.
🪄 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: 29e54836-c989-421c-8b80-a8f02748f61a
📒 Files selected for processing (31)
docs/fe-oss-apis/dsa.mdpython/cudnn/deepseek_sparse_attention/README.mdpython/cudnn/deepseek_sparse_attention/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_backward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/__init__.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_compressed_top_k_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_interface.pypython/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.pypython/cudnn/deepseek_sparse_attention/indexer_forward/api.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/api.pypython/cudnn/deepseek_sparse_attention/indexer_top_k/compress_top_k_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/api.pypython/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100.pypython/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100_mxfp8.pypython/cudnn/deepseek_sparse_attention/utils/runtime.pypython/cudnn/deepseek_sparse_attention/utils/sm100/gemm.pypython/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.pypython/cudnn/deepseek_sparse_attention/utils/sm100/mxfp8_scale_utils.pypython/cudnn/deepseek_sparse_attention/utils/tensor_conversion.pytest/python/fe_api/dsa/dsa_reference.pytest/python/fe_api/dsa/dsa_utils.pytest/python/fe_api/dsa/test_DSA_dense_score_recompute.pytest/python/fe_api/dsa/test_DSA_indexer_forward.pytest/python/fe_api/dsa/test_DSA_indexer_forward_compressed.pytest/python/fe_api/dsa/test_DSA_mxfp8_scale_utils.py
🚧 Files skipped from review as they are similar to previous changes (18)
- python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100_mxfp8.py
- python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm100.py
- python/cudnn/deepseek_sparse_attention/init.py
- python/cudnn/deepseek_sparse_attention/score_recompute/indexer_score_unified_sm100.py
- python/cudnn/deepseek_sparse_attention/utils/tensor_conversion.py
- python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py
- python/cudnn/deepseek_sparse_attention/utils/runtime.py
- test/python/fe_api/dsa/dsa_reference.py
- python/cudnn/deepseek_sparse_attention/indexer_backward/api.py
- test/python/fe_api/dsa/dsa_utils.py
- python/cudnn/deepseek_sparse_attention/score_recompute/api.py
- python/cudnn/deepseek_sparse_attention/utils/sm100/gemm.py
- python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm100.py
- python/cudnn/deepseek_sparse_attention/utils/sm100/mma_desc.py
- python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py
- python/cudnn/deepseek_sparse_attention/indexer_forward/_interface.py
- docs/fe-oss-apis/dsa.md
- python/cudnn/deepseek_sparse_attention/indexer_forward/api.py
| sW_dtype = Float32 if self.use_fp8_scales else self.w_dtype | ||
| sW_struct = cute.struct.Align[cute.struct.MemRange[sW_dtype, self.tile_m], 128] | ||
| sKScale_struct = cute.struct.Align[cute.struct.MemRange[Float32, self.tile_n * self.kv_stages], 128] | ||
| sLseReduce_struct = cute.struct.Align[cute.struct.MemRange[Float32, 32], 128] |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
LSE reduction scratch is hard-coded to 32 floats instead of being derived from the tile geometry, overflowing for qhead_per_kvhead=16. _write_lse_from_local_accum needs 8 slots per q token per stage with a stage base of 0/16, so q_tokens_per_stage > 2 (qhpkv 16) writes past the 32-element buffer.
python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py#L134-L134: sizesLseReduce_structas8 * self.q_tokens_per_tile(and derivelse_reduce_basefromq_stage_idx * self.q_tokens_per_stage * 8at the call sites), or rejectcompute_lsefor qhpkv 16.python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py#L1316-L1316: update theget_tensorlayout to the same geometry-derived extent rather than the literal(32,).
📍 Affects 1 file
python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py#L134-L134(this comment)python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py#L1316-L1316
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudnn/deepseek_sparse_attention/indexer_forward/indexer_fwd_sm90.py`
at line 134, Resize sLseReduce_struct at line 134 to 8 * self.q_tokens_per_tile
and derive lse_reduce_base at each _write_lse_from_local_accum call from
q_stage_idx * self.q_tokens_per_stage * 8, preserving the existing stage
indexing. At line 1316, update get_tensor’s layout extent to the same
geometry-derived size instead of (32,); both sites must support
qhead_per_kvhead=16 without out-of-bounds writes.
6 tasks
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
3 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 FP8/MXFP8 support to the DSA indexer score paths and introduces an
SM100 compressed indexer-forward + Top-K pipeline.
The SM100 indexer score implementation is shared by dense forward and dense
recompute, while the compressed path produces Top-K indices, selected logits,
optional fused softmax, and optional LSE without materializing the full dense
score tensor.
SM100 MXFP8 indexer paths support
qhead_per_kv_head ∈ {32, 64}. The QH32epilogue is compile-time specialized separately so the existing QH64 hot path
and launch topology remain unchanged.
Key changes
FP8/MXFP8 score paths
descales.
Q/K with packed E8M0 block scales.
indexer score recompute.
variable-length inputs.
Compressed indexer forward
indexer_forward_top_k_wrapper.default, and optional FP32 LSE.
output indices; global KV indices are returned by default.
deterministic=True, resolve exact-value ties at the K-th boundarytoward the smallest local KV indices. This makes the selected set
reproducible but does not sort the output slots.
LSE buffers.
candidate-buffer sizing helpers.
Runtime and backward integration
conversion on one caller-selected CUDA stream. The launch path does not add
cross-stream synchronization; callers own input readiness and tensor
lifetime.
prefixes to the host. Callers guarantee valid sequence, scale, causal-offset,
and candidate-prefix values, alignment, and backing storage.
index_score, avoiding indexer-score recompute and a separate softmax.B=1BSHDviews with global Top-K indices.
attn_scoreandindex_score;callers must copy either tensor if it must be preserved.
Validation
Validated on an NVIDIA B200 with CUDA 13.2, PyTorch 2.11, and cuDNN backend
9.20:
The full DSA suite completed with
70 passed, 7 skipped, 4 failed. The fourfailures reproduce identically on current
upstream/developin an independentworktree:
top_k=705;Formatting and static checks:
QH64 B200 performance was remeasured with the same inputs, 10 warmups, and 100
CUDA-event iterations:
Limitations / notes
compressed indexer paths currently require MQA (
H_kv = 1).head_dim = 128. SM100 MXFP8 additionally requiresqhead_per_kv_head ∈ {32, 64}andsf_vec_size = 32.MXFP8, LSE, or explicit
q_causal_offsets.-1,logit
-inf, and softmax0.1 <= top_k <= 2048, although they use different implementations.kernels require BF16 operands.
q_causal_offsetscomputes candidate offsets eagerly andis not CUDA-graph-capturable. THD graph capture requires explicit maximum
sequence lengths, caller-provided candidate offsets, and preallocated
outputs. Local indices are required only for strict zero-transient-allocation
capture; global conversion is capturable but records temporaries in the graph
pool.
Summary by CodeRabbit