Add paged MQA logits (attn_scores) kernels for Blackwell SM100 - #4365
Merged
Merged
Conversation
dhiraj113
requested review from
Aneureka,
aleozlx,
bkryu,
cyx-6,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
August 6, 2026 03:07
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:
📝 WalkthroughWalkthroughAdded SM100-only FP8 and FP4 paged MQA logits APIs. The change includes GPU scheduling, CuTe DSL kernels, package exports, trace templates, precompilation, and CUDA-gated correctness tests. ChangesPaged MQA logits
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
🧹 Nitpick comments (11)
flashinfer/attn_scores/attn_scores.py (1)
176-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winChain the exception.
Ruff reports B904 here. Use
raise ... from Noneto hide theKeyErrorcontext.♻️ Proposed change
except KeyError: - raise ValueError(f"Unsupported dtype for paged_mqa_logits: {dtype}") + raise ValueError( + f"Unsupported dtype for paged_mqa_logits: {dtype}" + ) from None🤖 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/attn_scores/attn_scores.py` around lines 176 - 180, Update the exception handling in _to_cutlass so the ValueError raised for an unsupported dtype explicitly suppresses the caught KeyError context by chaining it from None.Source: Linters/SAST tools
flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py (1)
600-604: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse integer arithmetic for the power-of-two round-up.
math.log2returns a float. The expression relies on the float result being exact for powers of two.int.bit_length()gives the same answer without any floating-point step and removes themathimport dependency for this line.♻️ Proposed change
- self.num_tmem_alloc_cols_total = max(1 << math.ceil(math.log2(raw_total)), 32) + self.num_tmem_alloc_cols_total = max( + 1 << (raw_total - 1).bit_length(), 32 + )🤖 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/attn_scores/kernels/fp4_paged_mqa_logits.py` around lines 600 - 604, Update the num_tmem_alloc_cols_total calculation in the relevant allocator initialization to use integer bit-length arithmetic for rounding raw_total up to the next power of two, while preserving the minimum of 32. Remove the now-unused math import if this is its only use.flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py (1)
212-216: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffUntested performance knobs on the public compile path.
remove_kv_wait_in_epilogue,early_tmem_copy,smem_subpartition_opt,max_kv_pipeline, andmax_umma_pipelineeach change barrier choreography or SMEM layout._compile_fp8_kernelinflashinfer/attn_scores/attn_scores.pynever sets them, and the added tests do not cover them. Each flag is therefore an unexercised code path in a warp-specialized kernel.Either add coverage for the flags you intend to keep, or mark them clearly as experimental in the class docstring.
🤖 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/attn_scores/kernels/fp8_paged_mqa_logits.py` around lines 212 - 216, Mark remove_kv_wait_in_epilogue, early_tmem_copy, smem_subpartition_opt, max_kv_pipeline, and max_umma_pipeline as experimental in the relevant public class docstring, since _compile_fp8_kernel does not set them and tests do not cover them. Clearly state that these performance knobs are unsupported or unvalidated on the public compile path.flashinfer/attn_scores/kernels/schedule_kernel.py (1)
120-148: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider a binary search instead of the fully unrolled linear scan.
cutlass.range_constexpr(kAligned)unrolls the inner scan completely, and it is nested inside thekMaxSmChunksloop.aligned_bis part of the compile key, so a large batch generates a very large kernel body and a long JIT compile.prefix_sumis non-decreasing, so a bounded binary search overlog2(kAligned)steps gives the same result with a constant instruction count.This is optional; keep the current form if the measured compile time stays acceptable for the batch sizes you target.
🤖 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/attn_scores/kernels/schedule_kernel.py` around lines 120 - 148, Replace the fully unrolled prefix_sum scan in the schedule metadata loop with a bounded binary search over the non-decreasing prefix_sum array, using approximately log2(kAligned) iterations to find the first value greater than seg_starts and preserve q_idx_out semantics. Keep the existing kv_split_idx calculation and metadata writes unchanged, and retain the linear scan only if compile-time measurements show it remains acceptable.flashinfer/attn_scores/__init__.py (1)
23-29: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnsorted
__all__in both new package files. Ruff reports RUF022 on each new__init__.py. The shared root cause is that both__all__lists follow import order instead of isort-style alphabetical order.
flashinfer/attn_scores/__init__.py#L23-L29: reorder toaligned_context_len,compute_paged_mqa_logits_schedule,fp4_paged_mqa_logits,fp8_paged_mqa_logits,precompile_paged_mqa_logits.flashinfer/attn_scores/kernels/__init__.py#L18-L18: reorder to["FP4MQALogitsKernel", "FP8MQALogitsKernel"].🤖 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/attn_scores/__init__.py` around lines 23 - 29, Sort the __all__ entries alphabetically in flashinfer/attn_scores/__init__.py lines 23-29: use aligned_context_len, compute_paged_mqa_logits_schedule, fp4_paged_mqa_logits, fp8_paged_mqa_logits, and precompile_paged_mqa_logits. Also reorder __all__ in flashinfer/attn_scores/kernels/__init__.py line 18 to FP4MQALogitsKernel followed by FP8MQALogitsKernel.Source: Linters/SAST tools
tests/attn_scores/test_attn_scores.py (3)
180-184: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the unsupported
remove_online_sf_transposecombination fail loudly.The helper applies the transpose only when
page_size == 128. If a caller passesremove_online_sf_transpose=Truewith another page size, the helper silently returns non-transposed scale factors while the kernel is told to skip the in-kernel transpose. The result is silently wrong instead of an error. Add an assertion so a future parametrization cannot hit this path.♻️ Proposed guard
sf_per_block = sf.view(num_blocks, page_size) - if remove_online_sf_transpose and page_size == 128: + if remove_online_sf_transpose: + assert page_size == 128, ( + "remove_online_sf_transpose requires page_size=128; " + f"got {page_size}" + ) sf_per_block = ( sf_per_block.reshape(num_blocks, 4, 32).transpose(-1, -2).contiguous().reshape(num_blocks, 128) )🤖 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/attn_scores/test_attn_scores.py` around lines 180 - 184, Add an assertion in the helper before the conditional transpose to require page_size == 128 whenever remove_online_sf_transpose is true, so unsupported combinations fail immediately while existing valid behavior remains unchanged.
255-261: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
_valid_causal_maskin the numerical tests.
_valid_causal_maskalready computes the causal-valid mask, but only the twoout=/schedule_meta=tests call it. Six tests recompute the identical expression inline: lines 350-353, 413-416, 460-463, 527-530, 592-595, and 652-655. Replace each inline block withneginf_mask = ~_valid_causal_mask(context_lens, next_n, max_model_len, device). This removes the duplication and keeps the masking rule in one place.🤖 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/attn_scores/test_attn_scores.py` around lines 255 - 261, Replace the six duplicated inline causal-mask computations in the numerical tests with calls to _valid_causal_mask(context_lens, next_n, max_model_len, device), assigning its negation to neginf_mask. Preserve the surrounding test logic and use the existing helper as the single source for the masking rule.
753-781: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the equivalent validation test for
fp4_paged_mqa_logits.This test covers the FP8 guard rails only. The FP4 API has a separate dtype allowlist that also admits bfloat16, its own
_validate_outcall, and thenext_n=4atom-split branch that rewritescontext_lensandblock_tablebefore validation. None of those error paths are covered. Add atest_fp4_input_validationthat mirrors the four cases here againstfp4_paged_mqa_logits.🤖 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/attn_scores/test_attn_scores.py` around lines 753 - 781, Add a test_fp4_input_validation alongside test_fp8_input_validation covering fp4_paged_mqa_logits with the same four guard-rail cases: undersized out using aligned_context_len, an unsupported output dtype while preserving FP4’s allowed bfloat16 behavior, CPU context_lens, and int64 block_table. Use an FP4-supported setup and exercise the next_n=4 atom-split path so validation also covers its context_lens and block_table rewrites.benchmarks/_profile_attn_overhead.py (1)
1-5: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a module docstring and record the private-API dependency.
This script imports
_compute_schedule_metadata,_compile_fp8_kernel, and_to_cutlassfromflashinfer.attn_scores.attn_scores. These are private symbols. If a later refactor renames them, this script breaks silently and nobody notices, because no test covers it.The other four benchmarks in this layer start with a docstring that states what is measured. Add one here. State that the script measures per-call host overhead and that it depends on private symbols.
🤖 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/_profile_attn_overhead.py` around lines 1 - 5, Add a module docstring at the start of benchmarks/_profile_attn_overhead.py stating that the script measures per-call host overhead and depends on private symbols from flashinfer.attn_scores.attn_scores. Keep the existing imports and benchmark behavior unchanged.benchmarks/bench_paged_mqa_logits_cudagraph.py (1)
93-97: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCall the imported
aligned_context_leninstead of recomputing the alignment.Line 14 imports
aligned_context_len. Line 97 does not call it. It recomputes the same value from the private_SPLIT_KVimported at line 21. Line 168 repeats the same expression.
flashinfer/attn_scores/attn_scores.pylines 406-416 definealigned_context_lenfor exactly this purpose, and its docstring shows thisout=pre-allocation pattern. If_SPLIT_KVchanges, the public helper follows the change and these two lines do not.Use the helper at both sites. The private
_SPLIT_KVimport at line 21 then becomes unnecessary.♻️ Proposed refactor for both loops
- aligned = ((max_ml + _SPLIT_KV - 1) // _SPLIT_KV) * _SPLIT_KV + aligned = aligned_context_len(max_ml)Apply the same change at line 168, then drop
_SPLIT_KVfrom the import at line 21.🤖 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_paged_mqa_logits_cudagraph.py` around lines 93 - 97, Replace the manual alignment expressions in both loops, including the `max_ml` calculation near `fp8_cfgs` and the corresponding site near line 168, with calls to the imported `aligned_context_len` helper. Remove the now-unused private `_SPLIT_KV` import while preserving the existing pre-allocation behavior.benchmarks/bench_paged_mqa_logits.py (1)
99-109: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBind the loop variables as closure defaults to clear Ruff B023.
Ruff reports B023 on every closure in both loops. The current code is correct at runtime, because
bench()consumes each closure inside the same iteration, before the loop rebinds the names. So this is a lint failure, not a behavior defect.If Ruff runs in CI for
benchmarks/, this file fails the check. Bind the captured names as default arguments to silence B023 and to make the capture explicit.Ruff also reports E702 for the semicolon-joined statements at lines 77, 109, 148, and 173.
♻️ Proposed refactor for the FP8 closures
- def fn_trt(): - compiled_trt(kv_flat, qu8, w2d, logits_trt, blk_tbl, cl, sched, total, B) + def fn_trt(compiled_trt=compiled_trt, kv_flat=kv_flat, qu8=qu8, w2d=w2d, + logits_trt=logits_trt, blk_tbl=blk_tbl, cl=cl, sched=sched, + total=total, B=B): + compiled_trt(kv_flat, qu8, w2d, logits_trt, blk_tbl, cl, sched, total, B) - def fn_fi_sched(): - fp8_paged_mqa_logits(q, kv_fused, w, cl, blk_tbl, max_ml) + def fn_fi_sched(q=q, kv_fused=kv_fused, w=w, cl=cl, blk_tbl=blk_tbl, max_ml=max_ml): + fp8_paged_mqa_logits(q, kv_fused, w, cl, blk_tbl, max_ml) - def fn_fi_cached(): - fp8_paged_mqa_logits(q, kv_fused, w, cl, blk_tbl, max_ml, schedule_meta=sched_fi) + def fn_fi_cached(q=q, kv_fused=kv_fused, w=w, cl=cl, blk_tbl=blk_tbl, + max_ml=max_ml, sched_fi=sched_fi): + fp8_paged_mqa_logits(q, kv_fused, w, cl, blk_tbl, max_ml, schedule_meta=sched_fi) # Warmup all - fn_trt(); fn_fi_sched(); fn_fi_cached(); torch.cuda.synchronize() + fn_trt() + fn_fi_sched() + fn_fi_cached() + 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_paged_mqa_logits.py` around lines 99 - 109, Update the benchmark closure definitions in both loops, including fn_trt, fn_fi_sched, and fn_fi_cached, to bind each captured loop variable through closure default arguments and clear Ruff B023. Replace the semicolon-joined warmup statements at the referenced locations with separate statements to clear E702, preserving execution order and behavior.Source: Linters/SAST tools
🤖 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/_profile_fp8_overhead.py`:
- Around line 13-17: Update the benchmark setup around block_table and kv_fused
so the block table has ceil(ctx / pbk) columns and the KV pool uses the
corresponding num_phys block count instead of the hardcoded 100. Initialize the
table consistently with the sibling benchmark before calling
compute_paged_mqa_logits_schedule and the compiled kernel.
In `@benchmarks/bench_paged_mqa_logits_eager_fair.py`:
- Line 40: Replace the hardcoded _SPLIT_KV constant with the imported
aligned_context_len() result, then update both out_trt and out_fi sizing sites
to use that value. Remove the _SPLIT_KV definition and ensure both line-78 and
line-152 allocation paths derive their trailing-column size from
aligned_context_len().
In `@benchmarks/bench_paged_mqa_logits.py`:
- Around line 31-33: Replace the hardcoded NUM_SMS value with the device’s
actual SM count so both benchmark paths use identical scheduling. In
benchmarks/bench_paged_mqa_logits.py lines 31-33, import get_device_sm_count
from flashinfer.utils and initialize NUM_SMS after DEVICE using
torch.device(DEVICE); make the same initialization change in
benchmarks/bench_paged_mqa_logits_cudagraph.py line 38, reusing its existing
import; and in benchmarks/bench_paged_mqa_logits_eager_fair.py line 36, add the
import and initialize NUM_SMS from torch.device(DEVICE).
- Line 14: Replace the hardcoded TRT-LLM path in
benchmarks/bench_paged_mqa_logits.py:14-14,
benchmarks/bench_paged_mqa_logits_cudagraph.py:10-10, and
benchmarks/bench_paged_mqa_logits_eager_fair.py:16-16 with a TRTLLM_PATH
environment-variable lookup, exiting clearly when unset and guarding TRT-LLM
imports so FlashInfer-only columns still run. In
benchmarks/_profile_attn_overhead.py:2-2 and
benchmarks/_profile_fp8_overhead.py:2-2, remove the sys.path.insert calls and
unnecessary sys imports so the installed flashinfer package is used.
In `@flashinfer/attn_scores/attn_scores.py`:
- Around line 465-467: Add a trace= template to the decorator configuration for
both public APIs, fp8_paged_mqa_logits and fp4_paged_mqa_logits, so fi_trace()
can generate benchmark-definition JSON for their tensor inputs and outputs. Keep
the existing supported_compute_capability and flashinfer_api decorators
unchanged.
- Around line 549-556: Update the output slicing in the attention-score path
around _validate_out so both provided and newly allocated outputs return exactly
B * next_n rows. Slice the row dimension before or alongside the existing
max_context_len column slice, including the corresponding path around the second
occurrence noted in the comment.
- Around line 654-702: Handle schedule_meta explicitly in the next_n == 4
expansion path: either reject a caller-provided schedule_meta for this mode with
a clear validation error, or document that it must be computed from the expanded
kernel_ctx_lens for 2B sequences. Update the surrounding function contract and
validation near the schedule_meta handling without changing the existing
automatic computation path.
In `@flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py`:
- Around line 420-422: Update the constructor logic around
remove_online_sf_transpose so it raises an error when the flag is requested with
phys_block_kv != 128, rather than resetting it to False. Preserve the requested
flag for valid configurations and ensure invalid requests fail before assigning
self.remove_online_sf_transpose.
In `@flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py`:
- Around line 273-285: Clamp the computed stage count in the max_kv_pipeline
branch to at least one after calculating (SMEM_BUDGET - qw_total) //
kv_scale_per_stage. Update the assignment to self.num_kv_stages while preserving
the existing calculation and the default three-stage path when max_kv_pipeline
is false.
In `@flashinfer/attn_scores/kernels/schedule_kernel.py`:
- Around line 103-107: Reformat the cute.arch.shuffle_sync_up call within the
scheduling loop so its arguments are wrapped according to the pre-commit
formatter’s style, without changing its behavior.
In `@tests/attn_scores/test_attn_scores.py`:
- Line 69: Run the project formatter via pre-commit on all files, then apply the
resulting formatting changes to tests/attn_scores/test_attn_scores.py, including
the affected lines around logits and the other reported locations, and commit
the formatter output without changing behavior.
- Around line 608-649: Update test_fp4_paged_mqa_logits_remove_sf_transpose so
every parametrized case exercises a distinct transpose configuration: restrict
phys_block_kv to 128, or add a separate 64-specific assertion that compares
meaningfully different behavior. Ensure the test does not compare identical
remove_online_sf_transpose=False setups.
- Around line 357-369: Add a finiteness assertion for kernel output in the valid
region before computing finite masks: in tests/attn_scores/test_attn_scores.py
at lines 357-369 (test_fp8_paged_mqa_logits), 419-424
(test_fp8_paged_mqa_logits_fp16), 466-471 (test_fp8_paged_mqa_logits_next_n4),
534-546 (test_fp4_paged_mqa_logits), and 598-603
(test_fp4_paged_mqa_logits_next_n4), assert
torch.isfinite(out.float()[~neginf_mask]).all(). Keep the existing comparison
logic after each assertion.
- Around line 292-296: Update the GPU schedule test around
compute_paged_mqa_logits_schedule to skip explicitly when _CUTE_DSL_AVAILABLE is
false, before invoking the helper or synchronizing CUDA. Keep the existing
CPU/GPU comparison unchanged when the CuTe DSL is available, ensuring the test
cannot pass via _compute_schedule_metadata.
---
Nitpick comments:
In `@benchmarks/_profile_attn_overhead.py`:
- Around line 1-5: Add a module docstring at the start of
benchmarks/_profile_attn_overhead.py stating that the script measures per-call
host overhead and depends on private symbols from
flashinfer.attn_scores.attn_scores. Keep the existing imports and benchmark
behavior unchanged.
In `@benchmarks/bench_paged_mqa_logits_cudagraph.py`:
- Around line 93-97: Replace the manual alignment expressions in both loops,
including the `max_ml` calculation near `fp8_cfgs` and the corresponding site
near line 168, with calls to the imported `aligned_context_len` helper. Remove
the now-unused private `_SPLIT_KV` import while preserving the existing
pre-allocation behavior.
In `@benchmarks/bench_paged_mqa_logits.py`:
- Around line 99-109: Update the benchmark closure definitions in both loops,
including fn_trt, fn_fi_sched, and fn_fi_cached, to bind each captured loop
variable through closure default arguments and clear Ruff B023. Replace the
semicolon-joined warmup statements at the referenced locations with separate
statements to clear E702, preserving execution order and behavior.
In `@flashinfer/attn_scores/__init__.py`:
- Around line 23-29: Sort the __all__ entries alphabetically in
flashinfer/attn_scores/__init__.py lines 23-29: use aligned_context_len,
compute_paged_mqa_logits_schedule, fp4_paged_mqa_logits, fp8_paged_mqa_logits,
and precompile_paged_mqa_logits. Also reorder __all__ in
flashinfer/attn_scores/kernels/__init__.py line 18 to FP4MQALogitsKernel
followed by FP8MQALogitsKernel.
In `@flashinfer/attn_scores/attn_scores.py`:
- Around line 176-180: Update the exception handling in _to_cutlass so the
ValueError raised for an unsupported dtype explicitly suppresses the caught
KeyError context by chaining it from None.
In `@flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py`:
- Around line 600-604: Update the num_tmem_alloc_cols_total calculation in the
relevant allocator initialization to use integer bit-length arithmetic for
rounding raw_total up to the next power of two, while preserving the minimum of
32. Remove the now-unused math import if this is its only use.
In `@flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py`:
- Around line 212-216: Mark remove_kv_wait_in_epilogue, early_tmem_copy,
smem_subpartition_opt, max_kv_pipeline, and max_umma_pipeline as experimental in
the relevant public class docstring, since _compile_fp8_kernel does not set them
and tests do not cover them. Clearly state that these performance knobs are
unsupported or unvalidated on the public compile path.
In `@flashinfer/attn_scores/kernels/schedule_kernel.py`:
- Around line 120-148: Replace the fully unrolled prefix_sum scan in the
schedule metadata loop with a bounded binary search over the non-decreasing
prefix_sum array, using approximately log2(kAligned) iterations to find the
first value greater than seg_starts and preserve q_idx_out semantics. Keep the
existing kv_split_idx calculation and metadata writes unchanged, and retain the
linear scan only if compile-time measurements show it remains acceptable.
In `@tests/attn_scores/test_attn_scores.py`:
- Around line 180-184: Add an assertion in the helper before the conditional
transpose to require page_size == 128 whenever remove_online_sf_transpose is
true, so unsupported combinations fail immediately while existing valid behavior
remains unchanged.
- Around line 255-261: Replace the six duplicated inline causal-mask
computations in the numerical tests with calls to
_valid_causal_mask(context_lens, next_n, max_model_len, device), assigning its
negation to neginf_mask. Preserve the surrounding test logic and use the
existing helper as the single source for the masking rule.
- Around line 753-781: Add a test_fp4_input_validation alongside
test_fp8_input_validation covering fp4_paged_mqa_logits with the same four
guard-rail cases: undersized out using aligned_context_len, an unsupported
output dtype while preserving FP4’s allowed bfloat16 behavior, CPU context_lens,
and int64 block_table. Use an FP4-supported setup and exercise the next_n=4
atom-split path so validation also covers its context_lens and block_table
rewrites.
🪄 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: 7acb3822-f84a-4b68-97b6-9f2842a74fdb
📥 Commits
Reviewing files that changed from the base of the PR and between e493ed8 and 20fdd8662ff2049d5b50af202d6c6711dbb22a6e.
📒 Files selected for processing (13)
benchmarks/_profile_attn_overhead.pybenchmarks/_profile_fp8_overhead.pybenchmarks/bench_paged_mqa_logits.pybenchmarks/bench_paged_mqa_logits_cudagraph.pybenchmarks/bench_paged_mqa_logits_eager_fair.pyflashinfer/__init__.pyflashinfer/attn_scores/__init__.pyflashinfer/attn_scores/attn_scores.pyflashinfer/attn_scores/kernels/__init__.pyflashinfer/attn_scores/kernels/fp4_paged_mqa_logits.pyflashinfer/attn_scores/kernels/fp8_paged_mqa_logits.pyflashinfer/attn_scores/kernels/schedule_kernel.pytests/attn_scores/test_attn_scores.py
dhiraj113
force-pushed
the
dhiraj113/indexer
branch
from
August 6, 2026 03:57
e86abc5 to
baa48f8
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py (2)
12-15: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a
next_n=4FP4 reference case.
fp4_paged_mqa_logitsatom-splitsnext_n=4and changes the effective batch,context_lens, andblock_table. The current cases use onlynext_n=1andnext_n=2, so this test does not compare that path with the reference.Add a
shape_kwargscase withnext_n=4.Confidence: high.
🤖 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_fp4_paged_mqa_logits_reference_correctness.py` around lines 12 - 15, Add a third shape_kwargs test case for the FP4 paged MQA logits reference comparison using next_n=4, with valid batch_size, max_context_len, and phys_block_kv values consistent with the existing cases, so the atom-split path is exercised.
3-9: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the redundant CUDA availability decorators.
The test environment assumes CUDA availability. Keep
_skip_if_not_sm100_or_103()as the architecture gate. Remove thepytestimport after removing the decorator.
tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py#L3-L9: removepytestand@pytest.mark.skipif(...).tests/trace/test_fp8_paged_mqa_logits_reference_correctness.py#L3-L9: removepytestand@pytest.mark.skipif(...).Confidence: high. As per coding guidelines, “Match established project style.” Based on learnings, tests assume CUDA and should avoid CPU-only guards.
🤖 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_fp4_paged_mqa_logits_reference_correctness.py` around lines 3 - 9, Remove the redundant pytest import and CUDA availability decorator from tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py lines 3-9, and make the same change in tests/trace/test_fp8_paged_mqa_logits_reference_correctness.py lines 3-9. Preserve _skip_if_not_sm100_or_103() as the architecture gate in both tests.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py`:
- Around line 12-15: Add a third shape_kwargs test case for the FP4 paged MQA
logits reference comparison using next_n=4, with valid batch_size,
max_context_len, and phys_block_kv values consistent with the existing cases, so
the atom-split path is exercised.
- Around line 3-9: Remove the redundant pytest import and CUDA availability
decorator from tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py
lines 3-9, and make the same change in
tests/trace/test_fp8_paged_mqa_logits_reference_correctness.py lines 3-9.
Preserve _skip_if_not_sm100_or_103() as the architecture gate in both tests.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 777ccb17-059d-4724-a0d9-822dc9b6c20d
📥 Commits
Reviewing files that changed from the base of the PR and between 20fdd8662ff2049d5b50af202d6c6711dbb22a6e and e86abc5d0d5ab0d6eb3f0620bdc6d471ddc0e402.
📒 Files selected for processing (8)
flashinfer/attn_scores/attn_scores.pyflashinfer/trace/templates/attn_scores.pytests/trace/example.pytests/trace/fi_trace_out/fp4_paged_mqa_logits_nn2_H64_Dp64_pbk64.jsontests/trace/fi_trace_out/fp8_paged_mqa_logits_nn2_H64_D128_pbk64.jsontests/trace/test_fi_trace_template_consistency.pytests/trace/test_fp4_paged_mqa_logits_reference_correctness.pytests/trace/test_fp8_paged_mqa_logits_reference_correctness.py
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (6)
flashinfer/trace/templates/attn_scores.py (1)
404-470: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winThe FP4 reference default dtype and the declared output dtype disagree.
Line 604 declares the trace output as
bfloat16, matching the API default atflashinfer/attn_scores/attn_scores.pyline 631._fp4_paged_mqa_logits_referencedefaultsoutput_dtype=torch.float32at line 413 and returnslogits.to(output_dtype)at line 470. The reference therefore produces float32 while the schema promises bfloat16.
_paged_mqa_logits_masked_checkcasts both sides with.float(), so the comparison still runs. The mismatch matters for anything that trusts the declared schema, such as generated benchmark fixtures. Align the reference default with the declared output dtype.♻️ Proposed change
block_table, max_context_len, - output_dtype=torch.float32, + output_dtype=torch.bfloat16, ):Also applies to: 601-607
🤖 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/trace/templates/attn_scores.py` around lines 404 - 470, Update the output_dtype default in _fp4_paged_mqa_logits_reference to bfloat16 so the reference output matches the declared trace schema and API default. Keep the existing logits.to(output_dtype) conversion unchanged.flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py (1)
212-217: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThese five tuning flags are not part of the compile cache key.
_compile_fp8_kernelinflashinfer/attn_scores/attn_scores.pypasses none ofremove_kv_wait_in_epilogue,early_tmem_copy,smem_subpartition_opt,max_kv_pipeline, ormax_umma_pipeline, so all five always hold their defaults and no collision occurs today. If a later change threads any of them through, two structurally different kernels will share one@functools.cacheentry and one tag.Either add them to the
_compile_fp8_kernelsignature and the tag now, or add a short comment recording that they must join the cache key before they become configurable.🤖 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/attn_scores/kernels/fp8_paged_mqa_logits.py` around lines 212 - 217, Document near the five tuning flags in the affected kernel configuration that remove_kv_wait_in_epilogue, early_tmem_copy, smem_subpartition_opt, max_kv_pipeline, and max_umma_pipeline must be added to _compile_fp8_kernel’s signature and cache tag before becoming configurable. Do not change their current defaults or behavior.flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py (1)
84-193: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThese five f16x2 helpers duplicate
fp8_paged_mqa_logits.pylines 93-193 verbatim.pack_f16x2,unpack_f16x2,fma_f16x2,max_f16x2, andadd_f16x2are byte-identical in both kernel files. Move them, together with the bf16 variants, into a shared module such asflashinfer/attn_scores/kernels/_packed_math.pyand import them from both kernels. A future fix to one copy will otherwise miss the other.🤖 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/attn_scores/kernels/fp4_paged_mqa_logits.py` around lines 84 - 193, Extract pack_f16x2, unpack_f16x2, fma_f16x2, max_f16x2, add_f16x2, and the corresponding bf16 helpers into a shared module such as _packed_math.py. Remove both kernel-local copies and import the shared symbols in fp4_paged_mqa_logits.py and fp8_paged_mqa_logits.py, preserving their existing behavior and signatures.tests/attn_scores/test_attn_scores.py (1)
753-782: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the matching validation test for the FP4 entry point.
test_fp8_input_validationcovers the undersizedout=, the unsupported output dtype, CPUcontext_lens, and an int64block_table.fp4_paged_mqa_logitsruns the same two validators with a different dtype set (_FP4_DTYPESatflashinfer/attn_scores/attn_scores.pyline 51) and its own_validate_outcall at line 730, and none of it is exercised.Add an equivalent
test_fp4_input_validation. One case is FP4-specific:output_dtype=torch.int8(or any dtype outside fp32/fp16/bf16) must raise, and the FP4 message text differs from the FP8 one.A second gap is worth covering in the same pass:
next_n=4with an explicitschedule_meta. The host expands the batch to2Band rewritescontext_lens, so a schedule built from the originalcontext_lensproduces wrong logits with no error today.🤖 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/attn_scores/test_attn_scores.py` around lines 753 - 782, Add a matching test function for fp4_paged_mqa_logits covering undersized out validation, FP4-specific unsupported output dtype messaging, CPU context_lens, and int64 block_table, using the FP4 case setup and validators. Also add coverage for next_n=4 with an explicit schedule_meta built from the original context_lens, verifying the expanded-batch schedule is rejected or produces the correct validated behavior rather than silently wrong logits.benchmarks/bench_paged_mqa_logits_cudagraph.py (1)
16-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFour of these five private imports are unused, and one can break the import. The script uses only
_SPLIT_KV._to_cutlassis defined inside theif _CUTE_DSL_AVAILABLE:block inflashinfer/attn_scores/attn_scores.py(line 185), so thisfrom ... import _to_cutlassraisesImportErroron any machine withoutnvidia-cutlass-dsl. Drop the unused names.Line 14 already imports the public
aligned_context_len. Use it at lines 97 and 168 instead of recomputing the alignment from_SPLIT_KV; then this import block can be removed entirely.♻️ Proposed change
-from flashinfer.attn_scores.attn_scores import ( - _compile_fp8_kernel, - _compile_fp4_kernel, - _compute_schedule_metadata, - _to_cutlass, - _SPLIT_KV, -) from flashinfer.utils import get_device_sm_countThen at lines 97 and 168:
- aligned = ((max_ml + _SPLIT_KV - 1) // _SPLIT_KV) * _SPLIT_KV + aligned = aligned_context_len(max_ml)🤖 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_paged_mqa_logits_cudagraph.py` around lines 16 - 22, Remove the entire private import block from the benchmark, retaining only the existing public aligned_context_len import and any other symbols actually used. In the benchmark code paths around the current alignment calculations, update both uses to call aligned_context_len instead of recomputing alignment from _SPLIT_KV; remove _SPLIT_KV if it is no longer referenced.flashinfer/attn_scores/attn_scores.py (1)
573-573: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
aligned_context_leninstead of repeating the formula. Line 441 defines the public helper, and lines 573 and 728 recompute the same expression. Three copies of one alignment rule invite drift.- aligned_ctx = ((max_context_len + _SPLIT_KV - 1) // _SPLIT_KV) * _SPLIT_KV + aligned_ctx = aligned_context_len(max_context_len)Apply the same change at line 728.
🤖 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/attn_scores/attn_scores.py` at line 573, Replace the repeated alignment formula in the affected code with the existing aligned_context_len helper defined near line 441, and apply the same reuse at the corresponding computation near line 728. Preserve the current inputs and resulting alignment behavior while eliminating both duplicate expressions.
🤖 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/attn_scores/attn_scores.py`:
- Around line 477-480: Validate a caller-provided out buffer in the scheduling
function before either _gpu_schedule or the CPU fallback writes to it, matching
the existing _validate_out behavior used by fp8_paged_mqa_logits. Enforce the
documented shape (num_sms+1, 2), int32 dtype, and CUDA/device placement; retain
the current allocation path when out is None.
In `@flashinfer/trace/templates/attn_scores.py`:
- Around line 180-195: Update the checker around the finite/valid mask so any
non-finite value in the actual kernel output within the causal region
immediately returns False, while preserving tolerance for non-finite reference
values caused by fp8/fp4 accumulation. Ensure the all-invalid case cannot return
True merely because kernel output was excluded, and keep both
fp8_paged_mqa_logits_trace and fp4_paged_mqa_logits_trace using the corrected
checker.
In `@tests/trace/fi_trace_out/fp4_paged_mqa_logits_nn2_H64_Dp64_pbk64.json`:
- Line 125: Update _paged_mqa_logits_masked_check so non-finite actual logits in
unmasked positions where the reference is finite cause the check to fail. Define
validity from the causal mask and finite reference values, while separately
reject any non-finite values in act at those positions; retain exclusion of
reference-non-finite and masked positions from numeric comparison.
---
Nitpick comments:
In `@benchmarks/bench_paged_mqa_logits_cudagraph.py`:
- Around line 16-22: Remove the entire private import block from the benchmark,
retaining only the existing public aligned_context_len import and any other
symbols actually used. In the benchmark code paths around the current alignment
calculations, update both uses to call aligned_context_len instead of
recomputing alignment from _SPLIT_KV; remove _SPLIT_KV if it is no longer
referenced.
In `@flashinfer/attn_scores/attn_scores.py`:
- Line 573: Replace the repeated alignment formula in the affected code with the
existing aligned_context_len helper defined near line 441, and apply the same
reuse at the corresponding computation near line 728. Preserve the current
inputs and resulting alignment behavior while eliminating both duplicate
expressions.
In `@flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py`:
- Around line 84-193: Extract pack_f16x2, unpack_f16x2, fma_f16x2, max_f16x2,
add_f16x2, and the corresponding bf16 helpers into a shared module such as
_packed_math.py. Remove both kernel-local copies and import the shared symbols
in fp4_paged_mqa_logits.py and fp8_paged_mqa_logits.py, preserving their
existing behavior and signatures.
In `@flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py`:
- Around line 212-217: Document near the five tuning flags in the affected
kernel configuration that remove_kv_wait_in_epilogue, early_tmem_copy,
smem_subpartition_opt, max_kv_pipeline, and max_umma_pipeline must be added to
_compile_fp8_kernel’s signature and cache tag before becoming configurable. Do
not change their current defaults or behavior.
In `@flashinfer/trace/templates/attn_scores.py`:
- Around line 404-470: Update the output_dtype default in
_fp4_paged_mqa_logits_reference to bfloat16 so the reference output matches the
declared trace schema and API default. Keep the existing logits.to(output_dtype)
conversion unchanged.
In `@tests/attn_scores/test_attn_scores.py`:
- Around line 753-782: Add a matching test function for fp4_paged_mqa_logits
covering undersized out validation, FP4-specific unsupported output dtype
messaging, CPU context_lens, and int64 block_table, using the FP4 case setup and
validators. Also add coverage for next_n=4 with an explicit schedule_meta built
from the original context_lens, verifying the expanded-batch schedule is
rejected or produces the correct validated behavior rather than silently wrong
logits.
🪄 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: 900eeae7-18f6-483e-b7f7-0f14624cfbaf
📥 Commits
Reviewing files that changed from the base of the PR and between e493ed8 and baa48f860c1b6a4be0d2df44cdc8b3135ca10fd3.
📒 Files selected for processing (20)
benchmarks/_profile_attn_overhead.pybenchmarks/_profile_fp8_overhead.pybenchmarks/bench_paged_mqa_logits.pybenchmarks/bench_paged_mqa_logits_cudagraph.pybenchmarks/bench_paged_mqa_logits_eager_fair.pyflashinfer/__init__.pyflashinfer/attn_scores/__init__.pyflashinfer/attn_scores/attn_scores.pyflashinfer/attn_scores/kernels/__init__.pyflashinfer/attn_scores/kernels/fp4_paged_mqa_logits.pyflashinfer/attn_scores/kernels/fp8_paged_mqa_logits.pyflashinfer/attn_scores/kernels/schedule_kernel.pyflashinfer/trace/templates/attn_scores.pytests/attn_scores/test_attn_scores.pytests/trace/example.pytests/trace/fi_trace_out/fp4_paged_mqa_logits_nn2_H64_Dp64_pbk64.jsontests/trace/fi_trace_out/fp8_paged_mqa_logits_nn2_H64_D128_pbk64.jsontests/trace/test_fi_trace_template_consistency.pytests/trace/test_fp4_paged_mqa_logits_reference_correctness.pytests/trace/test_fp8_paged_mqa_logits_reference_correctness.py
🚧 Files skipped from review as they are similar to previous changes (9)
- tests/trace/example.py
- tests/trace/test_fi_trace_template_consistency.py
- benchmarks/_profile_attn_overhead.py
- flashinfer/init.py
- benchmarks/_profile_fp8_overhead.py
- tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py
- tests/trace/test_fp8_paged_mqa_logits_reference_correctness.py
- tests/trace/fi_trace_out/fp8_paged_mqa_logits_nn2_H64_D128_pbk64.json
- flashinfer/attn_scores/kernels/schedule_kernel.py
dhiraj113
force-pushed
the
dhiraj113/indexer
branch
from
August 6, 2026 04:18
baa48f8 to
b50d2de
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (9)
tests/attn_scores/test_attn_scores.py (3)
294-316: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winSkip the test when CuTe DSL is unavailable.
compute_paged_mqa_logits_scheduletakes the GPU path only when_CUTE_DSL_AVAILABLEis True. If CuTe DSL is not installed, the helper falls back to_compute_schedule_metadata, and line 314 compares the CPU reference against itself. The test then passes without running the GPU schedule kernel.💚 Proposed fix
from flashinfer.attn_scores.attn_scores import ( + _CUTE_DSL_AVAILABLE, _cached_num_sms, _compute_schedule_metadata, compute_paged_mqa_logits_schedule, ) from flashinfer.utils import get_device_index + if not _CUTE_DSL_AVAILABLE: + pytest.skip("GPU schedule kernel requires nvidia-cutlass-dsl") + device = "cuda"🤖 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/attn_scores/test_attn_scores.py` around lines 294 - 316, Skip this GPU schedule test when CuTe DSL is unavailable by checking the module’s _CUTE_DSL_AVAILABLE flag before invoking compute_paged_mqa_logits_schedule. Keep the existing reference and GPU comparison unchanged when the flag is enabled, ensuring the test only exercises the GPU kernel path.
754-824: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winThe
phys_block_kv=64case does not test the transpose path.Line 795 and line 807 both evaluate
phys_block_kv == 128to False whenphys_block_kvis 64. The buffer at line 794 is then built the same way as the one at line 811, and both calls passremove_online_sf_transpose=False. For that parametrization the test compares one configuration against an identical configuration. Only the 128 case exercises the feature. Restrict the parametrization to 128, or add a distinct assertion for the 64 case.🤖 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/attn_scores/test_attn_scores.py` around lines 754 - 824, Update test_fp4_paged_mqa_logits_remove_sf_transpose so every parametrized case exercises distinct configurations: either restrict phys_block_kv to 128, where remove_online_sf_transpose is enabled for the primary and disabled for the cross-check, or add a separate meaningful 64-case assertion instead of comparing identical calls.
395-409: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFive tests drop non-finite kernel output from their assertions.
Each test builds
finite = torch.isfinite(out_m) & torch.isfinite(ref_m)and then removes those positions fromvalid. A NaN or inf that the kernel writes at a causally valid position is excluded fromassert_closeand zeroed before the cosine check, so the test passes. Add a finiteness assertion on the kernel output in the valid region at each site.
tests/attn_scores/test_attn_scores.py#L395-L409(this comment): asserttorch.isfinite(out.float()[~neginf_mask]).all()before computingfinite, intest_fp8_paged_mqa_logits.tests/attn_scores/test_attn_scores.py#L483-L492: add the same assertion intest_fp8_paged_mqa_logits_fp16.tests/attn_scores/test_attn_scores.py#L554-L563: add the same assertion intest_fp8_paged_mqa_logits_next_n4.tests/attn_scores/test_attn_scores.py#L649-L665: add the same assertion intest_fp4_paged_mqa_logits.tests/attn_scores/test_attn_scores.py#L740-L749: add the same assertion intest_fp4_paged_mqa_logits_next_n4.🤖 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/attn_scores/test_attn_scores.py` around lines 395 - 409, Add a finiteness assertion for kernel output over all causally valid positions before computing `finite` in each listed test: `test_fp8_paged_mqa_logits`, `test_fp8_paged_mqa_logits_fp16`, `test_fp8_paged_mqa_logits_next_n4`, `test_fp4_paged_mqa_logits`, and `test_fp4_paged_mqa_logits_next_n4`. Check `out.float()[~neginf_mask]` with `torch.isfinite(...).all()`, while preserving the existing comparison logic afterward.flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py (1)
428-430: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDo not silently ignore
remove_online_sf_transpose.When
phys_block_kv != 128, the constructor resets the flag to False._compile_fp4_kernelinflashinfer/attn_scores/attn_scores.pystill puts the requested value in the compile cache key and in the tag, so two different requests map to the same generated kernel under different tags. The caller also receives no signal that the requested layout was ignored. If the caller pre-arranged GMEM SF into UTCCP chunk layout, the kernel transposes already-transposed data and produces wrong results.Raise instead of downgrading.
🛠️ Proposed change
- if remove_online_sf_transpose and phys_block_kv != 128: - remove_online_sf_transpose = False + if remove_online_sf_transpose and phys_block_kv != 128: + raise ValueError( + "remove_online_sf_transpose requires phys_block_kv=128; " + f"got {phys_block_kv}." + )🤖 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/attn_scores/kernels/fp4_paged_mqa_logits.py` around lines 428 - 430, Update the constructor logic around remove_online_sf_transpose so it raises an error when the flag is requested with phys_block_kv != 128, rather than resetting it to False. Preserve the requested flag unchanged for supported configurations, ensuring callers are explicitly notified instead of allowing _compile_fp4_kernel to cache and tag an incompatible kernel.flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py (1)
272-289: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winClamp the computed stage counts to a positive minimum.
Two optional flags compute a stage count that can reach zero:
- Line 273: for a large
N,TMEM_COLS // (2 * self.N)is 0, sonum_umma_stagesbecomes 0.- Line 287: for a large
N * head_dim,qw_totalapproaches or exceedsSMEM_BUDGET, sonum_kv_stagesbecomes zero or negative.A non-positive stage count creates an invalid pipeline and an invalid SMEM layout. No current caller sets either flag, so this is latent.
🛠️ Proposed guard
if max_umma_pipeline: - self.num_umma_stages = min(2, TMEM_COLS // (2 * self.N)) + self.num_umma_stages = max(1, min(2, TMEM_COLS // (2 * self.N))) else: self.num_umma_stages = 1 if max_kv_pipeline: @@ self.num_kv_stages = (SMEM_BUDGET - qw_total) // kv_scale_per_stage + if self.num_kv_stages < 2: + raise ValueError( + f"max_kv_pipeline leaves no SMEM for KV stages " + f"(computed {self.num_kv_stages}); reduce N or head_dim." + )🤖 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/attn_scores/kernels/fp8_paged_mqa_logits.py` around lines 272 - 289, Clamp the computed stage counts in the max_umma_pipeline and max_kv_pipeline branches to a minimum of one. Update the assignments to self.num_umma_stages and self.num_kv_stages so large N or insufficient SMEM cannot produce zero or negative pipeline stages, while preserving the existing upper bound and default values.flashinfer/attn_scores/attn_scores.py (3)
694-744: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject or document
schedule_metafornext_n=4.For
next_n == 4, the function expands the batch to2Band rewritescontext_lensintokernel_ctx_lens. Whenschedule_metais None, the schedule comes fromkernel_ctx_lens, which is correct. When the caller passesschedule_meta, the function uses it unchanged. A caller who builds it withcompute_paged_mqa_logits_schedule(context_lens)produces a schedule for the originalBsequences. The kernel then indexesq_idxover[0, 2B)with wrong split boundaries and returns wrong logits without an error.Reject the combination, or state the requirement in the docstring.
🛠️ Proposed guard
if next_n == 4: + if schedule_meta is not None: + raise ValueError( + "next_n=4 uses an internal atom-split (batch 2B). Pass schedule_meta " + "computed from the split context_lens, or leave it as None." + ) exp_B = B * 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 `@flashinfer/attn_scores/attn_scores.py` around lines 694 - 744, Handle schedule_meta explicitly in the next_n == 4 expansion path: either reject caller-provided metadata with a clear validation error, or document that it must be computed for the expanded kernel_ctx_lens shape [2B] rather than the original context_lens. Keep automatic schedule computation based on kernel_ctx_lens unchanged.
574-576: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReturn exactly
B * next_nrows whenouthas more rows.
_validate_outacceptsout.shape[0] > rows. The slice at line 576 keeps every row, so the returned tensor can exceed the documented[B*next_n, max_context_len]shape. The extra rows hold uninitialized data. The same problem exists at line 731 infp4_paged_mqa_logits.♻️ Proposed change
- logits = out[:, :max_context_len] + logits = out[: B * next_n, :max_context_len]🤖 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/attn_scores/attn_scores.py` around lines 574 - 576, Update the `out` handling in the enclosing attention-score function to slice rows as well as columns, returning exactly `B * next_n` rows and `max_context_len` columns. Apply the same correction in `fp4_paged_mqa_logits`, preserving `_validate_out`'s ability to accept larger output buffers while excluding unused rows from the returned tensor.
476-487: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winValidate the caller-provided
outbuffer before the schedule kernel writes into it.The docstring requires
outto be[num_sms+1, 2]int32 on CUDA. The function does not check this._gpu_schedulepasses the buffer directly to the compiled kernel, which writesnum_sms+1rows. A smaller or wrongly-typed buffer causes an out-of-bounds device write. The CPU fallback at line 485 raises only on a size mismatch.
fp8_paged_mqa_logitsvalidates itsoutparameter with_validate_out. Apply the same check here.🛡️ Proposed guard
+ if out is not None: + if tuple(out.shape) != (num_sms + 1, 2): + raise ValueError( + f"out must have shape ({num_sms + 1}, 2); got {tuple(out.shape)}" + ) + if out.dtype != torch.int32 or out.device != device: + raise ValueError( + f"out must be an int32 tensor on {device}; got " + f"dtype={out.dtype}, device={out.device}" + ) + if use_gpu_kernel and _CUTE_DSL_AVAILABLE and context_lens.is_cuda:🤖 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/attn_scores/attn_scores.py` around lines 476 - 487, Validate a caller-provided `out` buffer before either scheduling path uses it, matching the existing `_validate_out` behavior used by `fp8_paged_mqa_logits`. In the schedule function around `_gpu_schedule` and `_compute_schedule_metadata`, enforce shape `[num_sms + 1, 2]`, `torch.int32` dtype, and CUDA placement as required by the docstring, while preserving allocation and return behavior when `out` is None.flashinfer/trace/templates/attn_scores.py (1)
178-195: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winThe checker discards non-finite kernel output instead of failing on it.
Line 180 builds
finitefrom both tensors, and line 181 removes those positions fromvalid. A NaN or inf that the kernel writes at a causally valid position is therefore excluded fromtorch.allcloseand zeroed at line 179 before the cosine fallback, so the function returns True. Line 182 makes this worse at the limit: if every valid position is non-finite,valid.any()is False and the check returns True with nothing compared.Both
fp8_paged_mqa_logits_traceandfp4_paged_mqa_logits_traceuse this checker, so neither can detect a kernel that emits NaN. Fail when the actual output is non-finite inside the causal region.💚 Proposed change
r = ref.float().masked_fill(neginf_mask, 0) a = act.float().masked_fill(neginf_mask, 0) - finite = torch.isfinite(r) & torch.isfinite(a) - valid = (~neginf_mask) & finite + # A non-finite actual output inside the causal region is a kernel failure. + if not torch.isfinite(a[~neginf_mask]).all(): + return False + # A non-finite reference is tolerated: fp8/fp4 accumulation order differs at extremes. + valid = (~neginf_mask) & torch.isfinite(r) if not valid.any(): return True🤖 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/trace/templates/attn_scores.py` around lines 178 - 195, Update the shared checker around the `finite`, `valid`, and `valid.any()` logic to fail immediately when `act` is non-finite at any causally valid position, rather than filtering those positions out or returning success when none remain. Continue masking invalid causal positions and comparing finite valid values as before so both `fp8_paged_mqa_logits_trace` and `fp4_paged_mqa_logits_trace` reject NaN or infinite kernel outputs.
🧹 Nitpick comments (2)
flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py (1)
1569-1571: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused accumulator tensor in each UMMA branch.
is_umma_warp_0buildstCtAcc_base_1but uses onlytCtAcc_base_0(line 1689).is_umma_warp_1at lines 1755-1758 buildstCtAcc_base_0but uses onlytCtAcc_base_1(line 1843). Drop the unused tensor in each branch to make the TMEM ownership explicit.🤖 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/attn_scores/kernels/fp4_paged_mqa_logits.py` around lines 1569 - 1571, Remove the unused accumulator tensor construction from each UMMA branch: delete tCtAcc_base_1 creation in is_umma_warp_0, and delete tCtAcc_base_0 creation in is_umma_warp_1. Preserve the branch-specific accumulator tensors that are consumed later.flashinfer/attn_scores/__init__.py (1)
23-29: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to satisfy RUF022.Ruff reports
__all__is not sorted. Apply isort-style ordering.♻️ Proposed change
__all__ = [ + "aligned_context_len", + "compute_paged_mqa_logits_schedule", + "fp4_paged_mqa_logits", "fp8_paged_mqa_logits", - "fp4_paged_mqa_logits", - "compute_paged_mqa_logits_schedule", - "aligned_context_len", "precompile_paged_mqa_logits", ]🤖 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/attn_scores/__init__.py` around lines 23 - 29, Sort the entries in the __all__ list alphabetically using isort-style ordering to satisfy RUF022, preserving all existing exports.Source: Linters/SAST tools
🤖 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/attn_scores/attn_scores.py`:
- Around line 698-699: Clamp the first split length in the ctx_pair construction
near kernel_ctx_lens so context_lens values below 2 produce 0 instead of a
negative value. Preserve the existing second split and reshape behavior,
ensuring the kernel receives non-negative split context lengths.
---
Duplicate comments:
In `@flashinfer/attn_scores/attn_scores.py`:
- Around line 694-744: Handle schedule_meta explicitly in the next_n == 4
expansion path: either reject caller-provided metadata with a clear validation
error, or document that it must be computed for the expanded kernel_ctx_lens
shape [2B] rather than the original context_lens. Keep automatic schedule
computation based on kernel_ctx_lens unchanged.
- Around line 574-576: Update the `out` handling in the enclosing
attention-score function to slice rows as well as columns, returning exactly `B
* next_n` rows and `max_context_len` columns. Apply the same correction in
`fp4_paged_mqa_logits`, preserving `_validate_out`'s ability to accept larger
output buffers while excluding unused rows from the returned tensor.
- Around line 476-487: Validate a caller-provided `out` buffer before either
scheduling path uses it, matching the existing `_validate_out` behavior used by
`fp8_paged_mqa_logits`. In the schedule function around `_gpu_schedule` and
`_compute_schedule_metadata`, enforce shape `[num_sms + 1, 2]`, `torch.int32`
dtype, and CUDA placement as required by the docstring, while preserving
allocation and return behavior when `out` is None.
In `@flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py`:
- Around line 428-430: Update the constructor logic around
remove_online_sf_transpose so it raises an error when the flag is requested with
phys_block_kv != 128, rather than resetting it to False. Preserve the requested
flag unchanged for supported configurations, ensuring callers are explicitly
notified instead of allowing _compile_fp4_kernel to cache and tag an
incompatible kernel.
In `@flashinfer/attn_scores/kernels/fp8_paged_mqa_logits.py`:
- Around line 272-289: Clamp the computed stage counts in the max_umma_pipeline
and max_kv_pipeline branches to a minimum of one. Update the assignments to
self.num_umma_stages and self.num_kv_stages so large N or insufficient SMEM
cannot produce zero or negative pipeline stages, while preserving the existing
upper bound and default values.
In `@flashinfer/trace/templates/attn_scores.py`:
- Around line 178-195: Update the shared checker around the `finite`, `valid`,
and `valid.any()` logic to fail immediately when `act` is non-finite at any
causally valid position, rather than filtering those positions out or returning
success when none remain. Continue masking invalid causal positions and
comparing finite valid values as before so both `fp8_paged_mqa_logits_trace` and
`fp4_paged_mqa_logits_trace` reject NaN or infinite kernel outputs.
In `@tests/attn_scores/test_attn_scores.py`:
- Around line 294-316: Skip this GPU schedule test when CuTe DSL is unavailable
by checking the module’s _CUTE_DSL_AVAILABLE flag before invoking
compute_paged_mqa_logits_schedule. Keep the existing reference and GPU
comparison unchanged when the flag is enabled, ensuring the test only exercises
the GPU kernel path.
- Around line 754-824: Update test_fp4_paged_mqa_logits_remove_sf_transpose so
every parametrized case exercises distinct configurations: either restrict
phys_block_kv to 128, where remove_online_sf_transpose is enabled for the
primary and disabled for the cross-check, or add a separate meaningful 64-case
assertion instead of comparing identical calls.
- Around line 395-409: Add a finiteness assertion for kernel output over all
causally valid positions before computing `finite` in each listed test:
`test_fp8_paged_mqa_logits`, `test_fp8_paged_mqa_logits_fp16`,
`test_fp8_paged_mqa_logits_next_n4`, `test_fp4_paged_mqa_logits`, and
`test_fp4_paged_mqa_logits_next_n4`. Check `out.float()[~neginf_mask]` with
`torch.isfinite(...).all()`, while preserving the existing comparison logic
afterward.
---
Nitpick comments:
In `@flashinfer/attn_scores/__init__.py`:
- Around line 23-29: Sort the entries in the __all__ list alphabetically using
isort-style ordering to satisfy RUF022, preserving all existing exports.
In `@flashinfer/attn_scores/kernels/fp4_paged_mqa_logits.py`:
- Around line 1569-1571: Remove the unused accumulator tensor construction from
each UMMA branch: delete tCtAcc_base_1 creation in is_umma_warp_0, and delete
tCtAcc_base_0 creation in is_umma_warp_1. Preserve the branch-specific
accumulator tensors that are consumed later.
🪄 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: b4497067-c163-4ebc-91d4-a0ed935db9b6
📥 Commits
Reviewing files that changed from the base of the PR and between e493ed8 and b50d2deb2dc963b425c078491eadb3586b019dcd.
📒 Files selected for processing (15)
flashinfer/__init__.pyflashinfer/attn_scores/__init__.pyflashinfer/attn_scores/attn_scores.pyflashinfer/attn_scores/kernels/__init__.pyflashinfer/attn_scores/kernels/fp4_paged_mqa_logits.pyflashinfer/attn_scores/kernels/fp8_paged_mqa_logits.pyflashinfer/attn_scores/kernels/schedule_kernel.pyflashinfer/trace/templates/attn_scores.pytests/attn_scores/test_attn_scores.pytests/trace/example.pytests/trace/fi_trace_out/fp4_paged_mqa_logits_nn2_H64_Dp64_pbk64.jsontests/trace/fi_trace_out/fp8_paged_mqa_logits_nn2_H64_D128_pbk64.jsontests/trace/test_fi_trace_template_consistency.pytests/trace/test_fp4_paged_mqa_logits_reference_correctness.pytests/trace/test_fp8_paged_mqa_logits_reference_correctness.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tests/trace/test_fp4_paged_mqa_logits_reference_correctness.py
- tests/trace/example.py
- tests/trace/test_fp8_paged_mqa_logits_reference_correctness.py
- tests/trace/test_fi_trace_template_consistency.py
- tests/trace/fi_trace_out/fp8_paged_mqa_logits_nn2_H64_D128_pbk64.json
- flashinfer/init.py
- tests/trace/fi_trace_out/fp4_paged_mqa_logits_nn2_H64_Dp64_pbk64.json
- flashinfer/attn_scores/kernels/schedule_kernel.py
dhiraj113
added
model: dsv4
op: misc
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/attn_scores/attn_scores.py`:
- Line 895: Close the unterminated multiline string surrounding the FP4 compile
call before precompile_paged_mqa_logits(), ensuring the
precompile_paged_mqa_logits() invocation is parsed as Python code and the module
has valid syntax.
In `@tests/attn_scores/test_attn_scores_adversarial.py`:
- Around line 496-497: Complete the truncated fp4_paged_mqa_logits() invocation
in the adversarial attention-score test by supplying the remaining arguments and
closing its parentheses, then ensure the file ends with a final newline so
pytest can collect the module.
🪄 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: d99cdfaf-e099-47ba-a144-1d9534fc2359
📥 Commits
Reviewing files that changed from the base of the PR and between b50d2deb2dc963b425c078491eadb3586b019dcd and 27e8095fafde41f04633fc6290b96744d2f4277b.
📒 Files selected for processing (4)
flashinfer/attn_scores/attn_scores.pyflashinfer/trace/templates/attn_scores.pytests/attn_scores/test_attn_scores.pytests/attn_scores/test_attn_scores_adversarial.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/attn_scores/test_attn_scores.py
- flashinfer/trace/templates/attn_scores.py
dhiraj113
force-pushed
the
dhiraj113/indexer
branch
from
August 6, 2026 04:55
27e8095 to
5e6398d
Compare
The validator checked context_lens.shape[0] but not its rank, while the very
next line checks both for block_table. A [B,1] tensor has the same shape[0] as
a [B] one, so it passed and failed later against the rank-1 compiled fake:
ValueError: Mismatched Tensor on argument #0 when calling:
`__call__(context_lens: Tensor([n0], int32 ...
That message never mentions rank, speaks the compiled kernel's vocabulary
rather than the API's, and only appears after JIT compilation -- so on a cold
cache the caller waits through a kernel build to learn about an extra
dimension. [B,1] is a real shape: SGLang's DeepGEMM path produces it.
Fold the rank check into the existing shape check so context_lens and
block_table are validated the same way.
The test matches on the message text rather than the exception type: the old
path also raised ValueError, so a bare pytest.raises would have passed against
unfixed code. Confirmed by reverting the fix -- the test fails, and the
pre-fix failure is still a ValueError.
Addresses the context_lens rank review thread on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
precompile_paged_mqa_logits(device=...) moved only one of the three things
that make a kernel specific to a GPU:
num_sms came from the device= argument
cache tag came from the process-wide current torch device
codegen came from ordinal 0, because CUTLASS's detect_gpu_arch calls
cuDeviceGet(0) unconditionally and ignores the CUDA context
Measured on a heterogeneous host (cuda:0 B100 sm_100a / cuda:1 L40S sm_89),
precompile(device="cuda:1") produced a kernel with the L40S's 142 SMs baked
in, compiled for sm_100a. With a worker's per-rank device set to cuda:1 it
would additionally have been filed under an sm89 cache tag -- writing an
sm100a binary into the sm89 directory of a shared, persistent, cross-process
cache, where a later genuine sm89 process would load and run it.
* Add _cached_gpu_arch(device_index), resolving from the requested device
via get_device_capability(device) and normalising through
CompilationContext so the string matches arch names used elsewhere.
* Pass --gpu-arch in the options string of all three cute.compile sites
(FP8, FP4, schedule). The cute.compile[(GPUArch(...),)] subscript form is
silently discarded when options= is also passed, and these kernels need
--enable-tvm-ffi; verified by falsification -- asking for sm_80 now
correctly refuses an sm100a-only kernel, and did not before.
* Key the in-process caches and the on-disk JIT tag on the arch string.
Not cute.GPUArch: its equality is identity-based, so GPUArch(x) ==
GPUArch(x) is False and every lookup would miss.
* precompile now defaults to the current device rather than a hard-coded
cuda:0, normalises a bare "cuda", and runs the build inside
torch.cuda.device(target) so JitSpecCuteDsl tags the cache for the
target rather than for whatever device happened to be current.
Adding arch to the JIT tag invalidates existing on-disk cache entries, so the
first run after this pays a full recompile. That is intended: the old entries
were tagged without arch, which is the ambiguity being fixed.
Only fp8/fp4 kernel sources are verbatim upstream ports and both are
unchanged; schedule_kernel.py has no upstream counterpart.
Verified on a heterogeneous two-GPU host: arch resolution no longer drifts
with the current device, and precompiling for a device that cannot run the
kernel now fails loudly instead of silently mis-tagging.
Addresses the precompile device/architecture review thread on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
output_dtype is part of the compile cache key, but precompile warmed only one
value per variant: float32 for FP8 and bfloat16 for FP4, each taken from that
API's default. FP8's default matches what callers use; FP4's does not. A
consumer whose logits ABI binds scores as C float must pass
output_dtype=torch.float32, which was a complete cache miss -- leaving exactly
the first-request JIT the helper exists to prevent.
* Add an output_dtypes parameter. When omitted, each variant warms its
common set: float32 for FP8, and both bfloat16 and float32 for FP4 -- the
API default plus the dtype a float logits ABI requires. An explicit tuple
builds only what the caller runs, which was not previously expressible.
* Validate the requested dtypes against each variant's supported set up
front, so an impossible build fails at the helper rather than from inside
a compile.
FP4 now builds 18 kernels rather than 9 by default, roughly 3s -> 6s of
deployment setup, in exchange for removing a first-request compile.
The regression test asserts cache hits for both warmed dtypes and requires an
unwarmed one to miss -- otherwise the hit assertions would hold even if
nothing were warmed. It clears the in-process cache first: that cache is
shared across the session and other tests in this file compile fp4 with every
supported output dtype, so without clearing, the miss assertion is decided by
test ordering rather than by what precompile did. Verified under both fixed
and randomised test order. Clearing costs +0.17s (+0.7%) on the suite,
measured over three runs each way: it drops only the in-memory functools
entries, and the on-disk CuTe-DSL cache still serves them, so repopulation is
a load rather than a rebuild.
Addresses the FP4 precompile output-dtype review thread on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two independent defects in the paged-MQA trace artifacts, plus one of my own.
The committed FP4 JSON was not standalone-runnable: its embedded init called
_pack_ue8m0_to_int() without defining it, and emitted _quantize_to_fp4_e2m1
twice instead. The live dependency tuple already listed the right helpers, so
the artifact was simply stale -- regenerating in a fresh interpreter emits both
correctly, and no change to the dependency renderer was needed.
Separately, the FP4 reference defaulted to float32 while both the public API
and this template's output schema default to bfloat16. The correctness test hid
the disagreement by forcing the kernel to float32, which made it match the
reference's wrong default. FP8 had no such split -- reference, schema and API
all say float32 -- so FP4's reference was the outlier.
* Regenerate both paged-MQA artifacts.
* Default the FP4 reference to bfloat16, so reference, schema, generated JSON
and API agree. The math is done in float regardless; only the final cast
follows the dtype. The docstring records why, so it cannot drift back.
* Stop forcing float32 in the correctness test and assert the dtype on the
schema, the kernel output and the reference output.
* Extend test_rendered_source_standalone.py to exec and invoke both committed
paged-MQA JSONs. That test is what catches this class: a rendered string
missing an inlined dependency still works at dump time because module
globals resolve it, and only breaks for a consumer that exec()s the JSON.
* Fix example.py's docstring, which still named the artifacts _pbk64.json
after the phys_block_kv -> block_size rename.
Regenerating only these two artifacts takes ~5s: auto-dump is per-call via
FLASHINFER_TRACE_DUMP, so the two APIs can be invoked directly rather than
running tests/trace/example.py, which rebuilds all 118 artifacts in ~15 min
(mostly JIT for GEMM kernels this PR does not touch) and churns files it does
not own.
Not done here: emitting a separate traced FP32 definition. output_dtype is not
an axis and the filename derives from name_prefix plus the const axes, so a
second dtype would collide with the existing artifact; it needs either a new
Const axis (renaming every paged-MQA artifact) or a second TraceTemplate with
its own prefix and an explicit dump. Left as a question on the review thread.
Addresses the trace-artifact and reference-dtype review threads on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The module and trace-template docstrings both stated
output[...] = relu(Σ_h w[h] · (Q @ Kᵀ)) · scale[pos]
but the kernels apply ReLU per head, before weighting and reduction:
output[...] = Σ_h w[h] · relu(Q @ Kᵀ)
The two disagree whenever head scores have mixed signs. Verified against the
FP8 kernel on a B200 with the reviewer's example -- head scores [1, -1] and
weights [1, 1] give 1.0, where relu-of-the-sum predicts 0.0. They also disagree
on range: relu-of-the-sum cannot be negative, but with a negative weight the
kernel returns a negative logit, so the old text asserted a bound the kernel
does not respect.
The scale factor also differed by variant and was documented as if shared. FP8
multiplies the result by the per-token KV scale; FP4 folds its per-(token,
K-group) UE8M0 scales into dequantizing Q and K, so it has no trailing scale
factor at all.
No code change: the implementation and both reference implementations were
always correct. Only the prose describing them was wrong, which is why the
correctness tests -- which compare kernel against reference -- stayed green.
A regression test pins both properties using inputs where the two candidate
formulas give different answers: mixed-sign head scores must yield the
per-head-ReLU result, and a negative weight must yield a negative logit.
Without it a documentation fix leaves nothing to stop the contract drifting
back, or the kernel being changed to match the old text.
The PR description needs no change: it describes the op in prose and carries
no formula.
Addresses the operation-contract review thread on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The scheduler correctly gives a ctx=0 row no work, but the task iterator
terminates only on exact coordinate equality with the CTA's end boundary. It
loads the next row's length, sees next_num_kv == 0, discards it, and steps onto
the empty row anyway -- taking a full pipeline turn including TMA loads, UMMA
and epilogue stores.
Shown without timing: prefill the out= backing buffer with a sentinel and run
context_lens=[128, 0]. Before this change the ctx=0 row had 256/256 elements
written; after, zero. Both FP8 and FP4.
Empty rows have no scheduled work, so no other CTA adopts them: the one CTA
holding the live row walks every empty row on its way to its endpoint while the
other SMs idle. Cost is linear in the number of trailing zero rows -- one active
row of ctx=4096 plus N zeros went 1.000x / 1.131x / 1.354x at N = 0 / 15 / 31,
and is now flat at 1.000x / 1.011x / 1.003x.
These are real inputs rather than a synthetic edge case: a CUDA-graph filler of
sequence length 1 becomes 0 after the C4 metadata's seq_len // 4, as do genuine
lengths 1-3, so an under-filled captured batch carries one zero row per unused
slot.
The skip stops at this CTA's end boundary. Overshooting it would make the
exact-equality termination unreachable and hang the persistent loop -- the same
failure mode as a stale schedule_meta. Applied byte-identically to all six
warp-role traversals in each kernel: the roles rendezvous through pipelines, so
any disagreement about which row is next would deadlock rather than merely
diverge.
Verified:
* outputs remain bitwise identical to TensorRT-LLM across 14 eager and
CUDA-graph configs. The kernels are now deliberately divergent from
upstream for zero-length rows only; for any input without them the results
still match exactly.
* no common-case regression: patched vs unpatched with no zero rows anywhere,
six configs to B=256, worst ratio 1.018x with three configs below 1.0.
* termination holds for six shapes where CTA boundaries land on and around
zero rows, including 31 trailing zeros.
The regression test landed first as xfail(strict=True) so the defect was tracked
before the fix existed; the marker is removed here, which strict mode would have
forced anyway.
Addresses the zero-length row review threads on flashinfer-ai#4365.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
MAX_NUM_W_IN_REG caps the per-slot weight count, but the register footprint is the product NUM_W_IN_REG * next_n. Since NUM_W_IN_REG is min(MAX_NUM_W_IN_REG, num_heads), the cap stops binding once num_heads <= it, and the footprint grows to num_heads * next_n == N -- which the API admits up to 256, past the register file. At num_heads=64 the cap binds and incidentally holds the product at <=160; at num_heads=32 it goes inert and next_n=6 asks for 192 registers per thread. Bound the product by 160 registers, the largest footprint the existing policy already requests at the shape it was tuned for (num_heads=64, next_n=4), so every num_heads=64 configuration keeps its current value. Over all 188 accepted shapes: 151 unchanged and 37 clamped for the fp32 epilogue, none changed for fp16, and nothing changes that was not already over budget. NUM_W_IN_REG is the split point between the register and SMEM epilogue paths, both unrolled 4 wide, so the clamp snaps to a multiple of 4. Off-multiple values make the two paths stop tiling the subtile and silently emit non-finite logits (measured: 1e+36 with hundreds of NaNs at NUM_W_IN_REG=21). All four existing constants are multiples of 4, so this was latent; an assert now pins it, mirroring _EPI_SUBTILE_UNROLL in attn_scores.py which applies the same granularity to num_heads // num_epi_subtiles. Measured on sm_100a (B200), min-of-250 over 4 alternating rounds: num_heads=32 gains 1.23x at next_n=6 and 1.20x at next_n=8. On num_heads=64 the A/B spread (2.5% max) sits inside the baseline-vs-baseline A/A noise floor (3.0%), as expected since those shapes compile to an identical NUM_W_IN_REG. Adds correctness coverage at num_heads=32 for next_n 6/8 (clamped) and 4 (control) -- a regime that previously had no tests. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
max_context_len is a Var axis and need not be page-aligned, but both references write in whole physical pages. FP4 assigns [:total_len] where total_len is n_blk * block_size, and FP8 assigns [s:e] per block. When the last page runs past max_context_len the destination slice is shorter than the right-hand side and the assignment raises. At max_context_len=257 with block_size=64 the last page spans 256..320 against a 257-wide output: FP4 reports (257) vs (320), FP8 reports (1) vs (64). Also reproduces at 300/64 and 513/128. Every width exercised so far was page-aligned, so the page-grid extent and the logical output width agreed by coincidence and the mismatch never surfaced. Clip both the destination and the right-hand side to max_context_len, and stop FP8 early once a page starts past the end. Verified against the kernels on causally-valid positions: fp8 max|d| 3.05e-05, fp4 0 to 9.8e-04, uniform across aligned and non-aligned widths. Aligned widths are unchanged, so the committed fi_trace_out JSONs remain valid and were not regenerated. Adds 257/64 and 513/128 to both reference-correctness tests. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Phase 3 of the schedule kernel finds, for each SM, the number of fully-assigned
sequences: count{ j < batch_size : prefix_sum[j] <= seg_starts }. prefix_sum is
non-decreasing -- the comment above the loop already said so -- which makes this
an upper_bound. It was computed with a linear scan over the whole array.
That scan cost O(kAligned) twice over. At runtime every lane walked the entire
prefix_sum. At compile time range_constexpr pastes the body once per iteration,
and the scan sits inside the kMaxSmChunks strip loop, so the emitted IR held
roughly 5 * kAligned copies of it. Instruction count grew linearly with the
batch size, and the compiler passes over the resulting single huge basic block
are superlinear, so compile time grew faster still (measured 2.7x per doubling).
Measured on sm_100a (B200) with FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 to force
real compiles rather than cache loads:
B compile before -> after per-launch before -> after
256 5.8s -> 0.5s 0.013ms -> 0.008ms
512 15.6s -> 0.5s 0.033ms -> 0.008ms
1024 41.6s -> 0.6s 0.069ms -> 0.008ms
2048 138.6s -> 0.8s 0.144ms -> 0.014ms
4096 - 2.3s - 0.033ms
8192 - 4.5s - 0.035ms
16384 - 12.6s - 0.085ms
B >= 4096 has no baseline column because it had no practical baseline:
extrapolating the measured growth puts those compiles at 450s and up.
The per-launch win matters as much as the compile win. At B=2048 the scheduler
took 0.144ms against a ~45us main kernel -- more than three times the attention
it exists to schedule. It is now roughly 30% of it.
Binary search keeps the trip count static at ceil(log2(kAligned)) + 1, so
constexpr unrolling remains appropriate: 13 iterations at B=2048 instead of
2048. Compile time is not constant afterwards -- phases 1 and 2 still unroll
over kAligned // 32 -- but the dominant term is gone, leaving a term about 80x
smaller.
Correctness is checked against the CPU numpy planner, an independent
implementation, for exact equality rather than tolerance. Six context_lens
distributions per batch size in the timing sweep, plus a separate tie-heavy
sweep: a zero-length row contributes num_segs == 0, which is the only way
prefix_sum repeats a value, and duplicates are exactly where a partition search
that advances on < instead of <= silently undercounts. Zero-length rows are
supported input, so this is a real case and not a synthetic one.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The four public entry points were absent from docs/api/*.rst, which the documentation check flags as missing. Add docs/api/attn_scores.rst and register it in the toctree. Records the two properties that are easy to get wrong and are not evident from the signatures: the rectifier is applied per head, before the weighted sum (not to the summed score), and the kernels write every position unconditionally, so the caller must mask positions beyond each request's context length. Also states the schedule-reuse invalidation rule, since reusing schedule_meta across a 256-token boundary is silently wrong unless FLASHINFER_VALIDATE_INPUTS is enabled. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Commit ac36997 added the zero-length-row skip to the six task-advance blocks but left TMA warp 0's producer lookahead at a raw q_idx + 1. The consumers therefore skip a ctx==0 row while the producer stages that row's Q and weights (and SF_Q on FP4), so the row after the gap reduces with the skipped row's query and weights. Stage counts still balance, so it does not hang -- it silently returns wrong logits. A trailing zero cannot expose this: the prefetch_next < end_q_idx guard suppresses the bad fetch. The zero-row test added in ac36997 used a trailing zero, which is why it passed. An interspersed zero, sized so one CTA's range spans the gap, does expose it. Measured on sm_100a with context_lens=[128, 0, num_sms*256]: before: max|d| 845.12 against the pure-torch reference after : 1.06e-07 relative, matching a zero-free control at 1.05e-07 The lookahead now performs the same bounded skip as the task iterators, so the producer and the consumers select the same next row. Adds test_adv_interspersed_zero_row_uses_correct_q for both variants, with distinct per-row weights so reducing with a neighbour's weights is unmissable. It fails on the pre-fix kernels and passes after. Found by adversarial review; the reviewer warned about exactly this coupling in r3824800923 and r3824808841. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Commit 96a5d1e pasted the check into the wrong body. fp8_paged_mqa_logits called _validate_schedule_meta_fresh twice -- the first with fn_name="fp4_paged_mqa_logits", so an FP8 caller debugging a stale schedule was pointed at the wrong public function -- and fp4_paged_mqa_logits never called it at all. The consequence was not cosmetic: a stale schedule_meta passed to the FP4 entry point reached the launch unchecked and hung the persistent kernel, which is the exact failure the check exists to prevent, and it hung even with FLASHINFER_VALIDATE_INPUTS=1. FP8 now checks once under its own name; FP4 checks under its own name. Verified that both raise ValueError instead of hanging. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two errors in the page added by ee3694e: - The shared formula carried FP8's trailing per-token KV scale, which FP4 does not have -- MXFP4 block scales are folded into the dequantised values, so FP4 has no per-position factor. Stating the per-variant behaviour explicitly instead of implying a common s_p. - "The kernels write every position unconditionally" was made false by the zero-length-row skip in ac36997: a context_lens[b] == 0 row is now never written at all. Scoped that statement to within a request's context, and documented that a zero-length row leaves its out= row untouched -- which matters because out= is typically uninitialised. The trace artifacts embed the reference implementation's SOURCE, not only its outputs, so f6d8965's reasoning that unchanged aligned-width results meant they needed no regeneration was wrong: both JSONs still carried the pre-fix reference that raises at a non-page-aligned max_context_len. Regenerated; the diff is the embedded reference string only. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The block_table width and max_context_len bounds were described as checked arguments -- "too few is an out-of-bounds read. FLASHINFER_VALIDATE_INPUTS=1 checks at runtime" -- which reads as though the environment variable closes the hole. It does not. _validate_paged_bounds returns immediately when the variable is unset (the default) and unconditionally during CUDA-graph capture, where the device-to-host copy it needs is illegal. Capture is precisely the deployment mode this API was hardened for, so on that path there is no check at any setting. Nothing about the behaviour changes; the docstrings now say what it is. Both preconditions are named as the caller's to satisfy, a violation is named as undefined behaviour rather than "an out-of-bounds read", and the environment variable is described as a development aid that does not make the kernel safe. Predicating the tail loads so the invariant holds capture-side remains the better fix and is deliberately deferred: it touches eight sites across two pipelined kernels and belongs with the broader CUDA-graph-safe validation pass rather than at the end of a review round. Raised in PR flashinfer-ai#4365 review r3824399380. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The FP8 and FP4 compilers route through build_and_load_cute_dsl_kernel, which
exports an object file and reloads it, so only the first process ever compiles a
given specialization. The schedule kernel was the one holdout: a bare
cute.compile behind functools.cache, which is process-local. Since _gpu_schedule
specializes on ceil(batch_size / 32) * 32, every worker recompiled every batch
bucket it touched.
The asymmetry was visible on disk rather than in timings: attn_scores_fp8 and
attn_scores_fp4 had accumulated 55 cached artifacts between them while no
schedule module directory existed at all.
Routing it through the same helper, measured on sm_100a with the module
directory wiped first so process 1 is a genuine cold compile:
bucket cold warm speedup
64 0.52s 0.04s 13.8x
256 0.31s 0.00s 138x
1024 0.45s 0.00s 194x
extra_key_files is deliberately just this module, not attn_scores.py as the
FP8/FP4 wrappers use. Those prepare shapes in the wrapper, so that file affects
their codegen; it affects nothing here, where the generated code is fully
determined by PagedMQALogitsScheduleKernel plus the four cache-key values.
Including it would invalidate every cached bucket on any unrelated edit. The
under-invalidation direction was checked: _SPLIT_KV reaches the kernel as the
split_kv parameter, which is in both the in-process key and the on-disk tag, so
changing it yields a different artifact rather than a stale hit.
precompile_paged_mqa_logits gains batch_sizes= to warm buckets, defaulting to
None. No default bucket list is invented: guessing typical batch sizes would
compile kernels most deployments never call. The docstring now states that no
schedule buckets are warmed unless asked, which turns a previously implicit
over-promise ("call this and startup is warm") into an explicit scope. Sizes are
deduplicated to their bucket first, so [1, 2, 4, 8, 16] builds one kernel.
Raised in PR flashinfer-ai#4365 review r3832370463.
AI-assisted (Claude Opus 5).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
cute.compile does not refuse a --gpu-arch the running device cannot execute. On sm_100a, targets sm_90a, sm_103a and sm_80 each compiled cleanly and raised only on first invocation. The failure therefore travelled from "the moment an impossible target was requested" to the middle of whichever request happened to call the kernel first. Two changes, doing different jobs. JitSpecCuteDsl.load() now returns the exported artifact whenever one exists, rather than the in-process cute.compile result on a cache miss. That removes a real divergence -- a cache *hit* already returned the reloaded artifact, so miss and hit were handing back different objects from the same call depending on cache state. It does NOT by itself fix the deferred failure: load_module performs no architecture check either, which was measured rather than assumed. The fix is _arch_for_launch(), used by the three compile sites that launch (fp8, fp4, and the schedule kernel). It rejects a target whose arch differs from the current device's, naming both devices, both archs, and the one-line remedy. precompile_paged_mqa_logits is deliberately exempt: building a foreign-arch artifact for a worker that will later run on that device is its documented job, and it never invokes what it builds. Only the arch is checked. A same-arch device mismatch (two identical GPUs, tensors on one and the current device the other) is a stream-placement problem rather than a codegen one -- the launch takes the current device's stream regardless of arch -- and is left alone here rather than folded in untested. Verified on a heterogeneous host (parley: L40S sm_89 at cuda:0, sm_100a at cuda:1). Tensors on cuda:1 with cuda:0 current now raise; precompile for cuda:1 from a cuda:0-current process still succeeds; same-device calls are unaffected. The rejection surfaces from compute_paged_mqa_logits_schedule, since the schedule is built before the main kernel -- which is why all three sites needed wiring, not just the two public entry points. Because load() is shared infrastructure, every consumer of build_and_load_cute_dsl_kernel was regressed: attn_scores 569, cute_dsl_core unit tests 28 on both sm_100a and sm_120, nvfp4_quantize 10065, fp4 quantize 13, topk_varlen 115, utils/test_topk 1495, sm12x moe kernel cache 128 on an RTX 5080, fused_kda_decode 80, and mm_fp4 cute-dsl 2376. Raised in PR flashinfer-ai#4365 review r3824968121. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
_paged_mqa_logits_masked_check is the oracle for both the FP8 and FP4 trace templates, and it returned True in two cases where it had compared nothing. An all-NaN kernel output passed. `finite` was built from BOTH sides and used to drop positions from `valid`, so every position where the kernel emitted NaN removed itself from the comparison; with enough of them `valid.any()` went False and the function short-circuited to True. A single NaN inside the causal region was likewise filtered out and the remaining elements compared clean. A nonzero actual against an all-zero reference also passed. With the reference norm at zero the relative error is undefined, and the code substituted 0.0 -- which no threshold can exceed -- so any actual was accepted. Now: a non-finite value in the ACTUAL output inside the causal region fails, checked before `valid` is built so it cannot mask its own detection. A non-finite REFERENCE stays tolerated, which is what the fp8/fp4 accumulation-order argument in the docstring actually covers; the original text conflated the two. An all-zero reference falls back to an absolute comparison against zero. An empty `valid` is a pass only when the causal region is itself empty -- a region that exists but whose reference is entirely non-finite yields no evidence, so claiming success there would repeat the same mistake. Adds tests/trace/test_paged_mqa_logits_checker.py, which drives the checker directly rather than through a kernel: 13 cases covering both defects plus guards against over-rejecting (matching output, quantisation noise, tolerated non-finite reference, non-causal region, systematic scale error, shape mismatch, sequence-wrapped outputs). Six of them fail against the previous checker, verified by stashing the fix. This mattered more than a test-only change usually would: with this oracle, neither template could detect a kernel emitting NaN. The producer-lookahead bug fixed in 6beb139 was caught only because it produced wrong finite values. Raised by coderabbitai and endorsed by the reviewer in PR flashinfer-ai#4365 review r3725824140. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The kernels take the TVM-FFI environment stream, which resolves to the current stream of the *current* device. Launching with tensors that live on another device onto that stream is wrong independently of architecture, so a non-current target was either silently mis-launched (same arch) or rejected by the arch guard added in 5b13ddb (different arch). That guard closed the codegen half and explicitly left this half open. _on_device() now makes the target current around both compile and launch at all three sites. Compile is inside the same scope because JitSpecCuteDsl tags its on-disk cache from the current device; launch is inside because that is what the stream resolves against. When the target is already current -- every single-GPU call -- it returns contextlib.nullcontext() rather than switching. This subsumes the arch mismatch rather than merely permitting it: with the target current during compile and launch, the arch matches by construction. _arch_for_launch() is therefore unreachable from the launch paths and is kept only as a backstop against a future caller that compiles for a device without entering it; its test is renamed and reworded to say so, since a passing test describing the old design would misrepresent the code. Verified on a heterogeneous host (parley: L40S sm_89 at cuda:0, sm_100a at cuda:1 and cuda:4). Both non-current cases now run and produce output BIT- IDENTICAL to the same call made with the target already current -- max|d| = 0 in each. The same-arch case is the one the reviewer raised; the different-arch case previously raised and now works, which is strictly better than refusing. precompile for a non-current device still succeeds. Cost on the ordinary path, measured rather than assumed: _on_device(current) is 1662 ns in isolation, but against a ~45us eager call the end-to-end A/B is 0.964x worst case against an A/A control of 0.975x on the same config -- inside the noise floor. Graph replay is unaffected (1.007-1.016x), as expected since the Python executes once at capture. Eager was timed with perf_counter, not CUDA events, because a host-side cost is invisible to device timing. Raised in PR flashinfer-ai#4365 review r3824968121. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The committed fixtures embed the `check` function's source as well as `reference`, so daca3e3 changed the live checker without updating what the artifacts carry. Both JSONs still had the pre-fix version: the symmetric non-finite filter that let an all-NaN kernel output pass, and the `if rnorm > 0 else 0.0` fallback that let any actual pass against an all-zero reference. Same oversight as f6d8965, which claimed the artifacts needed no regeneration because aligned-width *outputs* were unchanged. They store code, not just results, and that is true of `check` exactly as it was of `reference`. Regenerated; the diff is the two embedded source strings. Raised in PR flashinfer-ai#4365 review r3725824140. AI-assisted (Claude Opus 5). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
dhiraj113
force-pushed
the
dhiraj113/indexer
branch
from
August 24, 2026 19:14
5b03c7c to
82d3c57
Compare
Collaborator
Author
|
@flashinfer-bot run |
Collaborator
Author
|
/bot run tests/attn_scores |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #64334393: 16/16 executed test jobs passed |
leejnau
approved these changes
Aug 24, 2026
kahyunnam
approved these changes
Aug 25, 2026
|
/bot run tests/attn_scores |
Collaborator
|
@imisszxq is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
1 task
11 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
This change adds two new operations to FlashInfer,
fp8_paged_mqa_logitsandfp4_paged_mqa_logits. Both compute the attention-score "logits" that DeepSeekMLA's sparse-attention indexer uses to decide which KV tokens to keep. They run
only on Blackwell (SM100/SM103).
Given a query, a paged FP8/FP4 KV cache, per-head weights, and per-request
context lengths, each op produces, for every request and KV position:
The rectifier is applied per head, before the weighted sum. The FP8 variant
uses per-token FP32 KV scales; the FP4 variant uses MXFP4 with UE8M0 block
scale factors.
New code
flashinfer/attn_scores/— the two public functions, the CuTe-DSL kernels(
fp8_paged_mqa_logits.py,fp4_paged_mqa_logits.py), and an on-GPU schedulekernel. All are ported from TensorRT-LLM with attribution.
The ports are not byte-identical to upstream. Three deliberate divergences,
each measured and covered by tests:
them and runs a full pipeline turn (TMA + UMMA + epilogue) for no work
rather than per-slot count. Upstream's cap goes inert once
num_heads <= MAX_NUM_W_IN_REG, letting the cache reach 256 registers/threadscanning: 173× faster to compile and 10× cheaper per launch at batch 2048
Outputs stay bit-identical to upstream for any input without zero-length rows.
The public functions infer shapes (batch, heads, head dim, block size, next_n)
from the inputs, reshape them for the kernel, build the CTA schedule, and
launch. FP4 supports
next_nin 1..3; larger values are rejected at the APIboundary with an actionable error.
Both functions validate dtypes, shapes, and tensor devices. Checks that require
a device-to-host copy — block-table bounds, per-row context lengths, schedule
staleness — are opt-in behind
FLASHINFER_VALIDATE_INPUTS.Scheduling
The persistent kernel needs a per-call work assignment, computed on the GPU by
default (no host round-trip, capturable in a CUDA graph) with a CPU fallback.
compute_paged_mqa_logits_schedule()lets a caller precompute it once and passit back via
schedule_meta=.A reused schedule is valid only while
ceil(context_lens / 256)is unchangedfor every request. It must be recomputed whenever any context length crosses a
256-token boundary — including on every CUDA-graph replay, since a static
buffer provides stable storage but not stable contents. A stale schedule is
detected under
FLASHINFER_VALIDATE_INPUTS.padded_context_len()reports the required output width;precompile_paged_mqa_logits()warms the compile cache.out=andschedule_meta=exist so the whole call can be captured and replayed.Masking
The kernels write every position unconditionally and do not apply a causal
or context mask. Callers must mask out-of-context positions afterward. This
matches how the upstream TensorRT-LLM kernel is used.
Tracing and docs
Both functions are registered with fi_trace
(
flashinfer/trace/templates/attn_scores.py) with reference implementationsand generated benchmark-definition JSON.
docs/api/attn_scores.rstdocumentsthe public entry points.
Tests
A correctness suite comparing kernel output to a pure-torch reference across
dtypes,
next_n, block sizes, and context lengths; an adversarial suitecovering boundary, degenerate and skewed inputs, GPU-vs-CPU schedule equality
(to batch 12288, including zero-length rows), and determinism; and
trace-template consistency and reference tests.
🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit