[BugFix] Reserve the bonus query slot in DFlash scheduling budget - #51256
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: HF-001 <1670186653@qq.com>
| if self.parallel_drafting: | ||
| # For parallel drafting, we need one new slot per 'masked' token | ||
| slots_per_req = self.num_speculative_tokens - 1 | ||
| if self.use_dflash(): |
There was a problem hiding this comment.
nit: Could we include self.use_dflash() in the existing increment condition below? This keeps the one-slot adjustment in a single place:
if self.uses_draft_model() or self.use_dflash():
slots_per_req += 1
There was a problem hiding this comment.
@dreamer-89 Thank you for your suggestion. The modifications have been completed
| assert cfg.scheduler_config.async_scheduling is True | ||
|
|
||
|
|
||
| def test_dflash_max_num_new_slots_for_drafting(): |
There was a problem hiding this comment.
Could we parameterize this test to cover both ordinary parallel drafting (K - 1) and DFlash (K) to ensure the fix does not change slot accounting for other parallel drafters.
There was a problem hiding this comment.
@dreamer-89 Thank you for your suggestion. have added the p-eagle test case
Signed-off-by: HF-001 <1670186653@qq.com>
|
Thanks @HF-001, currently pre-commit check is failing and it needs a label to run the check. I will let the code owners/folks who have permission to add it. |
@dreamer-89 Thank you |
@dreamer-89 Hello, can you help me find someone to label it? Thank you |
|
@mgoin @robertgshaw2-redhat, could one of you add the required verified/ready label, or point us to someone with permission? |
|
@HF-001 Thanks for this fix! Would be great if you could update the max_num_new_slots_for_drafting and the test to cover all cases here: |
| ("eagle3", 7), | ||
| ("dflash", 8), |
There was a problem hiding this comment.
By default eagle3 has parallel_drafting=False, only p-eagle has it set to true so will return 7. Could you please add some test to cover all cases?
There was a problem hiding this comment.
@zixi-qi Thank you for your feedback. have added a wealth of test cases
| # Draft models do not slice the draft tokens, while DFlash adds a | ||
| # bonus query before the masked tokens. | ||
| slots_per_req += 1 | ||
| return slots_per_req |
There was a problem hiding this comment.
I think current max_num_new_slots_for_drafting logic is a bit convoluted, would be good to handle each case separately and add some examples of expected output for each algorithm as comment to make it clearer.
Reference implementation below(not requiring you to copy exactly)
• @property
def max_num_new_slots_for_drafting(self) -> int:
"""Return the maximum additional drafting slots per request.
Let K be ``num_speculative_tokens``. For the standard configurations:
Algorithm method parallel_drafting return
EAGLE3 eagle3 False 0
P-EAGLE eagle3 True K - 1
DFlash dflash True K
DSpark dspark True K - 1
MTP mtp False 0
N-gram ngram False 0
Draft model draft_model False 1
The scheduler already reserves one slot for the request's next token.
"""
num_draft_tokens = self.num_speculative_tokens
if self.use_dflash():
# DFlash uses one bonus query followed by K mask queries.
return num_draft_tokens
if self.parallel_drafting:
# P-EAGLE and default DSpark use K total query positions.
if self.uses_draft_model():
# Parallel draft-model inputs retain an additional unsliced token.
return num_draft_tokens
return num_draft_tokens - 1
if self.uses_draft_model():
# Serial draft-model inputs retain one unsliced token.
return 1
return 0
There was a problem hiding this comment.
@zixi-qi Thank you, have optimized according to your suggestion
Signed-off-by: HF-001 <1670186653@qq.com>
Signed-off-by: HF-001 <1670186653@qq.com>
|
✅ @HF-001, CI is now available for this PR.
|
|
/ci run |
|
✅ Triggered Buildkite CI #83565 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #83632 for commit |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #83632. |
test_dspark_sequential_sampling_writes_persistent_draft_logits called DSparkSpeculator.clear_runtime_draft_logits, which no longer exists, and the fixture missed enable_adaptive_verification because object.__new__ skips __init__. The missing method is not a lost fix -- I first read it that way and was wrong. fbbc8e7 added a reuse-and-clear scheme (assign base_logits into draft_logits, clear it afterwards); 7d9970d replaced that six days later with a persistent preallocated buffer, because DSpark drafting is CUDA-graph replayed and a Python-side reassignment does not run per replay. The clearing call became a `pass` and was later deleted. Restoring it would reintroduce something removed for cause. So the invariant is still worth asserting -- the draft-logits buffer must never be replaced -- but it has to be checked against something that still exists. The test now runs a second sampling pass and asserts buffer identity across it. Both this and the three test_mtp failures were red at 67f5de5 as well as at the merge head, so neither came from the 08-13 merge. What the merge did was move this one's failure to a later line, which is what made it visible. Remaining in tests/v1/spec_decode: test_dflash_drafter_window_reserves_bonus_token, whose fix is upstream vllm-project#51256 -- in the sixteen commits not yet merged.
test_dflash_drafter_window_reserves_bonus_token is upstream's, byte-identical to theirs, and it builds SimpleNamespace runner stubs carrying exactly the fields upstream's _input_fits_in_drafter reads. This fork's version reads two more: self.parallel_config, for the per-rank gate-off sentinel added with the TP drafter-gate work, and self._drafter_gate_off_logged, its log counter. So our production change broke their mock -- AttributeError from a stub that is correct for upstream and incomplete for us. Their test, our behaviour, our stub update. I expected upstream vllm-project#51256 (Reserve the bonus query slot in DFlash scheduling budget) to fix this, because the PR title matches the test name. It did not and could not: the failure was never upstream's. Verified after merging it -- same AttributeError at the same line. All five tests/v1/spec_decode failures identified earlier are now green: 3 x test_mtp (propose() signature drift), test_dspark_config (assertion on a method we deliberately removed), and this one.
…lm-project#51256) Signed-off-by: HF-001 <1670186653@qq.com>
4 participants
Purpose
The existing generic parallel-drafting calculation only reserves K - 1
slots for DFlash, leaving the scheduling budget short by one slot per
request.
For example, with:
max_num_batched_tokens = 2048
max_num_seqs = 256
num_speculative_tokens = 8
the previous calculation allowed:
max_num_scheduled_tokens = 2048 - 7 * 256 = 256
However, a full DFlash batch may require:
256 * (8 + 1) = 2304
query tokens, exceeding max_num_batched_tokens.
DFlash has K + 1 query tokens per request, so its net drafting expansion is K slots rather than the generic parallel-drafting value of K - 1.