Repository navigation
Conversation
Preview approval requiredThis pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit |
This comment has been minimized.
This comment has been minimized.
726adda to
6ccf68f
Compare
This comment has been minimized.
This comment has been minimized.
6ccf68f to
1da0f79
Compare
This comment has been minimized.
This comment has been minimized.
1da0f79 to
c439fd5
Compare
tejachavali
left a comment
There was a problem hiding this comment.
The mechanism is clear and the wiring looks careful. Padding the front of each paged-prefill query block up to the computed minimum keeps the average packed query length above the tile threshold, the output gather restores real-row order, and the padded indices reset on each call_begin_forward, so a batch that skips padding cannot reuse a stale plan. The MLA whitelist split is the right companion change too. The end-to-end radix-cache CI case plus the cached_tokens == 0 guard is a genuine regression test.
What I would want settled before merge: is FlashInfer's tail alignment of a query block to its KV segment guaranteed across versions, or observed on the probed version? The whole fix rests on real rows keeping their causal positions. Can a prefill batch contain a zero-length request, and if so, does it drag the batch average back below the threshold even though every real request is padded? And the PR body shows all three CI slots failing, is that a known flake or something this needs to address?
|
@tejachavali Thanks for the careful read. All three questions, answered: Tail alignment: This is a kernel invariant, not an observation. Zero-length requests: The scheduler keeps them out of a prefill batch. Prefix matching is capped at Red CI slots: Known gate behavior, not flakes. |
c439fd5 to
9f80d4a
Compare
This comment has been minimized.
This comment has been minimized.
832757c to
b02bd26
Compare
Validated on A100 (SM80) — padding restores bit-exactness; plus one follow-up bugI tested this change on a single A100-SXM4-80GB (PCIe, SM80) to check whether the fix holds on Ampere, where the deterministic path resolves to different backends than on Hopper/Blackwell. Cherry-picked on top of sglang main I re-checked against your current head Method. For each prompt: flush cache → request A (full recompute) → request B (radix hit), compare token ids and the per-step top-k logprobs. To keep the prefix cache enabled and still get logprobs, I used
Notes that may be useful for triage:
What the flips look like. All 7 flipped prompts diverge at a point where at least one side has Follow-up bug: padding is incompatible with the full prefill CUDA graph. With this PR + Full mode allocates a static q buffer for the captured Performance on A100, for reference. fa3 prefill is faster than flashinfer+padding on this GPU: 9.4k tokens 783 ms vs 856 ms, 4.7k 381 ms vs 430 ms, and 8-concurrent 9.4k 12,660 vs 11,256 tok/s. I'm not suggesting changing the SM80 default (SM80/SM90 → Also, prefill CUDA graph modes look numerically neutral here: |
b02bd26 to
48787aa
Compare
|
@Woody-Hu Thank you for the thorough A100 validation and for the full-prefill-graph report; that was a real bug. Root cause confirmed: the plain-EXTEND branch of Fix in 48787aa: Regression test: One caveat on my side: this box is SM120 with only hybrid Qwen3.5 weights, so I could not re-run your |
|
Correction to my previous comment. "replays are already tile-invariant across prefix lengths" is only true within one captured bucket. Under the full prefill CUDA graph the tile comes from the bucket's static token count, so it does not move with the prefix length inside a bucket, but a radix hit (a few-token extend) and the full recompute of the same prompt will usually land in different buckets, and different buckets can pick different tiles. So this PR makes no bit-exactness claim for On skip vs. explicit error: I went with skipping because the configuration itself is legitimate and padding on a cuda-graph wrapper is neither needed nor possible, so an error would only reject users for no gain. The real limitation is the one above, and it belongs in the docs rather than in a startup error; happy to add a sentence to I have tightened the same sentence in the PR body. |
With --enable-deterministic-inference, FlashInfer prefill was bit-exact only when the prompt was uncached. Once the radix cache turned a prompt into a prefix hit, only the tail (q_len = seq_len - prefix_len) went to BatchPrefillWithPagedKVCacheWrapper. That changed avg_packed_qo_len and made FA2DetermineCtaTileQ pick a different cta_tile_q bucket than the full prefill; different tile sizes split the reduction differently, so cached and uncached runs diverged even at temperature 0. Under deterministic inference, pad the front of each paged-prefill q block with zero rows up to deterministic_prefill_min_q(num_qo_heads, num_kv_heads, head_dim), a per-model length that keeps avg_packed_qo_len above every bucket threshold. FlashInfer aligns the q block to the tail of the KV segment, so real rows keep their causal positions; padded rows are discarded by gathering the output back to the real rows. Padding is disabled for ragged prefill, sliding-window KV pools, custom KV indices/masks, multi-item scoring, piecewise CUDA graph capture and cuda-graph prefill wrappers: the full prefill CUDA graph plans against a static q buffer sized at capture and picks its tile from the captured max packed qo length, so raising qo_indptr there broke the q.shape[0] == qo_indptr[-1] invariant at capture time. Also: - allow flashinfer in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND and split the absorbed-MLA whitelist into MLA_SUPPORTED_... so flashinfer stays rejected for MLA models - make the radix_cache mode in test_deterministic fail when the cache did not actually hit - add TestFlashinferDeterministicRadixCache to CUDA CI; it disables the prefill CUDA graph because the breakable graph adds its own non-determinism - register the padding unit test in base-b and add a kernel-level test in base-b-kernel-unit proving a padded short extend matches the full prefill bit for bit, for single and mixed batches, and that a cuda-graph wrapper keeps the real qo_indptr - document radix cache as supported for FlashInfer Hybrid GDN models (e.g. Qwen3.5) still show cached/uncached drift for q_len <= 32 with every backend; that is a pre-existing property of the hybrid path, not of this change. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Co-Authored-By: Claude Code <noreply@anthropic.com>
48787aa to
be80cac
Compare
Environment
Server arguments:
Request parameters:
ReproductionI tested 5 different prompts. For each prompt, I sent one cold-cache request, followed by 4 repeated requests to test cache-hit behavior. I compared the generated token sequences and their corresponding logprobs between cold-cache and cache-hit requests. Results1. Default configuration (Radix Cache enabled) I observed the same behavior in two independent runs:
2. Radix Cache disabled I repeated the test with
3. Deterministic inference enabled (Radix Cache still enabled) I repeated the test with
This suggests that enabling deterministic inference does not resolve the inconsistency between cold-cache and cache-hit execution paths when Radix Cache is enabled. Initial FindingsBased on these results, the issue appears to be closely related to Radix Cache. With Radix Cache disabled, both token sequences and logprobs remain consistent. With it enabled, the cold-cache and cache-hit paths produce different numerical results, sometimes leading to different generated tokens. Enabling I suspect this may be related to the same class of issues discussed in #41351 and #46351, and that the numerical consistency of cached computation deserves further investigation. I've attached a minimal reproduction script along with the raw experimental evidence. I'm happy to test any proposed fix on my current RTX 4090 + Qwen3.5-4B setup. Additional note: So far, both execution paths still correctly answer the key question (needle) in each prompt. Therefore, this currently appears to be a determinism and consistency issue rather than a demonstrated answer-quality regression. I have not observed evidence of incorrect answers in these tests. |
|
Status: reproduced on current FlashInfer main; this PR's padding works, and a FlashInfer-native fix exists in flashinfer-ai/flashinfer#2830 (needs a rebase). What we found: in the eager plan, FA2 paged prefill picks Next steps:
The hybrid GDN (Qwen3.5) drift for q_len <= 32 is a separate path, covered by neither fix. |
Motivation
With
--enable-deterministic-inference, the FlashInfer backend produced bit-identical logprobs only while a prompt was uncached. As soon as the radix cache served a prefix hit, the same prompt attemperature=0could return different logprobs and even a different first token. This makes deterministic inference unusable with the radix cache, which is on by default.The deterministic-inference docs listed radix cache as unsupported for FlashInfer and the server disabled it automatically. This PR removes that limitation.
Root cause
For a paged prefill, FlashInfer sizes its CTA tile from
through
FA2DetermineCtaTileQ. A cached request only sends the uncached tail to attention (q_len = seq_len - prefix_len), soavg_packed_qo_lenshrinks and the kernel can pick a differentcta_tile_qbucket than the full prefill of the same prompt. Different tile sizes split the softmax reduction differently, and the results stop being bit-exact.The reduction order is otherwise batch invariant; only the tile choice differs.
Fix
Under deterministic inference, pad the front of each paged-prefill q block with zero rows up to
which keeps
avg_packed_qo_lenabove the bucket threshold for every batch shape (GQA groups and head dim included). FlashInfer aligns a q block to the tail of its KV segment, so the real rows keep their original causal positions; the padded rows produce garbage that is discarded by gathering the output back to the real rows. The padded q rows are built inside the attention layer, so the rest of the model still sees the real token count.Padding is skipped for ragged prefill, sliding-window KV pools, custom KV indices/masks, multi-item scoring, piecewise CUDA graph capture and cuda-graph prefill wrappers, where the assumption above does not hold or the extra copy is not worth it. The full prefill CUDA graph plans against a q buffer sized at capture, and flashinfer's
enable_cuda_graphplan derives the tile from the bucket's static token count rather than the per-batch average, so padding there is neither needed nor possible and raisingqo_indptronly broke theq.shape[0] == qo_indptr[-1]invariant (reported by Woody-Hu on A100). A radix hit and the full recompute of the same prompt usually land in different buckets, so this PR makes no bit-exactness claim for--cuda-graph-backend-prefill full; it only stops the crash there. The bit-exactness claims below are for eagerly planned attention, which is what the new CI case runs.Enabling radix cache for FlashInfer also required splitting the deterministic backend whitelist: the absorbed-MLA check reused
RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND, so addingflashinferthere would have silently allowed an unsupported MLA combination. The MLA check now uses its ownMLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND.Validation
BatchPrefillWithPagedKVCacheWrapper, cold vs radix hit, Qwen3-8B shapes 32/8/128,head_dim=128,min_q=17): with paddingmax|diff| = 0.0(bit-exact); without paddingmax|diff| = 0.001953125.1.335e-5).sglang.test.test_deterministic --test-mode radix_cachepasses, and the mode now fails ifcached_tokens == 0, so it cannot silently pass on a cache miss.qwen3.5-9bon one GPU with FlashInfer + deterministic inference, 16 concurrent requests, median of 3 bursts: the worst case where every request is a 1-4 token extend behind a 1024-token shared prefix goes from 65.7 ms to 70.7 ms per burst (~7.5%), while a cold full prefill of 16 x 1028 tokens is unchanged (1382.7 ms -> 1383.2 ms, within noise). The padding row copy only happens when a request is shorter thanmin_q; regular prefill never enters that path.max|delta| = 0.0.Tests
test/registered/unit/layers/attention/test_flashinfer_deterministic.py: 7/7.test/registered/kernels/ops/attention/test_flashinfer_host_plan.py: 5/5.test/registered/unit/test_model_overrides.py: 115/115, including a new regression test that FlashInfer + deterministic inference keeps the radix cache.TestFlashinferDeterministicRadixCacheruns the radix_cache mode on CUDA with--disable-prefill-cuda-graph.test/registered/kernels/ops/attention/test_flashinfer_deterministic_prefill.py(base-b-kernel-unit): padded 1-token and 3-token extends, and a mixed batch that keeps a long request next to a padded one, all bit-exact against their full prefill, plus a cuda-graph regression: drivingcall_begin_forwardwith a realuse_cuda_graph=Truewrapper keeps the realqo_indptrand attaches no gather indices, while the eager wrapper pads tomin_q; 6/6 locally.Caveat
Hybrid GDN models (Qwen3.5 family) still show a cached/uncached difference when the uncached chunk is very short (
q_len <= 32), with both FlashInfer and Triton, so it is a separate pre-existing property of the hybrid linear-attention path rather than of this padding. Forq_len >= 33both backends are bit-exact on those models as well. The tests use the regular Qwen3 model, where the fix is exact.The prefill CUDA graph is another independent source of non-determinism: with the default
cuda_graph_config.prefill.backend = "breakable"the same request differs between graph replay and eager execution, so the new CI case disables it. That is orthogonal to this PR and is worth a separate look. On A100, Woody-Hu measureddisabled/breakable/fullas bit-exact with each other on both fa3 and padded FlashInfer, so the drift above is so far an SM120 observation.Related work
#43612 aligns chunked-prefill chunk starts to FlashInfer's 128-row q-tile grid from the scheduler, and explicitly leaves radix-cache hits, tiny last chunks and batches of 1-token extends unaddressed. This PR covers the radix-cache-hit case from the attention side, enabled only under deterministic inference. The two are complementary: chunk alignment avoids off-grid starts when the scheduler is free to move them, while padding makes short extends match the full-prefill tile when the start is fixed by the prefix length.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #40807132297
Latest PR Test (Extra): ❌ Run #40807132207
Latest PR Test (AMD ROCm 10): ❌ Run #40807132296