buffer size insuffient Dspark sd for FlashInfer MNNVL allreduce - #50932
buffer size insuffient Dspark sd for FlashInfer MNNVL allreduce#50932khushali9 wants to merge 12 commits into
Conversation
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@ilmarkov Can you please verify and report your findings, it would be super helpful. |
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
|
Thank you for the PR! The PR makes sense to me. I suggest cleaning up the PR and only adding the minimal changes to add |
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
|
please fix precommit |
|
@khushali9 My bad. no worries then. |
|
@wzhao18 then can you add those tags to run all the tests, as you reviewed the PR. |
mgoin
added
bug
|
/ci run |
|
✅ @khushali9, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #83463 for commit |
|
/ci retry |
|
✅ Queued 2 failed job(s) for retry in Buildkite CI #83463. |
|
/ci run |
|
✅ Triggered Buildkite CI #83572 for commit |
|
/ci retry |
|
@khushali9 no worry about unrelated CI failure. @mgoin can help with force merge. |
|
@wzhao18 that would be great, thank you so much. |
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #83759 for commit |
5 participants
fixes #50877
Purpose
FlashInferAllReduce.should_use_fi_ar gates on:
max_workspace_size is the size of the whole MNNVL allocation (2 MB for TP8). But
the MNNVL backend is Lamport-based and rotates through NUM_LAMPORT_BUFFERS=3
buffers, so only ~1/3 of the budget backs any single all-reduce:
budget 2,097,152 bytes → gate says 146 tokens
per-buffer 698,368 bytes → real capacity is 48 tokens
Any tensor between 49–146 tokens passes the gate and then aborts inside the kernel.
The drafter produces B × (1 + num_speculative_tokens) tokens per step — e.g. 14
requests × 4 = 56 tokens — landing exactly in this window. Without a drafter,
decode batches produce B tokens and stay below 48.
The same defect exists in fused_allreduce_gemma_rms_norm._can_use_flashinfer,
which its docstring notes "mirrors FlashInferAllReduce".
Fix
After the existing budget bound (kept as a cheap pre-filter), call the workspace's
own is_buffer_size_sufficient(tp_size, num_tokens, hidden_dim, dtype) method rather
than reimplementing the Lamport arithmetic. This is on the abstract base class
(flashinfer-python==0.6.15.post1, the pinned version), correct for both mnnvl and
trtllm, and version-proof.
What this does NOT fix
AllReduceFusionPass has the same too-permissive bound
but self.max_token_num there also drives the workspace allocation request — fixing it
requires separating those two concerns and verifying the one-shot/two-shot split,
which needs GB200 hardware. Leaving that for a follow-up.
Workaround (no code change needed)
VLLM_ALLREDUCE_USE_FLASHINFER=0
Disables the eager FlashInferAllReduce path; the compiled fusion pass is unaffected.
Test Plan
I have added test file and you can see output below,
every assertion in the capacity-accounting logic changed ran and succeeded.
pytest tests/compile/passes/distributed/test_fusion_all_reduce.py -v , already existing test.
On GB200
python -m vllm.entrypoints.openai.api_server
--model --speculative-config '{"model": ""}' --tensor-parallel-size 8
if reporter can test on GB200 that would be perfect.
Test Result
Cannot test as I do not have GB200, if @ilmarkov can test it, that would be best.
AI assistance
This fix was developed with Claude (claude-opus-4-6) assistance. I have reviewed
every changed line and understand the root cause and the design.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)