Repository navigation
feat(attention): add SM120 sparse MLA kernels - #3395
Merged
Merged
Conversation
lucifer1004
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jimmyzho,
kahyunnam,
nv-yunzheq,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
May 23, 2026 10:44
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:
📝 WalkthroughWalkthroughImplements a full SM120 Sparse-MLA paged-attention backend: SM120 model/scale/kv traits, PTX/CUDA primitives and helpers, DSV3.2/DSV4 decode and prefill kernels (including dual-cache), C++ orchestrator and TVM bindings, Python JIT/API/autotuner/wrapper, tests, traces, and a microbenchmark. ChangesSparse-MLA SM120 Paged Attention
Sequence Diagram (high-level interaction): sequenceDiagram
participant PythonAPI
participant TVM_FFI
participant CppDispatch
participant CUDA_Kernels
PythonAPI->>TVM_FFI: call sparse_mla_sm120_paged_attention(...)
TVM_FFI->>CppDispatch: validate & forward raw pointers/params
CppDispatch->>CUDA_Kernels: launch prefill or decode kernels
CUDA_Kernels-->>CppDispatch: write output/out_lse
CppDispatch-->>PythonAPI: in-place outputs available
Estimated code review effort: Possibly related issues:
Possibly related PRs:
Suggested labels: Suggested reviewers:
Poem:
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces sparse-MLA paged attention support for SM120 (Blackwell) GPUs, implementing optimized CUDA kernels for the DSv3.2 and DSv4 model families. The changes include a full integration with JIT compilation, TVM-FFI bindings, and a Python API featuring a wrapper class for CUDA graph compatibility. Feedback primarily focuses on enhancing the performance of the latency-sensitive decode path by pre-allocating scratch buffers to avoid runtime overhead. Additional suggestions involve refining the SM count fallback logic for broader hardware compatibility and resolving an inconsistency between shared memory layout documentation and the prefill kernel's implementation of rope weight staging.
Comment on lines
+206
to
+215
| mid_out = torch.empty( | ||
| (num_tokens, num_heads, num_splits, d_v), | ||
| dtype=torch.bfloat16, | ||
| device=q.device, | ||
| ) | ||
| mid_lse = torch.empty( | ||
| (num_tokens, num_heads, num_splits), | ||
| dtype=torch.float32, | ||
| device=q.device, | ||
| ) |
Contributor
There was a problem hiding this comment.
Allocating mid_out and mid_lse scratch buffers using torch.empty inside the functional API's implementation (_paged_attention) is a performance bottleneck for the latency-sensitive decode path (num_tokens <= 64). These allocations occur on every call, which introduces Python overhead and is incompatible with CUDA graph capture. These buffers should be passed as optional arguments to the functional API, allowing the Wrapper class to pre-allocate and reuse them.
Comment on lines
+416
to
+421
| self._out_lse = torch.empty( | ||
| (max_num_tokens, max_num_heads), | ||
| dtype=torch.float32, | ||
| device=self._device, | ||
| ) | ||
|
|
Contributor
There was a problem hiding this comment.
The BatchSparseMLAPagedAttentionWrapper should pre-allocate the mid_out and mid_lse scratch tensors used by the decode kernels during its __init__. Currently, it only pre-allocates out_lse, but the underlying kernels for num_tokens <= 64 require additional workspace. Pre-allocating these for the maximum decode batch size (64) and maximum splits (32) would make the run() method truly allocation-free and safe for CUDA graphs.
| } else { | ||
| int sm_count = 0; | ||
| CUDA_CHECK_BOOL(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0)); | ||
| if (sm_count <= 0) sm_count = 188; |
Contributor
There was a problem hiding this comment.
The fallback value for sm_count is hardcoded to 188. While this is a reasonable default for high-end Blackwell GPUs like B200, it might be sub-optimal for other SM120/SM121 variants with different SM counts. Consider using a more conservative default or ensuring that cudaDeviceGetAttribute is reliable in the target environments.
Comment on lines
+123
to
+126
| // Scratch area: q_rope staging (O2 overlap) and xv_rope weight staging | ||
| // q_rope needs: N_HG * HPB * D_ROPE * sizeof(bf16) = 4096 bytes | ||
| // xv_rope needs: HPB * BI * sizeof(bf16) = 2048 bytes | ||
| static constexpr size_t SMEM_SCRATCH = N_HG * HPB * D_ROPE * sizeof(bf16); |
Contributor
There was a problem hiding this comment.
The comment indicates that SMEM_SCRATCH is used for xv_rope weight staging, but the prefill kernel implementation in prefill_kernel.cuh (e.g., line 450) appears to use sm.w_fp8 for this purpose. Please verify if the scratch area is actually used for xv_rope weights and update the comment or implementation to be consistent.
![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: 6
🧹 Nitpick comments (7)
tests/attention/test_sparse_mla_sm120.py (1)
336-360: ⚡ Quick winAssert
out_lseagainst the reference as well, not justoutput.These tests already compute
_ref_lse, but never validateout_lse. Adding that check will catch regressions in sink-merge and masking semantics thatoutput-only checks can miss.Proposed fix (apply similarly to the other test cases below)
- ref_out, _ref_lse = _ref_sparse_attn( + ref_out, _ref_lse = _ref_sparse_attn( q, kv_dequant, indices, sm_scale, d_v, attn_sink=attn_sink ) @@ torch.testing.assert_close(output, ref_out, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(out_lse, _ref_lse, atol=5e-2, rtol=5e-2)🤖 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 `@tests/attention/test_sparse_mla_sm120.py` around lines 336 - 360, Add an assertion comparing out_lse to the computed reference _ref_lse after calling flashinfer.sparse_mla_sm120_paged_attention: use torch.testing.assert_close(out_lse, _ref_lse, atol=<same_tol>, rtol=<same_tol>) (match the tolerances used for output or use appropriate float32 tolerances) so the test validates the LSE sink-merge/masking semantics; locate the variables out_lse and _ref_lse and the call to flashinfer.sparse_mla_sm120_paged_attention in the test and insert the assertion immediately after the existing output check.tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json (1)
113-123: ⚡ Quick winUse concrete dtypes for optional inputs in trace spec
topk_lengthandattn_sinkcurrently use"dtype": "unknown"(Lines 113 and 121). If this JSON is consumed for validation or benchmark generation, concrete dtypes are safer and reduce ambiguity.Suggested fix
"topk_length": { @@ - "dtype": "unknown", + "dtype": "int32", @@ "attn_sink": { @@ - "dtype": "unknown", + "dtype": "float32",🤖 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 `@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json` around lines 113 - 123, The trace spec entries "topk_length" and "attn_sink" currently declare "dtype": "unknown"; change them to concrete dtypes: set "topk_length" to an integer dtype (e.g., "int32") and set "attn_sink" to a floating dtype (e.g., "float32") so validators/benchmarks have unambiguous types; locate the entries by the keys "topk_length" and "attn_sink" in the JSON and update their "dtype" fields accordingly while leaving their "shape"/"optional"/"description" values intact.include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh (1)
76-93: 💤 Low valueConsider adding a clarifying comment for the index clamping behavior.
The clamping of negative indices to 0 (line 78) silently reads data from slot 0 for invalid entries. While the QK masking ensures these don't affect the final result (softmax weights are ~0), a brief comment here would help future maintainers understand this intentional design choice and prevent confusion.
🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh` around lines 76 - 93, The loop clamps negative indices to zero (idx = (idx >= 0) ? idx : 0;) which intentionally causes invalid entries to read slot 0; add a concise inline comment next to this line explaining that negative indices are mapped to 0 to avoid OOB reads and that downstream QK masking/softmax weights ensure these dummy reads do not affect results (refer to symbols: indices, idx, BI, IO_THREADS, KV::SCALE_IN_KV_SMEM, PAGE_BLOCK_SIZE, cp_async_bulk_g2s*). Keep the comment short and state the rationale and safety guarantees so future maintainers understand the design choice.csrc/sparse_mla_sm120_prefill.cu (2)
61-65: 💤 Low valueStatic configuration flag is technically racy but benign.
The
static bool configuredpattern can race if multiple CPU threads call the same template instantiation concurrently. However, sincecudaFuncSetAttributeis idempotent for the same value, the worst case is redundant calls. This is a common pattern in CUDA launchers.Consider using
std::call_onceif thread-safety guarantees are desired:`#include` <mutex> static std::once_flag configured_flag; std::call_once(configured_flag, [&]() { if (smem_bytes > 48 * 1024) { cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); } });Also applies to: 103-107
🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 61 - 65, The static flag configured used around the cudaFuncSetAttribute call (checking smem_bytes and calling cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)) can race if multiple CPU threads instantiate/launch the same kernel; replace the ad-hoc static bool configured with a thread-safe one-time initialization using std::call_once and a std::once_flag (e.g., configured_flag) so the cudaFuncSetAttribute call is executed exactly once per process when smem_bytes > 48*1024; apply the same change to the other occurrence around lines 103-107.
225-226: 💤 Low valueDead code: unreachable return statement.
Line 225 is unreachable because all branches in the preceding if-else chain (lines 212-221) return. The compiler should optimize this away, but consider removing it or adding an assertion for clarity.
♻️ Proposed fix
else if (topk == 2048) DISPATCH_BY_NH_CM(FP8, 2048); - else - return false; `#undef` DISPATCH_BY_NH_CM `#undef` DISPATCH_MG_CM - return false; // unreachable + return false; // topk not in {128, 512, 1024, 2048}🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 225 - 226, The trailing "return false;" at the end of the function in csrc/sparse_mla_sm120_prefill.cu is unreachable because every branch above returns; remove the dead return or replace it with an explicit assertion/unreachable hint (e.g., assert(false) or a compiler intrinsic like __builtin_unreachable() / unreachable()) to make intent clear and help static analysis; update the function's final line accordingly (modify the trailing return in the function that contains the unreachable return).include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh (1)
636-650: 💤 Low valueOptional: Consider extracting duplicated tile-routing lambdas.
The lambdas
tile_kv_ptr,tile_stride, andtile_idx_ptrare duplicated between IO warps (lines 636-650) and math warps (lines 715-727). While functionally correct, extracting these to device-side helper functions or a single definition at kernel scope could reduce maintenance burden.Also applies to: 715-727
🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh` around lines 636 - 650, Extract the duplicated per-tile routing logic by replacing the three lambdas tile_kv_ptr, tile_stride, and tile_idx_ptr with shared device-visible helpers (e.g., inline __device__ or __device__ __forceinline__ template functions) declared at kernel scope so both IO warps and math warps call the same implementation; these helpers should accept the tile index (ti) plus the necessary captured variables (KV_cache, KV_cache_extra, stride_kv_block, stride_kv_block_extra, idx_base, indices_extra, s_i, and template parameter TOPK_EXTRA or branch on it) and return the same types (const uint8_t*, size_t, const int32_t*), preserving the existing logic for (ti < NI) and the if constexpr TOPK_EXTRA == 0 special case. Ensure function names match the original intent (e.g., tile_kv_ptr_helper, tile_stride_helper, tile_idx_ptr_helper) and update call sites in both locations (previous lambda uses at lines 636-650 and 715-727) to call these helpers so maintenance and duplication are removed.include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh (1)
300-302: 💤 Low valueStale comment: warp count is 8, not 4.
The comment says "warp_id < DSV4_N_WARPS = 4" but
DSV4_N_WARPSis defined as 8 at line 34.📝 Suggested fix
// ────────────────────────────────────────────────────────────── - // Math warps branch (warp_id < DSV4_N_WARPS = 4) + // Math warps branch (warp_id < DSV4_N_WARPS = 8) // ──────────────────────────────────────────────────────────────🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh` around lines 300 - 302, The comment above the "Math warps branch" incorrectly states "warp_id < DSV4_N_WARPS = 4" while DSV4_N_WARPS is actually 8; update the comment text near the Math warps branch to reflect the correct warp count (e.g., "warp_id < DSV4_N_WARPS = 8" or simply "warp_id < DSV4_N_WARPS") so it matches the DSV4_N_WARPS constant used in the kernel.
🤖 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 `@benchmarks/bench_sparse_mla_sm120.py`:
- Around line 310-317: The printed DSv3.2 metadata is incorrect: update the
message strings that mention "page_block_size=1" in the DSv3.2 section so they
reflect the actual benchmark setting page_block_size=64; specifically change the
print call that outputs "DSv3.2 decode-dsv3_2 path (topk=2048,
page_block_size=1):" to "page_block_size=64" and ensure any related header text
or metadata emitted alongside dsv3_2_configs (the list comprehension variable)
is consistent with the actual page_block_size used earlier in the script.
In `@csrc/sparse_mla_sm120_decode_dsv3_2.cu`:
- Around line 79-81: The code calls cudaDeviceGetAttribute(&sm_count,
cudaDevAttrMultiProcessorCount, 0) which always queries device 0 and can be
wrong on multi‑GPU setups; update the logic in this block (around sm_count and
its use for chunks_per_block) to first determine the current device (e.g., via
cudaGetDevice or cudaGetDeviceProperties) and then call cudaDeviceGetAttribute
or inspect properties for that device id, keep the existing DSV3_2_CUDA_CHECK
error handling, and preserve the fallback default (sm_count = 188) only if the
device query fails.
In `@csrc/sparse_mla_sm120_decode_dsv4.cu`:
- Around line 81-83: The current code hardcodes device 0 when calling
cudaDeviceGetAttribute to populate sm_count; instead query the active device and
use it: call cudaGetDevice (or accept a device id parameter) to obtain the
current device id and pass that id to cudaDeviceGetAttribute (the call wrapped
by CUDA_CHECK_BOOL), then keep the existing fallback logic (if sm_count <= 0 set
default 188). Update references around sm_count and the
CUDA_CHECK_BOOL(cudaDeviceGetAttribute(..., 0)) invocation so the attribute call
uses the obtained device variable rather than literal 0.
In `@csrc/sparse_mla_sm120_jit_binding.cu`:
- Around line 114-127: The code accesses extra_indices.value() inside the
extra_kv_cache.has_value() block which can dereference an empty optional; update
the logic in the block that computes extra_topk_arg/stride/pbs (the code
referencing extra_kv_cache, extra_indices, extra_topk_arg, pbs_extra_arg,
stride_extra_kv_block) to first verify extra_indices.has_value() and handle the
missing case (e.g., early return/error, or compute a safe default) before
calling extra_indices.value(); ensure any control path that relies on
extra_indices uses the validated/guarded value.
In `@flashinfer/sparse_mla_sm120.py`:
- Around line 712-717: The try/except currently swallows all errors when calling
AutoTuner.get().load_configs(str(path)), making failures invisible; change the
except to capture the exception (e.g., except Exception as e) and log a warning
including the path, mtime and the exception message so the failure is visible
(use the module logger if available or Python's logging) while still leaving
_decode_dsv4_cache_mtime unchanged so the next cold call retries. Ensure you
reference AutoTuner.get().load_configs, the variables path and mtime, and
preserve the existing retry behavior.
In
`@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json`:
- Line 56: The constraint for head_dim_qk currently allows (512, 576) which
conflicts with the concrete trace where head_dim_qk is fixed to 512; update the
constraint for the symbol head_dim_qk to only permit 512 (e.g., replace "(512,
576)" with "(512)" or otherwise remove 576) so the allowed values match the
fixed axis value used in this trace artifact.
---
Nitpick comments:
In `@csrc/sparse_mla_sm120_prefill.cu`:
- Around line 61-65: The static flag configured used around the
cudaFuncSetAttribute call (checking smem_bytes and calling
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_bytes)) can race if multiple CPU threads instantiate/launch the same
kernel; replace the ad-hoc static bool configured with a thread-safe one-time
initialization using std::call_once and a std::once_flag (e.g., configured_flag)
so the cudaFuncSetAttribute call is executed exactly once per process when
smem_bytes > 48*1024; apply the same change to the other occurrence around lines
103-107.
- Around line 225-226: The trailing "return false;" at the end of the function
in csrc/sparse_mla_sm120_prefill.cu is unreachable because every branch above
returns; remove the dead return or replace it with an explicit
assertion/unreachable hint (e.g., assert(false) or a compiler intrinsic like
__builtin_unreachable() / unreachable()) to make intent clear and help static
analysis; update the function's final line accordingly (modify the trailing
return in the function that contains the unreachable return).
In `@include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh`:
- Around line 76-93: The loop clamps negative indices to zero (idx = (idx >= 0)
? idx : 0;) which intentionally causes invalid entries to read slot 0; add a
concise inline comment next to this line explaining that negative indices are
mapped to 0 to avoid OOB reads and that downstream QK masking/softmax weights
ensure these dummy reads do not affect results (refer to symbols: indices, idx,
BI, IO_THREADS, KV::SCALE_IN_KV_SMEM, PAGE_BLOCK_SIZE, cp_async_bulk_g2s*). Keep
the comment short and state the rationale and safety guarantees so future
maintainers understand the design choice.
In `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh`:
- Around line 300-302: The comment above the "Math warps branch" incorrectly
states "warp_id < DSV4_N_WARPS = 4" while DSV4_N_WARPS is actually 8; update the
comment text near the Math warps branch to reflect the correct warp count (e.g.,
"warp_id < DSV4_N_WARPS = 8" or simply "warp_id < DSV4_N_WARPS") so it matches
the DSV4_N_WARPS constant used in the kernel.
In `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh`:
- Around line 636-650: Extract the duplicated per-tile routing logic by
replacing the three lambdas tile_kv_ptr, tile_stride, and tile_idx_ptr with
shared device-visible helpers (e.g., inline __device__ or __device__
__forceinline__ template functions) declared at kernel scope so both IO warps
and math warps call the same implementation; these helpers should accept the
tile index (ti) plus the necessary captured variables (KV_cache, KV_cache_extra,
stride_kv_block, stride_kv_block_extra, idx_base, indices_extra, s_i, and
template parameter TOPK_EXTRA or branch on it) and return the same types (const
uint8_t*, size_t, const int32_t*), preserving the existing logic for (ti < NI)
and the if constexpr TOPK_EXTRA == 0 special case. Ensure function names match
the original intent (e.g., tile_kv_ptr_helper, tile_stride_helper,
tile_idx_ptr_helper) and update call sites in both locations (previous lambda
uses at lines 636-650 and 715-727) to call these helpers so maintenance and
duplication are removed.
In `@tests/attention/test_sparse_mla_sm120.py`:
- Around line 336-360: Add an assertion comparing out_lse to the computed
reference _ref_lse after calling flashinfer.sparse_mla_sm120_paged_attention:
use torch.testing.assert_close(out_lse, _ref_lse, atol=<same_tol>,
rtol=<same_tol>) (match the tolerances used for output or use appropriate
float32 tolerances) so the test validates the LSE sink-merge/masking semantics;
locate the variables out_lse and _ref_lse and the call to
flashinfer.sparse_mla_sm120_paged_attention in the test and insert the assertion
immediately after the existing output check.
In
`@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json`:
- Around line 113-123: The trace spec entries "topk_length" and "attn_sink"
currently declare "dtype": "unknown"; change them to concrete dtypes: set
"topk_length" to an integer dtype (e.g., "int32") and set "attn_sink" to a
floating dtype (e.g., "float32") so validators/benchmarks have unambiguous
types; locate the entries by the keys "topk_length" and "attn_sink" in the JSON
and update their "dtype" fields accordingly while leaving their
"shape"/"optional"/"description" values intact.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: a1b0ad7b-3323-48fd-b83f-d8426595d635
📥 Commits
Reviewing files that changed from the base of the PR and between bff85f3 and caa907ef47ba3266e493f7a46414fd7097a50935.
📒 Files selected for processing (34)
.gitignorebenchmarks/bench_sparse_mla_sm120.pycsrc/sparse_mla_sm120.cucsrc/sparse_mla_sm120_decode_dsv3_2.cucsrc/sparse_mla_sm120_decode_dsv4.cucsrc/sparse_mla_sm120_jit_binding.cucsrc/sparse_mla_sm120_prefill.cuflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/sparse_mla_sm120.pyflashinfer/sparse_mla_sm120.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/sparse_mla_sm120/arch/barrier.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/common.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/cp_async.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/ldmatrix_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/mma_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/d2_load_b.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/online_softmax.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/xv_rope_mma.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/model_type.hinclude/flashinfer/attention/sparse_mla_sm120/model/scale_convert.cuhinclude/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuhtests/attention/test_sparse_mla_sm120.pytests/trace/example.pytests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json
Comment on lines
+310
to
+317
| # DSv3.2 decode-dsv3_2: topk fixed at 2048, page_block_size=1, 656 B/token. | ||
| # Sweep is num_heads × num_tokens (with_sink off; sink is a per-head | ||
| # epilogue scale, negligible on the bandwidth-bound critical path). | ||
| dsv3_2_configs = [(h, t) for h in (8, 16, 32, 64, 128) for t in (1, 16, 32, 64)] | ||
|
|
||
| print() | ||
| print("DSv3.2 decode-dsv3_2 path (topk=2048, page_block_size=1):") | ||
| print(header) |
Contributor
There was a problem hiding this comment.
Fix DSv3.2 page-block-size text to match the actual benchmark config.
Line 310 and Line 316 still report page_block_size=1, but the benchmark uses page_block_size = 64 (Line 199). This makes printed benchmark metadata misleading.
Proposed fix
- # DSv3.2 decode-dsv3_2: topk fixed at 2048, page_block_size=1, 656 B/token.
+ # DSv3.2 decode-dsv3_2: topk fixed at 2048, page_block_size=64, 656 B/token.
@@
- print("DSv3.2 decode-dsv3_2 path (topk=2048, page_block_size=1):")
+ print("DSv3.2 decode-dsv3_2 path (topk=2048, page_block_size=64):")🧰 Tools
🪛 Ruff (0.15.13)
[warning] 311-311: Comment contains ambiguous × (MULTIPLICATION SIGN). Did you mean x (LATIN SMALL LETTER X)?
(RUF003)
🤖 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 `@benchmarks/bench_sparse_mla_sm120.py` around lines 310 - 317, The printed
DSv3.2 metadata is incorrect: update the message strings that mention
"page_block_size=1" in the DSv3.2 section so they reflect the actual benchmark
setting page_block_size=64; specifically change the print call that outputs
"DSv3.2 decode-dsv3_2 path (topk=2048, page_block_size=1):" to
"page_block_size=64" and ensure any related header text or metadata emitted
alongside dsv3_2_configs (the list comprehension variable) is consistent with
the actual page_block_size used earlier in the script.
Comment on lines
+712
to
+717
| try: | ||
| AutoTuner.get().load_configs(str(path)) | ||
| _decode_dsv4_cache_mtime = mtime | ||
| except Exception: | ||
| # Keep mtime unchanged so the next cold call retries. | ||
| pass |
Contributor
There was a problem hiding this comment.
Silent exception swallowing hinders debugging.
The static analysis tool correctly flagged this try-except-pass pattern. When the disk cache fails to load (corrupt JSON, version mismatch), silently ignoring the error makes debugging difficult. Consider logging a warning.
🔧 Proposed fix to add logging
+import logging
+
+_logger = logging.getLogger(__name__)
+
def _decode_dsv4_maybe_load_cache() -> None:
...
try:
AutoTuner.get().load_configs(str(path))
_decode_dsv4_cache_mtime = mtime
- except Exception:
+ except Exception as e:
# Keep mtime unchanged so the next cold call retries.
- pass
+ _logger.debug("Failed to load decode-dsv4 cache from %s: %s", path, e)🧰 Tools
🪛 Ruff (0.15.13)
[error] 715-717: try-except-pass detected, consider logging the exception
(S110)
[warning] 715-715: Do not catch blind exception: Exception
(BLE001)
🤖 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 `@flashinfer/sparse_mla_sm120.py` around lines 712 - 717, The try/except
currently swallows all errors when calling
AutoTuner.get().load_configs(str(path)), making failures invisible; change the
except to capture the exception (e.g., except Exception as e) and log a warning
including the path, mtime and the exception message so the failure is visible
(use the module logger if available or Python's logging) while still leaving
_decode_dsv4_cache_mtime unchanged so the next cold call retries. Ensure you
reference AutoTuner.get().load_configs, the variables path and mtime, and
preserve the existing retry behavior.
lucifer1004
force-pushed
the
sparse-mla-sm120
branch
3 times, most recently
from
May 23, 2026 11:48
17a6d1f to
348826f
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
csrc/sparse_mla_sm120_prefill.cu (1)
61-65: ⚖️ Poor tradeoffStatic
configuredflag is not thread-safe.The
static bool configuredpattern for one-timecudaFuncSetAttributeis racy when multiple host threads concurrently launch kernels. While the worst outcome is redundant attribute setting (not incorrect behavior), consider usingstd::call_onceorstd::atomic_flagfor correctness.🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 61 - 65, The static bool configured used before calling cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes) is racy across host threads; replace this pattern with a thread-safe one (e.g., use a std::once_flag + std::call_once or std::atomic_flag) so the attribute is set only once safely when smem_bytes > 48*1024; update the code paths referencing configured and keep the condition on smem_bytes and the call to cudaFuncSetAttribute(kernel, ...) unchanged while guarding that call with the chosen thread-safe one-time initialization mechanism.tests/attention/test_sparse_mla_sm120.py (1)
336-359: ⚡ Quick winAssert
out_lseas well.These cases compute
ref_lseand then discard it. That leaves the merged-LSE contract untested, so a regression in sink handling or split/prefill normalization can still pass while returning the wrongout_lse.Example change
- ref_out, _ref_lse = _ref_sparse_attn( + ref_out, ref_lse = _ref_sparse_attn( q, kv_dequant, indices, sm_scale, d_v, attn_sink=attn_sink ) ... torch.testing.assert_close(output, ref_out, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(out_lse, ref_lse, atol=5e-2, rtol=5e-2)Also applies to: 409-429, 484-504, 562-582, 675-697
🤖 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 `@tests/attention/test_sparse_mla_sm120.py` around lines 336 - 359, The test computes a reference LSE (_ref_lse) but never asserts the kernel-produced out_lse, leaving merged-LSE/sink behavior untested; after calling flashinfer.sparse_mla_sm120_paged_attention add an assertion comparing out_lse to _ref_lse using torch.testing.assert_close (same dtype/precision expectations as the output check) with appropriate atol/rtol (e.g., atol=5e-2, rtol=5e-2) so failures in sink handling or split/prefill normalization are caught.
🤖 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 `@flashinfer/aot.py`:
- Around line 548-549: The Sparse-MLA AOT registration is incorrectly placed
inside the add_moe gate so passing --add-moe=false disables it; move the
gen_sparse_mla_sm120_module() registration out of the add_moe conditional by
removing it from the block that checks add_moe and appending
jit_specs.append(gen_sparse_mla_sm120_module()) unconditionally (or into the
broader registration path) so the function is always called regardless of the
add_moe flag.
In `@flashinfer/sparse_mla_sm120.py`:
- Around line 422-423: The decorators on the public methods __init__ and run()
(and the similar occurrence at the other location) use `@flashinfer_api` without a
trace= argument which breaks tracing; update the decorator usage on the wrapper
APIs (the __init__ method and the run method in this module, and the other
decorated API noted) to include an explicit trace= template (e.g., a short
unique string or format that includes class/name and key params) such that
fi_trace() can identify them, or if these wrappers are intentionally internal
remove `@flashinfer_api`; ensure you modify the `@flashinfer_api`(...) call to pass
trace="module.ClassName.__init__" or trace="module.ClassName.run" (or an
equivalent descriptive template) so tracing and auto-dump JSON generation work.
- Around line 494-506: The current wrapper only squeezes 4-D output when q.dim()
== 4, so a call with 3-D q and 4-D output (e.g., output shape [T,1,H,D]) will be
passed through incorrectly to sparse_mla_sm120_paged_attention; update the logic
to normalize output independently of q by checking output.dim() == 4 and
validating output.size(1) == 1 (raising the same ValueError message if not) and
then squeezing output (same as the existing output.squeeze(1)) outside or
separate from the q.dim() == 4 branch so both cases are handled correctly.
- Around line 226-294: topk_length and extra_topk_length must be clamped to the
actual index tensor widths to avoid kernels reading past indices; before
computing num_splits or calling any of the dispatch/kernels (e.g.,
sparse_mla_sm120_decode_dsv4, module.sparse_mla_sm120_decode_dsv3_2,
module.sparse_mla_sm120_paged_attention) compute safe_topk = min(topk_length,
indices.shape[-1]) and safe_extra_topk = min(extra_topk_length,
extra_indices.shape[-1]) (guard extra_indices is not None), use those safe
values when deriving num_splits/num_splits_main/num_splits_extra and pass
safe_topk/safe_extra_topk into downstream logic and kernel args instead of the
original potentially-too-large values; this change affects where num_splits is
computed and the arguments passed to the decode/paged calls and to
_decode_scratch_views.
In `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh`:
- Around line 129-145: The issue is that io_bulk_gather_tile(...)
writes/advances sm.mbar_kv (waking math warps) before io_gather_scales(...)
finishes, so on DSV4 math may read sm.kv_scale_bufs[...] from the new slot too
early; to fix it, reorder the operations so that io_gather_scales(...) for a
given slot is issued and completed before the corresponding
io_bulk_gather_tile(...) that updates sm.mbar_kv for that same slot (i.e., call
io_gather_scales for sm.kv_scale_bufs[(ti+1)&1] prior to calling
io_bulk_gather_tile that writes sm.mbar_kv + ((ti+1)&1)), and keep the
__threadfence_block() after both gathers to ensure visibility; apply the same
change in the prologue (tile 0) and inside the for loop that uses actual_ni so
scales are always loaded before the barrier flip.
---
Nitpick comments:
In `@csrc/sparse_mla_sm120_prefill.cu`:
- Around line 61-65: The static bool configured used before calling
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_bytes) is racy across host threads; replace this pattern with a thread-safe
one (e.g., use a std::once_flag + std::call_once or std::atomic_flag) so the
attribute is set only once safely when smem_bytes > 48*1024; update the code
paths referencing configured and keep the condition on smem_bytes and the call
to cudaFuncSetAttribute(kernel, ...) unchanged while guarding that call with the
chosen thread-safe one-time initialization mechanism.
In `@tests/attention/test_sparse_mla_sm120.py`:
- Around line 336-359: The test computes a reference LSE (_ref_lse) but never
asserts the kernel-produced out_lse, leaving merged-LSE/sink behavior untested;
after calling flashinfer.sparse_mla_sm120_paged_attention add an assertion
comparing out_lse to _ref_lse using torch.testing.assert_close (same
dtype/precision expectations as the output check) with appropriate atol/rtol
(e.g., atol=5e-2, rtol=5e-2) so failures in sink handling or split/prefill
normalization are caught.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: ec2d7b35-611b-4fe9-862a-0a8c81a9444c
📥 Commits
Reviewing files that changed from the base of the PR and between caa907ef47ba3266e493f7a46414fd7097a50935 and 348826ff82d4ec82c5e2f633a661a0ed3e8d9e52.
📒 Files selected for processing (34)
.gitignorebenchmarks/bench_sparse_mla_sm120.pycsrc/sparse_mla_sm120.cucsrc/sparse_mla_sm120_decode_dsv3_2.cucsrc/sparse_mla_sm120_decode_dsv4.cucsrc/sparse_mla_sm120_jit_binding.cucsrc/sparse_mla_sm120_prefill.cuflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/sparse_mla_sm120.pyflashinfer/sparse_mla_sm120.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/sparse_mla_sm120/arch/barrier.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/common.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/cp_async.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/ldmatrix_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/mma_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/d2_load_b.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/online_softmax.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/xv_rope_mma.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/model_type.hinclude/flashinfer/attention/sparse_mla_sm120/model/scale_convert.cuhinclude/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuhtests/attention/test_sparse_mla_sm120.pytests/trace/example.pytests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json
✅ Files skipped from review due to trivial changes (1)
- tests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.json
Comment on lines
+422
to
+423
| @supported_compute_capability([120, 121]) | ||
| @flashinfer_api |
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win
Add explicit trace metadata to the decorated wrapper APIs.
__init__ and run() are public methods decorated with @flashinfer_api, but neither supplies trace=. Please either add wrapper-specific trace templates or drop the decorator here; decorating without trace= breaks the tracing contract for public APIs. As per coding guidelines flashinfer/*.py: "Every public API decorated with @flashinfer_api should carry a trace= argument so that fi_trace() works and auto-dump produces a benchmark definition JSON."
Also applies to: 468-468
🤖 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 `@flashinfer/sparse_mla_sm120.py` around lines 422 - 423, The decorators on the
public methods __init__ and run() (and the similar occurrence at the other
location) use `@flashinfer_api` without a trace= argument which breaks tracing;
update the decorator usage on the wrapper APIs (the __init__ method and the run
method in this module, and the other decorated API noted) to include an explicit
trace= template (e.g., a short unique string or format that includes class/name
and key params) such that fi_trace() can identify them, or if these wrappers are
intentionally internal remove `@flashinfer_api`; ensure you modify the
`@flashinfer_api`(...) call to pass trace="module.ClassName.__init__" or
trace="module.ClassName.run" (or an equivalent descriptive template) so tracing
and auto-dump JSON generation work.
lucifer1004
force-pushed
the
sparse-mla-sm120
branch
from
May 23, 2026 12:09
348826f to
6574a54
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json (1)
56-56:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAlign the constraint with the fixed trace axis value.
Line 56 conflicts with the fixed axis value (
head_dim_qk = 512) and makes this artifact internally inconsistent.Suggested fix
- "head_dim_qk in (512, 576)", + "head_dim_qk == 512",🤖 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 `@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json` at line 56, The constraint "head_dim_qk in (512, 576)" conflicts with the fixed axis value head_dim_qk = 512; update the constraint to match the fixed axis by replacing "head_dim_qk in (512, 576)" with a constraint that only allows 512 (e.g., "head_dim_qk in (512)" or "head_dim_qk = 512") so the trace artifact is internally consistent.
🧹 Nitpick comments (3)
include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh (1)
141-160: 💤 Low valueNamespace-scope constants shadow model-specific traits.
D_ROPEandD_Vare defined both as namespace-scope constants (lines 143-144) and within eachKVCacheTraitsspecialization. Currently both models share these values, but if a future model type differs, the namespace-scope constants could cause confusion. Consider documenting that these are universal across all supported models or removing the duplication.🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh` around lines 141 - 160, Namespace-scope D_ROPE and D_V duplicate per-model KVCacheTraits and can shadow model-specific values; either remove the namespace-scope constants and use KVCacheTraits::<D_ROPE>/<D_V> everywhere, or keep them but rename to DEFAULT_D_ROPE/DEFAULT_D_V and add a clear comment plus static_asserts (e.g., static_assert(DEFAULT_D_ROPE == KVCacheTraits<Model>::D_ROPE) for each supported model) to enforce/declare they are intended to be universal; update all references to use the chosen source of truth (KVCacheTraits or DEFAULT_*), and add a short explanatory comment near KVCacheTraits and the constants to avoid future confusion.csrc/sparse_mla_sm120_prefill.cu (1)
61-65: 💤 Low valueThread-safety note on static
configuredflag.The
static bool configuredpattern is not thread-safe across concurrent kernel launches with different template instantiations. However, sincecudaFuncSetAttributeis idempotent when setting the same value, the worst case is redundant API calls without correctness impact. Consider usingstd::call_onceif strict single-initialization semantics are desired.🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 61 - 65, The static bool configured used around cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes) is not thread-safe; replace it with a std::once_flag and std::call_once (or protect with a mutex) so the cudaFuncSetAttribute call happens only once safely across threads and template instantiations: add a static std::once_flag (include <mutex>), then use std::call_once(flag, [&]{ if (smem_bytes > 48 * 1024) cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); }); referencing the existing kernel and smem_bytes identifiers.tests/attention/test_sparse_mla_sm120.py (1)
292-757: ⚡ Quick winAdd one regression for the
topk_lengthpaths.These tests cover
-1masking, but none of them exercisetopk_length/extra_topk_lengthbeing smaller than the backing index width. A single decode case and a single dual-cache case would lock down the length-truncation logic that the new kernels and bindings rely on.🤖 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 `@tests/attention/test_sparse_mla_sm120.py` around lines 292 - 757, Add two regression cases that exercise the topk_length / extra_topk_length truncation: (1) a decode-style test variant of test_sparse_mla_sm120_decode_dsv4 that constructs indices with some entries >= topk_length (and some -1), passes a smaller topk_length argument to flashinfer.sparse_mla_sm120_paged_attention so the kernel must truncate the index width, and asserts output matches the _ref_sparse_attn reference computed using only the truncated columns; (2) a dual-cache variant of test_sparse_mla_sm120_prefill_dsv4_dual that similarly creates extra_idx values beyond extra_topk_length (and -1), calls flashinfer.sparse_mla_sm120_paged_attention with extra_kv_cache and extra_indices plus extra_topk_length/extra_pbs length args so the extra_topk_length truncation path is exercised, and compares output/out_lse to the _ref_sparse_attn reference built using the truncated virtual_idx. Ensure both tests use flashinfer.sparse_mla_sm120_paged_attention (the same symbol used elsewhere) and preserve the existing masking convention.
🤖 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 `@flashinfer/sparse_mla_sm120.py`:
- Around line 171-198: In _decode_scratch_views validate that any provided
mid_out and mid_lse have shapes large enough for (num_tokens, num_heads,
num_splits, d_v) (for mid_out) and (num_tokens, num_heads, num_splits) (for
mid_lse) before returning sliced views; if a provided tensor is undersized,
either allocate a new buffer of the required shape (matching dtype/device) or
raise a clear ValueError mentioning mid_out/mid_lse and the expected dims;
update the function to check mid_out.shape and mid_lse.shape against num_tokens,
num_heads, num_splits, d_v and only return slices when the sizes are sufficient.
In `@flashinfer/trace/templates/attention.py`:
- Around line 1128-1241: The new TraceTemplate sparse_mla_sm120_paged_trace only
models the single-cache inputs but the SM120 public APIs accept
extra_kv_cache/extra_indices/extra_topk_length for dual-cache (decode/prefill)
paths; add a corresponding TraceTemplate (or extend
sparse_mla_sm120_paged_trace) to include inputs "extra_kv_cache",
"extra_indices", and "extra_topk_length" with matching shapes/dtypes and update
constraints to reflect relationships (e.g., extra_indices.shape[0] ==
num_tokens, extra_indices.shape[-1] == topk, extra_kv_cache.shape[1] ==
page_block_size), and ensure tags/description note dual-cache behavior so
fi_trace() can serialize both single- and dual-cache callsites used by
BatchSparseMLAPagedAttentionWrapper.run()/sparse_mla_sm120_paged_attention.
---
Duplicate comments:
In
`@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json`:
- Line 56: The constraint "head_dim_qk in (512, 576)" conflicts with the fixed
axis value head_dim_qk = 512; update the constraint to match the fixed axis by
replacing "head_dim_qk in (512, 576)" with a constraint that only allows 512
(e.g., "head_dim_qk in (512)" or "head_dim_qk = 512") so the trace artifact is
internally consistent.
---
Nitpick comments:
In `@csrc/sparse_mla_sm120_prefill.cu`:
- Around line 61-65: The static bool configured used around
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_bytes) is not thread-safe; replace it with a std::once_flag and
std::call_once (or protect with a mutex) so the cudaFuncSetAttribute call
happens only once safely across threads and template instantiations: add a
static std::once_flag (include <mutex>), then use std::call_once(flag, [&]{ if
(smem_bytes > 48 * 1024) cudaFuncSetAttribute(kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); }); referencing the
existing kernel and smem_bytes identifiers.
In `@include/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuh`:
- Around line 141-160: Namespace-scope D_ROPE and D_V duplicate per-model
KVCacheTraits and can shadow model-specific values; either remove the
namespace-scope constants and use KVCacheTraits::<D_ROPE>/<D_V> everywhere, or
keep them but rename to DEFAULT_D_ROPE/DEFAULT_D_V and add a clear comment plus
static_asserts (e.g., static_assert(DEFAULT_D_ROPE ==
KVCacheTraits<Model>::D_ROPE) for each supported model) to enforce/declare they
are intended to be universal; update all references to use the chosen source of
truth (KVCacheTraits or DEFAULT_*), and add a short explanatory comment near
KVCacheTraits and the constants to avoid future confusion.
In `@tests/attention/test_sparse_mla_sm120.py`:
- Around line 292-757: Add two regression cases that exercise the topk_length /
extra_topk_length truncation: (1) a decode-style test variant of
test_sparse_mla_sm120_decode_dsv4 that constructs indices with some entries >=
topk_length (and some -1), passes a smaller topk_length argument to
flashinfer.sparse_mla_sm120_paged_attention so the kernel must truncate the
index width, and asserts output matches the _ref_sparse_attn reference computed
using only the truncated columns; (2) a dual-cache variant of
test_sparse_mla_sm120_prefill_dsv4_dual that similarly creates extra_idx values
beyond extra_topk_length (and -1), calls
flashinfer.sparse_mla_sm120_paged_attention with extra_kv_cache and
extra_indices plus extra_topk_length/extra_pbs length args so the
extra_topk_length truncation path is exercised, and compares output/out_lse to
the _ref_sparse_attn reference built using the truncated virtual_idx. Ensure
both tests use flashinfer.sparse_mla_sm120_paged_attention (the same symbol used
elsewhere) and preserve the existing masking convention.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 01adcc98-febf-4a7d-b52e-a45a37939b6c
📥 Commits
Reviewing files that changed from the base of the PR and between 348826ff82d4ec82c5e2f633a661a0ed3e8d9e52 and 6574a54d4d36f9680c1be34c895bfdee34f47a20.
📒 Files selected for processing (34)
.gitignorebenchmarks/bench_sparse_mla_sm120.pycsrc/sparse_mla_sm120.cucsrc/sparse_mla_sm120_decode_dsv3_2.cucsrc/sparse_mla_sm120_decode_dsv4.cucsrc/sparse_mla_sm120_jit_binding.cucsrc/sparse_mla_sm120_prefill.cuflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/sparse_mla_sm120.pyflashinfer/sparse_mla_sm120.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/sparse_mla_sm120/arch/barrier.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/common.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/cp_async.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/ldmatrix_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/mma_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/d2_load_b.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/online_softmax.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/xv_rope_mma.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/model_type.hinclude/flashinfer/attention/sparse_mla_sm120/model/scale_convert.cuhinclude/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuhtests/attention/test_sparse_mla_sm120.pytests/trace/example.pytests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json
Comment on lines
+430
to
+550
| local_sum[1] += p[nt][2] + p[nt][3]; | ||
| } | ||
| #pragma unroll | ||
| for (int s = 2; s >= 1; s >>= 1) { | ||
| local_sum[0] += __shfl_xor_sync(0xffffffff, local_sum[0], s); | ||
| local_sum[1] += __shfl_xor_sync(0xffffffff, local_sum[1], s); | ||
| } | ||
|
|
||
| // Cross-warp reduce. | ||
| if (tid == 0) { | ||
| sm_warp_max[warp_id * HPB + gid] = local_max[0]; | ||
| sm_warp_max[warp_id * HPB + gid + 8] = local_max[1]; | ||
| sm_warp_sum[warp_id * HPB + gid] = local_sum[0]; | ||
| sm_warp_sum[warp_id * HPB + gid + 8] = local_sum[1]; | ||
| } | ||
| bar_sync_t<3, DSV4_MATH_THREADS>(); | ||
| if (threadIdx.x < VALID_HPB) { | ||
| const int h = threadIdx.x; | ||
| float wmax[DSV4_N_WARPS], wsum[DSV4_N_WARPS]; | ||
| #pragma unroll | ||
| for (int w = 0; w < DSV4_N_WARPS; w++) { | ||
| wmax[w] = sm_warp_max[w * HPB + h]; | ||
| wsum[w] = sm_warp_sum[w * HPB + h]; | ||
| } | ||
| float bmax = -1e30f; | ||
| #pragma unroll | ||
| for (int w = 0; w < DSV4_N_WARPS; w++) bmax = fmaxf(bmax, wmax[w]); | ||
| float bsum = 0.f; | ||
| #pragma unroll | ||
| for (int w = 0; w < DSV4_N_WARPS; w++) bsum += wsum[w] * exp2f(wmax[w] - bmax); | ||
| sm_warp_max[h] = bmax; | ||
| sm_warp_sum[h] = bsum; | ||
| } | ||
| bar_sync_t<3, DSV4_MATH_THREADS>(); | ||
|
|
||
| const float block_local_max0 = sm_warp_max[gid]; | ||
| const float block_local_max1 = sm_warp_max[gid + 8]; | ||
| const float block_local_sum0 = sm_warp_sum[gid]; | ||
| const float block_local_sum1 = sm_warp_sum[gid + 8]; | ||
|
|
||
| // Online softmax update. | ||
| float new_gmax0 = fmaxf(global_max[0], block_local_max0); | ||
| float new_gmax1 = fmaxf(global_max[1], block_local_max1); | ||
| const float alpha0 = (global_max[0] > -1e29f) ? exp2f(global_max[0] - new_gmax0) : 0.f; | ||
| const float alpha1 = (global_max[1] > -1e29f) ? exp2f(global_max[1] - new_gmax1) : 0.f; | ||
| // Two distinct rescale factors: | ||
| // block_rescale = exp(block_local_max - new_gmax) — rescales the | ||
| // block-wide sum (already weighted by per-warp local_max during | ||
| // cross-warp reduction) into the new global frame. Used for the | ||
| // global_sum update. | ||
| // warp_rescale = exp(local_max[w] - new_gmax) — rescales THIS | ||
| // warp's p (computed in the warp's own local_max frame) into the | ||
| // new global frame. CRITICAL for correctness when a warp covers | ||
| // only invalid candidates: the post-mask qk == -1e30 * sm_scale * | ||
| // LOG2E ≈ -6.38e28 becomes the warp's local_max[0], and softmax | ||
| // gives p ≡ 1 (exp2(qk - local_max) = exp2(0)). Without the | ||
| // per-warp factor, these spurious 1s would leak into sm_p_full | ||
| // and corrupt the RoPE MMA. The exp(local_max - new_gmax) factor | ||
| // drives them to ~0. | ||
| const float block_rescale0 = exp2f(block_local_max0 - new_gmax0); | ||
| const float block_rescale1 = exp2f(block_local_max1 - new_gmax1); | ||
| const float warp_rescale0 = exp2f(local_max[0] - new_gmax0); | ||
| const float warp_rescale1 = exp2f(local_max[1] - new_gmax1); | ||
|
|
||
| if (chunk_idx > chunk_lo) { | ||
| #pragma unroll | ||
| for (int vc = 0; vc < N_V_CHUNKS; vc++) { | ||
| #pragma unroll | ||
| for (int nt = 0; nt < NT_PER_WARP_XV; nt++) { | ||
| acc_nope[vc][nt][0] *= alpha0; | ||
| acc_nope[vc][nt][1] *= alpha0; | ||
| acc_nope[vc][nt][2] *= alpha1; | ||
| acc_nope[vc][nt][3] *= alpha1; | ||
| } | ||
| } | ||
| #pragma unroll | ||
| for (int nt = 0; nt < ROPE_N_TILES; nt++) { | ||
| acc_rope[nt][0] *= alpha0; | ||
| acc_rope[nt][1] *= alpha0; | ||
| acc_rope[nt][2] *= alpha1; | ||
| acc_rope[nt][3] *= alpha1; | ||
| } | ||
| global_sum[0] = global_sum[0] * alpha0 + block_local_sum0 * block_rescale0; | ||
| global_sum[1] = global_sum[1] * alpha1 + block_local_sum1 * block_rescale1; | ||
| } else { | ||
| global_sum[0] = block_local_sum0 * block_rescale0; | ||
| global_sum[1] = block_local_sum1 * block_rescale1; | ||
| } | ||
| global_max[0] = new_gmax0; | ||
| global_max[1] = new_gmax1; | ||
|
|
||
| // Stage 2.75: sm_p_full = p * warp_rescale. Each warp uses its OWN | ||
| // local_max-based rescale to kill all-invalid-warp contributions. | ||
| float w_pre[DSV4_QK_N_TILES][4]; | ||
| #pragma unroll | ||
| for (int nt = 0; nt < DSV4_QK_N_TILES; nt++) { | ||
| w_pre[nt][0] = p[nt][0] * warp_rescale0; | ||
| w_pre[nt][1] = p[nt][1] * warp_rescale0; | ||
| w_pre[nt][2] = p[nt][2] * warp_rescale1; | ||
| w_pre[nt][3] = p[nt][3] * warp_rescale1; |
Contributor
There was a problem hiding this comment.
All-invalid chunks still contribute clamped slot-0 values.
When every candidate in a half-head is masked, local_max becomes the mask sentinel, so p = exp2(qk - local_max) collapses to 1s instead of 0s. If the whole block is invalid, warp_rescale also stays 1, and the later XV accumulations consume the slot-0 KV data that was gathered for the clamped -1 indices. Track an explicit “any valid candidate” bit per half-head and force local_max=-inf, local_sum=0, and w_pre=0 when none are valid.
lucifer1004
force-pushed
the
sparse-mla-sm120
branch
from
May 23, 2026 13:17
6574a54 to
19591cc
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 13
♻️ Duplicate comments (2)
csrc/sparse_mla_sm120_decode_dsv3_2.cu (1)
79-81:⚠️ Potential issue | 🟠 Major | ⚡ Quick winQuery the active CUDA device before reading SM count.
This still samples device 0 unconditionally, so
chunks_per_blockcan be tuned for the wrong GPU in multi-device runs.🔧 Minimal fix
- int sm_count = 0; - DSV3_2_CUDA_CHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0)); + int device = 0; + DSV3_2_CUDA_CHECK(cudaGetDevice(&device)); + int sm_count = 0; + DSV3_2_CUDA_CHECK(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/sparse_mla_sm120_decode_dsv3_2.cu` around lines 79 - 81, The code reads SM count for device 0 unconditionally; call cudaGetDevice (or equivalent) to obtain the active device ID first and then pass that device ID into cudaDeviceGetAttribute instead of hardcoding 0 (i.e., replace the cudaDeviceGetAttribute(..., cudaDevAttrMultiProcessorCount, 0) call with one that uses the retrieved device variable), keeping the existing sm_count fallback logic intact.include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh (1)
430-550:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAll-invalid splits still accumulate the clamped slot-0 KV.
The
warp_rescalefix only helps when some other warp/chunk contributes a larger valid max. If the whole split is invalid,new_gmax == local_max,warp_rescalestays 1, and the clamped slot-0 values still flow into XV.🔧 Minimal fix direction
+ bool any_valid0 = false, any_valid1 = false; ... - if (abs_c0 >= split_cand_end || idx0 < 0) { + if (abs_c0 >= split_cand_end || idx0 < 0) { qk[nt][0] = -1e30f; qk[nt][2] = -1e30f; + } else { + any_valid0 = true; + any_valid1 = true; } ... + if (!any_valid0) { + local_max[0] = -1e30f; + local_sum[0] = 0.f; + } + if (!any_valid1) { + local_max[1] = -1e30f; + local_sum[1] = 0.f; + }The key is to carry an explicit validity bit per half-head and force both the local sum and the staged weights to zero when no valid candidate exists.
🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh` around lines 430 - 550, The fix: add and propagate a per-half-head validity flag so fully-invalid splits cannot contribute; after computing per-warp local_max[0/1] and local_sum[0/1] (and before storing sm_warp_max/sm_warp_sum and before computing p/w_pre), compute valid0 = (local_sum[0] > 0 || any qk entries > -1e29) and valid1 likewise (or derive from comparing local_max to the clamp sentinel), then if a half-head is not valid set local_sum[*]=0 and zero out p[nt][*] (and later w_pre[nt][*]) for all nt so that sm_warp_sum and block/global accumulators and acc_nope/acc_rope do not get clamped slot-0 contributions; propagate the validity through the cross-warp reduction (store/read sm_warp_valid alongside sm_warp_max/sm_warp_sum) and use it when computing block_rescale/warp_rescale and alpha updates so invalid halves remain zeroed.
🧹 Nitpick comments (1)
tests/trace/test_fi_trace_template_consistency.py (1)
391-395: ⚡ Quick winAdd a targeted fi_trace completeness test for
sparse_mla_sm120_paged.The skip is reasonable for the auto-generated path, but it removes template completeness coverage unless a dedicated test exercises the dual-cache inputs (
extra_kv_cache/extra_indices) and resolvesextra_*axes.🤖 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 `@tests/trace/test_fi_trace_template_consistency.py` around lines 391 - 395, Add a dedicated fi_trace completeness test that targets the "sparse_mla_sm120_paged" template by creating a test function (e.g., test_fi_trace_sparse_mla_sm120_paged_completeness) in tests/trace/test_fi_trace_template_consistency.py which supplies the dual-cache inputs extra_kv_cache and extra_indices with concrete shapes/values so the extra_* axes (extra_topk, extra_page_block_size) are resolved; invoke the same fi_trace/template completeness assertions used by other tests (reuse the existing helper/assertion logic) and register the new test so the template coverage includes the paged dual-cache path.
🤖 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 `@csrc/sparse_mla_sm120_decode_dsv3_2.cu`:
- Around line 111-119: The merge kernel can overflow the shared sm_lse buffer
when num_splits > DSV4_MAX_SPLITS; before calling
sparse_mla_decode_dsv4_merge_kernel (the merge_kernel invocation that takes
mid_out, mid_lse, output, out_lse, attn_sink, num_tokens, num_splits) add a
guard that rejects or clamps invalid inputs by checking if num_splits >
DSV4_MAX_SPLITS and returning/erroring (or logging and returning) instead of
launching the kernel; ensure the check uses the constant DSV4_MAX_SPLITS and
that no GPU kernel launch occurs when the check fails.
In `@csrc/sparse_mla_sm120_jit_binding.cu`:
- Around line 54-55: The current check using TVM_FFI_ICHECK_GE(kv_cache.ndim(),
2) is too permissive; replace it with a strict validation that accepts only 2-D
paged caches or exactly the 4-D paged layout you expect. Concretely, change the
guards around kv_cache (the places using TVM_FFI_ICHECK_GE(kv_cache.ndim(), 2))
to require either kv_cache.ndim() == 2 OR (kv_cache.ndim() == 4 &&
kv_cache.shape(2) == 1 && kv_cache.shape(1) == <PAGE_BLOCK_SIZE>), and emit a
clear error message referencing the expected shapes (e.g., "kv_cache must be 2D
[num_blocks, page_bytes] or 4D [num_blocks, page_block_size=64, 1, bpt]"). Apply
the same tightened check in all affected spots (the other uses near the blocks
you noted).
In `@csrc/sparse_mla_sm120_prefill.cu`:
- Around line 312-330: The code computes stride_kv_block from the runtime
page_block_size which can diverge from the compiled kernel block size (64) used
by kernels dispatched via dispatch_v32, dispatch_dsv4_single and
dispatch_dsv4_dual; validate page_block_size before using it and fail early if
it != 64 (or adjust to the compile-time constant) so the runtime stride can't
desynchronize kernel indexing. Locate the computation of stride_kv_block and the
extra branch that checks extra_KV_cache, add a guard that checks page_block_size
== 64 (or set stride_kv_block = 64 * stride_kv_row) and return false or log an
error if mismatched, ensuring the same validation applies before calling
dispatch_v32, dispatch_dsv4_single and dispatch_dsv4_dual.
- Around line 316-321: The dual-cache fast path currently branches on
extra_KV_cache != nullptr but the kernels assume extra_indices is valid; update
the guard before calling dispatch_dsv4_dual to ensure both extra_KV_cache and
extra_indices are non-null (e.g., change the if to require extra_KV_cache !=
nullptr && extra_indices != nullptr or explicitly return false when
extra_KV_cache is non-null but extra_indices is null) so dispatch_dsv4_dual and
its kernels never receive a half-populated dual-cache.
In `@csrc/sparse_mla_sm120.cu`:
- Around line 99-101: Ensure both KV cache tensors are 8-bit element storage
before reinterpreting them as raw bytes: add a runtime check after
CHECK_CUDA(kv_cache) / CHECK_LAST_DIM_CONTIGUOUS(kv_cache) that verifies
kv_cache.scalar_type() is either at::kByte or at::kChar (or that
element_size()/itemsize() == 1) and fail with a clear error message if not;
apply the same validation to the other KV cache checks around the 141-148 region
to prevent effective_stride_kv_row() and the kernel from reading the wrong byte
layout.
- Around line 106-156: Add explicit rank and shape validation for companion
tensors: verify `indices`, `output`, and `out_lse` have expected ndim and sizes
matching num_tokens and num_heads (e.g., indices.ndim() == 3 and
indices.size(0)==num_tokens, indices.size(1)==num_heads, indices.size(2)==topk)
and that `output` and `out_lse` match their expected shapes; check
`topk_length`/`extra_topk_length` when present are 1-D with length num_tokens
and that `topk` equals indices.size(-1); for the optional extra cache validate
`extra_indices`/`extra_kv_cache` ranks and that
`extra_topk==extra_indices.size(-1)` and `extra_page_block_size==ekv.size(-3)`;
use existing helpers like CHECK_INPUT_AND_TYPE, CHECK_CUDA,
CHECK_LAST_DIM_CONTIGUOUS and TVM_FFI_ICHECK_EQ/GT to emit host-side failures
and include the variable names (`indices`, `output`, `out_lse`, `topk_length`,
`extra_indices`, `extra_kv_cache`, `extra_topk_length`) so mismatches are caught
before the kernel launch.
In `@flashinfer/sparse_mla_sm120.py`:
- Around line 171-212: In _decode_scratch_views validate that d_v equals the
fixed kernel width (512) and raise a clear ValueError if not; similarly add the
same check where similar scratch sizing occurs (the other decode-path handlers
mentioned) so callers cannot pass a smaller d_v that would underflow the SM120
decode kernels — e.g., in _decode_scratch_views and the corresponding
scratch-resolver functions referenced in the review, check if d_v != 512 and
raise ValueError("d_v must be 512 for SM120 decode kernels") before allocating
or slicing buffers.
In `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh`:
- Around line 89-94: Clamp the per-token top-k length before computing chunk
counts: when reading topk_length_ptr at t_idx (used to set topk_len), bound it
to the compiled TOPK (and >=0) so that topk_len = min(max(topk_length_ptr ?
topk_length_ptr[t_idx] : TOPK, 0), TOPK). Then use that clamped topk_len for
computing num_chunks_total, chunk_lo, and chunk_hi (symbols: topk_length_ptr,
topk_len, TOPK, num_chunks_total, chunk_lo, chunk_hi, DSV3_2_CAND_WINDOW, t_idx)
to prevent num_chunks_total from exceeding the allocated [num_tokens, TOPK]
indices buffer.
In `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh`:
- Around line 92-103: Clamp the runtime top-k lengths before computing chunk
counts: when reading topk_length_ptr or extra_topk_length_ptr into
topk_len/extra_topk_len, bound them to the maximum available index sizes (TOPK
for the main indices and extra_topk for the extra_KV_cache path) to avoid
walking past index buffers; then use those clamped values when computing
num_orig_chunks, num_extra_chunks and num_chunks_total (refer to the variables
topk_len, extra_topk_len, TOPK, extra_topk, num_orig_chunks, num_extra_chunks).
In `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh`:
- Around line 106-107: Clamp the runtime topk_length to the compiled TOPK before
deriving tile counts: when computing topk_len and actual_ni in the MG kernel
(same pattern as the SG kernel), read cold.topk_length, set topk_len =
min((int)TOPK, cold.topk_length ? __ldg(cold.topk_length + s_i) : TOPK) (or
equivalent clamp) and then compute actual_ni = (topk_len + BI - 1) / BI; this
prevents actual_ni from exceeding TOPK and avoids walking past the end of
indices in the SG and single-cache MG paths.
- Around line 487-493: The epilogue unconditionally reads cold.attn_sink +
h_start + gid + 8 and computes s1/denom1/il1 which is invalid when there are
only 8 public heads; update the block around cold.attn_sink, s1, denom1 and il1
so the second-half load and math only run when those indices are valid (e.g.,
guard with a compile/runtime check like NUM_HEADS > 8 or (h_start + gid + 8) <
total_heads), and set il1 = 0.f otherwise; change references to s1/denom1 and
il1 in this scope to follow that guard so no out-of-bounds read occurs (look for
cold.attn_sink, h_start, gid, s1, denom1, il1, sm.l_smem and sm.m_smem to locate
the code).
In `@tests/attention/test_sparse_mla_sm120.py`:
- Line 232: Remove the unused num_heads binding by replacing the three-name
unpack with an underscore placeholder; in the context where q is unpacked (in
_ref_sparse_attn / the test), change "num_tokens, num_heads, d_qk = q.shape" to
"num_tokens, _, d_qk = q.shape" so Ruff no longer flags the unused variable.
In `@tests/trace/example.py`:
- Around line 432-457: The new example calls
flashinfer.sparse_mla_sm120_paged_attention in tests/trace/example.py but the
generated fi_trace_out JSON artifact for that call wasn’t committed; rerun the
trace generator by executing tests/trace/example.py to produce the updated
fi_trace_out artifact (which includes the paged-attention trace for
flashinfer.sparse_mla_sm120_paged_attention), then add and commit the
regenerated JSON so the example and checked-in traces remain in sync.
---
Duplicate comments:
In `@csrc/sparse_mla_sm120_decode_dsv3_2.cu`:
- Around line 79-81: The code reads SM count for device 0 unconditionally; call
cudaGetDevice (or equivalent) to obtain the active device ID first and then pass
that device ID into cudaDeviceGetAttribute instead of hardcoding 0 (i.e.,
replace the cudaDeviceGetAttribute(..., cudaDevAttrMultiProcessorCount, 0) call
with one that uses the retrieved device variable), keeping the existing sm_count
fallback logic intact.
In `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh`:
- Around line 430-550: The fix: add and propagate a per-half-head validity flag
so fully-invalid splits cannot contribute; after computing per-warp
local_max[0/1] and local_sum[0/1] (and before storing sm_warp_max/sm_warp_sum
and before computing p/w_pre), compute valid0 = (local_sum[0] > 0 || any qk
entries > -1e29) and valid1 likewise (or derive from comparing local_max to the
clamp sentinel), then if a half-head is not valid set local_sum[*]=0 and zero
out p[nt][*] (and later w_pre[nt][*]) for all nt so that sm_warp_sum and
block/global accumulators and acc_nope/acc_rope do not get clamped slot-0
contributions; propagate the validity through the cross-warp reduction
(store/read sm_warp_valid alongside sm_warp_max/sm_warp_sum) and use it when
computing block_rescale/warp_rescale and alpha updates so invalid halves remain
zeroed.
---
Nitpick comments:
In `@tests/trace/test_fi_trace_template_consistency.py`:
- Around line 391-395: Add a dedicated fi_trace completeness test that targets
the "sparse_mla_sm120_paged" template by creating a test function (e.g.,
test_fi_trace_sparse_mla_sm120_paged_completeness) in
tests/trace/test_fi_trace_template_consistency.py which supplies the dual-cache
inputs extra_kv_cache and extra_indices with concrete shapes/values so the
extra_* axes (extra_topk, extra_page_block_size) are resolved; invoke the same
fi_trace/template completeness assertions used by other tests (reuse the
existing helper/assertion logic) and register the new test so the template
coverage includes the paged dual-cache path.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: bc8a63e0-3135-4693-9d62-aafb7d6ccf86
📥 Commits
Reviewing files that changed from the base of the PR and between 6574a54d4d36f9680c1be34c895bfdee34f47a20 and 19591ccd8b61048a701f1ca9a066967fe6c6f463.
📒 Files selected for processing (35)
.gitignorebenchmarks/bench_sparse_mla_sm120.pycsrc/sparse_mla_sm120.cucsrc/sparse_mla_sm120_decode_dsv3_2.cucsrc/sparse_mla_sm120_decode_dsv4.cucsrc/sparse_mla_sm120_jit_binding.cucsrc/sparse_mla_sm120_prefill.cuflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/sparse_mla_sm120.pyflashinfer/sparse_mla_sm120.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/sparse_mla_sm120/arch/barrier.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/common.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/cp_async.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/ldmatrix_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/mma_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/d2_load_b.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/online_softmax.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/xv_rope_mma.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/model_type.hinclude/flashinfer/attention/sparse_mla_sm120/model/scale_convert.cuhinclude/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuhtests/attention/test_sparse_mla_sm120.pytests/trace/example.pytests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/test_fi_trace_template_consistency.py
✅ Files skipped from review due to trivial changes (1)
- tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json
Comment on lines
+312
to
+330
| const size_t stride_kv_block = (size_t)page_block_size * (size_t)stride_kv_row; | ||
| const size_t stride_kv_block_extra = | ||
| (extra_KV_cache != nullptr) ? (size_t)extra_page_block_size * (size_t)extra_stride_kv_row : 0; | ||
|
|
||
| if (extra_KV_cache != nullptr) { | ||
| if (mt != ModelType::DSV4) return false; | ||
| return dispatch_dsv4_dual(num_heads, topk, topk_extra, extra_page_block_size, Q, KV_cache, | ||
| indices, extra_KV_cache, extra_indices, attn_sink, output, out_lse, | ||
| sm_scale, num_tokens, stride_kv_block, stride_kv_block_extra, | ||
| topk_length, extra_topk_length, stream); | ||
| } | ||
|
|
||
| switch (mt) { | ||
| case ModelType::DSV3_2: | ||
| return dispatch_v32(num_heads, topk, Q, KV_cache, indices, attn_sink, output, out_lse, | ||
| sm_scale, num_tokens, stride_kv_block, topk_length, stream); | ||
| case ModelType::DSV4: | ||
| return dispatch_dsv4_single(num_heads, topk, Q, KV_cache, indices, attn_sink, output, out_lse, | ||
| sm_scale, num_tokens, stride_kv_block, topk_length, stream); |
Contributor
There was a problem hiding this comment.
Validate the main page_block_size against the compiled envelope.
Every main-cache kernel instantiated here is <..., 64>, but stride_kv_block is still derived from the runtime page_block_size. Accepting any other value silently desynchronizes the stride from the kernel's block/local index math.
🔧 Minimal fix
bool sparse_mla_prefill_dispatch(ModelType mt, int num_heads, int topk, int page_block_size,
int topk_extra, int extra_page_block_size, const bf16* Q,
const uint8_t* KV_cache, const int32_t* indices,
const uint8_t* extra_KV_cache, const int32_t* extra_indices,
bf16* output, float* out_lse, float sm_scale, int num_tokens,
int stride_kv_row, int extra_stride_kv_row, const float* attn_sink,
const int* topk_length, const int* extra_topk_length,
cudaStream_t stream) {
+ if (page_block_size != 64) return false;
+
const size_t stride_kv_block = (size_t)page_block_size * (size_t)stride_kv_row;🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 312 - 330, The code computes
stride_kv_block from the runtime page_block_size which can diverge from the
compiled kernel block size (64) used by kernels dispatched via dispatch_v32,
dispatch_dsv4_single and dispatch_dsv4_dual; validate page_block_size before
using it and fail early if it != 64 (or adjust to the compile-time constant) so
the runtime stride can't desynchronize kernel indexing. Locate the computation
of stride_kv_block and the extra branch that checks extra_KV_cache, add a guard
that checks page_block_size == 64 (or set stride_kv_block = 64 * stride_kv_row)
and return false or log an error if mismatched, ensuring the same validation
applies before calling dispatch_v32, dispatch_dsv4_single and
dispatch_dsv4_dual.
Comment on lines
+316
to
+321
| if (extra_KV_cache != nullptr) { | ||
| if (mt != ModelType::DSV4) return false; | ||
| return dispatch_dsv4_dual(num_heads, topk, topk_extra, extra_page_block_size, Q, KV_cache, | ||
| indices, extra_KV_cache, extra_indices, attn_sink, output, out_lse, | ||
| sm_scale, num_tokens, stride_kv_block, stride_kv_block_extra, | ||
| topk_length, extra_topk_length, stream); |
Contributor
There was a problem hiding this comment.
Reject half-populated dual-cache inputs.
Entering the dual-cache path only checks extra_KV_cache != nullptr, but the kernels dereference extra_indices unconditionally once they're on that path.
🔧 Minimal fix
+ const bool has_extra_kv = extra_KV_cache != nullptr;
+ const bool has_extra_indices = extra_indices != nullptr;
+ if (has_extra_kv != has_extra_indices) return false;
+
- if (extra_KV_cache != nullptr) {
+ if (has_extra_kv) {
if (mt != ModelType::DSV4) return false;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (extra_KV_cache != nullptr) { | |
| if (mt != ModelType::DSV4) return false; | |
| return dispatch_dsv4_dual(num_heads, topk, topk_extra, extra_page_block_size, Q, KV_cache, | |
| indices, extra_KV_cache, extra_indices, attn_sink, output, out_lse, | |
| sm_scale, num_tokens, stride_kv_block, stride_kv_block_extra, | |
| topk_length, extra_topk_length, stream); | |
| const bool has_extra_kv = extra_KV_cache != nullptr; | |
| const bool has_extra_indices = extra_indices != nullptr; | |
| if (has_extra_kv != has_extra_indices) return false; | |
| if (has_extra_kv) { | |
| if (mt != ModelType::DSV4) return false; | |
| return dispatch_dsv4_dual(num_heads, topk, topk_extra, extra_page_block_size, Q, KV_cache, | |
| indices, extra_KV_cache, extra_indices, attn_sink, output, out_lse, | |
| sm_scale, num_tokens, stride_kv_block, stride_kv_block_extra, | |
| topk_length, extra_topk_length, stream); |
🤖 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 `@csrc/sparse_mla_sm120_prefill.cu` around lines 316 - 321, The dual-cache fast
path currently branches on extra_KV_cache != nullptr but the kernels assume
extra_indices is valid; update the guard before calling dispatch_dsv4_dual to
ensure both extra_KV_cache and extra_indices are non-null (e.g., change the if
to require extra_KV_cache != nullptr && extra_indices != nullptr or explicitly
return false when extra_KV_cache is non-null but extra_indices is null) so
dispatch_dsv4_dual and its kernels never receive a half-populated dual-cache.
Comment on lines
+92
to
+103
| const int topk_len = topk_length_ptr ? __ldg(topk_length_ptr + t_idx) : TOPK; | ||
| const int extra_topk_len = | ||
| (extra_KV_cache != nullptr) | ||
| ? (extra_topk_length_ptr ? __ldg(extra_topk_length_ptr + t_idx) : extra_topk) | ||
| : 0; | ||
|
|
||
| // Chunk range this block owns. Total chunks = main + extra (extra layout | ||
| // is concatenated immediately after main; per-chunk dispatch in IO + math | ||
| // routes to the right source). | ||
| const int num_orig_chunks = (topk_len + DSV4_CAND_WINDOW - 1) / DSV4_CAND_WINDOW; | ||
| const int num_extra_chunks = (extra_topk_len + DSV4_CAND_WINDOW - 1) / DSV4_CAND_WINDOW; | ||
| const int num_chunks_total = num_orig_chunks + num_extra_chunks; |
Contributor
There was a problem hiding this comment.
Clamp runtime top-k lengths to the actual index buffers.
These lengths feed num_orig_chunks / num_extra_chunks directly. If a caller passes topk_length > TOPK or extra_topk_length > extra_topk, the gather path walks past the end of the corresponding indices tensor.
🔧 Minimal fix
- const int topk_len = topk_length_ptr ? __ldg(topk_length_ptr + t_idx) : TOPK;
- const int extra_topk_len =
+ const int topk_len_raw = topk_length_ptr ? __ldg(topk_length_ptr + t_idx) : TOPK;
+ const int topk_len = max(0, min(topk_len_raw, TOPK));
+ const int extra_topk_len_raw =
(extra_KV_cache != nullptr)
? (extra_topk_length_ptr ? __ldg(extra_topk_length_ptr + t_idx) : extra_topk)
: 0;
+ const int extra_topk_len = max(0, min(extra_topk_len_raw, extra_topk));🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuh` around
lines 92 - 103, Clamp the runtime top-k lengths before computing chunk counts:
when reading topk_length_ptr or extra_topk_length_ptr into
topk_len/extra_topk_len, bound them to the maximum available index sizes (TOPK
for the main indices and extra_topk for the extra_KV_cache path) to avoid
walking past index buffers; then use those clamped values when computing
num_orig_chunks, num_extra_chunks and num_chunks_total (refer to the variables
topk_len, extra_topk_len, TOPK, extra_topk, num_orig_chunks, num_extra_chunks).
| Honors ``attn_sink`` (FlashMLA V4 sink-merge convention) and | ||
| ``topk_length`` (per-token valid-length mask). | ||
| """ | ||
| num_tokens, num_heads, d_qk = q.shape |
Contributor
There was a problem hiding this comment.
Drop the unused num_heads binding.
Ruff's RUF059 will keep flagging this unpack because num_heads is never read in _ref_sparse_attn.
🔧 Minimal fix
- num_tokens, num_heads, d_qk = q.shape
+ num_tokens, _num_heads, d_qk = q.shape📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| num_tokens, num_heads, d_qk = q.shape | |
| num_tokens, _num_heads, d_qk = q.shape |
🧰 Tools
🪛 Ruff (0.15.13)
[warning] 232-232: Unpacked variable num_heads is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/attention/test_sparse_mla_sm120.py` at line 232, Remove the unused
num_heads binding by replacing the three-name unpack with an underscore
placeholder; in the context where q is unpacked (in _ref_sparse_attn / the
test), change "num_tokens, num_heads, d_qk = q.shape" to "num_tokens, _, d_qk =
q.shape" so Ruff no longer flags the unused variable.
lucifer1004
force-pushed
the
sparse-mla-sm120
branch
from
May 23, 2026 15:39
19591cc to
3609a6c
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
♻️ Duplicate comments (12)
tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json (1)
56-56:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winKeep the constraint aligned with the fixed
head_dim_qk.This trace fixes
head_dim_qkto512, but the constraint still permits576, so a validator can accept a shape this concrete artifact never produces.Suggested fix
- "head_dim_qk in (512, 576)", + "head_dim_qk == 512",🤖 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 `@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json` at line 56, Update the constraint "head_dim_qk in (512, 576)" so it matches the fixed head_dim_qk value used in the trace: remove 576 and restrict the allowed values to only 512; locate the constraint string "head_dim_qk in (512, 576)" in the trace JSON and replace it with a single-value constraint that only permits 512 to prevent validators from accepting shapes the artifact cannot produce.csrc/sparse_mla_sm120.cu (2)
99-101:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject non-8-bit cache tensors before byte reinterpretation.
Both cache paths are cast to
const uint8_t*, but neither verifies 8-bit scalar storage. A wider dtype changes the stride math and makes the kernel read the packed KV layout incorrectly.Suggested fix
CHECK_CUDA(kv_cache); CHECK_LAST_DIM_CONTIGUOUS(kv_cache); + TVM_FFI_ICHECK_EQ(kv_cache.dtype().bits, 8) + << "kv_cache must use 8-bit storage"; + TVM_FFI_ICHECK_EQ(kv_cache.dtype().lanes, 1) + << "kv_cache must be a scalar 8-bit tensor"; CHECK_INPUT_AND_TYPE(indices, dl_int32); CHECK_INPUT_AND_TYPE(output, dl_bfloat16); CHECK_INPUT_AND_TYPE(out_lse, dl_float32); @@ CHECK_CUDA(ekv); CHECK_LAST_DIM_CONTIGUOUS(ekv); + TVM_FFI_ICHECK_EQ(ekv.dtype().bits, 8) + << "extra_kv_cache must use 8-bit storage"; + TVM_FFI_ICHECK_EQ(ekv.dtype().lanes, 1) + << "extra_kv_cache must be a scalar 8-bit tensor"; CHECK_INPUT_AND_TYPE(eidx, dl_int32);Also applies to: 142-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 `@csrc/sparse_mla_sm120.cu` around lines 99 - 101, The cache pointer is later reinterpreted as const uint8_t* but the code only checks CUDA device and last-dim contiguity (CHECK_CUDA(kv_cache); CHECK_LAST_DIM_CONTIGUOUS(kv_cache);), so add a guard that rejects non-8-bit element types (or verifies element size == 1 byte) before any byte reinterpretation; update the same pattern around the other cache checks (the block near lines 142-149) so both cache paths validate scalar width is 8-bit (e.g., via a CHECK or ASSERT like CHECK_ELEMENT_SIZE(kv_cache,1) or CHECK_SCALAR_TYPE(kv_cache, uint8_t)) to prevent incorrect stride math when the tensor uses wider dtypes.
106-156:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate the non-
qtensor ranks and leading dimensions up front.
num_tokensandnum_headscome fromq, butindices,output,out_lse,topk_length, and the optional extra tensors are used without checking that they match. A mismatch here becomes device-side OOB instead of a clear host-side error.Suggested fix
CHECK_DIM(3, q); + CHECK_DIM(2, indices); + CHECK_DIM(3, output); + CHECK_DIM(2, out_lse); const int num_tokens = static_cast<int>(q.size(0)); const int num_heads = static_cast<int>(q.size(1)); const int d_qk = static_cast<int>(q.size(2)); const int topk = static_cast<int>(indices.size(-1)); const int page_block_size = static_cast<int>(kv_cache.size(-3)); + + TVM_FFI_ICHECK_EQ(indices.size(0), num_tokens) + << "indices must be [num_tokens, topk]"; + TVM_FFI_ICHECK_EQ(output.size(0), num_tokens); + TVM_FFI_ICHECK_EQ(output.size(1), num_heads); + TVM_FFI_ICHECK_EQ(out_lse.size(0), num_tokens); + TVM_FFI_ICHECK_EQ(out_lse.size(1), num_heads); @@ if (extra_kv_cache.has_value()) { TVM_FFI_ICHECK(extra_indices.has_value()) << "extra_kv_cache requires extra_indices"; const auto& ekv = extra_kv_cache.value(); const auto& eidx = extra_indices.value(); + CHECK_DIM(2, eidx); + TVM_FFI_ICHECK_EQ(eidx.size(0), num_tokens) + << "extra_indices must be [num_tokens, extra_topk]"; @@ } + + if (topk_length.has_value()) { + const auto& tl = topk_length.value(); + CHECK_INPUT_AND_TYPE(tl, dl_int32); + CHECK_DIM(1, tl); + TVM_FFI_ICHECK_EQ(tl.size(0), num_tokens) + << "topk_length must be [num_tokens]"; + } + if (extra_topk_length.has_value()) { + const auto& etl = extra_topk_length.value(); + CHECK_INPUT_AND_TYPE(etl, dl_int32); + CHECK_DIM(1, etl); + TVM_FFI_ICHECK_EQ(etl.size(0), num_tokens) + << "extra_topk_length must be [num_tokens]"; + }🤖 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 `@csrc/sparse_mla_sm120.cu` around lines 106 - 156, The code reads sizes from q (num_tokens, num_heads) but does not validate that tensors used later (indices, output, out_lse, topk_length, extra_indices, extra_kv_cache, extra_topk_length) share the expected ranks and leading dimensions, which can cause device OOB; add host-side checks after computing num_tokens/num_heads/topk/page_block_size that assert indices.ndim()==3 (or expected rank), indices.size(0)==num_tokens, indices.size(1)==num_heads and indices.size(-1)==topk, output.ndim()/out_lse.ndim() and their leading dims match num_tokens and num_heads, topk_length (if present) has size(num_tokens), and for extra_inputs ensure extra_indices.size(0)==num_tokens, extra_indices.size(1)==num_heads, extra_topk==extra_indices.size(-1), and extra_kv_cache.size(-3)==extra_page_block_size; use the same TVM_FFI_ICHECK_EQ/ICHECK_* helpers and CHECK_* macros already in this file to fail fast on the host (refer to symbols: indices, output, out_lse, topk_length, extra_indices, extra_kv_cache, extra_topk_length, attn_sink).csrc/sparse_mla_sm120_decode_dsv3_2.cu (2)
110-118:⚠️ Potential issue | 🔴 CriticalReject oversized
num_splitsbefore launching the DSV4 merge kernel.This path reuses
sparse_mla_decode_dsv4_merge_kernelbut does not mirror the DSV4 entrypoint'snum_splits > DSV4_MAX_SPLITSguard. Larger values can overflow the merge kernel's fixed split staging.🤖 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 `@csrc/sparse_mla_sm120_decode_dsv3_2.cu` around lines 110 - 118, Before launching sparse_mla_decode_dsv4_merge_kernel, add the same guard used in the DSV4 entrypoint to reject or clamp oversize splits: check num_splits against DSV4_MAX_SPLITS (or the equivalent constant) and return/log/error out if num_splits > DSV4_MAX_SPLITS, or set num_splits = DSV4_MAX_SPLITS before calling merge_kernel; update the call site around merge_kernel/sparse_mla_decode_dsv4_merge_kernel to use the validated/clamped value so the merge kernel cannot overflow its fixed split staging.
79-80:⚠️ Potential issue | 🟠 MajorUse the active device for the SM-count heuristic.
This still queries
cudaDevAttrMultiProcessorCounton device0, sochunks_per_blockcan be chosen from the wrong SM count in multi-GPU runs.🤖 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 `@csrc/sparse_mla_sm120_decode_dsv3_2.cu` around lines 79 - 80, The code fetches SM count with cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0) which always queries device 0; call cudaGetDevice(&device) (or cudaGetDevice(¤tDevice)) first and pass that device id into cudaDeviceGetAttribute so sm_count reflects the active device used by chunks_per_block selection (update usages around sm_count, cudaDeviceGetAttribute, and chunks_per_block).csrc/sparse_mla_sm120_decode_dsv4.cu (1)
81-82:⚠️ Potential issue | 🟠 MajorUse the active device for the SM-count heuristic.
This still queries
cudaDevAttrMultiProcessorCounton device0, sochunks_per_blockis tuned against the wrong GPU whenever the launch runs on another device.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/sparse_mla_sm120_decode_dsv4.cu` around lines 81 - 82, The code reads cudaDevAttrMultiProcessorCount for device 0, which can be wrong if another GPU is active; change the call to query the active device by first obtaining the current device (e.g., via cudaGetDevice) then pass that device ID to cudaDeviceGetAttribute when populating sm_count (used by the chunks_per_block heuristic). Update the CUDA_CHECK_BOOL call around cudaDeviceGetAttribute to use the retrieved device ID so sm_count and any downstream logic (chunks_per_block) are tuned for the active GPU.csrc/sparse_mla_sm120_jit_binding.cu (2)
114-117:⚠️ Potential issue | 🟠 MajorValidate
extra_indicesbefore dereferencing it.This block still assumes
extra_indicesis present wheneverextra_kv_cacheis set and unconditionally callsextra_indices.value(). A malformed call will crash inside the binding instead of failing with a clean validation error.🤖 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 `@csrc/sparse_mla_sm120_jit_binding.cu` around lines 114 - 117, The code assumes extra_indices is present whenever extra_kv_cache is set and dereferences extra_indices.value() directly; add a validation guard so that inside the block that checks extra_kv_cache.has_value() you first verify extra_indices.has_value() and emit a clear validation error (or return/fail) if it's missing before using extra_indices.value() to compute extra_topk_arg or any sizes; keep the existing logic that reads ekv = extra_kv_cache.value() and checks ekv.ndim() only after extra_indices presence is confirmed.
53-55:⚠️ Potential issue | 🟠 MajorTighten decode KV-cache rank/layout validation.
Both decode bindings still accept any
kv_cache.ndim() >= 2, but the non-2D branch immediately treats the tensor as the 4-D paged layout. That means 3-D/5-D tensors, or 4-D tensors with the wrong page dimension, can slip through and be launched underpage_block_size=64assumptions.Also applies to: 117-125, 151-153
flashinfer/sparse_mla_sm120.py (3)
436-437: 🛠️ Refactor suggestion | 🟠 MajorAdd explicit trace metadata to the wrapper decorators.
BatchSparseMLAPagedAttentionWrapper.__init__andBatchSparseMLAPagedAttentionWrapper.runare public@flashinfer_apisurfaces, but both decorators still omittrace=. That breaks the tracing contract for wrapper calls. As per coding guidelines "Every public API decorated with@flashinfer_apishould carry atrace=argument so thatfi_trace()works and auto-dump produces a benchmark definition JSON."Also applies to: 482-483
🤖 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 `@flashinfer/sparse_mla_sm120.py` around lines 436 - 437, The public methods BatchSparseMLAPagedAttentionWrapper.__init__ and BatchSparseMLAPagedAttentionWrapper.run are decorated with `@flashinfer_api` but lack the required trace= metadata; update the decorators on both definitions (the lines with `@flashinfer_api` above __init__ and run) to include a trace argument (e.g., trace="BatchSparseMLAPagedAttentionWrapper.__init__" and trace="BatchSparseMLAPagedAttentionWrapper.run") so the fi_trace() machinery and auto-dump can produce proper benchmark JSONs; keep the existing supported_compute_capability decorators intact.
327-334:⚠️ Potential issue | 🟠 MajorReject non-512
d_von the SM120 public surface.The decode launchers never receive
d_v; they always execute the fixed-width 512-value kernels. Allowing callers to sizeoutputand_mid_outwith a smallerd_vcan under-allocate the decode path and turnnum_tokens <= 64calls into out-of-bounds writes.Also applies to: 443-470
🤖 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 `@flashinfer/sparse_mla_sm120.py` around lines 327 - 334, The public SM120 entry points must reject any d_v != 512 to prevent under-allocation of decode kernels; add an explicit check at the start of the public functions that accept d_v (look for the function(s) with parameters named d_v, topk_length, attn_sink, extra_kv_cache, extra_indices, extra_topk_length, mid_out, mid_lse and the analogous overload later) and raise a clear exception (e.g., ValueError) if d_v != 512, and also validate that output/_mid_out tensors are sized for 512 elements per head; apply the same check to the other matching public surface (the second block noted in the review) so decode launchers always assume fixed-width 512.
240-308:⚠️ Potential issue | 🟠 MajorClamp runtime top-k lengths before dispatching.
topk_lengthandextra_topk_lengthare still forwarded unchecked. If either exceeds its paired index width, the decode/prefill kernels can walk past the providedindices/extra_indicesbuffers.🤖 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 `@flashinfer/sparse_mla_sm120.py` around lines 240 - 308, Clamp the runtime top-k lengths (topk_length and extra_topk_length) to the actual index widths before any dispatch/FFI call to avoid out-of-bounds reads: compute clipped_topk = min(topk_length, indices.shape[-1]) and clipped_extra_topk = min(extra_topk_length, extra_indices.shape[-1] if extra_indices is not None else 0) and use these clipped values in the calls to sparse_mla_sm120_decode_dsv4, module.sparse_mla_sm120_decode_dsv3_2, and module.sparse_mla_sm120_paged_attention (and in any dispatch predicate logic that depends on topk_length/extra_topk_length) so kernels never see a length larger than the corresponding indices buffer.include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh (1)
489-495:⚠️ Potential issue | 🟠 MajorGuard the second-half sink load when
NUM_HEADS == 8.This SG epilogue still unconditionally reads
cold.attn_sink[h_start + gid + 8]. For the small-TP path the public sink buffer only has 8 elements, so this reads past the caller allocation even though heads 8..15 only exist as padding in shared memory.🤖 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 `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh` around lines 489 - 495, The epilogue unconditionally reads cold.attn_sink[h_start + gid + 8] even when only 8 public heads exist, causing an out-of-bounds access for the small-TP path; fix it by guarding the second-half load and its dependent computations (the reads of cold.attn_sink + h_start + gid + 8, denom1 and il1) so they only execute when the build/packet actually contains the upper 8 heads (e.g. use a compile-time guard like `#if` NUM_HEADS > 8 / `#endif` or an equivalent runtime check) and leave the first-half (s0/denom0/il0) unchanged — target the expressions around cold.attn_sink, h_start, gid, denom1, il1 and sm.m_smem to locate the code to change.
🧹 Nitpick comments (1)
tests/trace/test_fi_trace_template_consistency.py (1)
391-394: ⚡ Quick winAdd a targeted
fi_tracecompleteness test before skipping this template.Skipping
sparse_mla_sm120_pagedhere leaves the new trace covered only by structural checks and the non-test example script, so regressions in resolved dtypes or required Const-axis values on the common single-cache path will no longer be caught in CI. A focused end-to-end test for the single-cache form would keep coverage without needing to model the optional extra-cache axes.🤖 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 `@tests/trace/test_fi_trace_template_consistency.py` around lines 391 - 394, Add a focused end-to-end fi_trace completeness test for the single-cache form of the "sparse_mla_sm120_paged" template before it is skipped: in tests/trace/test_fi_trace_template_consistency.py locate the skip/blacklist entry for "sparse_mla_sm120_paged" and insert a small test that constructs inputs that exercise the single-cache path (omit or pass None for optional extra_kv_cache/extra_indices axes or set them to empty/absent shapes), call the same fi_trace/trace-completeness helper used elsewhere in this file, and assert the trace resolves dtypes and Const-axis values (e.g., extra_topk/extra_page_block_size) as expected; this keeps CI coverage for resolved dtypes/Const-axis without modeling the optional extra-cache axes.
🤖 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 `@csrc/sparse_mla_sm120_prefill.cu`:
- Around line 61-65: The code currently sets a function-local static configured
= true regardless of cudaFuncSetAttribute(...) success, which can make failures
sticky and is not per-device; change the logic in the prefill kernel launcher(s)
that reference configured, kernel and smem_bytes so that you check the return
value of cudaFuncSetAttribute(...) (cudaError_t rc = cudaFuncSetAttribute(...))
and only set the cached flag to true when rc == cudaSuccess, logging/propagating
the error otherwise; also replace the single function-local static configured
with a per-device cache (e.g. a small array or unordered_map keyed by
cudaGetDevice() or use cudaGetDeviceProperties device id) so the configured flag
is tracked per CUDA device for the kernel in csrc/sparse_mla_sm120_prefill.cu.
In `@flashinfer/trace/templates/attention.py`:
- Around line 1292-1407: The template sparse_mla_sm120_decode_dsv4_trace
advertises "optional dual-cache decode" but its TraceTemplate inputs/constraints
don't include the dual-cache arguments, so fi_trace cannot model callsites that
pass extra_kv_cache/extra_indices/extra_topk_length; update the TraceTemplate to
add optional inputs named "extra_kv_cache" (same shape/dtype as "kv_cache"),
"extra_indices" (same shape/dtype as "indices"), and "extra_topk_length"
(optional int32 like "topk_length"), mark them optional in the inputs dict of
sparse_mla_sm120_decode_dsv4_trace, and add corresponding constraints (e.g.,
extra_indices.shape[0] == num_tokens and extra_indices.shape[-1] == extra_topk
or extra_topk_length handling) so traces can represent both single- and
dual-cache callsites when TraceTemplate is used by fi_trace().
In `@include/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuh`:
- Around line 365-466: The bug is that fully-masked half-heads (local_max ==
sentinel) still produce p==1 and nonzero sums; fix by detecting sentinel
local_max values (e.g., local_max[0] or local_max[1] <= -1e29f) after the
per-warp reduction and zeroing-out that half-head's contributions: set
p[nt][0/1] = 0 (for half 0) or p[nt][2/3] = 0 (for half 1) for all nt and ensure
corresponding local_sum elements are 0 (or set warp_rescale0/1 = 0) before
writing sm_warp_sum/sm_warp_max and before computing block_local_sum0/1 so
fully-invalid half-heads don't contribute the clamped slot-0 entry; update code
around local_max, p, local_sum, warp_rescale, and the place that writes
sm_warp_sum (use the existing symbols local_max, p, local_sum, warp_rescale0/1,
sm_warp_sum, block_local_sum0/1).
In `@include/flashinfer/attention/sparse_mla_sm120/model/model_type.h`:
- Around line 36-41: The BF16 dispatch note in the ComputeMode comment is stale;
update the comment near ComputeMode to state that BF16 is now also used for the
DSV4 prefill path (specifically for topk == 128 and for dual-cache variants) —
mention ComputeMode::BF16, the DSV4 prefill path, and the topk==128 / dual-cache
cases so readers know BF16 is exercised by current dispatch sites.
In `@tests/attention/test_sparse_mla_sm120.py`:
- Around line 55-58: The test suite is incorrectly skipped on SM121 because the
pytestmark only checks is_sm120a_supported(torch.device("cuda")); update the
gate so the suite runs on both SM120a and SM121a by changing the condition to
allow either architecture (e.g., use is_sm120a_supported(...) or
is_sm121a_supported(...)) or replace with a compute-capability check via
get_compute_capability() to accept compute capability for SM120/SM121; modify
the pytestmark reference that currently uses is_sm120a_supported and
torch.device("cuda") accordingly so the file is not skipped on SM121.
In `@tests/trace/example.py`:
- Around line 459-503: The new example call to
flashinfer.sparse_mla_sm120_decode_dsv4 in tests/trace/example.py was added but
the corresponding fi_trace_out JSON artifact was not committed; rerun the
example script (execute tests/trace/example.py) so it generates the matching
trace output for the sparse_mla_sm120_decode_dsv4 call, add the produced
fi_trace_out JSON to the repo alongside the existing paged-attention trace, and
commit that generated file so the code example and checked-in traces remain
aligned.
---
Duplicate comments:
In `@csrc/sparse_mla_sm120_decode_dsv3_2.cu`:
- Around line 110-118: Before launching sparse_mla_decode_dsv4_merge_kernel, add
the same guard used in the DSV4 entrypoint to reject or clamp oversize splits:
check num_splits against DSV4_MAX_SPLITS (or the equivalent constant) and
return/log/error out if num_splits > DSV4_MAX_SPLITS, or set num_splits =
DSV4_MAX_SPLITS before calling merge_kernel; update the call site around
merge_kernel/sparse_mla_decode_dsv4_merge_kernel to use the validated/clamped
value so the merge kernel cannot overflow its fixed split staging.
- Around line 79-80: The code fetches SM count with
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0) which
always queries device 0; call cudaGetDevice(&device) (or
cudaGetDevice(¤tDevice)) first and pass that device id into
cudaDeviceGetAttribute so sm_count reflects the active device used by
chunks_per_block selection (update usages around sm_count,
cudaDeviceGetAttribute, and chunks_per_block).
In `@csrc/sparse_mla_sm120_decode_dsv4.cu`:
- Around line 81-82: The code reads cudaDevAttrMultiProcessorCount for device 0,
which can be wrong if another GPU is active; change the call to query the active
device by first obtaining the current device (e.g., via cudaGetDevice) then pass
that device ID to cudaDeviceGetAttribute when populating sm_count (used by the
chunks_per_block heuristic). Update the CUDA_CHECK_BOOL call around
cudaDeviceGetAttribute to use the retrieved device ID so sm_count and any
downstream logic (chunks_per_block) are tuned for the active GPU.
In `@csrc/sparse_mla_sm120_jit_binding.cu`:
- Around line 114-117: The code assumes extra_indices is present whenever
extra_kv_cache is set and dereferences extra_indices.value() directly; add a
validation guard so that inside the block that checks extra_kv_cache.has_value()
you first verify extra_indices.has_value() and emit a clear validation error (or
return/fail) if it's missing before using extra_indices.value() to compute
extra_topk_arg or any sizes; keep the existing logic that reads ekv =
extra_kv_cache.value() and checks ekv.ndim() only after extra_indices presence
is confirmed.
In `@csrc/sparse_mla_sm120.cu`:
- Around line 99-101: The cache pointer is later reinterpreted as const uint8_t*
but the code only checks CUDA device and last-dim contiguity
(CHECK_CUDA(kv_cache); CHECK_LAST_DIM_CONTIGUOUS(kv_cache);), so add a guard
that rejects non-8-bit element types (or verifies element size == 1 byte) before
any byte reinterpretation; update the same pattern around the other cache checks
(the block near lines 142-149) so both cache paths validate scalar width is
8-bit (e.g., via a CHECK or ASSERT like CHECK_ELEMENT_SIZE(kv_cache,1) or
CHECK_SCALAR_TYPE(kv_cache, uint8_t)) to prevent incorrect stride math when the
tensor uses wider dtypes.
- Around line 106-156: The code reads sizes from q (num_tokens, num_heads) but
does not validate that tensors used later (indices, output, out_lse,
topk_length, extra_indices, extra_kv_cache, extra_topk_length) share the
expected ranks and leading dimensions, which can cause device OOB; add host-side
checks after computing num_tokens/num_heads/topk/page_block_size that assert
indices.ndim()==3 (or expected rank), indices.size(0)==num_tokens,
indices.size(1)==num_heads and indices.size(-1)==topk,
output.ndim()/out_lse.ndim() and their leading dims match num_tokens and
num_heads, topk_length (if present) has size(num_tokens), and for extra_inputs
ensure extra_indices.size(0)==num_tokens, extra_indices.size(1)==num_heads,
extra_topk==extra_indices.size(-1), and
extra_kv_cache.size(-3)==extra_page_block_size; use the same
TVM_FFI_ICHECK_EQ/ICHECK_* helpers and CHECK_* macros already in this file to
fail fast on the host (refer to symbols: indices, output, out_lse, topk_length,
extra_indices, extra_kv_cache, extra_topk_length, attn_sink).
In `@flashinfer/sparse_mla_sm120.py`:
- Around line 436-437: The public methods
BatchSparseMLAPagedAttentionWrapper.__init__ and
BatchSparseMLAPagedAttentionWrapper.run are decorated with `@flashinfer_api` but
lack the required trace= metadata; update the decorators on both definitions
(the lines with `@flashinfer_api` above __init__ and run) to include a trace
argument (e.g., trace="BatchSparseMLAPagedAttentionWrapper.__init__" and
trace="BatchSparseMLAPagedAttentionWrapper.run") so the fi_trace() machinery and
auto-dump can produce proper benchmark JSONs; keep the existing
supported_compute_capability decorators intact.
- Around line 327-334: The public SM120 entry points must reject any d_v != 512
to prevent under-allocation of decode kernels; add an explicit check at the
start of the public functions that accept d_v (look for the function(s) with
parameters named d_v, topk_length, attn_sink, extra_kv_cache, extra_indices,
extra_topk_length, mid_out, mid_lse and the analogous overload later) and raise
a clear exception (e.g., ValueError) if d_v != 512, and also validate that
output/_mid_out tensors are sized for 512 elements per head; apply the same
check to the other matching public surface (the second block noted in the
review) so decode launchers always assume fixed-width 512.
- Around line 240-308: Clamp the runtime top-k lengths (topk_length and
extra_topk_length) to the actual index widths before any dispatch/FFI call to
avoid out-of-bounds reads: compute clipped_topk = min(topk_length,
indices.shape[-1]) and clipped_extra_topk = min(extra_topk_length,
extra_indices.shape[-1] if extra_indices is not None else 0) and use these
clipped values in the calls to sparse_mla_sm120_decode_dsv4,
module.sparse_mla_sm120_decode_dsv3_2, and
module.sparse_mla_sm120_paged_attention (and in any dispatch predicate logic
that depends on topk_length/extra_topk_length) so kernels never see a length
larger than the corresponding indices buffer.
In `@include/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuh`:
- Around line 489-495: The epilogue unconditionally reads cold.attn_sink[h_start
+ gid + 8] even when only 8 public heads exist, causing an out-of-bounds access
for the small-TP path; fix it by guarding the second-half load and its dependent
computations (the reads of cold.attn_sink + h_start + gid + 8, denom1 and il1)
so they only execute when the build/packet actually contains the upper 8 heads
(e.g. use a compile-time guard like `#if` NUM_HEADS > 8 / `#endif` or an equivalent
runtime check) and leave the first-half (s0/denom0/il0) unchanged — target the
expressions around cold.attn_sink, h_start, gid, denom1, il1 and sm.m_smem to
locate the code to change.
In
`@tests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.json`:
- Line 56: Update the constraint "head_dim_qk in (512, 576)" so it matches the
fixed head_dim_qk value used in the trace: remove 576 and restrict the allowed
values to only 512; locate the constraint string "head_dim_qk in (512, 576)" in
the trace JSON and replace it with a single-value constraint that only permits
512 to prevent validators from accepting shapes the artifact cannot produce.
---
Nitpick comments:
In `@tests/trace/test_fi_trace_template_consistency.py`:
- Around line 391-394: Add a focused end-to-end fi_trace completeness test for
the single-cache form of the "sparse_mla_sm120_paged" template before it is
skipped: in tests/trace/test_fi_trace_template_consistency.py locate the
skip/blacklist entry for "sparse_mla_sm120_paged" and insert a small test that
constructs inputs that exercise the single-cache path (omit or pass None for
optional extra_kv_cache/extra_indices axes or set them to empty/absent shapes),
call the same fi_trace/trace-completeness helper used elsewhere in this file,
and assert the trace resolves dtypes and Const-axis values (e.g.,
extra_topk/extra_page_block_size) as expected; this keeps CI coverage for
resolved dtypes/Const-axis without modeling the optional extra-cache axes.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 45221b76-5edb-4862-9922-8fd1e0faa517
📥 Commits
Reviewing files that changed from the base of the PR and between 19591ccd8b61048a701f1ca9a066967fe6c6f463 and 3609a6c.
📒 Files selected for processing (35)
.gitignorebenchmarks/bench_sparse_mla_sm120.pycsrc/sparse_mla_sm120.cucsrc/sparse_mla_sm120_decode_dsv3_2.cucsrc/sparse_mla_sm120_decode_dsv4.cucsrc/sparse_mla_sm120_jit_binding.cucsrc/sparse_mla_sm120_prefill.cuflashinfer/__init__.pyflashinfer/aot.pyflashinfer/jit/sparse_mla_sm120.pyflashinfer/sparse_mla_sm120.pyflashinfer/trace/templates/attention.pyinclude/flashinfer/attention/sparse_mla_sm120/arch/barrier.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/common.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/cp_async.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/ldmatrix_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/arch/mma_sm120.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/d2_load_b.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/fp8_quant.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/online_softmax.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/q_rope.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/smem_layout.cuhinclude/flashinfer/attention/sparse_mla_sm120/common/xv_rope_mma.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv3_2_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/decode_dsv4_kernel.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/kv_cache_traits.cuhinclude/flashinfer/attention/sparse_mla_sm120/model/model_type.hinclude/flashinfer/attention/sparse_mla_sm120/model/scale_convert.cuhinclude/flashinfer/attention/sparse_mla_sm120/prefill_kernel.cuhtests/attention/test_sparse_mla_sm120.pytests/trace/example.pytests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/fi_trace_out/sparse_mla_sm120_paged_h32_dqk512_dv512_topk512_ps64_kvb584.jsontests/trace/test_fi_trace_template_consistency.py
✅ Files skipped from review due to trivial changes (2)
- .gitignore
- tests/trace/fi_trace_out/sparse_mla_sm120_decode_dsv4_h32_dqk512_dv512_topk512_ps64_kvb584.json
Collaborator
|
[FAILED] Pipeline #54602527: 9/20 passed |
5 tasks done
sekfung
pushed a commit
to sekfung/sglang
that referenced
this pull request
Jul 2, 2026
Integrate FlashInfer's SM120 sparse_mla_sm120 decode kernel for
DeepSeek-V4-Flash on RTX PRO 6000 (SM120). Decode ITL improves
1.84x vs Triton (59ms vs 109ms, ISL=8K OSL=1K BS=1).
Key challenge: SGLang SWA KV cache uses page_size=256 with footer
layout (scales at page end), but FlashInfer decode_dsv4 requires
page_block_size=64. Solution: page-split at call site — split
256-token pages into 4 virtual 64-token pages with correct footer
layout and remap indices before calling FlashInfer.
- Auto-detect FlashInfer SM120 via try-import, fallback to Triton
- Page-split + index remap in _flash_mla_flashinfer()
- Uses BatchSparseMLAPagedAttentionWrapper.run() for decode dispatch
- Extra cache (C4/C128) passed as-is (native dual-cache support)
- GSM8K 10/10 = 100%, CUDA graph BS=1,2,4,8 all pass
- Env var SGLANG_SM120_FLASHMLA_BACKEND={auto,flashinfer,triton}
Requires: flashinfer >= 0.6.12 with sparse-mla-sm120 support
(lucifer1004/flashinfer@sparse-mla-sm120, PR flashinfer-ai/flashinfer#3395)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
aldakata
added a commit
to aldakata/flashinfer
that referenced
this pull request
Jul 15, 2026
quantize_kv_model1 viewed the packed buffer as NHD, rejected by _check_sm120_dsv4_kv_cache_layout in the public API since flashinfer-ai#3395; metadata-only change
25 of 68 tasks
yyihuang
added a commit
that referenced
this pull request
Sep 27, 2026
… (padded Q, separate SWA/compressed tables, caller-owned workspace, length offsets) and regenerate the SM100/SM103 programs (#5591) ## Summary Hardens the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a) for the asks in #4671 and re-exports the generated programs: - padded query rows beyond the metadata token count are ignored (no metadata read past `num_tokens`, `-1` never dereferenced, padded output rows untouched); - separate SWA and compressed index tables (`extra_sparse_indices`, `extra_sparse_topk_lens`) are consumed in-kernel, the combined table keeps working; - caller-owned workspace: `get_cake_dsv4_workspace_bytes(...)`, no device allocation on any call path (MTP, CUDA graph capture), split counters self-reset; - raw lengths plus `sparse_topk_lens_offset`; - many-token (MTP) rows route to the bodies that beat trtllm-gen (16-512 tokens): BF16/H128 SWA and SWA+topk4x rows use the persistent prefill body from 64 tokens (`_BF16_H128_PREFILL_MIN_TOKENS`), BF16/H128 topk128x width 260/388 rows above 16 tokens use the row-first single-owner program (`_BF16_TOPK128X_SPLIT_MAX_TOKENS`); FP8/H64 rows admitted to the persistent body with >= 128 tokens use the H64-specific single-CTA M64 program (`fp8_h64_prefill_source_persistent_m64`, `_FP8_H64_M64_MIN_TOKENS`); BF16/H32 topk4x rows share the retained-KV topk128x program (`bf16_h32_topk128x_early_v47`, last-arriver split merge); - BF16/H128 SWA+topk4x rows whose persistent-grid tail clusters are the majority launch the boustrophedon work-feed program of the prefill body (`bf16_h128_prefill_v42_snake`, `_bf16_h128_prefill_uses_snake_feed`); - the boustrophedon program `bf16_h128_prefill_v42_snake` is regenerated from a K-reuse body (the PV operand is served from the gathered K tile, 25 % fewer L2 bytes per KV tile); the striped program is byte-identical to the previous export; - the FP8/H128 persistent programs rescale the O accumulator lazily (the running max only advances when a tile raises it by more than 8 in the exp2 domain, so most tiles skip the TMEM rescale on the serial PV chain; P stays inside e4m3, O / sum unchanged); - the BF16/H128 persistent prefill program drains O with 32-byte stores (half the LSU instruction count of the epilogue tail); - FP8 H8/H16 canonical 12-token rows (batch 3, ragged q_len 5, widths 192/256 topk4x or 260 topk128x) use the FP8 port of the 1-CTA SwapsAb body (`fp8_h8_h16_source_exact`, heads on the MMA N side), 1.24-1.49x vs the default path on both arches; - 12-token decode routing follows the Cake seeds: FP8 low-head widths up to 384 -> one-partition producer, FP8/H64 rows with >= 3 complete sparse tiles -> the persistent FP8 body, BF16/H128 topk128x width 260/388 -> the two-stage split3/split4 programs on both targets. Also hardens the default trtllm-gen host path (#4671 P1/P4, see `flashinfer/mla/_core.py`). ## Validation Export r9 (Cake 421f1a5671a, this branch at f63041c + the regenerated programs), same-session paired CUPTI (active-union median, 500 warmup / 3000 calls per arm, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 | 134/134 | | **94 canonical rows faster than the default path** | **94/94, geomean 1.24x, min 1.006x** | **94/94, geomean 1.23x, min 1.013x** | | 40 hardening rows faster than the default path | **40/40, geomean 1.28x, min 1.016x** | **40/40, geomean 1.25x, min 1.029x** | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, on the exported programs) | 292 passed | 292 passed | | `tests/mla/test_cake_dsv4*.py` (CPU, this host) | 208 passed | 208 passed | The default trtllm-gen path itself is hardened for #4671 P1/P4 in `flashinfer/mla/_core.py`. ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21, `flashinfer/decode.py`, `flashinfer/mla/_core.py`, DSv4 TRTLLM-GEN kernels); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16, `flashinfer/mla/cake_dsv4.py`, `flashinfer/jit/cake_dsv4.py`, `csrc/cake_dsv4/`), regenerated by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py` (DSv4 RopeQuant exposure #4918, 27d5b02, 2026-09-07). - **Branch merge-base** with upstream `main` at measurement: d3af75a. The branch has since merged upstream `main` c0ea2a1; of the upstream commits in between, the only ones touching `flashinfer/mla/_core.py` (#4557, #5552, #5577) change `_trtllm_batch_decode_with_kv_cache_mla_impl` (other `backend="cake"` MLA decode paths) and leave the DSv4 default path and `trtllm_batch_decode_sparse_mla_dsv4` unchanged, so the comparison above still describes the merged head. ## CI `tests/mla/test_cake_dsv4.py` and `tests/mla/test_cake_dsv4_hardening.py` (CPU and GPU parametrizations). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
added a commit
that referenced
this pull request
Sep 27, 2026
…epilogue stores in the DSv4 sparse-MLA Cake programs (SM100/SM103) (#5610) ## Summary Round-2 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5591: - **FP8/H128 persistent prefill program (prefill-style rows with 128+ query tokens, batch 2-48 items x 16 KV tiles):** the two K and two V loader warps issued their TMA `gather4` with lane-divergent coordinates, so ptxas serialised every gather into a 32-iteration per-lane loop and the loaders spent ~1.0 us of every 1.12-us KV tile issuing. They now issue from one elected lane with warp-uniform coordinates (`shfl.sync` broadcast, one unrolled 16-quartet site per warp with the KV descriptor selected per tile), the same form the default trtllm-gen kernels use. Same barriers, pipeline depth and data movement; only the issuing lane changes. The loop is now paced by the tensor pipe (16 QK + 8 PV MMAs per tile, ~0.8 us). prefill-style-000089 / 000093: GB300 1.02x / 1.01x -> **1.24x / 1.21x**, B200 1.05x / 1.06x -> **1.19x / 1.18x** vs the default path. The persistent seed compiles with `-Xptxas=--register-usage-level=10` so the exported build keeps the softmax schedule of the source build (without it ptxas re-orders the exp2/convert chains across the P-publication fence: 3-5 % slower on these rows). - **BF16/H128 K-reuse (snake) prefill program (`bf16_h128_prefill_v42_snake`, hardening rows 27 / 37):** the epilogue drains O with 32-byte stores like the striped program already did. hardening-000037: GB300 1.016x -> **1.107x**, B200 1.028x -> **1.100x**; hardening-000027: 1.175x -> 1.190x, 1.172x -> 1.173x. - Every other program is a byte-identical regeneration (renamed for the new producer revision). - New CPU test: registered programs build with public nvcc/ptxas options only, and the FP8/H128 persistent variants pin the register-usage level. The one-KV-tile SWA-only work items of the BF16/H128 persistent prefill body (hardening rows 23 / 29 / 35) were re-attributed and left unchanged: the item is a per-SM chain bound by the SM's L2<->SM port (256 KB per CTA per item at the measured ~66 GB/s = the measured 3.9 / 4.1 us period); the byte levers (K-reuse body, four issuing warps, DSM exchange of the V halves) measured negative or closed by arithmetic. Those rows stay at 1.03-1.16x. ## Validation Export r11 (Cake `db0609184cd`, this branch at 8589d49 + the regenerated programs), same-session paired CUPTI (active-union median, 500 warmup / 3000 calls per arm, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 | 134/134 | | **94 canonical rows faster than the default path** | **94/94, geomean 1.243x, min 1.054x** (#5591: 1.237x, min 1.006x) | **94/94, geomean 1.232x, min 1.041x** (#5591: 1.226x, min 1.013x) | | 40 hardening rows faster than the default path | **40/40, geomean 1.30x, min 1.058x** (#5591: 1.28x, min 1.016x) | **40/40, geomean 1.25x, min 1.037x** (#5591: 1.25x, min 1.029x) | | exported program vs Cake source build (agreement gate, 134 rows) | geomean 1.002x, every row within 3 % | geomean 1.004x, every row within 3 % | | compute-sanitizer synccheck + memcheck on the two changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, on the exported programs) | 292 passed | 292 passed | | `tests/mla/test_cake_dsv4*.py` (CPU, this host) | 210 passed | 210 passed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21, `flashinfer/decode.py`, `flashinfer/mla/_core.py`, DSv4 TRTLLM-GEN kernels); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16) as hardened and regenerated by #5591 (8589d49, 2026-09-27), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py` (DSv4 RopeQuant exposure #4918, 27d5b02, 2026-09-07). - **Branch base** = upstream `main` 8589d49 (the #5591 merge); no upstream commits in between. ## CI `tests/mla/test_cake_dsv4.py` and `tests/mla/test_cake_dsv4_hardening.py` (CPU and GPU parametrizations). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved correctness of CAKE DSV4 attention kernels by refining shared-memory synchronization, sparse gather operations, and pipeline phase handling. - Updated output processing to use more efficient vectorized stores in selected kernels. - Corrected kernel registration and dispatch across supported GPU architectures, including updates to available variants and compile settings. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This was referenced Sep 28, 2026
yyihuang
added a commit
that referenced
this pull request
Sep 28, 2026
…03 (persistent-item epilogue, index ring, retained-KV decode body, SwapsAb V alias, routing) (#5630) ## Summary Round-3 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5610: - **FP8/H128 persistent prefill program (128+ query-token rows):** the steady KV tile already runs at the tensor-pipe floor; the per-item epilogue drained the 64x256 fp32 accumulator with 16-byte stores (4096 half-sector transactions per CTA per item). It now uses 32-byte stores. prefill-style-000089 / 000093: GB300 1.24x / 1.23x -> **1.33x / 1.31x**, B200 1.17x / 1.18x -> **1.25x / 1.25x** vs the default path. - **Sub-128-token FP8/H128 uniform-gather program:** compiled without `--register-usage-level=10` (kept on the lane-gather program): -5.8 % SASS, GB300 decode rows 007/010/031/034 1.24x -> **1.26x**. - **BF16/H128 striped prefill program:** the sparse-index block of the next work item is staged during the current item (two-slot ring) and loaded batched. prefill-style-000088 / 000092: GB300 1.23x -> **1.34x / 1.35x**, B200 1.13x -> **1.21x / 1.22x**; hardening rows 21-37 all faster. - **BF16/H32 retained-KV decode program (rows 73/74/76/77):** hoisted scale loads, one epilogue release fence, eight warps issue the gathers: GB300 1.08-1.10x -> **1.30-1.34x**, B200 1.09-1.12x -> **1.32-1.33x**. - **BF16/FP8 H8/H16 decode programs (16 rows):** PV reads the V quarter from the staged K tile instead of re-gathering it: GB300 FP8 rows 1.36x -> **1.47x** (geomean), BF16 1.23x -> **1.31x**; B200 1.39x -> **1.52x**, 1.26x -> **1.35x**. - **Routing:** width-260 BF16/H128 topk128x rows run the four-owner two-stage program (GB300 1.23x -> **1.44x**, B200 1.11x -> **1.29x**); dense (not ragged) BF16/H64 rows follow the ragged rules (fixed-Q guard row: GB300 1.03x -> **1.21x**, B200 1.03x -> **1.11x**); the three-owner topk128x program and the fixed-Q programs are no longer reachable and leave the registry. - Every other program is a byte-identical regeneration (renamed for the new producer revision). ## Validation Export r12 (this branch: host and tests at `84b04ff02`, regenerated programs as published here), same-session paired CUPTI (active-union median, 500 warmup / 3000 calls per arm, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 | 134/134 | | **94 canonical rows faster than the default path** | 94/94, geomean **1.300x**, min 1.098x (#5610: 1.243x, min 1.054x) | 94/94, geomean **1.280x**, min 1.027x (#5610: 1.232x, min 1.041x) | | 40 hardening rows faster than the default path | 40/40, geomean 1.330x, min 1.071x (#5610: 1.30x, min 1.058x) | 40/40, geomean 1.307x, min 1.065x (#5610: 1.25x, min 1.037x) | | exported program vs source build (agreement gate, 134 rows) | 134/134, geomean 1.0018 | 134/134, geomean 1.0026 | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed | 296 passed | | `tests/mla/test_cake_dsv4*.py` (CPU, this host) | 212 passed, 84 skipped | 212 passed, 84 skipped | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e, 2026-09-27), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Updated CAKE DSv4 routing across supported GPU variants, including revised handling for dense BF16/H64 queries and selected BF16/H128 input widths. * Expanded coverage for sparse-index staging and attention processing across value-head groups. * **Performance** * Optimized output packing and stores, and adjusted kernel staging and memory allocation for updated execution paths. * **Compatibility** * Removed obsolete kernel variants and registrations; workloads that previously used those routes may now follow different supported routes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
added a commit
that referenced
this pull request
Sep 28, 2026
…0/SM103 (stacked on #5630) (#5644) ## Summary Round-4 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5630 (this PR is stacked on #5630; host and tests unchanged): - **BF16/FP8 H8/H16 decode programs (16 rows):** the FP8 body issues its K gathers back-to-back after hoisting the page-offset loads and runs an 8-slot K ring; both bodies get a deeper K/V alias ring so a 3-tile item's second tile no longer waits for the first PV. GB300 FP8 rows 1.40x -> **1.70x** (geomean), B200 1.52x -> **1.73x**; 3-tile BF16 rows -0.2..-0.5 us (GB300 decode-000062 1.50x -> 1.64x). - **BF16/H128 persistent prefill programs (128+ query-token rows and the K-reuse rows):** FA4-style lazy softmax rescale (a row keeps its running max until the true max grows by more than 2^8; the deferred factor is applied then). prefill-style-000088: GB300 1.34x -> **1.47x**, B200 1.22x -> **1.37x**; 000092 1.35x -> **1.41x** / 1.22x -> **1.33x**. - **BF16 one-tile SWA decode programs (18 rows):** the sparse-index load is issued before the Q TMA (statement order only; bit-identical output): -0.1..-0.26 us per row. - Every other program is a byte-identical regeneration (renamed for the new producer revision). ## Validation Export r13 (regenerated programs as published here; host and tests at #5630's `76d586ff6`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 | 134/134 | | **94 canonical rows faster than the default path** | 94/94, geomean **1.343x**, min 1.130x (#5630: 1.300x, min 1.098x) | 94/94, geomean **1.319x**, min 1.048x (#5630: 1.280x, min 1.027x) | | 40 hardening rows faster than the default path | 40/40, geomean **1.353x**, min 1.085x (#5630: 1.330x, min 1.071x) | 40/40, geomean **1.310x**, min 1.069x (#5630: 1.307x, min 1.065x) | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed | 296 passed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (76d586f), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance Improvements** * Updated attention processing to handle more key/value data stages per pipeline, potentially improving throughput. * Sparse-index lookups now occur earlier in the loading process, helping streamline attention data preparation. * **Numerical Improvements** * Refined softmax accumulation to avoid unnecessary correction for small score changes while preserving correction for larger changes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
added a commit
that referenced
this pull request
Sep 29, 2026
…0/SM103 (exact numerics; stacked on #5644) (#5686) ## Summary Round-5 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5644, plus one host-side binding line. Every regenerated program computes bit-identical output to its #5644 predecessor: this round only re-times data movement and epilogues, with no change to the arithmetic. - **BF16 H8/H16 decode programs (16 rows):** eight warps issue the K gathers (both arches); the FP8 H8/H16 programs do the same on SM100a. - **BF16/H128 persistent prefill programs (SM103a):** the O epilogue loads the next tile's TMEM accumulator ahead of the store (−1.5..−4.3 %). - **BF16/H64 split-K decode programs (9 rows):** owner-CTA form with hoisted partial select, zero-fill and 256-bit stores: −1.7..−2.9 us per row (GB300 class geomean 1.25x -> **1.51x**, B200 1.18x -> **1.45x**). - **FP8/H128 persistent prefill programs:** TMA-store O epilogue (−0.5..−0.6 us on the 9-tile rows, −7.3 % on prefill-style-000093). - **BF16/H32 retained-KV decode programs (4 rows):** the softmax publishes P before the row-sum exchange (separate stats barrier), −0.3..−1.2 %. - Every other program is a byte-identical regeneration (renamed for the new producer revision). - **Host (`flashinfer/mla/cake_dsv4.py`, first commit of this branch):** the FP8 persistent bodies now store O through a TMA descriptor (`tmap_o`); `_TMA_SOURCE_ALIASES` binds it to the same `[tokens, heads, 512]` output rows the plain `O` argument sees. Without it every FP8 launch on those routes raised "argument 'tmap_o' (tma_buffer) has no host value" (28 failures in the hardening tests on the regenerated programs alone). `test_registered_arg_plans_use_known_names` now also requires every TMA descriptor to alias a tensor the host actually places in its value table, so a vocabulary-only match cannot pass again. ## Validation Exports of the regenerated programs as published here (B200: export r15 against #5644's `6d479bea1`; GB300: export r16 against this branch's host-fix commit `fd2131236`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 correct; 133/134 export gates (row 085: baseline-arm order sensitivity, see doc § Round-over-round) | 134/134 | | **94 canonical rows faster than the default path** | 94/94 > 1, geomean **1.370x**, min 1.129x (#5644: 1.343x, min 1.130x) | 94/94, geomean **1.356x**, min 1.103x (#5644: 1.319x, min 1.048x) | | 40 hardening rows faster than the default path | 40/40 > 1, geomean **1.385x**, min 1.096x (#5644: 1.353x, min 1.085x) | 40/40 > 1, geomean **1.333x**, min 1.045x (#5644: 1.310x, min 1.069x) | | output bit-identical to the #5644 programs on every changed row | yes | yes | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed, 0 failed | 296 passed, 0 failed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (ebe0efb), round 4 #5644 (6d479be), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added output tensor-map support for selected CAKE DSV4 attention paths, enabling outputs to be staged and transferred through optimized tensor operations. * Updated selected kernels to process larger output groups and adjust work distribution for supported GPU architectures. * Refreshed kernel-family registration and dispatch across supported architectures. * **Bug Fixes** * Improved handling of invalid key/value indices and synchronization of output statistics and transfers. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
added a commit
that referenced
this pull request
Sep 30, 2026
…0/SM103 (exact numerics; stacked on #5686) (#5700) ## Summary Round-6 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5686. No host change. Every regenerated program computes bit-identical output to its #5686 predecessor: this round only re-times data movement, store widths, pipeline depths and the kernel ABI, with no change to the arithmetic. - **BF16/H128 split-K programs, classes E / F / B / C (18 rows):** reducer family — 256-bit partial stores, batched partial loads, and a leaner 5-split reducer on SM103a; class E additionally gets a per-arch owner (SM100a: 4-stage K ring; SM103a: new 8-slot K/PV ring body). E rows −2.7..−3.7 us (GB300) / −2.3..−2.5 us (B200); F rows −0.4..−0.5 us; B S=3 rows −0.65..−0.70 us. - **FP8/H64 single-CTA SWA programs, class K1 (11 rows):** 256-bit O stores in the lane-per-row epilogue, −0.06..−0.27 us. - **BF16 one-tile SWA programs, classes A / D / G (17 rows) + the H32 retained-KV and H8/H16 guards decode programs (SM103a):** tensor maps passed by value (`__grid_constant__`) instead of through a pointer, −0.26..−1.03 us per row (−4.6..−9.0 % on A/D/G). - **FP8/H64 wide one-partition programs, class K4 (2 rows):** third 32 KiB K stage, −0.06..−0.16 us. - **FP8/H128 persistent programs, classes K2 / L (33 rows, SM103a):** epilogue TMEM load-ahead of 4, −0.6..−3.3 % on the H128 rows. - **BF16/H128 striped prefill programs, class M (7 rows):** V rows issued ahead (LDS order only), −1.5..−4.7 % (B200), −0.3..−1.6 % (GB300). - **BF16 H8/H16 decode programs, class H (8 rows, SM100a):** the never-read padded gather rows of the last item are skipped, −0.45..−1.02 us on the padded rows. - Every other program is a byte-identical regeneration (renamed for the new producer revision). ## Validation Exports of the regenerated programs as published here (B200: export r16; GB300: export r17, both from Cake tree `939659882ba` against a fresh clone of #5686's `10c80d75b`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 correct; 134/134 export gates | 134/134 correct; 134/134 export gates (decode-000002 and hardening-000036 in a second measurement round, see doc § Gate table) | | **94 canonical rows faster than the default path** | 94/94 > 1, geomean **1.411x**, min 1.116x (#5686: 1.370x, min 1.129x) | 94/94 > 1, geomean **1.352x**, min 1.075x (#5686: 1.356x, min 1.103x) | | 40 hardening rows faster than the default path | 40/40 > 1, geomean **1.390x**, min 1.092x (#5686: 1.385x, min 1.096x) | 40/40 > 1, geomean **1.354x**, min 1.085x (#5686: 1.333x, min 1.045x) | | output bit-identical to the #5686 programs on every changed row | yes (92 changed rows) | yes (48 changed rows) | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed, 0 failed | 296 passed, 0 failed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (ebe0efb), round 4 #5644 (6d479be), round 5 #5686 (10c80d7), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optimized attention-kernel variants for newer NVIDIA GPU architectures, including improved handling of split outputs and KV data. * Removed the need for callers to provide a separate tensor-map workspace for several kernel variants. * **Performance Improvements** * Optimized data loading, pipelining, and output conversion across multiple GPU variants. * Added a specialized path for combining three partial outputs. * **Bug Fixes** * Improved bounds handling for gathered KV rows, including out-of-range indices. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
andakai
added a commit
to andakai/vllm
that referenced
this pull request
Oct 3, 2026
FlashInfer's SM120 sparse MLA kernels (flashinfer-ai/flashinfer#3395 through 0.7.0.post1) are built for 64-token pages only, so the 256-row branch could never be served. Re-page the flat rows at 64 and drop the single-use SM120_PAGE_ROWS constant. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Dakai An <dakaian108@gmail.com>
This was referenced Nov 5, 2026
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
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
Sparse-MLA paged attention for SM120a / SM121a, covering both the DSv4 family
(
d_qk=512) and the DSv3.2 / GLM-5.1 family (d_qk=576). Adds decode(warp-spec, split-K) and prefill (TMA + warp-spec) kernels in the SM120 sparse
MLA JIT/AOT module, with dispatch through the existing MLA Python APIs.
Surfaces
The user-facing surface is limited to extending existing MLA entry points with
SM120/SM121 sparse backend coverage. SM120-specific dispatchers, runners, and
direct kernel wrappers remain internal implementation details; users should
continue to call the public
flashinfer.mlaAPIs below. The public APIs keeptheir existing decode-oriented names, while the SM120 sparse backend internally
selects decode or prefill kernels from the flattened token count.
Public Python APIs
flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4(...)accepts BF16 query tensors, a packed uint8 SWA KV cache with 584 bytes per
token, optional attention sinks, and an optional second packed KV cache via
extra_sparse_indices/extra_sparse_topk_lens.swa_topk_lensdescribes the active SWA segment. When the optionalcompressed segment is present,
extra_sparse_topk_lensdescribes its activelength. Both sparse segments share one online-softmax denominator.
flashinfer.mla.trtllm_batch_decode_with_kv_cache_mla(...)DSv3.2 / GLM-style packed sparse MLA (
d_qk=576, 656 bytes per token).backend="auto"on SM120/SM121 andsparse_mla_top_k > 0, the callselects the sparse backend.
backend="sparse"selects it explicitly.kv_scale_formatcontrols the packed-cache scale convention:"auto"/"pow2_fp32"for DSv3.2-style power-of-two FP32 scales, and"arbitrary_fp32"for GLM-style arbitrary FP32 scales.Internal / JIT Implementation Details
flashinfer.mla._sparse_mla_sm120._sparse_mla_sm120_paged_attention(...)isan internal dispatcher that shares validation and buffer handling between
decode and prefill shapes. It is registered as the custom op
flashinfer::sparse_mla_sm120_paged_attention._SparseMLAPagedAttentionRunneris an internal helper used by tests,benchmarks, and the public API implementation to own LSE and decode scratch
buffers.
flashinfer.mla._sparse_mla_sm120.sparse_mla_sm120_decode_dsv4(...)and theDSv3.2 decode wrapper are internal direct kernel wrappers around the JIT
module and AutoTuner
chunks_per_blocktactic cache.flashinfer.jit.mla.gen_sparse_mla_sm120_module()wires the new SM120 sparseMLA CUDA sources into the JIT/AOT build.
Performance
RTX PRO 6000 Blackwell Server Edition (SM120a). KV pool sized > L2 (>=600 MB)
so analytical numbers reflect DRAM-bound production.
attn_TFLOPs = 2 * T * num_heads * topk * (d_qk + d_v) / latency.DSv4 prefill, single cache (584 B/slot footer, FP8 paged):
DSv4 prefill, dual cache (main SWA cache
topk=128plus secondary C4A /C128A cache, 584 B/slot footer, FP8 paged):
DSv4 decode (T=1):
DSv3.2 / GLM-5.1 (656 B/slot inline, FP8 paged,
topk = 2048):Testing
122 tests covering decode + prefill, single + dual cache, attention sink, top-k
length masking, and full DSv4 / DSv3.2 / GLM-5.1 shape coverage. Tests
compute-sanitizer --tool racecheckclean on representative shapes.End-to-end validated by serving GLM-5.1 and DSv4-Flash through these
kernels.