Skip to content

[Fix] attention: Make FlashInfer prefill bit-exact on radix-cache hits - #46351

Open
jbymy2000 wants to merge 1 commit into
sgl-project:mainfrom
jbymy2000:fix/flashinfer-deterministic-radix-cache
Open

jbymy2000 wants to merge 1 commit into
sgl-project:mainfrom
jbymy2000:fix/flashinfer-deterministic-radix-cache

Conversation

@jbymy2000

@jbymy2000 jbymy2000 commented Oct 25, 2026 •

Copy link
Copy Markdown

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 at temperature=0 could 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

avg_packed_qo_len = sum(q_len_i * gqa_group) / batch_size

through FA2DetermineCtaTileQ. A cached request only sends the uncached tail to attention (q_len = seq_len - prefix_len), so avg_packed_qo_len shrinks and the kernel can pick a different cta_tile_q bucket 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

min_q = deterministic_prefill_min_q(num_qo_heads, num_kv_heads, head_dim)

which keeps avg_packed_qo_len above 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_graph plan 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 raising qo_indptr only broke the q.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 adding flashinfer there would have silently allowed an unsupported MLA combination. The MLA check now uses its own MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND.

Validation

  • FlashInfer kernel probe (BatchPrefillWithPagedKVCacheWrapper, cold vs radix hit, Qwen3-8B shapes 32/8/128, head_dim=128, min_q=17): with padding max|diff| = 0.0 (bit-exact); without padding max|diff| = 0.001953125.
  • Same probe on short queries q=2/4/6/37: padded runs are bit-exact against the full prefill; unpadded runs differ.
  • End-to-end tiny dense model with radix cache enabled: padded prefill gives identical logprobs for cached and uncached requests; disabling the padding reproduces the divergence (1.335e-5).
  • Interactive sglang.test.test_deterministic --test-mode radix_cache passes, and the mode now fails if cached_tokens == 0, so it cannot silently pass on a cache miss.
  • Prefill overhead, qwen3.5-9b on 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 than min_q; regular prefill never enters that path.
  • Independent reproduction on A100 (SM80) by Woody-Hu (see the PR thread): Qwen3-4B, 3.7k-token prompts, recompute vs radix hit. Whitelist-only (no padding) flips 5/8 greedy outputs with step-0 top-k drift up to 0.107; with the padding 0/8 flips and max|delta| = 0.0.

Tests

  • New test/registered/unit/layers/attention/test_flashinfer_deterministic.py: 7/7.
  • Existing test/registered/kernels/ops/attention/test_flashinfer_host_plan.py: 5/5.
  • Existing test/registered/unit/test_model_overrides.py: 115/115, including a new regression test that FlashInfer + deterministic inference keeps the radix cache.
  • New CI case TestFlashinferDeterministicRadixCache runs the radix_cache mode on CUDA with --disable-prefill-cuda-graph.
  • New kernel-level test 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: driving call_begin_forward with a real use_cuda_graph=True wrapper keeps the real qo_indptr and attaches no gather indices, while the eager wrapper pads to min_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. For q_len >= 33 both 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 measured disabled / breakable / full as 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

@mintlify

mintlify Bot commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit be80cac.

Approve preview

@github-actions

This comment has been minimized.

@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch from 726adda to 6ccf68f Compare October 25, 2026 02:32
@github-actions

This comment has been minimized.

@jbymy2000 jbymy2000 changed the title [Deterministic] Support radix cache with the FlashInfer attention backend [Fix] attention: Make FlashInfer prefill bit-exact on radix-cache hits Oct 25, 2026
@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch from 6ccf68f to 1da0f79 Compare October 25, 2026 02:35
@github-actions

This comment has been minimized.

@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch from 1da0f79 to c439fd5 Compare October 25, 2026 08:25

@tejachavali tejachavali left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@jbymy2000

Copy link
Copy Markdown
Author

@tejachavali Thanks for the careful read. All three questions, answered:

Tail alignment: This is a kernel invariant, not an observation. prefill.cuh maps query row i to sequence position q_idx + kv_len - qo_len and masks with kv_idx > idx_in_original_seq; scheduler.cuh even comments the layout as "right aligned". SGLang's existing prefix-extend path already relies on the same offset (flashinfer_backend.py notes the causal offsets use seq_len - q_len). After padding, a real row j still maps to j + prefix_len, exactly as without padding. It is not a documented per-version guarantee, but the new end-to-end CI case is the guard: any change to the alignment breaks cached vs uncached bit-exactness there.

Zero-length requests: The scheduler keeps them out of a prefill batch. Prefix matching is capped at input_len - 1 (_compute_max_prefix_len, to enable logprob computation) and travels into the radix key as its limit; the chunked path rejects extend_len <= 0; retracted requests re-enter through the same capped match with at least the last generated token to recompute. Your concern is real though -- if one ever did reach the batch it would count in batch_size and drag the average down -- so the helper deliberately leaves q_len == 0 at zero instead of padding it, and does not fabricate a query row without a KV segment.

Red CI slots: Known gate behavior, not flakes. pr-gate requires the run-ci label, which only maintainers can add, so it fails fast and every test stage is skipped (run 40176548704: check-changes passed, all base-* stages skipped). The only real failures were the PR title format and one ruff-format line, both fixed at head c439fd58e3 (title and lint now pass). Once the label lands, the new radix-cache case will run.

@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch from c439fd5 to 9f80d4a Compare October 26, 2026 07:30
@github-actions

This comment has been minimized.

@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch 2 times, most recently from 832757c to b02bd26 Compare October 27, 2026 08:10
@Woody-Hu

Copy link
Copy Markdown
Contributor

Validated on A100 (SM80) — padding restores bit-exactness; plus one follow-up bug

I 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 3c07ff2654 (the base of the image I use: torch 2.13.0+cu130, triton 3.7.1, flashinfer 0.6.18). Model Qwen3-4B (36 layers, 8 KV heads, head_dim 128, BF16), greedy, 3.7k-token prompts, 64 generated tokens, 8 prompts per arm, server restarted between arms.

I re-checked against your current head b02bd26bcb: flashinfer_deterministic.py is byte-identical to what I tested and the _can_pad_prefill_q / padding call site are unchanged, so these numbers still apply.

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 return_logprob: true with logprob_start_len = -1 — the default 0 clamps max_prefix_len to min(input_len - 1, logprob_start_len), which silently disables the hit (cached_tokens == 0).

arm source flags backend used hits greedy flips max|Δtop-k| at step 0
A0 baseline none flashinfer 8/8 2/8 0.0026–0.083
A1 baseline --enable-deterministic-inference fa3 (SM80 fallback) 8/8 0/8 0.0
A2 baseline + --attention-backend flashinfer flashinfer 0/8 n/a —
A3 baseline + 1 line (add flashinfer to RADIX_SUPPORTED_..., no padding) as A2 flashinfer 8/8 5/8 0.0026–0.107
B1 this PR --enable-deterministic-inference fa3 8/8 0/8 0.0
B2 this PR + --attention-backend flashinfer flashinfer (padded) 8/8 0/8 0.0

Notes that may be useful for triage:

  • A0 shows the defect is not specific to --enable-deterministic-inference: the default server on A100 also flips 2/8 between recompute and a radix hit.
  • A2 is not "still broken", it is not tested at all: on the baseline, handle_deterministic_inference does not list flashinfer as radix-compatible and sets disable_radix_cache=True, so cached_tokens is 0 for all 8 prompts. Your whitelist split is what makes the combination usable at all.
  • A3 vs B2 differ only in the padding for a non-MLA model (the whitelist/MLA_SUPPORTED_... part is unreachable for Qwen3-4B), so 5/8 → 0/8 is attributable to the padding.
  • A1/B1 explain why nobody hits this on A100 by default: _deterministic_attention_backend maps SM80/SM90 to fa3, and fa3 sets num_splits = 1 if deterministic else 0 (flashattention_backend.py), so there is no q-length-dependent reduction split to begin with. Bit-exactness was already true there; forcing flashinfer was the untested path.

