feat(cake_msa): add Blackwell minimax sparse attention source kernels - #4355
Conversation
|
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:
📝 WalkthroughWalkthroughBlackwell MSA support now covers SM100a and SM103a sparse prefill, decode, and top-k operations. The change adds CUDA kernels and bindings, architecture-aware dispatch, JIT/AOT registration, workspace handling, route manifests, benchmarks, documentation, and extensive validation tests. ChangesBlackwell MSA backend
Estimated code review effort: 5 (Critical) | ~180 minutes Merge Risk: 🔴 Critical · up to The new Blackwell sparse-attention routes still contain unresolved issues that can reject valid CUDA Graph workloads, read invalid tensor regions, or return incorrect attention and LSE results; merge should be blocked until the input validation, descriptor bounds, and reduction correctness issues are fixed. Sequence Diagram(s)sequenceDiagram
participant PublicAPI
participant DeviceDispatch
participant BlackwellJIT
participant FFI
participant CUDAKernel
PublicAPI->>DeviceDispatch: select Blackwell path
DeviceDispatch->>BlackwellJIT: load target-specific MSA module
BlackwellJIT->>FFI: resolve variant entry point
FFI->>CUDAKernel: validate tensors and launch kernel
CUDAKernel-->>PublicAPI: write attention outputs and LSE
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🧹 Nitpick comments (5)
benchmarks/bench_cake_msa_sm100.py (3)
72-152: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the comparable/unsupported shape counts from
SHAPES.The literal counts
5and1appear at Line 834, Line 974, Line 1029, and Line 1030. If someone adds a row toSHAPES, the script fails with a count error that does not name the cause. Compute the counts once next toSHAPES.♻️ Proposed derived counts
SHAPES_BY_LABEL = {shape["label"]: shape for shape in SHAPES} +COMPARABLE_SHAPE_COUNT = sum(1 for shape in SHAPES if shape["q_dtype"] != "float16") +UNSUPPORTED_SHAPE_COUNT = len(SHAPES) - COMPARABLE_SHAPE_COUNTThen use
COMPARABLE_SHAPE_COUNTat Lines 834, 974, and 1029, andUNSUPPORTED_SHAPE_COUNTat Line 1030.🤖 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_cake_msa_sm100.py` around lines 72 - 152, Derive COMPARABLE_SHAPE_COUNT and UNSUPPORTED_SHAPE_COUNT once next to SHAPES from the shape metadata, rather than hard-coding 5 and 1. Update the assertions or checks at the referenced uses around lines 834, 974, 1029, and 1030 to use the corresponding constants, preserving the existing comparable-versus-unsupported classification.
679-722: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated baseline-import verification.
The same seven-line check appears at Lines 680-686 and Lines 715-721. Move it into one helper and call it from both branches.
♻️ Proposed helper
+def _import_baseline(baseline_root: Path) -> Any: + imported_baseline = importlib.import_module("fmha_sm100") + imported_baseline_root = Path(imported_baseline.__file__).resolve().parents[2] + if imported_baseline_root != baseline_root: + raise RuntimeError( + f"expected fmha_sm100 from {baseline_root}, " + f"imported {imported_baseline_root}" + ) + return imported_baselineThen replace both blocks with
_import_baseline(baseline_root).🤖 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_cake_msa_sm100.py` around lines 679 - 722, Extract the duplicated baseline module import and root validation into a helper named _import_baseline, accepting baseline_root and returning the imported module. Replace both inline verification blocks in the verify branch and the non-flashinfer branch with calls to _import_baseline(baseline_root), preserving the existing mismatch error behavior.
507-525: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the failure record carry the same keys as the success record.
The early-return failure record omits
reference,candidate_nonfinite_count, andbaseline_nonfinite_count, which the success record at Lines 554 and 567-568 includes. A shape or dtype mismatch therefore produces a JSON record with a different schema. No consumer reads it today because_run_parentraises first, so this is schema hygiene for the reported artifact.♻️ Proposed key alignment
return { "status": "failed", "passed": False, + "reference": "pinned_public_fmha_sm100_sparse_atten_func", "candidate_public_api": candidate_api, "baseline_public_api": baseline_api,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/bench_cake_msa_sm100.py` around lines 507 - 525, Update the early-return failure record in the shape/dtype mismatch branch to include the same schema keys as the success record, specifically reference, candidate_nonfinite_count, and baseline_nonfinite_count, using the appropriate values or nulls for this failure path. Preserve the existing failure status and diagnostic fields.flashinfer/msa_ops/_cake_sm100.py (1)
130-158: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBound
_eager_decode_dummiesgrowth.The cache key includes
stream_ptr, and entries are never removed. Each entry holds a128 x 2 x 128device tensor, about 64 KiB for bf16. A process that creates many short-lived CUDA streams accumulates one entry per distinct stream pointer, and the memory is never released.
_flat_kv_route_cachealready implements a size cap with dead-reference sweeping at Lines 443-450. Apply the same bound here, or dropstream_ptrfrom the key if the dummy buffer does not need per-stream isolation.🤖 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/msa_ops/_cake_sm100.py` around lines 130 - 158, Bound the global _eager_decode_dummies cache used by _decode_tma_dummy so short-lived stream pointers cannot accumulate device tensors indefinitely. Reuse the existing _flat_kv_route_cache size-cap and dead-reference sweeping pattern, or remove stream_ptr from the key if per-stream isolation is unnecessary, while preserving workspace-backed behavior.csrc/cake_msa/cake_msa_decode_fp8_flat.cu (1)
3836-3841: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBoth kernels emit an unreachable paged gather path. Each file assigns
num_n_blocks_3 = 0immediately before two gather4 loops, so roughly 100 lines of TMA gather code never execute in either variant. Do not hand-edit these frozen generated files; change the CAKE generator instead.
csrc/cake_msa/cake_msa_decode_fp8_flat.cu#L3836-L3841: suppress the gather4 loops for flat FP8 decode, where the native FP8 path above already stages KV.csrc/cake_msa/cake_msa_decode_fp8_paged.cu#L3857-L3862: suppress the same loops for paged FP8 decode.🤖 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/cake_msa/cake_msa_decode_fp8_flat.cu` around lines 3836 - 3841, The CAKE generator emits unreachable gather4 loops because num_n_blocks_3 is set to zero before them. Update the generator to suppress these loops in both csrc/cake_msa/cake_msa_decode_fp8_flat.cu lines 3836-3841 and csrc/cake_msa/cake_msa_decode_fp8_paged.cu lines 3857-3862, rather than editing the generated files; preserve the native FP8 KV staging path in the flat variant.
🤖 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_cake_msa_sm100.py`:
- Around line 913-918: Update the process-order selection around the
comparable-row branch to alternate using a dedicated count of comparable rows
processed, rather than SHAPES index parity. Increment the count only for
comparable rows, including indices 0, 1, 2, 3, and 5 while excluding the FP16
row, and preserve the existing minimax-first/flashinfer-first alternation.
In `@csrc/cake_msa/cake_msa_decode_bf16_flat_binding.cu`:
- Around line 498-537: Update the scalar validation in both BF16 flat and paged
decode bindings so num_requests, num_kv_heads, and record_tasks are required to
be greater than zero in addition to fitting the i32 range. Add positivity checks
alongside the existing range checks before kernel dispatch, preserving the
current error-reporting style and messages for invalid inputs.
In `@csrc/cake_msa/cake_msa_decode_fp16_flat_binding.cu`:
- Around line 504-506: Update the scalar validation for num_kv_heads and
record_tasks in both flat and paged binding files to require positive i32
values, not merely values within the i32 range. Add the checks beside the
existing range checks for the binding argument validation symbols, preserving
the current range-error behavior while rejecting zero and negative divisors
before kernel launch. If these files are generated, make the equivalent change
in the generator template.
In `@csrc/cake_msa/cake_msa_decode_fp8_flat_binding.cu`:
- Around line 394-399: Both Run bindings validate arg_Q’s device ID before
confirming it is CUDA, causing CPU inputs to report the wrong error. In
csrc/cake_msa/cake_msa_decode_fp8_flat_binding.cu lines 394-399 and
csrc/cake_msa/cake_msa_decode_fp8_paged_binding.cu lines 392-397, move
CheckCudaTensor(arg_Q, "Q") before ffi::CUDADeviceGuard and CheckCakeMsaTarget,
preserving the remaining validation order.
In `@csrc/cake_msa/cake_msa_decode_fp8_paged_binding.cu`:
- Around line 205-215: Update the K TMA global-dimension validation to require
d2 to be at least 128, matching the two producer loads at y coordinates 0 and
64; preserve the existing box dimensions and error handling. Apply the same
minimum page/sequence extent to the EncodeTma_V counterpart and the flat binding
checks identified in the diff so all CAKE MSA paged decode paths reject
undersized dimensions.
In `@csrc/cake_msa/cake_msa_decode_m16_bf16_paged_binding.cu`:
- Around line 205-215: Verify the Python dispatch path for the paged CAKE MSA
route and determine whether it restricts page_size to 128. If smaller pages can
reach this kernel, update the global-dimension validation in EncodeTma_K and
EncodeTma_V to require at least 128 tokens for global_dim[1], matching the two
64-row TMA loads; preserve the existing validation and dispatch behavior when
page_size is already constrained to 128.
In `@csrc/cake_msa/cake_msa_prefill_m128_fp8_paged_binding.cu`:
- Around line 165-179: Add an innermost extent validation to both EncodeTma_k
and EncodeTma_v, requiring d1 == 128 before constructing the TMA descriptor.
Keep the existing global_dim and stride definitions unchanged, and use
TVM_FFI_CHECK with ValueError so smaller or mismatched FP8 head dimensions are
rejected.
In `@csrc/cake_msa/cake_msa_topk_binding.cu`:
- Around line 131-161: Update Run to validate tensor sizes against the launch
geometry before dispatch: require arg_num_valid_pages to be no greater than
arg_max_k_tiles, and require arg_output.numel() to be at least grid_x multiplied
by 16. Use the existing TVM_FFI_CHECK validation style and reference the
relevant scalar/tensor values in clear error messages, alongside the current
checks in Run.
In `@flashinfer/msa_ops/__init__.py`:
- Around line 24-30: Ensure the module containing supports_packed_kv is
importable on the project’s declared minimum Python version by adding the
annotations future import when that floor is below Python 3.10; otherwise verify
and preserve the declared 3.10+ requirement. Keep the torch.device | str
annotation and supports_packed_kv behavior unchanged.
In `@flashinfer/msa_ops/_cake_sm100.py`:
- Around line 838-859: Inspect the decode kernel writes and binding signatures
for partial_o, partial_m, and partial_d, then update the allocations in the
decode path so each slot contains only the extent the kernel actually indexes
rather than an unnecessary _BLOCK_SIZE tile. Keep partial_slots as the
token/KV-head/split count, ensure all three buffer shapes and downstream
indexing remain consistent, and preserve workspace reuse so the eager path does
not allocate oversized buffers on every invocation.
In `@tests/jit/test_cake_msa_jit.py`:
- Around line 94-99: Ensure gen_cake_msa_module’s cache is cleared during
teardown as well as setup, preferably via an autouse fixture scoped to these
tests; then remove the explicit cache_clear calls from the affected test bodies.
Apply the same cleanup to both patched TARGET_CUDA_ARCHS test cases so cached
specs cannot leak after monkeypatch restores the shared context.
In `@tests/msa_ops/test_cake_msa_sm100.py`:
- Around line 1063-1088: The test named test_cuda_graph_decode_route_uses_m128
does not exercise CUDA-graph routing because _select_decode_route receives
capturing=False. Set capturing=True in this test call so the graph-capture
branch is covered, while preserving the expected ("m128", False, None) result.
---
Nitpick comments:
In `@benchmarks/bench_cake_msa_sm100.py`:
- Around line 72-152: Derive COMPARABLE_SHAPE_COUNT and UNSUPPORTED_SHAPE_COUNT
once next to SHAPES from the shape metadata, rather than hard-coding 5 and 1.
Update the assertions or checks at the referenced uses around lines 834, 974,
1029, and 1030 to use the corresponding constants, preserving the existing
comparable-versus-unsupported classification.
- Around line 679-722: Extract the duplicated baseline module import and root
validation into a helper named _import_baseline, accepting baseline_root and
returning the imported module. Replace both inline verification blocks in the
verify branch and the non-flashinfer branch with calls to
_import_baseline(baseline_root), preserving the existing mismatch error
behavior.
- Around line 507-525: Update the early-return failure record in the shape/dtype
mismatch branch to include the same schema keys as the success record,
specifically reference, candidate_nonfinite_count, and baseline_nonfinite_count,
using the appropriate values or nulls for this failure path. Preserve the
existing failure status and diagnostic fields.
In `@csrc/cake_msa/cake_msa_decode_fp8_flat.cu`:
- Around line 3836-3841: The CAKE generator emits unreachable gather4 loops
because num_n_blocks_3 is set to zero before them. Update the generator to
suppress these loops in both csrc/cake_msa/cake_msa_decode_fp8_flat.cu lines
3836-3841 and csrc/cake_msa/cake_msa_decode_fp8_paged.cu lines 3857-3862, rather
than editing the generated files; preserve the native FP8 KV staging path in the
flat variant.
In `@flashinfer/msa_ops/_cake_sm100.py`:
- Around line 130-158: Bound the global _eager_decode_dummies cache used by
_decode_tma_dummy so short-lived stream pointers cannot accumulate device
tensors indefinitely. Reuse the existing _flat_kv_route_cache size-cap and
dead-reference sweeping pattern, or remove stream_ptr from the key if per-stream
isolation is unnecessary, while preserving workspace-backed behavior.
🪄 Autofix
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 Plus
Run ID: 88e53856-f261-49c3-8352-8b954f1b1468
📒 Files selected for processing (50)
benchmarks/bench_cake_msa_sm100.pycsrc/cake_msa/cake_msa_decode_bf16_flat.cucsrc/cake_msa/cake_msa_decode_bf16_flat_binding.cucsrc/cake_msa/cake_msa_decode_bf16_paged.cucsrc/cake_msa/cake_msa_decode_bf16_paged_binding.cucsrc/cake_msa/cake_msa_decode_fp16_flat.cucsrc/cake_msa/cake_msa_decode_fp16_flat_binding.cucsrc/cake_msa/cake_msa_decode_fp16_paged.cucsrc/cake_msa/cake_msa_decode_fp16_paged_binding.cucsrc/cake_msa/cake_msa_decode_fp8_flat.cucsrc/cake_msa/cake_msa_decode_fp8_flat_binding.cucsrc/cake_msa/cake_msa_decode_fp8_paged.cucsrc/cake_msa/cake_msa_decode_fp8_paged_binding.cucsrc/cake_msa/cake_msa_decode_m16_bf16_flat.cucsrc/cake_msa/cake_msa_decode_m16_bf16_flat_binding.cucsrc/cake_msa/cake_msa_decode_m16_bf16_paged.cucsrc/cake_msa/cake_msa_decode_m16_bf16_paged_binding.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_flat.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_flat_binding.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_gqa16_flat.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_gqa16_flat_binding.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_gqa16_paged.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_gqa16_paged_binding.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_paged.cucsrc/cake_msa/cake_msa_prefill_m128_bf16_paged_binding.cucsrc/cake_msa/cake_msa_prefill_m128_fp16_flat.cucsrc/cake_msa/cake_msa_prefill_m128_fp16_flat_binding.cucsrc/cake_msa/cake_msa_prefill_m128_fp16_paged.cucsrc/cake_msa/cake_msa_prefill_m128_fp16_paged_binding.cucsrc/cake_msa/cake_msa_prefill_m128_fp8_flat.cucsrc/cake_msa/cake_msa_prefill_m128_fp8_flat_binding.cucsrc/cake_msa/cake_msa_prefill_m128_fp8_paged.cucsrc/cake_msa/cake_msa_prefill_m128_fp8_paged_binding.cucsrc/cake_msa/cake_msa_prefill_m64_bf16_flat.cucsrc/cake_msa/cake_msa_prefill_m64_bf16_flat_binding.cucsrc/cake_msa/cake_msa_topk.cucsrc/cake_msa/cake_msa_topk_binding.cudocs/api/sparse.rstflashinfer/aot.pyflashinfer/jit/__init__.pyflashinfer/jit/cake_msa.pyflashinfer/msa_ops/__init__.pyflashinfer/msa_ops/_cake_sm100.pyflashinfer/msa_ops/sparse_decode.pyflashinfer/msa_ops/sparse_prefill.pyflashinfer/msa_ops/sparse_topk_select.pytests/jit/test_cake_msa_jit.pytests/msa_ops/test_cake_msa_sm100.pytests/msa_ops/test_cake_msa_source.pytests/msa_ops/test_packed_kv.py
| TVM_FFI_CHECK(cuda_stream >= 0, ValueError) << "cuda_stream must be non-negative"; | ||
| ffi::CUDADeviceGuard device_guard(arg_Q.device().device_id); | ||
| CheckCakeMsaTarget(arg_Q.device().device_id); | ||
| cudaStream_t stream = reinterpret_cast<cudaStream_t>(static_cast<uintptr_t>(cuda_stream)); | ||
| CheckCudaTensor(arg_Q, "Q"); | ||
| CheckDtype(arg_Q, "Q", 4, 16, 1); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Both bindings read Q.device().device_id before they confirm Q is a CUDA tensor. Each Run constructs ffi::CUDADeviceGuard and calls CheckCakeMsaTarget on that id, and only afterwards calls CheckCudaTensor(arg_Q, "Q"). A CPU tensor reports device_id == 0, so the caller receives a compute-capability error instead of the accurate device-type error.
csrc/cake_msa/cake_msa_decode_fp8_flat_binding.cu#L394-L399: moveCheckCudaTensor(arg_Q, "Q")above theffi::CUDADeviceGuardconstruction.csrc/cake_msa/cake_msa_decode_fp8_paged_binding.cu#L392-L397: apply the same reordering.
📍 Affects 2 files
csrc/cake_msa/cake_msa_decode_fp8_flat_binding.cu#L394-L399(this comment)csrc/cake_msa/cake_msa_decode_fp8_paged_binding.cu#L392-L397
🤖 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/cake_msa/cake_msa_decode_fp8_flat_binding.cu` around lines 394 - 399,
Both Run bindings validate arg_Q’s device ID before confirming it is CUDA,
causing CPU inputs to report the wrong error. In
csrc/cake_msa/cake_msa_decode_fp8_flat_binding.cu lines 394-399 and
csrc/cake_msa/cake_msa_decode_fp8_paged_binding.cu lines 392-397, move
CheckCudaTensor(arg_Q, "Q") before ffi::CUDADeviceGuard and CheckCakeMsaTarget,
preserving the remaining validation order.
| uint64_t global_dim[3] = {(uint64_t)(128), (uint64_t)(d2), (uint64_t)(outer2)}; | ||
| TVM_FFI_CHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] > 0, ValueError) | ||
| << "TMA descriptor for 'K' resolved a non-positive global dim"; | ||
| TVM_FFI_CHECK(128u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError) | ||
| << "TMA box (128, 64, 1) exceeds resolved global dims for 'K'"; | ||
| uint64_t global_strides[2] = { | ||
| (uint64_t)((d1 * 8) / 8), | ||
| (uint64_t)(((d2 * d1) * 8) / 8), | ||
| }; | ||
| uint32_t box_dim[3] = {128u, 64u, 1u}; | ||
| uint32_t elem_strides[3] = {1u, 1u, 1u}; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The K bound check allows a page dimension that the kernel over-reads.
Line 208 requires only 64u <= global_dim[1], where global_dim[1] = d2 is the page dimension. The box extent along that axis is 64, so the check matches the box.
The kernel does not stop at one box. In cake_msa_decode_fp8_paged.cu the producer issues two loads per stage at Lines 3704-3705, with y coordinates 0 and 64. The second load therefore needs d2 >= 128. If a caller passes d2 == 64, this check passes, the TMA unit zero-fills the out-of-range tile, and the upper 64 tokens of every KV tile silently become zeros. The kernel produces wrong attention output and reports no error.
Raise the bound to the maximum coordinate the kernel actually reaches. EncodeTma_V at Line 208's counterpart, Line 283, has the same gap, and the flat binding repeats the pattern against the sequence dimension.
#!/bin/bash
# Description: Determine the page sizes the CAKE MSA SM100 paged decode path accepts and whether any layer enforces page_size >= 128.
set -euo pipefail
fd -t f -e py . flashinfer/msa_ops --exec rg -n -C 6 'page_size|page_shape|msa_max_pages' {} \;
fd -t f -e py . tests --exec rg -n -C 4 'page_size' {} \; | rg -n -i 'cake|msa' || true🔧 Proposed tightening of the bound checks
- TVM_FFI_CHECK(128u <= global_dim[0] && 64u <= global_dim[1] && 1u <= global_dim[2], ValueError)
- << "TMA box (128, 64, 1) exceeds resolved global dims for 'K'";
+ // The decode producer issues two boxes per stage at y = 0 and y = 64, so the
+ // page dimension must cover 128 rows, not just one box extent.
+ TVM_FFI_CHECK(128u <= global_dim[0] && 128u <= global_dim[1] && 1u <= global_dim[2], ValueError)
+ << "TMA source 'K' page dim " << global_dim[1]
+ << " must be at least 128 to cover both decode boxes (y = 0 and y = 64)";🤖 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/cake_msa/cake_msa_decode_fp8_paged_binding.cu` around lines 205 - 215,
Update the K TMA global-dimension validation to require d2 to be at least 128,
matching the two producer loads at y coordinates 0 and 64; preserve the existing
box dimensions and error handling. Apply the same minimum page/sequence extent
to the EncodeTma_V counterpart and the flat binding checks identified in the
diff so all CAKE MSA paged decode paths reject undersized dimensions.
| CheckCudaTensor(arg_max_score, "max_score"); | ||
| CheckDtype(arg_max_score, "max_score", 2, 32, 1); | ||
| CheckContiguous(arg_max_score, "max_score"); | ||
| CheckCudaTensor(arg_output, "output"); | ||
| CheckDtype(arg_output, "output", 0, 32, 1); | ||
| CheckContiguous(arg_output, "output"); | ||
| TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) | ||
| << "scalar 'num_heads' value " << arg_num_heads | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| TVM_FFI_CHECK(arg_max_k_tiles >= -2147483648LL && arg_max_k_tiles <= 2147483647LL, ValueError) | ||
| << "scalar 'max_k_tiles' value " << arg_max_k_tiles | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| TVM_FFI_CHECK(arg_total_q >= -2147483648LL && arg_total_q <= 2147483647LL, ValueError) | ||
| << "scalar 'total_q' value " << arg_total_q | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| TVM_FFI_CHECK(arg_num_valid_pages >= -2147483648LL && arg_num_valid_pages <= 2147483647LL, | ||
| ValueError) | ||
| << "scalar 'num_valid_pages' value " << arg_num_valid_pages | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| TVM_FFI_CHECK(arg_force_begin_blocks >= -2147483648LL && arg_force_begin_blocks <= 2147483647LL, | ||
| ValueError) | ||
| << "scalar 'force_begin_blocks' value " << arg_force_begin_blocks | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| TVM_FFI_CHECK(arg_force_end_blocks >= -2147483648LL && arg_force_end_blocks <= 2147483647LL, | ||
| ValueError) | ||
| << "scalar 'force_end_blocks' value " << arg_force_end_blocks | ||
| << " is outside i32 range [-2147483648, 2147483647]"; | ||
| CheckSameCudaDevice(arg_output, arg_max_score, "output", "max_score"); | ||
| TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) | ||
| << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " | ||
| << grid_z << ")"; |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Validate the tensor sizes against the scalar arguments.
Run checks dtype, contiguity, device, scalar ranges, and grid positivity. It does not relate the tensor sizes to the scalars.
The kernel indexes max_score[(head * max_k_tiles + block) * total_q + query] for block < num_valid_pages, and writes output[row * 16 + lane_1] for row = blockIdx.x. Two unvalidated conditions produce out-of-bounds device accesses:
num_valid_pages > max_k_tilesdrives the read index past the intended extent ofmax_score.arg_output.numel() < grid_x * 16drives the write past the end ofoutput.
This kernel has no TMA encoder, so it receives none of the geometry validation that the other CAKE MSA bindings get. Add the two checks. Both are local and use data already available.
🛡️ Proposed fix to bound the kernel indices
CheckSameCudaDevice(arg_output, arg_max_score, "output", "max_score");
TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError)
<< "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", "
<< grid_z << ")";
+ TVM_FFI_CHECK(arg_num_valid_pages >= 0 && arg_num_valid_pages <= arg_max_k_tiles, ValueError)
+ << "num_valid_pages " << arg_num_valid_pages << " must be in [0, max_k_tiles="
+ << arg_max_k_tiles << "]";
+ TVM_FFI_CHECK(arg_max_score.numel() >= arg_num_heads * arg_max_k_tiles * arg_total_q, ValueError)
+ << "max_score has " << arg_max_score.numel() << " elements, but the kernel indexes up to "
+ << (arg_num_heads * arg_max_k_tiles * arg_total_q);
+ TVM_FFI_CHECK(arg_output.numel() >= grid_x * 16, ValueError)
+ << "output has " << arg_output.numel() << " elements, but the kernel writes " << (grid_x * 16);📝 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.
| CheckCudaTensor(arg_max_score, "max_score"); | |
| CheckDtype(arg_max_score, "max_score", 2, 32, 1); | |
| CheckContiguous(arg_max_score, "max_score"); | |
| CheckCudaTensor(arg_output, "output"); | |
| CheckDtype(arg_output, "output", 0, 32, 1); | |
| CheckContiguous(arg_output, "output"); | |
| TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) | |
| << "scalar 'num_heads' value " << arg_num_heads | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_max_k_tiles >= -2147483648LL && arg_max_k_tiles <= 2147483647LL, ValueError) | |
| << "scalar 'max_k_tiles' value " << arg_max_k_tiles | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_total_q >= -2147483648LL && arg_total_q <= 2147483647LL, ValueError) | |
| << "scalar 'total_q' value " << arg_total_q | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_num_valid_pages >= -2147483648LL && arg_num_valid_pages <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'num_valid_pages' value " << arg_num_valid_pages | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_force_begin_blocks >= -2147483648LL && arg_force_begin_blocks <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'force_begin_blocks' value " << arg_force_begin_blocks | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_force_end_blocks >= -2147483648LL && arg_force_end_blocks <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'force_end_blocks' value " << arg_force_end_blocks | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| CheckSameCudaDevice(arg_output, arg_max_score, "output", "max_score"); | |
| TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) | |
| << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " | |
| << grid_z << ")"; | |
| CheckCudaTensor(arg_max_score, "max_score"); | |
| CheckDtype(arg_max_score, "max_score", 2, 32, 1); | |
| CheckContiguous(arg_max_score, "max_score"); | |
| CheckCudaTensor(arg_output, "output"); | |
| CheckDtype(arg_output, "output", 0, 32, 1); | |
| CheckContiguous(arg_output, "output"); | |
| TVM_FFI_CHECK(arg_num_heads >= -2147483648LL && arg_num_heads <= 2147483647LL, ValueError) | |
| << "scalar 'num_heads' value " << arg_num_heads | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_max_k_tiles >= -2147483648LL && arg_max_k_tiles <= 2147483647LL, ValueError) | |
| << "scalar 'max_k_tiles' value " << arg_max_k_tiles | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_total_q >= -2147483648LL && arg_total_q <= 2147483647LL, ValueError) | |
| << "scalar 'total_q' value " << arg_total_q | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_num_valid_pages >= -2147483648LL && arg_num_valid_pages <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'num_valid_pages' value " << arg_num_valid_pages | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_force_begin_blocks >= -2147483648LL && arg_force_begin_blocks <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'force_begin_blocks' value " << arg_force_begin_blocks | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| TVM_FFI_CHECK(arg_force_end_blocks >= -2147483648LL && arg_force_end_blocks <= 2147483647LL, | |
| ValueError) | |
| << "scalar 'force_end_blocks' value " << arg_force_end_blocks | |
| << " is outside i32 range [-2147483648, 2147483647]"; | |
| CheckSameCudaDevice(arg_output, arg_max_score, "output", "max_score"); | |
| TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) | |
| << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " | |
| << grid_z << ")"; | |
| TVM_FFI_CHECK(arg_num_valid_pages >= 0 && arg_num_valid_pages <= arg_max_k_tiles, ValueError) | |
| << "num_valid_pages " << arg_num_valid_pages << " must be in [0, max_k_tiles=" | |
| << arg_max_k_tiles << "]"; | |
| TVM_FFI_CHECK(arg_max_score.numel() >= arg_num_heads * arg_max_k_tiles * arg_total_q, ValueError) | |
| << "max_score has " << arg_max_score.numel() << " elements, but the kernel indexes up to " | |
| << (arg_num_heads * arg_max_k_tiles * arg_total_q); | |
| TVM_FFI_CHECK(arg_output.numel() >= grid_x * 16, ValueError) | |
| << "output has " << arg_output.numel() << " elements, but the kernel writes " << (grid_x * 16); |
🤖 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/cake_msa/cake_msa_topk_binding.cu` around lines 131 - 161, Update Run to
validate tensor sizes against the launch geometry before dispatch: require
arg_num_valid_pages to be no greater than arg_max_k_tiles, and require
arg_output.numel() to be at least grid_x multiplied by 16. Use the existing
TVM_FFI_CHECK validation style and reference the relevant scalar/tensor values
in clear error messages, alongside the current checks in Run.
| def supports_packed_kv(device: torch.device | str) -> bool: | ||
| """Return whether MSA accepts packed paged K/V views on ``device``.""" | ||
|
|
||
| normalized_device = torch.device(device) | ||
| return normalized_device.type == "cuda" and get_compute_capability( | ||
| normalized_device | ||
| ) in {(12, 0), (12, 1)} |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
Confirm the declared minimum Python version supports the torch.device | str annotation at runtime.
This module has no from __future__ import annotations, so the parameter annotation on Line 24 is evaluated when the module is imported. PEP 604 unions between two classes require Python 3.10 or newer. If the project still declares support for Python 3.9, this import raises TypeError: unsupported operand type(s) for |.
flashinfer/msa_ops/_cake_sm100.py uses the same syntax but adds from __future__ import annotations at its Line 19, so it is unaffected. Either add the same future import here or confirm the declared floor is 3.10+.
#!/bin/bash
# Description: Read the project's declared Python floor and check the future-import status of the changed module.
set -euo pipefail
echo "=== declared requires-python / target-version ==="
rg -n 'requires-python|target-version|python_requires' pyproject.toml setup.py setup.cfg 2>/dev/null || true
echo "=== classifiers ==="
rg -n 'Programming Language :: Python' pyproject.toml 2>/dev/null || true
echo "=== future import in the changed module ==="
rg -n 'from __future__ import annotations' flashinfer/msa_ops/__init__.py || echo "ABSENT in flashinfer/msa_ops/__init__.py"🛡️ Proposed fix if the floor is below 3.10
+from __future__ import annotations
+
import torch🤖 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/msa_ops/__init__.py` around lines 24 - 30, Ensure the module
containing supports_packed_kv is importable on the project’s declared minimum
Python version by adding the annotations future import when that floor is below
Python 3.10; otherwise verify and preserve the declared 3.10+ requirement. Keep
the torch.device | str annotation and supports_packed_kv behavior unchanged.
| partial_slots = int(total_q) * int(num_kv_heads) * max_splits | ||
| partial_o = _workspace_buffer( | ||
| workspace, | ||
| "decode_partial_o", | ||
| (partial_slots, _BLOCK_SIZE, _HEAD_DIM), | ||
| dtype=torch.float32, | ||
| device=q.device, | ||
| ) | ||
| partial_m = _workspace_buffer( | ||
| workspace, | ||
| "decode_partial_m", | ||
| (partial_slots, _BLOCK_SIZE), | ||
| dtype=torch.float32, | ||
| device=q.device, | ||
| ) | ||
| partial_d = _workspace_buffer( | ||
| workspace, | ||
| "decode_partial_d", | ||
| (partial_slots, _BLOCK_SIZE), | ||
| dtype=torch.float32, | ||
| device=q.device, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Verify the partial_o slot shape; the current sizing can allocate multiple gigabytes per decode call.
partial_slots counts one slot per (token, kv-head, split). Each partial_o slot then holds a full 128 x 128 float32 tile, which is 64 KiB. For total_q=1024, num_kv_heads=8, and topk=16, max_splits is 8, so partial_slots is 65536 and partial_o alone requires about 4 GiB.
The _BLOCK_SIZE row dimension looks redundant with the per-token slot index. partial_m and partial_d use (partial_slots, _BLOCK_SIZE), which is only consistent if the kernel treats every slot as a 128-row tile. Confirm what the decode kernel actually stores per slot.
The cost repeats on every call. When workspace is None, _workspace_buffer allocates fresh tensors at Line 205 instead of reusing them, so the eager decode path pays this allocation on each invocation.
#!/bin/bash
# Description: Determine the per-slot extent the decode kernels expect for the split partial buffers.
set -euo pipefail
echo "=== decode binding signatures: partial buffer params ==="
fd -e cu 'cake_msa_decode_.*_binding\.cu$' csrc --exec rg -n -C4 'partial_o|partialO|partial_m|partial_d|split_completion'
echo "=== decode kernel writes into the partial buffers ==="
fd -e cu 'cake_msa_decode_' csrc -E '*_binding.cu' --exec rg -n -C4 'partial_o|partialO'🤖 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/msa_ops/_cake_sm100.py` around lines 838 - 859, Inspect the decode
kernel writes and binding signatures for partial_o, partial_m, and partial_d,
then update the allocations in the decode path so each slot contains only the
extent the kernel actually indexes rather than an unnecessary _BLOCK_SIZE tile.
Keep partial_slots as the token/KV-head/split count, ensure all three buffer
shapes and downstream indexing remain consistent, and preserve workspace reuse
so the eager path does not allocate oversized buffers on every invocation.
| monkeypatch.setattr( | ||
| jit_core.current_compilation_context, | ||
| "TARGET_CUDA_ARCHS", | ||
| {target_arch}, | ||
| ) | ||
| cake_msa.gen_cake_msa_module.cache_clear() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Clear the module cache at teardown, not only at setup.
gen_cake_msa_module is keyed on (variant, target) and does not include TARGET_CUDA_ARCHS in the key. The test patches the shared jit_core.current_compilation_context singleton, clears the cache, then populates it with specs built from the fake arch set. monkeypatch restores TARGET_CUDA_ARCHS at teardown, but it does not clear the cache. Every later caller of gen_cake_msa_module("topk", "sm100f") in the same process receives the spec built under the patched arch set. The result is order-dependent test pollution.
Clear the cache after the test as well. The same problem exists at Lines 146-154.
🛠️ Proposed fix using an autouse fixture
+@pytest.fixture(autouse=True)
+def _reset_cake_msa_module_cache():
+ cake_msa.gen_cake_msa_module.cache_clear()
+ yield
+ cake_msa.gen_cake_msa_module.cache_clear()
+
+
`@pytest.mark.parametrize`(
("target", "target_arch", "expected_flag", "expected_define", "forbidden"),With the fixture in place, remove the explicit cache_clear() calls at Lines 99 and 151.
📝 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.
| monkeypatch.setattr( | |
| jit_core.current_compilation_context, | |
| "TARGET_CUDA_ARCHS", | |
| {target_arch}, | |
| ) | |
| cake_msa.gen_cake_msa_module.cache_clear() | |
| `@pytest.fixture`(autouse=True) | |
| def _reset_cake_msa_module_cache(): | |
| cake_msa.gen_cake_msa_module.cache_clear() | |
| yield | |
| cake_msa.gen_cake_msa_module.cache_clear() | |
| monkeypatch.setattr( | |
| jit_core.current_compilation_context, | |
| "TARGET_CUDA_ARCHS", | |
| {target_arch}, | |
| ) | |
| cake_msa.gen_cake_msa_module.cache_clear() |
🤖 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/jit/test_cake_msa_jit.py` around lines 94 - 99, Ensure
gen_cake_msa_module’s cache is cleared during teardown as well as setup,
preferably via an autouse fixture scoped to these tests; then remove the
explicit cache_clear calls from the affected test bodies. Apply the same cleanup
to both patched TARGET_CUDA_ARCHS test cases so cached specs cannot leak after
monkeypatch restores the shared context.
| def test_cuda_graph_decode_route_uses_m128() -> None: | ||
| device = _require_supported_gpu() | ||
| from flashinfer.msa_ops import MSASparseAttentionWorkspace | ||
| from flashinfer.msa_ops._cake_sm100 import _select_decode_route | ||
|
|
||
| workspace = MSASparseAttentionWorkspace(device) | ||
| q = torch.empty((4, 16, HEAD_DIM), dtype=torch.bfloat16, device=device) | ||
| k = torch.empty((4096, 1, HEAD_DIM), dtype=torch.bfloat16, device=device) | ||
| cu_k = torch.tensor([0, 4096], dtype=torch.int32, device=device) | ||
| kv_lens = torch.tensor([4096], dtype=torch.int32, device=device) | ||
|
|
||
| route, persistent_unsplit, path_force_fused = _select_decode_route( | ||
| q=q, | ||
| k=k, | ||
| cu_k=cu_k, | ||
| kv_lens=kv_lens, | ||
| group_size=16, | ||
| seqlen_q=4, | ||
| paged=False, | ||
| force_fused=True, | ||
| workspace=workspace, | ||
| route_key=("graph-stable-m128",), | ||
| capturing=False, | ||
| ) | ||
|
|
||
| assert (route, persistent_unsplit, path_force_fused) == ("m128", False, None) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
The test name claims CUDA-graph routing, but the call sets capturing=False.
test_cuda_graph_decode_route_uses_m128 passes capturing=False at Line 1085. The graph-capture branch of _select_decode_route is therefore never exercised. The two sibling tests at Lines 993 and 1033 also pass capturing=False, so this test only repeats the eager path with a workspace attached.
If the intent is to cover capture-time routing, pass capturing=True. If the intent is to cover workspace-attached eager routing, rename the test to match.
#!/bin/bash
# Inspect the routing helper to confirm what `capturing` and `workspace` control.
fd -t f '_cake_sm100.py' flashinfer --exec rg -n -C 45 'def _select_decode_route' {}🤖 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/msa_ops/test_cake_msa_sm100.py` around lines 1063 - 1088, The test
named test_cuda_graph_decode_route_uses_m128 does not exercise CUDA-graph
routing because _select_decode_route receives capturing=False. Set
capturing=True in this test call so the graph-capture branch is covered, while
preserving the expected ("m128", False, None) result.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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_cake_msa_sm100.py`:
- Around line 252-273: Remove or replace the TopK32 benchmark row identified by
stable_id official_decode_bf16_b64_q8_kv65536_h64_hkv4_k32_paged, preserving the
required coverage with topk no greater than 16 and keeping the manifest at 12
rows. In benchmarks/bench_cake_msa_sm100.py lines 441-442, restrict accepted
topk values to {4, 8, 16}; in tests/msa_ops/test_cake_msa_benchmark_manifest.py
lines 12-20, remove the TopK32 stable ID; and in lines 63-64, assert only the
supported TopK values.
🪄 Autofix
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 Plus
Run ID: bf05140a-28c5-40cc-a098-83c7e4c438cd
📒 Files selected for processing (2)
benchmarks/bench_cake_msa_sm100.pytests/msa_ops/test_cake_msa_benchmark_manifest.py
| stable_id="official_decode_bf16_b64_q8_kv65536_h64_hkv4_k32_paged", | ||
| tier="official_coverage", | ||
| source="minimax_official_sparse_decode_benchmark", | ||
| provenance=MINIMAX_BENCHMARK_PROVENANCE, | ||
| selection_rationale=( | ||
| "Official long-KV decode coordinate with the documented TopK32 " | ||
| "option; covers B64, KV65536, and the largest supported TopK." | ||
| ), | ||
| operation="sparse_decode", | ||
| batch_size=64, | ||
| seqlen_q=8, | ||
| seqlen_kv=65536, | ||
| q_dtype="bfloat16", | ||
| kv_dtype="bfloat16", | ||
| kv_layout="paged", | ||
| num_q_heads=64, | ||
| num_kv_heads=4, | ||
| topk=32, | ||
| causal=True, | ||
| force_fused=True, | ||
| seed=67, | ||
| baseline_mode="minimax_public", |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep the benchmark matrix within the TopK16 and 12-row contracts.
The backend contract supports exact selection of at most 16 pages. This row creates a 32-wide q2k tensor and passes it to the candidate public API. It can fail validation or kernel dispatch. It also makes the manifest contain 13 rows instead of the stated 12.
benchmarks/bench_cake_msa_sm100.py#L252-L273: remove or replace the TopK32 row while preserving required coverage with TopK16 or less.benchmarks/bench_cake_msa_sm100.py#L441-L442: restrict acceptedtopkvalues to{4, 8, 16}.tests/msa_ops/test_cake_msa_benchmark_manifest.py#L12-L20: remove the TopK32 stable ID.tests/msa_ops/test_cake_msa_benchmark_manifest.py#L63-L64: require only supported TopK values.
📍 Affects 2 files
benchmarks/bench_cake_msa_sm100.py#L252-L273(this comment)benchmarks/bench_cake_msa_sm100.py#L441-L442tests/msa_ops/test_cake_msa_benchmark_manifest.py#L12-L20tests/msa_ops/test_cake_msa_benchmark_manifest.py#L63-L64
🤖 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_cake_msa_sm100.py` around lines 252 - 273, Remove or replace
the TopK32 benchmark row identified by stable_id
official_decode_bf16_b64_q8_kv65536_h64_hkv4_k32_paged, preserving the required
coverage with topk no greater than 16 and keeping the manifest at 12 rows. In
benchmarks/bench_cake_msa_sm100.py lines 441-442, restrict accepted topk values
to {4, 8, 16}; in tests/msa_ops/test_cake_msa_benchmark_manifest.py lines 12-20,
remove the TopK32 stable ID; and in lines 63-64, assert only the supported TopK
values.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/msa_ops/test_cake_msa_benchmark_manifest.py (1)
145-145: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSelect the base shape by stable ID instead of list index.
SHAPE_MANIFEST[4]binds this test to manifest ordering. If a row is inserted before index 4, this test silently changes meaning, and the failure message will not explain why. Use the stable-ID lookup that the module already exposes.♻️ Proposed change
- base = benchmark.SHAPE_MANIFEST[4] + base = benchmark.SHAPES_BY_ID[_FP16_SHAPE_ID]🤖 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/msa_ops/test_cake_msa_benchmark_manifest.py` at line 145, Update the base-shape assignment in the test to use the module’s existing stable-ID lookup instead of indexing SHAPE_MANIFEST by position, preserving the intended shape selection when manifest ordering changes.tests/msa_ops/test_cake_msa_source.py (1)
182-194: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider scanning binding sources as well.
The loop reads only
_device_name(variant). Inline PTX can also appear in a binding translation unit. Including_binding_name(variant)in the same scan costs one extraread_textper variant and closes the gap.Note: the ast-grep hint on Line 186 reports XPath injection. That is a false positive. The code performs a regex search over local file text and uses no XPath API.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/msa_ops/test_cake_msa_source.py` around lines 182 - 194, Extend test_cake_msa_sources_use_cuda_12_8_global_vector_widths to scan both _device_name(variant) and _binding_name(variant) source files for _CUDA_12_8_ILLEGAL_GLOBAL_VECTOR matches, recording each offending source name in illegal_sites while preserving the existing failure message and regex-based scanning.Source: Linters/SAST tools
benchmarks/bench_cake_msa_sm100.py (1)
1011-1012: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCompute the reference before calling the candidate, or snapshot the inputs.
_verify_candidate_referencecalls the candidate first and then builds the reference from the sameinputstensors. If any candidate kernel writes intoq,k,v, orq2k, the reference silently consumes mutated data and the check becomes self-confirming. Compute the reference first to remove the ordering dependency.♻️ Proposed change
- candidate_output = _primary_output(candidate_call()) - reference_output = _candidate_reference_output(torch, shape, inputs) + reference_output = _candidate_reference_output(torch, shape, inputs) + candidate_output = _primary_output(candidate_call()) torch.cuda.synchronize()🤖 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_cake_msa_sm100.py` around lines 1011 - 1012, In _verify_candidate_reference, compute reference_output via _candidate_reference_output before invoking candidate_call, then obtain candidate_output afterward. Preserve the existing comparison logic while ensuring the reference uses the original q, k, v, and q2k inputs.
🤖 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_cake_msa_sm100.py`:
- Around line 1025-1039: Update the early failure return in the shape/dtype
mismatch path inside _verify_public_outputs() and _verify_candidate_reference()
so it includes the same parity keys as the success path, especially the
correctness payload expected by _validate_correctness_metadata and
validate_json(). Keep the existing failure diagnostics intact, but make the
returned row schema consistent across success and failure cases to avoid
missing-key lookups for report consumers.
In `@tests/msa_ops/test_cake_msa_source.py`:
- Line 125: Update _CUDA_12_8_ILLEGAL_GLOBAL_VECTOR to match ld.global and
st.global instructions with optional intervening qualifiers before v8.b32,
including forms such as .nc, .ca, and .L2::128B. Escape literal dots in the
regular expression so only the intended PTX syntax is matched.
---
Nitpick comments:
In `@benchmarks/bench_cake_msa_sm100.py`:
- Around line 1011-1012: In _verify_candidate_reference, compute
reference_output via _candidate_reference_output before invoking candidate_call,
then obtain candidate_output afterward. Preserve the existing comparison logic
while ensuring the reference uses the original q, k, v, and q2k inputs.
In `@tests/msa_ops/test_cake_msa_benchmark_manifest.py`:
- Line 145: Update the base-shape assignment in the test to use the module’s
existing stable-ID lookup instead of indexing SHAPE_MANIFEST by position,
preserving the intended shape selection when manifest ordering changes.
In `@tests/msa_ops/test_cake_msa_source.py`:
- Around line 182-194: Extend
test_cake_msa_sources_use_cuda_12_8_global_vector_widths to scan both
_device_name(variant) and _binding_name(variant) source files for
_CUDA_12_8_ILLEGAL_GLOBAL_VECTOR matches, recording each offending source name
in illegal_sites while preserving the existing failure message and regex-based
scanning.
🪄 Autofix
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 Plus
Run ID: a48e65bf-056d-42ae-a9e8-e492bbbde128
📒 Files selected for processing (11)
benchmarks/bench_cake_msa_sm100.pycsrc/cake_msa/cake_msa_decode_bf16_flat.cucsrc/cake_msa/cake_msa_decode_bf16_paged.cucsrc/cake_msa/cake_msa_decode_fp16_flat.cucsrc/cake_msa/cake_msa_decode_fp16_paged.cucsrc/cake_msa/cake_msa_decode_fp8_flat.cucsrc/cake_msa/cake_msa_decode_fp8_paged.cucsrc/cake_msa/cake_msa_decode_m16_bf16_flat.cucsrc/cake_msa/cake_msa_decode_m16_bf16_paged.cutests/msa_ops/test_cake_msa_benchmark_manifest.pytests/msa_ops/test_cake_msa_source.py
🚧 Files skipped from review as they are similar to previous changes (8)
- csrc/cake_msa/cake_msa_decode_fp16_paged.cu
- csrc/cake_msa/cake_msa_decode_fp16_flat.cu
- csrc/cake_msa/cake_msa_decode_fp8_paged.cu
- csrc/cake_msa/cake_msa_decode_fp8_flat.cu
- csrc/cake_msa/cake_msa_decode_m16_bf16_paged.cu
- csrc/cake_msa/cake_msa_decode_bf16_flat.cu
- csrc/cake_msa/cake_msa_decode_bf16_paged.cu
- csrc/cake_msa/cake_msa_decode_m16_bf16_flat.cu
| r"\binline CUtensorMap EncodeTma_([A-Za-z0-9_]+)\(" | ||
| ) | ||
| _DIRECT_TENSOR_MAP_LAUNCH_ARG = re.compile(r"&h_([A-Za-z0-9_]+)(?=,|})") | ||
| _CUDA_12_8_ILLEGAL_GLOBAL_VECTOR = re.compile(r"\b(?:ld|st)\.global\.v8\.b32\b") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Broaden the pattern to cover qualified ld.global/st.global forms.
The pattern requires global to be followed immediately by v8. PTX permits qualifiers between them, for example ld.global.nc.v8.b32, ld.global.ca.v8.b32, and ld.global.L2::128B.v8.b32. These forms are equally illegal under the 128-bit limit, but the current pattern does not match them, so the guard can pass on a regressing source. Unescaped . also matches any character; escape it for precision.
🐛 Proposed pattern
-_CUDA_12_8_ILLEGAL_GLOBAL_VECTOR = re.compile(r"\b(?:ld|st)\.global\.v8\.b32\b")
+_CUDA_12_8_ILLEGAL_GLOBAL_VECTOR = re.compile(
+ r"\b(?:ld|st)\.global(?:\.[A-Za-z0-9:]+)*\.v8\.b32\b"
+)📝 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.
| _CUDA_12_8_ILLEGAL_GLOBAL_VECTOR = re.compile(r"\b(?:ld|st)\.global\.v8\.b32\b") | |
| _CUDA_12_8_ILLEGAL_GLOBAL_VECTOR = re.compile( | |
| r"\b(?:ld|st)\.global(?:\.[A-Za-z0-9:]+)*\.v8\.b32\b" | |
| ) |
🤖 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/msa_ops/test_cake_msa_source.py` at line 125, Update
_CUDA_12_8_ILLEGAL_GLOBAL_VECTOR to match ld.global and st.global instructions
with optional intervening qualifiers before v8.b32, including forms such as .nc,
.ca, and .L2::128B. Escape literal dots in the regular expression so only the
intended PTX syntax is matched.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/msa_ops/test_cake_msa_long_decode_route.py`:
- Around line 15-18: Update _long_decode_tensors and the sibling synthetic
route-policy tensor fixture at tests/msa_ops/test_cake_msa_long_decode_route.py
lines 15-18 and 68-70 to create tensors on torch.device("meta"). Preserve their
existing shapes and dtypes so route selection continues inspecting the same
metadata without allocating full BF16 storage.
🪄 Autofix
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 Plus
Run ID: 19724c06-e9e7-4eaa-ac30-d5888d58fee0
📒 Files selected for processing (2)
flashinfer/msa_ops/_cake_sm100.pytests/msa_ops/test_cake_msa_long_decode_route.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/msa_ops/_cake_sm100.py
| def _long_decode_tensors(): | ||
| q = torch.empty((512, 64, HEAD_DIM), dtype=torch.bfloat16) | ||
| k = torch.empty((32768, 4, BLOCK_SIZE, HEAD_DIM), dtype=torch.bfloat16) | ||
| return q, k |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file list =="
git ls-files | rg '(^|/)test_cake_msa_long_decode_route\.py$|conftest\.py$' || true
echo "== target file excerpt =="
if [ -f tests/msa_ops/test_cake_msa_long_decode_route.py ]; then
wc -l tests/msa_ops/test_cake_msa_long_decode_route.py
sed -n '1,110p' tests/msa_ops/test_cake_msa_long_decode_route.py
fi
echo "== allocator/oom references in target and conftest =="
rg -n "empty\\(|_long_decode_tensors|neighboring_q4_keeps_folded_m128_route|long_decode|meta|flashinfer|get_compute_capability|is_sm90a_supported|OOM|oom|skip" tests/conftest.py tests/msa_ops/test_cake_msa_long_decode_route.py || true
echo "== shape/dtype use in relevant route selector? =="
rg -n "def .*route|_.*route|route|decode|q.*shape|k.*shape|device|dtype" tests/msa_ops/test_cake_msa_long_decode_route.py . -g '*.py' | head -200Repository: flashinfer-ai/flashinfer
Length of output: 29233
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate cake sm100 dispatcher =="
fd -a '_cake_sm100|cake_sm100' . || true
rg -n "class .*Cake|def _is_long_paged_gqa16_direct_decode|def _select_decode_route|_cake_sm100|flashinfer\.msa_ops" flashinfer tests -g '*.py' | head -200
echo "== route-policy implementation files =="
files=$(rg -l "_is_long_paged_gqa16_direct_decode|_select_decode_route" -g '*.py' .)
echo "$files"
for f in $files; do
echo "--- $f ---"
wc -l "$f"
sed -n '1,220p' "$f"
doneRepository: flashinfer-ai/flashinfer
Length of output: 24184
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '930,1025p' flashinfer/msa_ops/_cake_sm100.pyRepository: flashinfer-ai/flashinfer
Length of output: 2839
🌐 Web query:
PyTorch meta device torch.empty meta device shape dtype behavior no memory allocation
💡 Result:
The meta device in PyTorch is a specialized, abstract device type designed to hold tensor metadata (shape, data type, and device) without allocating or storing any actual numerical data [1][2]. When you use torch.empty(..., device='meta'), PyTorch creates a tensor object with the specified shape and dtype but assigns it a meta device context instead of a standard CPU or CUDA device [3][1]. Because no underlying data is allocated, this operation is extremely lightweight in terms of memory usage and compute time, regardless of how large the requested tensor is [4]. Key behavioral characteristics include: Memory and Compute: No memory is allocated for the tensor's data, and no compute is performed during creation [1][4]. This makes meta tensors highly efficient for applications such as shape inference in large models, model initialization on meta device before materialization, and performing abstract analysis of operations [1][4]. Unsupported Operations: Because meta tensors contain no actual data, operations that require access to that data—such as torch.nonzero(), item(), or other data-dependent calculations—will fail [1][2]. Interaction with Real Devices: You cannot directly convert a meta tensor to a CPU or CUDA tensor via standard methods because the system has no data to initialize [1]. To materialize a tensor originally created on the meta device, you typically use torch.nn.Module.to_empty() (for modules) or other factory functions that explicitly manage initialization [1]. The to_empty() method allows you to move a module to a target device (like 'cuda') while leaving the parameters uninitialized, after which you are expected to perform the actual parameter initialization manually [1][2]. You can use the torch.device('meta') context manager to force all tensor constructions within a block to default to the meta device, which is particularly useful for creating model architectures without incurring the memory overhead of actual parameter data [1][4].
Citations:
- 1: https://docs.pytorch.org/docs/stable/meta.md
- 2: https://docs.pytorch.org/docs/2.3/meta.html
- 3: Make meta a device (getting rid of empty_meta) pytorch/pytorch#53143
- 4: https://github.com/pytorch/tutorials/blob/main/recipes_source/recipes/reasoning_about_shapes.py
Use meta tensors for route-policy policy tests.
The route selectors only inspect shape, dtype, and option metadata, but the fixtures allocate full BF16 KV tensors. Use device=torch.device("meta") for the synthetic route-policy tensors to keep these tests from materializing large synthetic memory.
📍 Affects 1 file
tests/msa_ops/test_cake_msa_long_decode_route.py#L15-L18(this comment)tests/msa_ops/test_cake_msa_long_decode_route.py#L68-L70
🤖 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/msa_ops/test_cake_msa_long_decode_route.py` around lines 15 - 18,
Update _long_decode_tensors and the sibling synthetic route-policy tensor
fixture at tests/msa_ops/test_cake_msa_long_decode_route.py lines 15-18 and
68-70 to create tensors on torch.device("meta"). Preserve their existing shapes
and dtypes so route selection continues inspecting the same metadata without
allocating full BF16 storage.
Source: Coding guidelines
yyihuang
changed the title
0d7811f to
a97441a
Compare
|
Can you post some performance numbers? |
yyihuang
changed the title
# Conflicts: # .pre-commit-config.yaml
|
@flashinfer-bot run |
|
@flashinfer-bot run |
|
/bot run tests/msa_ops |
yyihuang
changed the title
4 participants
Summary
This pull request adds source-distributed Blackwell MSA kernels behind the existing
flashinfer.msa_opsAPIs. It covers sparse prefill, sparse decode, exact TopK-16 selection, architecture-specific routing for compute capability 10.0 and 10.3, JIT/AOT registration, workspace and CUDA Graph support, tests, and documentation.The final source inventory contains 75 units: 38 for SM100a and 37 for SM103a. The exported sources pass native replay, public API, direct reducer, exact-route, formatting, hash, and compute-sanitizer gates on both architectures.
Related to #4254. SGLang integration: sgl-project/sglang#35846.
Validation
Kernel A/B
One outer round per architecture in one process and at the same shapes. Every arm visit used 10 warmups and 30 CUPTI cold-L2 samples.
NVIDIA B200 (
sm_100a)Over the 11 baseline-comparable rows, baseline/source geomean is 2.408351x and baseline/export geomean is 2.452001x. Across all 13 rows, the export/source latency geomean is 0.984027, or a 1.016232x export speedup. Physical benchmark turnaround was 2,306 s; the Slurm runtime was 2,305 s.
NVIDIA GB300 (
sm_103a)Over the 11 baseline-comparable rows, baseline/source geomean is 2.508831x and baseline/export geomean is 2.563379x. Across all 13 rows, the export/source latency geomean is 0.980892, or a 1.019480x export speedup. Physical benchmark turnaround was 2,450 s; the Slurm runtime was 2,442 s.
SGLang end-to-end correctness
MiniMax-M3-MXFP8 was evaluated on 4x NVIDIA GB300 with TP4 using real
/v1/chat/completionsrequests, seed 20260819, temperature 0, top-p 1, one request thread, and a fresh server and cache for every arm. The three full-suite arms used the same frozen GPQA Diamond 198 and balanced LongBench-v2 100 inputs: a no-MSA Triton control, the standalone source kernel, and the FlashInfer export. TP ranks 0--3 were audited asmain_attn=tritonwith MSA disabled,main_attn=fmha_sm100, andmain_attn=flashinfer, respectively. Every arm completed 298/298 requests with no failures, retries, or measured-window compilation events.The full-suite source-to-export GPQA delta misses the existing threshold of no worse than -1 question; LongBench-v2 passes its no-worse-than -0.02 threshold. Because response-level churn was substantial, a preregistered fresh three-arm replay was then run on the correctness-discordant union: 32 GPQA and 22 LongBench-v2 examples, in canonical order, with identical request-contract hashes across arms. An independent final repetition used the same selected examples and request contract. These selected-subset counts are not full-suite accuracy:
In the first replay, source and export tied on both evaluations. GPQA control-to-source and control-to-export were both -2 questions; LongBench-v2 control-to-source and control-to-export were both -1 question. In the independent final replay, GPQA export was two questions below source but five above control, while LongBench-v2 export was four below both source and control. Correctness changed relative to the original selected responses 31/33 times in the first replay and 32/25 times in the final replay for GPQA/LongBench-v2. Every arm in both replays used identical cross-arm request hashes, completed 54/54 requests, and passed the required TP0--3 route, fixed-parity, fresh-lifecycle, and measured-window audits.
The first strict paired analysis classified GPQA as
original_export_signal_not_reproduced_generation_variance, LongBench-v2 ascommon_msa_ordering_with_generation_churn, and the overall result aspr_specific_regression_not_reproduced_generation_variance. The independent final replay again did not reproduce the original GPQA export-below-both ordering; it introduced a new LongBench-v2 export-only ordering that was absent from both the full suite and the first replay. The direction of the export-specific signal therefore does not reproduce across the full suite and the two permitted fresh replays. The full-suite GPQA threshold failure and the mixed final-replay LongBench-v2 result remain recorded, but the combined controlled evidence does not support a reproducible accuracy regression caused specifically by the exported PR kernel.Full E2E validation used FlashInfer commit
73d19992657ac89542f596d1e620bebb56ca9e00. The merged head144f12e333bf1179f730d4c9574dd96f0f7276a5preserves the same MSA payload, but did not receive a separate exact-head full E2E rerun.SGLang fixed-request serving A/B
The same model and hardware ran one fresh-server source-to-export round using 8,192 input plus 1,024 output tokens, 256 prompts at each concurrency, and no measured-call warmup. Both arms completed 1,024/1,024 requests without failure and passed startup, fixed-parity, route, and measured-window audits.
Here
Export speedupis export output throughput divided by standalone-source output throughput; it is not a comparison against the Triton control. Both arms use the public SGLang integration in sgl-project/sglang#35846. Serving measurements used FlashInfer commite8173d70a42f50b98b4c1aadfbfc7b291fcade67.Measured serving runtime was 2,960.09 s for source and 2,885.19 s for export. Physical turnaround was 2:41:36; the Slurm runtime was 2:41:12.
Summary by CodeRabbit