[Perf][DSV4.1] Segment FlashInfer FP8 prefill queries for faster attention tiles

[WoosukKwon]

Purpose

Bound FP8 prefill query segments to 1024 tokens in the DeepSeek-V4.1 FlashInfer sparse-MLA backend. The existing TRTLLM-GEN launcher selects a faster query tile for these shapes. An 8192-query request is presented as eight segments in one attention call, using the same query/KV/output storage and sparse indices.

Preserve absolute positions by giving each segment sequence length original_seq_len - original_query_len + segment_end. Query starts and segment sequence lengths use persistent buffers; the sequence-length values are refreshed from device metadata on each build. Decode dispatch and BF16 query grouping remain unchanged. This works with the existing FlashInfer API and requires no FlashInfer source patch.

This is a vLLM-side workaround for FlashInfer's tile-selection heuristic. The 1024-token cap should be reassessed when that heuristic changes; a future FlashInfer dispatch fix could remove this metadata adaptation.

Duplicate-work check

Checked #56217/comments, open PRs referencing that issue, and FlashInfer/prefill/segmentation/sparse-MLA searches. No open PR implements this query segmentation. Related SM120 backend fixes, SM90 prefill additions, and indexer chunk-planning changes target different paths. This is independent of #57204 and #57206.

Test Plan

.venv/bin/python -m pytest tests/kernels/attention/test_flashmla_sparse.py -k segments_preserve_positions_and_replay_updates -q
pre-commit run --files vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py tests/kernels/attention/test_flashmla_sparse.py
pre-commit run mypy-3.12 --files vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py tests/kernels/attention/test_flashmla_sparse.py --hook-stage manual

The tests cover FP8 segmentation and unchanged BF16 grouping, mixed decode/prefill batches, padded query offsets, preserved positions, persistent buffer addresses, and CUDA graph replay after sequence lengths change.

Test Result

Standalone commit 9890a68a89, based on 6ca2b23e22: 12 tests passed on GB200 (10 deselected). Applicable pre-commit hooks, including mypy 3.10, and explicit mypy 3.12 passed. The newer upstream output-projection fix is preserved.

Earlier four-GB200 matched profiles showed Q16 attention tiles replacing Q8 for 8K prefill, reducing total attention kernel time from 35.53 to 20.74 ms. Exploratory isolated full-forward latency reductions were 10.55% for 8K prefill/~8K KV, 8.50% for 4K prefill/~100K KV, and 9.84% for an 8K prefill plus 31 decode requests/~100K KV. Full-model parity probes observed zero output error.

Those isolated timing probes predate a benchmark metadata-reset correction affecting two SWA mapping calls in both arms; they are not fresh standalone-branch end-to-end measurements. The later corrected full matrix validated the combined configuration, rather than isolating this patch's contribution.

Real-weight evaluation previously completed with this segmentation in the combined optimized configuration: deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a40aa45a94ad051997016db3960a90277, TP4/EP4, native MegaMoE/shared fusion, FlashInfer sparse attention. Full GSM8K, lm-eval 0.4.12, five-shot multi-turn chat, thinking disabled, greedy generation, 1024 output tokens, seed 123:

Configuration	Correct / 1319	Accuracy
Auto-MoE/FlashMLA FP8-indexer baseline, three runs	1270–1274	96.2851–96.5883%
Combined optimized configuration, FP8 indexer	1269	96.2092%
Combined optimized configuration, MXFP4 sparse indexer	1266	95.9818%

These combined evaluations include separate MoE and indexer changes and do not isolate segmentation's quality effect. The 1M model limit was used; these GSM8K evaluations do not test 1M-token prompts.

AI assistance

Codex assisted with implementation, validation, and this description. Woosuk Kwon requested this PR after reviewing the segmentation approach, measured gains, and maintenance tradeoff.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.