Repository navigation
Conversation
… not the split-KV tile Deterministic flashinfer prefill is bit-exact across chunkings only when every chunk starts on the fa2 kernel's 128-row q-tile grid. The scheduler was aligning chunks to SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE (4096), which is the split-KV tile and has no bearing on this; it also forced any --chunked-prefill-size below 4096 to be rejected at startup. Read a PREFILL_Q_TILE_TOKENS = 128 constant from the flashinfer backend instead, leave the split tile a pure kernel knob, and pin the invariant with a kernel test (exact on the 128 grid, not exact off it). The Kimi-Linear DCP PD test moves to --chunked-prefill-size 128. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
kpham-sgl
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
hebiao064,
hnyls2002,
ispobock,
merrymercy and
xiezhq-hermann
as code owners
October 10, 2026 21:19
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The MLA wrapper has no deterministic mode, so aligning its chunks buys nothing; the Kimi-Linear test keeps its 64-token chunk unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The scheduler no longer matches backend names: it takes prefill_chunk_align off the live backend, which deterministic flashinfer MHA sets to its q-tile width and wrappers forward from their dense child. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
kpham-sgl
requested review from
hanming-lu,
yizhang2077 and
yuan-luo
as code owners
October 10, 2026 23:58
4 of 5 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Alternative to #43598. Two findings behind it, both measured on 4x B300 (flashinfer 0.7.0.post1):
SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE(4096), which is the split-KV tile. ProbingBatchPrefillWithPagedKVCacheWrapper(backend="fa2", fixed_split_size=4096)directly, a chunked prefill reproduces the unchunked one bit for bit iff every chunk starts on the kernel's 128-row q-tile grid (128 / gqa_grouptokens); the split tile value (4096 / 1024 / None) does not change which boundaries are exact.--attention-backend flashinferresolves toFlashInferMLAAttnBackend, which has no deterministic mode (no fixed split, ragged prefill for no-prefix extends,BatchMLAPagedAttentionWrapperfor chunked ones). A Kimi-Linear TP2 sweep shows chunked vs unchunked input logprobs differ by up to 4.2 and even the unchunked server is not batch-invariant, with or without alignment.So the backend now declares the grid instead of the scheduler matching names:
AttentionBackend.prefill_chunk_aligndefaults toNone, deterministicFlashInferAttnBackendsets it to its q-tile width (PREFILL_Q_TILE_TOKENS = 128), the hybrid wrappers forward their dense child's value, and the scheduler'sinit_prefill_chunk_alignjust reads it. The MLA backend leaves itNone, so the Kimi test is unchanged and starts on main as is; the env var stays a pure kernel knob.A new kernel test,
test_flashinfer_deterministic_chunk_invariance.py, pins the MHA invariant in both directions: exact on the 128 grid, not exact off it.Verified on the final revision: kernel test 2/2;
test_deterministic_chunk_align2/2; Qwen3-8B deterministic flashinfer with--chunked-prefill-size 256matches the unchunked server bit for bit on every prompt, sequential and concurrent, while--chunked-prefill-size 100is rejected at startup with "aligned to 128 tokens";TestKimiLinearDCPDisaggregationOK with the untouched test file (292 s). The TP arm fails on B300 withcached_tokens == 0both before and after this change, because B300 defaults Kimi-Linear to flashinfer, which disables the radix cache under deterministic inference; CI's H100 defaults to fa3.Not addressed here: deterministic inference for the flashinfer MLA path, and the MHA-path holes where a tiny last chunk, a chunk batched with many 1-token extends, or an off-grid radix-cache hit lands on a different fa2 tile.
CI States
Latest PR Test (Base): ❌ Run #38096840821
Latest PR Test (Extra): ❌ Run #38096840575
Latest PR Test (AMD ROCm 10): ❌ Run #38096840800