[Model] Add FlashInfer PrimTS QSA backend for Qwen3.8-Flash-Next - #56240
Draft
PerkzZheng wants to merge 2 commits into
Draft
PerkzZheng wants to merge 2 commits into
PerkzZheng wants to merge 2 commits into
Conversation
Consume compact indexer blocks through FlashInfer's prepared plan/run API. Support BF16 and FP8 KV caches, request-safe packed prefill, fixed MTP decode groups, and Triton fallback. Preserve reused MTP tail positions and share graph-safe workspace across ordered attention layers. Release cache-dependent plans and derived views on cache unbind/rebind. Keep the public diff focused on production integration; additional tests, benchmarks and integration notes remain on the local validation backup. Assisted-by: OpenAI Codex Signed-off-by: PerkzZheng <67892460+PerkzZheng@users.noreply.github.com>
Cache packed CPU/device query offsets in shared forward metadata. Reuse their device storage through the model-scoped weak buffer pool while FlashInfer plans retain it, avoiding per-layer copies and repeated plan preparation for unchanged query layouts. Keep route identity tied to request boundaries, padded token count, group size and device. Fixed-layout decode remains offset-free, and K/V plans and selected-block metadata remain layer-specific. Validated with 19 focused checks, preserved/current Qwen and MTP suites, changed-file hooks, and five prefill cases each under memcheck/initcheck. Validation artifacts remain local, outside the public production diff. Assisted-by: OpenAI Codex Signed-off-by: PerkzZheng <67892460+PerkzZheng@users.noreply.github.com>
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
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
1 participant
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.
Purpose
Integrate FlashInfer's QToken-KvBlock-Sparse-Attention backend for Qwen3.8-Flash-Next (Qwen4Exp). Depends on flashinfer-ai/flashinfer#4996 and a compatible CUTLASS DSL 4.7 installation.
VLLM_QSA_ATTENTION_BACKEND=auto|triton|prims_ts.autoselects PrimTS when supported and available, otherwise Triton. Fixed decode supports CUDA graphs.Related work checked: #55430 targets Triton SM121 prefill, #55065 targets dense SM120 prefill, and #55557/#54846 target KV quantization. This PR adds the FlashInfer SM100/SM103 sparse prefill-and-MTP backend rather than duplicating those implementations.
Test Plan
.venv/bin/python -m pytest tests/models/qwen4_exp/{test_qsa_reference,test_qsa_pre_indexer,test_config,test_ple}.py -qAdditional local integration, hook, and sanitizer gate:
bash qsa_bench/shared_qo_indptr_20260910/run_final.sh. This harness is retained outside the public source diff.Test Result
SM103: 143 upstream tests, 19 focused route/ownership/CUDA checks, 95 preserved QSA cases, and 10 ownership/MTP checks passed (overlapping suites). BF16/FP8 reference and decode-graph checks passed. Five prefill cases passed each of memcheck and initcheck with zero errors; changed-file hooks passed.
Prior frozen-revision model evaluation, TP2 / FP8 KV / MTP3:
GPQA includes two unfinished PrimTS answers and remains an accuracy follow-up, not an equivalence pass. These model evaluations were not rerun after the latest offset-sharing change. SM100 runtime validation remains pending.
E2E performance: all-layer GPU time
Prior matched Nsight measurements on GB300/SM103, TP2, at the frozen revisions in the linked report; not rerun after the latest offset-sharing change. Prefill measures request four after three warmups, with prefix caching disabled. Decode uses full CUDA graphs and 64 resident requests with MTP3 (256 target-query tokens), averaged over 33 scheduler iterations including draft work.
Speedup = Triton / PrimTS. These are warm, pure-stage GPU wall times across all layers, not API TTFT, queue latency, cold-L2 standalone timing, or time per accepted token. Sparse-path speedups include metadata/index expansion, attention, and reduction, with PDL overlap counted once. These are single matched captures, not confidence intervals. Both BF16 prefill backends use the same documented image-MoE padding fix; resident BS256 did not fit this TP2 cache configuration.
Detailed prior accuracy and pure-stage performance results. Standalone kernel benchmark suites.
AI assistance: OpenAI Codex assisted with implementation, validation, and this description. Kept as a draft for human review and dependency/validation follow-up.
Essential Elements of an Effective PR Description Checklist