iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV offload] Store interior chunk-boundary blocks under MTP/Eagle (#46972)

Browse Source 
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
drakosha
2026-07-07 13:16:52 +03:00
committed by GitHub
co-authored by Claude Or Ozeri
parent cbb5f045be
commit 3354dba381
2 changed files with 131 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+93 -20
View File
@@ -2216,15 +2216,17 @@ class TestEagle:
# -------------------------------------------------------------------
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_full_attn_store_excludes_trailing_block(
def test_full_attn_store_excludes_trailing_decode_block(
self, request_runner, async_scheduling: bool
):
"""Eagle full-attention group stores all blocks except the trailing
one.
"""Eagle full-attention group excludes the trailing block only while
decoding.
Setup: 2 groups — group 0 is normal full-attention, group 1 is
eagle full-attention. With a 3-block prompt, group 1 should store
only blocks 0 and 1, skipping block 2 (the volatile tail).
eagle full-attention. With a 3-block prompt, group 1 stores all 3
prompt blocks at the end of prefill (the trailing prompt block is
stable), but skips block 3 once it fills with decoded tokens (its
draft-layer KV is volatile until the next block starts).
"""
block_size = 4
block_size_factor = 1
@@ -2270,23 +2272,27 @@ class TestEagle:
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
# 4 decoded tokens fill block 3 entirely with decode tokens (one
# extra token so the block is stored under async scheduling too).
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
decoded_tokens=[1, 1, 1, 1, 1, EOS_TOKEN_ID],
expected_stored=(
(0, 0),
(0, 1),
(0, 2),
(0, 3),
(1, 0),
(1, 1),
(1, 2),
),
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_sw_store_excludes_trailing_block(
def test_sw_store_excludes_trailing_decode_block(
self, request_runner, async_scheduling: bool
):
"""Eagle sliding-window group stores all blocks except the trailing
one."""
"""Eagle sliding-window group stores all prompt blocks but excludes
the trailing block while decoding."""
block_size = 4
sliding_window = 8
num_gpu_blocks = 100
@@ -2321,15 +2327,18 @@ class TestEagle:
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
# 4 decoded tokens fill block 3 entirely with decode tokens.
runner.run(
decoded_tokens=[EOS_TOKEN_ID],
expected_stored=((0, 0), (0, 1)),
decoded_tokens=[1, 1, 1, 1, EOS_TOKEN_ID],
expected_stored=((0, 0), (0, 1), (0, 2)),
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_single_block_nothing_stored(self, request_runner, async_scheduling: bool):
"""An eagle group with only one block stores nothing: that block is
the tail."""
def test_single_block_stored_at_end_of_prefill(
self, request_runner, async_scheduling: bool
):
"""An eagle group with a single-block prompt stores it at the end of
prefill: prompt blocks are stable, so no tail is held back."""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
@@ -2360,17 +2369,80 @@ class TestEagle:
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=())
runner.manager.prepare_store.assert_not_called()
runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=((0, 0),))
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_multichunk_store_no_interior_holes(
self, request_runner, async_scheduling: bool
):
"""Eagle store must not drop interior blocks across prefill chunks.
Regression: the trailing-block exclusion (num_blocks - 1) was applied
when collecting keys, but next_stored_block_idx advanced by the
non-decremented count, so the trailing block of every chunked-prefill
chunk was skipped and never re-considered. With the harness chunk budget
(1000 tokens) and block_size 4, a prompt longer than one chunk lost the
block at the chunk boundary, leaving a permanent gap that caps prefix
reuse at the first hole. Only the trailing decode block may be held
back; all other blocks must be stored exactly once (no duplicates from
next_stored_block_idx regressing at the prefill->decode transition).
"""
block_size = 4
block_size_factor = 1
offloaded_block_size = block_size * block_size_factor
num_gpu_blocks = 1000
kv_cache_groups = [
KVCacheGroupSpec(
["layer0"],
FullAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
is_eagle_group=True,
),
]
runner = request_runner(
block_size=block_size,
num_gpu_blocks=num_gpu_blocks,
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
block_size_factor=block_size_factor,
)
assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group
# Prompt spans more than one prefill chunk (chunk budget 1000 tokens).
num_blocks = 256
runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
# Decode a few non-EOS tokens so prefill completes across both chunks
# before the request finishes.
runner._run([1, 1, 1, 1, EOS_TOKEN_ID], complete_transfers=True)
offsets = sorted(
b.request_block_offset
for t in runner.completed_stores
for b in t.gpu_blocks
)
# The stored blocks must be contiguous from 0: no interior block is
# dropped at a chunk boundary. (The bug left a gap at offloaded block
# 249, the tail of the first 1000-token chunk.)
assert offsets == list(range(len(offsets))), (
f"interior hole in stored blocks: {offsets}"
)
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):
"""Eagle group constrains load: convergence tightens both groups.
Store 3 offloaded blocks per group (eagle group skips tail → stores
2). Then a new request loads from CPU. The eagle group's post-pop hit
(2) does not tighten below group 0's hit (3), so both groups load
normally.
Store 3 offloaded blocks per group (all prompt blocks, so the eagle
group stores all 3 as well). Then a new request loads from CPU. The
eagle group pops its trailing hit block on load, tightening the hit
to 2 blocks for both groups.
"""
block_size = 4
block_size_factor = 1
@@ -2419,6 +2491,7 @@ class TestEagle:
(0, 2),
(1, 0),
(1, 1),
(1, 2),
),
)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+38 -7
View File
@@ -298,12 +298,39 @@ class RequestOffloadState:
for group_state, new_blocks in zip(self.group_states, new_block_id_groups):
group_state.block_ids.extend(new_blocks)
def storable_blocks(
self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int
) -> int:
"""Number of leading offloaded blocks eligible for store.
For eagle/MTP groups the volatile trailing block of the offloadable
range is excluded while decoding: the draft-layer KV of the last
accepted position may be rewritten after spec-token rejection. During
prefill the trailing block is stable (the draft input for a chunk's
last position is the next prompt token), so it is stored immediately.
The exclusion must be applied consistently everywhere
``next_stored_block_idx`` is derived: otherwise the trailing block of
each step is skipped on collection but jumped over by
``next_stored_block_idx``, so it is never re-considered and a
permanent hole breaks prefix-reuse lookup.
"""
num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens
if group_config.is_eagle_group and is_decoding:
num_blocks = max(0, num_blocks - 1)
return num_blocks
def advance_stored_idx(self, num_offloadable_tokens: int) -> None:
# max(): at the prefill->decode transition of a block-aligned prompt,
# storable_blocks drops by one (the eagle exclusion kicks in), and the
# index must not move backwards past already-stored blocks.
for group_config, group_state in zip(
self.config.kv_group_configs, self.group_states
):
num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
group_state.next_stored_block_idx = num_blocks
group_state.next_stored_block_idx = max(
group_state.next_stored_block_idx,
self.storable_blocks(group_config, num_offloadable_tokens),
)
def update_num_hit_blocks(self, num_cached_tokens: int) -> None:
for group_config, group_state in zip(
@@ -903,9 +930,9 @@ class OffloadingConnectorScheduler:
for group_config, group_state in zip(
self.config.kv_group_configs, req_status.group_states
):
num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
if group_config.is_eagle_group:
num_blocks = max(0, num_blocks - 1)
num_blocks = req_status.storable_blocks(
group_config, num_offloadable_tokens
)
start_block_idx = group_state.next_stored_block_idx
if num_blocks <= start_block_idx:
@@ -977,7 +1004,9 @@ class OffloadingConnectorScheduler:
is_sliding_window = (
group_config.sliding_window_size_in_blocks is not None
)
num_blocks = num_offloadable_tokens // group_config.offloaded_block_size
num_blocks = req_status.storable_blocks(
group_config, num_offloadable_tokens
)
start_block_idx = group_state.next_stored_block_idx
block_ids = group_state.block_ids
num_group_blocks = 0
@@ -1010,7 +1039,9 @@ class OffloadingConnectorScheduler:
group_sizes.append(num_group_blocks)
block_indices.append(start_gpu_block_idx or 0)
group_state.next_stored_block_idx = num_blocks
group_state.next_stored_block_idx = max(
group_state.next_stored_block_idx, num_blocks
)
src_spec = GPULoadStoreSpec(
src_block_ids, group_sizes=group_sizes, block_indices=block_indices