Repository navigation
Conversation
…tch-invariance tests The existing cases in tests/attention/test_batch_invariant_fa2.py use uniform batches (every request has the same qo_len and kv_len) and compare a prefix of the batch against a smaller batch of the same shape, so they never exercise a request whose batch-mates have different lengths. This adds three parametrised cases with heterogeneous batches: tensor-core decode of a 3000-token request alone versus with short or long batch-mates under `disable_split_kv` or `fixed_split_size` (bitwise equal today), FA2 prefill of a 37-token request alone versus with decode-like, mid-size or long prefill mates under `disable_split_kv` and under vLLM's batch-invariant flags, and a prefix-cache tail (only the uncached tail of the query against the full KV) versus the cold prefill of the whole sequence. The decode-like-mates prefill case and the tails whose packed length is at most 64 are strict xfails: FA2 picks `cta_tile_q` from the batch-average packed query length (flashinfer-ai#2424), so the tile and the bits change with the batch composition; flashinfer-ai#2830 adds the opt-in fixed tile. Tails whose cached prefix length is not a multiple of 128 are a non-strict xfail: on B200 the LSE differs by one fp32 ulp in most such cases and isolated output rows flip one ulp even with an unchanged tile, while tails whose prefix length is a multiple of 128 are bitwise equal in every probed case, so those two are asserted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… test Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
kahyunnam
requested review from
Anerudhan,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
jimmyzho,
nv-yunzheq,
qsang-nv,
saltyminty,
yyihuang and
yzh119
as code owners
November 3, 2026 21:51
Contributor
|
Warning Review limit reachedYou've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Next included review available in 22 minutes. View limit details
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. CodeRabbit Verify
|
Member
Author
|
@flashinfer-bot run tests/attention/test_batch_invariant_fa2.py |
Member
Author
|
/bot run tests/attention/test_batch_invariant_fa2.py |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #79003457: 29/29 executed test jobs passed |
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.
📌 Description
tests/attention/test_batch_invariant_fa2.pyonly checks that a request's output is unchanged when the batch is truncated; nothing asserts bitwise batch invariance when a request shares a batch with mates of a different shape, which is the case serving engines (vLLM batch-invariant mode, SGLang radix cache) actually hit. This PR adds three parametrised tests (22 cases, fp16, 32 q heads / 8 kv heads, head_dim 128, page 16) that compareoandlsewithtorch.equalagainst the same request run alone.test_batch_decode_tensor_cores_batch_composition(4 cases): a 3000-token tensor-core decode request alone vs batched with 16 short or 16 long mates, underdisable_split_kv=Trueorfixed_split_size=2048. All pass on main.test_batch_prefill_tensor_cores_batch_composition(6 cases): a 37-token causal FA2 prefill over 3000 KV alone vs with decode-like mates (qo 1), mid mates (qo 64) or long prefill mates (qo 1500), underdisable_split_kv=Trueand under vLLM's batch-invariant flags (fixed_split_size=4096, disable_split_kv=True). Mid and long mates pass; decode-like mates are a strict xfail because FA2 pickscta_tile_qfrom the batch-average packed query length (#2424), which #2830'sfixed_cta_tile_qis meant to fix.test_batch_prefill_tensor_cores_prefix_cache_tail(12 cases): a cold prefill of L tokens vs a tail-only query of t rows against the same L-long KV, the radix-cache-hit pattern from SGLang #46351. Tails 1, 4 and 16 change the tile and are strict xfails (#2424). Tails 37 and 200 keep the 128 tile but are a non-strict xfail: on B200 the tail is bit-exact wheneverL - tis a multiple of 128 (36/36 probed configurations across fp16/bf16, group size 1 and 4, three seeds) and otherwise shows a one-fp32-ulplsedifference and a few output rows flipping one ulp, data dependent. Two aligned-prefix cases, (1024, 256) and (3700, 116), are asserted bit-exact and pass. The misaligned-prefix drift is not covered by #2830 and is left as a non-strict xfail so it is recorded without blocking CI; I will file it separately once the kv-loop bound inprefill.cuhis confirmed as the cause.🔍 Related Issues
Refs #7937, #2424, #2830. Refs sgl-project/sglang#46351.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).pytest tests/attention/test_batch_invariant_fa2.pyon B200 (CUDA 13.0): 2026 passed, 10 xfailed, 2 xpassed (the twotail_len=37non-strict cases, which happen to be bit-exact with seed 0). The new cases alone: 10 passed, 10 xfailed, 2 xpassed. Not run on SM80/SM90/SM120; the aligned-prefix assertions and the strictness of the tile-rule xfails are verified on B200 only.Reviewer Notes
The strict xfails are intentional: once #2830 lands and vLLM passes
fixed_cta_tile_q, they should turn into assertions. If a non-Blackwell CI lane reports an XPASS on a strict case, flipping it tostrict=Falseis the right first move rather than loosening the comparison.🤖 Generated with Claude Code