DSA: Add FP8/MXFP8 and compressed Top-K indexer paths#370
DSA: Add FP8/MXFP8 and compressed Top-K indexer paths#370jiayus-nvidia wants to merge 4 commits into
Conversation
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.
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 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 Stage1GEMM
participant Stage2TopK
Caller->>TopKWrapper: Call compressed indexer forward
TopKWrapper->>IndexerInterface: Dispatch compressed logits
IndexerInterface->>Stage1GEMM: Build compact candidate buffer
Stage1GEMM-->>IndexerInterface: Return compact candidates
IndexerInterface->>Stage2TopK: Select Top-K and optional softmax
Stage2TopK-->>TopKWrapper: Return indices, logits, softmax, and LSE
TopKWrapper-->>Caller: Return TupleDict
Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 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
|
@cudnn-ci-bot run |
|
🚀 Running mirror pipeline Branch: cudnn-gh/pr-370-8c3ea50 |
6bda0b9 to
d7579a8
Compare
|
WIP for adding deterministic and e2e test |
3 participants
Summary
This PR extends the DSA indexer score path with FP8/MXFP8 support and adds an
SM100 compact forward + Top-K API.
It shares the SM100 indexer score implementation between dense forward and
dense recompute, and supports Top-K indices, logits, fused softmax, and optional
LSE without materializing the full dense score tensor.
Key changes
indexer score recompute.
indexer_forward_top_k_wrapper, which returns:backward
index_score, avoiding indexer-score recompute and a separatesoftmax.
B=1BSHD views and using global Top-K indices.attn_scoreandindex_score;callers that need either tensor afterward must pass a copy.
indexer forward and Top-K launches.
merge path.
test coverage.
Testing
Validated on an NVIDIA B200 with: