Skip to content

v0.28.0.dev20260908053955

v0.28.0.dev20260908053955: [Bugfix] Keep draft KV in sync when dynamic K is zero (#716)
@POWEHI2023 POWEHI2023 tagged this 08 Sep 05:36 
## Summary

- Keep the draft model KV cache synchronized when dynamic speculative
  decoding selects zero speculative tokens for a batch.
- Perform cache ingestion before returning from `propose()` when `K ==
0`.
- Preserve the existing behavior of skipping draft token generation for
  zero-token batches.

## Problem

Dynamic speculative decoding may assign different speculative-token
counts
based on batch size. For example:

- batch size 1: K = 3
- batch size 2: K = 0

Previously, `DraftModelProposer.propose()` returned immediately when K
was
zero. This skipped not only draft generation, but also the
cache-ingestion
step for committed target tokens.

When the batch later returned to K > 0, the draft model resumed from a
stale
KV-cache state. This reduced speculative-token acceptance even though
the
target-model output remained correct.

## Fix

Move the K=0 return until after:

- scheduler plan collection;
- committed-token ingestion into the draft model;
- draft sequence-length updates.

Draft generation is still skipped when K is zero, so the change only
keeps
the draft cache synchronized.

## Reproduction

Configuration:

- target model: Qwen/Qwen3-0.6B
- draft model: Qwen/Qwen3-0.6B
- two 384-token requests
- max batched tokens: 128
- batch size 1: K = 3
- batch size 2: K = 0

Before the change:

- K=0 proposer calls: 1
- cache-ingestion plans during the K=0 call: 0
- cold acceptance: 28/57 (49.12%)
- subsequent acceptance: 36/39, due to history masking the stale-cache
issue

After the change:

- K=0 proposer calls: 1
- cache-ingestion plans during the K=0 call: 2
- acceptance: 36/39 across all three runs
- generated token IDs are unchanged

## Validation

- [x] `scripts/lint.sh`
- [x] Relevant proposer/cache tests: 167 passed
- [x] Full non-slow test suite: 1986 passed, 15 skipped
- [x] Dynamic-K reproduction: K=0 produced two cache-ingestion plans
- [x] Generated token IDs remained unchanged
- [x] Full `scripts/test.sh` smoke run: 1986 passed, 15 skipped

---------

Signed-off-by: POWEHI2023 <zxe6b885@163.com>
Signed-off-by: Ranran Haoran Zhang <ranranhaoranzhang@gmail.com>
Co-authored-by: Ranran Haoran Zhang <ranranhaoranzhang@gmail.com>
Assets 2
Loading