[Bugfix] Fix uniform_batch for same-sized prefills - #39945
benchislett wants to merge 8 commits into
Conversation
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a prefill guard within the _is_uniform_decode method to prevent the misclassification of prefill batches as uniform decode batches. While the current implementation checks for zero computed tokens, review feedback highlights that this is insufficient for chunked prefills. It is recommended to instead compare computed tokens against prompt tokens to reliably determine if a request has finished prefilling.
mgoin
added
bug
MatthewBonanni
left a comment
MatthewBonanni
left a comment
There was a problem hiding this comment.
Gemini's suggestion is reasonable to me. Probably worth adding a chunked prefill edge case to the test too.
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
Do we know the root cause? This feels a bit like an attention backend problem if it can properly handle this; mostly just curious
We do have backends that treat small chunked prefills as decodes since it more efficient; so Im a bit resistant to global changes like this without better understanding the motivation
It does seem like we may need a better way for attention backends to disable full-CGs (e.g. cascade attention could use this); but there is bit of an ordering problem here since we want to pad for CGs before building attention metadata
|
The problem is the graph selection: when a prefill has exactly the wrong size, the dispatcher thinks it's a uniform_decode batch and issues a FULL CUDA graph which replays the decode kernel and not the prefill kernel. AFAIK this incongruity is causing it to fail to populate the KV cache for all the prefill slots. Not sure if it's supposed to be able to handle this case just by running the decode kernel. Probably being a hybrid model is also complicating things, could be that the FULL graph isn't handling the mamba state properly or some such |
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Port vllm-project#39945 onto the MiMo Black Benediction branch. The fix keeps same-sized prefill batches from being misclassified by the uniform_batch path, which DFlash relies on during draft proposal.
Port vllm-project#39945 onto the MiMo Black Benediction branch. The fix keeps same-sized prefill batches from being misclassified by the uniform_batch path, which DFlash relies on during draft proposal.
njhill
added
the
mrv1-only
Derive unfinished-prefill state from the current scheduler output so both the normal model-runner path and the early PP+SP path reject shape-aliased prefills. Keep the override hybrid-only and preserve CUDA graph capture semantics. Based-on: vllm-project/vllm#39945 Based-on: vllm-project/vllm#47123 Based-on: vllm-project/vllm#53059 Co-authored-by: Katsumi Takeuchi <contact@recutita.com> Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com>
5 participants
Purpose
FIX #39928
Testing
Added some more tests for the modified code. Can reproduce #39928 but it is resolved with this patch. No additional testing planned.