What the flips look like. All 7 flipped prompts diverge at a point where at least one side has top1 - top2 ∈ {0.0, 0.0625, 0.125} (4 of them are exact ties, e.g. 1687 = 281 to the last bit), while the two paths differ by 0.068–0.107 logprob at that same step. So these are tie-breaks caused by a small prefill numerical difference, not corrupted cache content: the padded path collapses the difference to exactly 0.0, including on the prompts where an exact tie exists. (Note on magnitudes: single-step drift is ≤0.107; when ids stay equal the difference keeps compounding over 64 decode steps to ~1–2.7, which is why I quote the step-0 number here.)

Follow-up bug: padding is incompatible with the full prefill CUDA graph. With this PR + --cuda-graph-backend-prefill full, capture fails at startup:

File ".../flashinfer_backend.py", line 1430, in forward_extend
    o = prefill_wrapper_paged.forward(...)
ValueError: q.shape[0] (16) does not match qo_indptr[-1] (17). For paged prefill,
q must have shape [total_tokens, num_heads, head_dim] where total_tokens = qo_indptr[-1].

Full mode allocates a static q buffer for the captured num_tokens (16 here) and precomputes cta_tile_q from max packed qo len, so raising qo_indptr to min_q = 17 breaks the invariant. Without the patch the same configuration captures fine (51.6 s). The PR body lists piecewise (breakable) graphs as skipped — _can_pad_prefill_q only checks not is_in_tc_piecewise_cuda_graph(), and the full prefill graph is a different backend not covered by that predicate. One option is to also skip padding when ForwardBatch.full_prefill_cuda_graph_capture is set (see runner/prefill_cuda_graph_runner.py:1641), or to reject the combination with an explicit error; silently crashing at capture time is the worst outcome.

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 → fa3 is both correct and faster); the A100 relevance of this PR is the explicitly-forced --attention-backend flashinfer case.

Also, prefill CUDA graph modes look numerically neutral here: disabled / breakable / full are bit-exact with each other on both fa3 and flashinfer+padded (6/6 prompts), and give no throughput gain on A100 (no PDL). If that holds more broadly it may be worth relaxing the "prefill graph is an independent non-determinism source" caveat for Ampere specifically — but I did not instrument whether replay actually fired for every request, so I'd treat that as a hint rather than a conclusion.

@jbymy2000
jbymy2000 force-pushed the fix/flashinfer-deterministic-radix-cache branch from b02bd26 to 48787aa Compare October 27, 2026 12:19
@jbymy2000

Copy link
Copy Markdown
Author

@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 init_forward_metadata_out_graph plans through the full-graph wrappers (full_cg_prefill_wrappers, built with use_cuda_graph=True), and call_begin_forward raised qo_indptr to min_q there as well, while the PrefillMetadata for that branch never attached the padded gather indices, so q kept its 16 captured rows against a 17-row plan. Padding has no business on a cuda-graph wrapper in the first place: its q buffer is sized at capture, and in the enable_cuda_graph branch of flashinfer's scheduler.cuh the tile is derived from the capture-time max_qo_len (total_num_rows - batch_size + 1), not from the per-batch average, so replays are already tile-invariant across prefix lengths without any help from the padding.

Fix in 48787aa: _can_pad_prefill_q now takes the paged wrapper and returns False when wrapper_paged.is_cuda_graph_enabled, which covers both capture and replay of the full prefill graph (and any other graph-captured prefill wrapper). The eager path is unchanged.

Regression test: test_flashinfer_deterministic_prefill.py now drives FlashInferIndicesUpdaterPrefill.call_begin_forward with a real use_cuda_graph=True wrapper and a 1-token extend behind a 100-token prefix, asserting qo_indptr[-1] stays 1 with no gather indices; the same call against an eager wrapper asserts it is padded to 17. 6/6 locally, plus the existing unit, host-plan and model-override suites.

One caveat on my side: this box is SM120 with only hybrid Qwen3.5 weights, so I could not re-run your --cuda-graph-backend-prefill full arm end to end. If you get a chance to re-run it on the new head, that would close the loop. The logprob_start_len = -1 note is a good catch for anyone reproducing with logprobs on. On prefill-graph neutrality: the breakable-graph drift I measured was on SM120, so I will keep that caveat scoped to what I observed and add your Ampere result to the body.

@jbymy2000

Copy link
Copy Markdown
Author

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 --cuda-graph-backend-prefill full: after 48787aa that configuration simply no longer crashes, and whether a radix hit is bit-exact there is a separate question that, as far as I can tell, nobody has measured yet. The padding only applies when attention is planned eagerly, and the bit-exactness claims in this PR are for that path, which is what the new CI case runs (--disable-prefill-cuda-graph).

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 deterministic_inference.mdx if maintainers prefer it stated there.

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>
@ListenJ

ListenJ commented Nov 2, 2026

Copy link
Copy Markdown

upstream-bundle.zip

Environment

  • SGLang: 0.5.23 (latest version on pip; latest GitHub release is v0.5.22)
  • Model: Qwen/Qwen3.5-4B (Hybrid GDN)
  • GPU: NVIDIA RTX 4090 24GB
  • NVIDIA Driver: 580.173.02
  • CUDA: 12.8

Server arguments:

  • --mem-fraction-static 0.85
  • --trust-remote-code
  • Radix Cache enabled by default.

Request parameters:

  • temperature=0
  • seed=0
  • max_tokens=192
  • logprobs=true

Reproduction

I 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.

Results

1. Default configuration (Radix Cache enabled)

I observed the same behavior in two independent runs:

  • 2 out of 5 prompts produced different token sequences, with the first divergence occurring at token positions 60 and 113, respectively.
  • The maximum absolute logprob difference between cold-cache and cache-hit requests ranged from 0.0642 to 1.3649.
  • All 4 cache-hit repetitions produced identical results for each prompt.

2. Radix Cache disabled

I repeated the test with --disable-radix-cache:

  • No token differences were observed across all 5 prompts.
  • The maximum absolute logprob difference between cold-cache and cache-hit requests was 0.0.

3. Deterministic inference enabled (Radix Cache still enabled)

I repeated the test with --enable-deterministic-inference while keeping Radix Cache enabled:

  • 3 out of 5 prompts still produced different token sequences, with divergences at positions 22, 45, and 142.
  • The maximum absolute logprob difference ranged from 0.0691 to 1.3964.

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 Findings

Based 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 --enable-deterministic-inference does not eliminate the issue.

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.

@kahyunnam

Copy link
Copy Markdown
Contributor

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 cta_tile_q from the batch's average packed q length (FA2DetermineCtaTileQ): 128 above 64, 64 above 16, else 16. A radix hit sends only the uncached tail, so for Qwen3-8B shapes (32/8 heads, head_dim 128) any tail of 16 tokens or fewer gets a different tile and reduction order. On A100 (bf16, causal, disable_split_kv=True): without padding, every tail of 16 or fewer differed from a cold prefill by up to 9.8e-4 (8 of 12 cases); with this PR's padding to min_q=17, 12 of 12 were bit-exact. A 512-token request alone versus with 63 one-row batch-mates differed by 3.9e-3. With flashinfer-ai/flashinfer#2830, which adds a caller-pinned fixed_cta_tile_q, every tile value (16, 64, 128) made all cases bit-exact with no padding. That PR conflicts with main and its allowed-value check needs rework for head_dim 512. Whether a pinned tile also settles the full-prefill CUDA-graph bucket question is untested.

Next steps:

  1. On the FlashInfer side, feat: Add Option For Fixed cta_tile_q flashinfer-ai/flashinfer#2830 (fixed query tile) is the fix and awaits a rebase and review; the opt-in versus always-invariant decision is [RL] Fixed prefill query tile (cta_tile_q) so prefill bits do not depend on batch-mates (#2830) flashinfer-ai/flashinfer#7924.
  2. FlashInfer maintainers: add a radix-hit and batch-mate bit-exactness test once the API lands.
  3. SGLang: once that ships, fixed_cta_tile_q=128 (64 or 16 for head_dim >= 256) can replace the q padding.

The hybrid GDN (Qwen3.5) drift for q_len <= 32 is a separate path, covered by neither fix.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants