@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4a806d08ee

[Bugfix] OffloadingConnector: stop zeroing offload hits under MTP/EAGLE spec decode (#52771)

Signed-off-by: Kam Basra <kameldipbasra@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Kam Basra committed Sep 7, 2026 at 13:16 UTC 4a806d08ee94b35356dd749bf814ee91ee93f8d0

3 files changed +297 -21

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+261 -12

@@ -50,6 +50,7 @@ from vllm.v1.kv_cache_interface import (

50 FullAttentionSpec,

51 KVCacheConfig,

52 KVCacheGroupSpec,

53 + MambaSpec,

54 SlidingWindowSpec,

55 )

56 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

@@ -3070,8 +3071,9 @@ class TestEagle:

3071

3072 Groups: 0=non-eagle full-attn, 1=eagle full-attn.

3073 Group 0 has only 1 hit (out of 3 keys) → max_hit tightens to 4.

3073 - This clears eagle_verified. Group 1 runs with max_hit=4 → only 1

3074 - key queried, 1 hit, pop to 0 → returns 0.

3074 + This clears eagle_verified. Group 1 runs with max_hit=4 → widened

3075 + query of 2 keys, 2 hits, pop to 1 → the confirmed boundary holds

3076 + at 4 tokens.

3077 """

3078 block_size = 4

3079 groups = [

@@ -3116,9 +3118,12 @@ class TestEagle:

3118 offload_keys_per_group=[[10, 11, 12], [1, 2, 3]],

3119 )

3120 # Group 0 (non-eagle FA): prefix finds 1 hit → max_hit=4, num_hit=4

3119 - # Group 1 (eagle FA): max_hit=4 → num_blocks=1, keys=[1].

3120 - # Finds 1 hit, pop to 0 → new_num_hit = 0 < block_size → return 0

3121 - assert sched._lookup(req_status) == 0

3121 + # Group 1 (eagle FA): max_hit=4, widened query → keys=[1, 2].

3122 + # Finds 2 hits, pop to 1 → boundary stays at 4 tokens. Before the

3123 + # #52735 fix the query was not widened for full-attention groups,

3124 + # so the pop landed on the only queried chunk and zeroed the

3125 + # whole request.

3126 + assert sched._lookup(req_status) == 4

3127

3128 def test_eagle_verified_survives_eagle_tighten(self, request_runner):

3129 """Eagle group tightening does NOT clear eagle_verified.

@@ -3234,10 +3239,12 @@ class TestEagle:

3239 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

3240 generate_store_output(keys)

3241 )

3237 - # 4 decoded tokens fill block 3 entirely with decode tokens (one

3238 - # extra token so the block is stored under async scheduling too).

3242 + # 4 decoded tokens fill block 3; the extra decode steps let its store

3243 + # job complete within this run under both scheduling modes. The

3244 + # eagle group holds back its volatile trailing block (1, 3) while the

3245 + # request is still decoding, while the normal group stores (0, 3).

3246 runner.run(

3240 - decoded_tokens=[1, 1, 1, 1, 1, EOS_TOKEN_ID],

3247 + decoded_tokens=[1, 1, 1, 1, 1, 1, 1],

3248 expected_stored=(

3249 (0, 0),

3250 (0, 1),

@@ -3248,6 +3255,13 @@ class TestEagle:

3255 (1, 2),

3256 ),

3257 )

3258 + # Once the request finishes, no spec rejection can rewrite the tail,

3259 + # so the exclusion is lifted (issue #52735): the held-back (1, 3) and

3260 + # the just-completed block 4 are stored for both groups.

3261 + runner.run(

3262 + decoded_tokens=[EOS_TOKEN_ID],

3263 + expected_stored=((1, 3),),

3264 + )

3265

3266 @pytest.mark.parametrize("async_scheduling", [True, False])

3267 def test_sw_store_excludes_trailing_decode_block(

@@ -3289,11 +3303,17 @@ class TestEagle:

3303 runner.manager.prepare_store.side_effect = lambda keys, req_context: (

3304 generate_store_output(keys)

3305 )

3292 - # 4 decoded tokens fill block 3 entirely with decode tokens.

3306 + # 4 decoded tokens fill block 3 entirely with decode tokens. The

3307 + # eagle group holds back its volatile trailing block while decoding.

3308 runner.run(

3294 - decoded_tokens=[1, 1, 1, 1, EOS_TOKEN_ID],

3309 + decoded_tokens=[1, 1, 1, 1],

3310 expected_stored=((0, 0), (0, 1), (0, 2)),

3311 )

3312 + # Finish lifts the exclusion; the tail block is stored (issue #52735).

3313 + runner.run(

3314 + decoded_tokens=[EOS_TOKEN_ID],

3315 + expected_stored=((0, 3),),

3316 + )

3317

3318 @pytest.mark.parametrize("async_scheduling", [True, False])

3319 def test_single_block_stored_at_end_of_prefill(

@@ -3569,8 +3589,8 @@ class TestEagle:

3589 )

3590 spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))

3591 scheduler = OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)

3572 - scheduler.manager.prepare_store.side_effect = (

3573 - lambda keys, req_context: generate_store_output(keys)

3592 + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (

3593 + generate_store_output(keys)

3594 )

3595

3596 request = MagicMock()

@@ -4171,3 +4191,232 @@ def test_retention_interval_zero_stores_only_replay_boundary(

4191 (1, 11),

4192 ),

4193 )

4194 +

4195 +

4196 +def _shared_kv_mtp_config():

4197 + """Speculative config for a shared-group MTP model: eagle-family method

4198 + whose drafter layer merges into a target KV-cache group, so no group

4199 + self-identifies as a drafter group (issue #52735)."""

4200 + spec = MagicMock(name="shared_kv_mtp_spec")

4201 + spec.use_eagle.return_value = True

4202 + spec.use_eagle_block_drop.return_value = True

4203 + spec.use_multi_module_mtp.return_value = False

4204 + spec.num_speculative_tokens_per_batch_size = None

4205 + spec.max_num_new_slots_for_drafting = 0

4206 + spec.num_speculative_tokens = 1

4207 + return spec

4208 +

4209 +

4210 +class TestSharedGroupMTPOffload:

4211 + """Regression tests for issue #52735: OffloadingConnector must keep

4212 + serving when speculative decoding is enabled but no KV-cache group is

4213 + annotated as a drafter group (shared-group MTP models)."""

4214 +

4215 + def test_no_annotation_marks_no_groups(self, request_runner):

4216 + """Spec decode on + zero annotated groups must NOT mark every group

4217 + as a drafter group; the full store->load roundtrip must match the

4218 + non-speculative behavior of test_two_groups_full_and_sliding_window."""

4219 + block_size = 4

4220 + kv_cache_groups = [

4221 + KVCacheGroupSpec(

4222 + ["layer0"],

4223 + FullAttentionSpec(

4224 + block_size=block_size,

4225 + num_kv_heads=1,

4226 + head_size=1,

4227 + dtype=torch.float32,

4228 + ),

4229 + ),

4230 + KVCacheGroupSpec(

4231 + ["layer1"],

4232 + SlidingWindowSpec(

4233 + block_size=block_size,

4234 + num_kv_heads=1,

4235 + head_size=1,

4236 + dtype=torch.float32,

4237 + sliding_window=8,

4238 + ),

4239 + ),

4240 + ]

4241 + runner = request_runner(

4242 + block_size=block_size,

4243 + num_gpu_blocks=100,

4244 + async_scheduling=True,

4245 + kv_cache_groups=kv_cache_groups,

4246 + speculative_config=_shared_kv_mtp_config(),

4247 + )

4248 + kv_group_configs = runner.connector_scheduler.config.kv_group_configs

4249 + assert [c.is_eagle_group for c in kv_group_configs] == [False, False]

4250 +

4251 + runner.new_request(token_ids=[0] * block_size * 3)

4252 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

4253 + generate_store_output(keys)

4254 + )

4255 + runner.run(decoded_tokens=[0])

4256 + runner.run(

4257 + decoded_tokens=[0] * (block_size * 3 + 2),

4258 + expected_stored=(0, 1, 2, 3, 4, 5),

4259 + )

4260 + runner.run(decoded_tokens=[EOS_TOKEN_ID])

4261 +

4262 + runner.scheduler.reset_prefix_cache()

4263 +

4264 + runner.new_request(token_ids=[0] * (block_size * 3 + 1))

4265 + runner.manager.lookup.return_value = LookupResult.HIT

4266 + runner.run(

4267 + decoded_tokens=[EOS_TOKEN_ID],

4268 + expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),

4269 + )

4270 +

4271 +

4272 +class TestMambaHybridOffloadServing:

4273 + """Store->finish->lookup flows on a full-attention + mamba-align hybrid,

4274 + with a manager that only HITs keys that were actually stored. Guards the

4275 + two collapse routes of issue #52735."""

4276 +

4277 + BLOCK = 4

4278 + MAMBA_BLOCK = 16

4279 + PROMPT_TOKENS = 28 # 7 hash blocks; 1 full mamba chunk

4280 +

4281 + def _make_scheduler(self, speculative_config, mamba_eagle=False):

4282 + vllm_config = _make_vllm_config(

4283 + extra_config={"self_describing_kv_events": True}

4284 + )

4285 + vllm_config.cache_config.block_size = self.BLOCK

4286 + vllm_config.cache_config.prefix_match_unit = self.BLOCK

4287 + vllm_config.speculative_config = speculative_config

4288 + vllm_config.kv_events_config = KVEventsConfig(

4289 + enable_kv_cache_events=True, publisher="null"

4290 + )

4291 + kv_cache_config = KVCacheConfig(

4292 + num_blocks=64,

4293 + kv_cache_tensors=[],

4294 + kv_cache_groups=[

4295 + KVCacheGroupSpec(

4296 + ["full_layer"],

4297 + FullAttentionSpec(

4298 + block_size=self.BLOCK,

4299 + num_kv_heads=1,

4300 + head_size=1,

4301 + dtype=torch.float32,

4302 + ),

4303 + is_eagle_group=mamba_eagle,

4304 + ),

4305 + KVCacheGroupSpec(

4306 + ["mamba_layer"],

4307 + MambaSpec(

4308 + block_size=self.MAMBA_BLOCK,

4309 + shapes=((1, 1),),

4310 + dtypes=(torch.float32,),

4311 + mamba_cache_mode="align",

4312 + ),

4313 + ),

4314 + ],

4315 + )

4316 + spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))

4317 + return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)

4318 +

4319 + def _roundtrip_served_tokens(self, scheduler) -> int:

4320 + """Request A: prompt-only store, finish; request A2: served tokens."""

4321 + stored: set = set()

4322 + scheduler.manager.lookup.side_effect = lambda key, ctx: (

4323 + LookupResult.HIT if key in stored else LookupResult.MISS

4324 + )

4325 + scheduler.manager.prepare_store.side_effect = lambda keys, ctx: (

4326 + generate_store_output(list(keys))

4327 + )

4328 +

4329 + def complete_all_jobs():

4330 + job_ids = list(scheduler._jobs.keys())

4331 + for jid in job_ids:

4332 + stored.update(scheduler._jobs[jid].keys)

4333 + if job_ids:

4334 + scheduler.update_connector_output(

4335 + KVConnectorOutput(

4336 + kv_connector_worker_meta=OffloadingWorkerMetadata(

4337 + completed_jobs={jid: 1 for jid in job_ids}

4338 + )

4339 + )

4340 + )

4341 +

4342 + def make_request(req_id):

4343 + request = MagicMock()

4344 + request.request_id = req_id

4345 + request.kv_transfer_params = None

4346 + request.num_prompt_tokens = self.PROMPT_TOKENS

4347 + request.num_tokens = self.PROMPT_TOKENS

4348 + request.num_computed_tokens = 0

4349 + request.block_hashes = [

4350 + BlockHash(f"h{i}".encode())

4351 + for i in range(self.PROMPT_TOKENS // self.BLOCK)

4352 + ]

4353 + request.all_token_ids = list(range(self.PROMPT_TOKENS))

4354 + request.lora_request = None

4355 + request.is_finished.return_value = False

4356 + request.status = None

4357 + scheduler.on_new_request(request)

4358 + return request

4359 +

4360 + req = make_request("A")

4361 + req_status = scheduler._req_status["A"]

4362 + req_status.num_locally_computed_tokens = 0

4363 + req_status.update_offload_keys()

4364 + assert scheduler._lookup(req_status) == 0 # cold

4365 +

4366 + n_full = self.PROMPT_TOKENS // self.BLOCK

4367 + n_mamba = self.PROMPT_TOKENS // self.MAMBA_BLOCK

4368 + req_status.group_states[0].block_ids[:] = list(range(1, n_full + 1))

4369 + req_status.group_states[1].block_ids[:] = list(range(101, 101 + n_mamba + 1))

4370 +

4371 + out = SimpleNamespace(

4372 + num_scheduled_tokens={"A": self.PROMPT_TOKENS},

4373 + finished_req_ids=set(),

4374 + kv_connector_block_state=KVConnectorBlockState(

4375 + block_ids={},

4376 + boundary_state_offloads={"A": [(1, 101, self.MAMBA_BLOCK)]},

4377 + ),

4378 + )

4379 + scheduler._build_partial_tail_store_jobs(out)

4380 + scheduler._build_store_jobs(out)

4381 + complete_all_jobs()

4382 +

4383 + req.num_computed_tokens = self.PROMPT_TOKENS

4384 + req.num_tokens = self.PROMPT_TOKENS + 1

4385 + req.is_finished.return_value = True

4386 + req_status.update_offload_keys()

4387 + scheduler.request_finished(req)

4388 + out2 = SimpleNamespace(num_scheduled_tokens={}, finished_req_ids={"A"})

4389 + scheduler._build_store_jobs(out2)

4390 + complete_all_jobs()

4391 +

4392 + make_request("A2")

4393 + rs2 = scheduler._req_status["A2"]

4394 + rs2.num_locally_computed_tokens = 0

4395 + rs2.update_offload_keys()

4396 + served = scheduler._lookup(rs2)

4397 + return served if served is not None else 0

4398 +

4399 + def test_shared_group_mtp_serves_like_non_speculative(self):

4400 + """Route 1 of #52735: with spec decode on and no drafter annotation,

4401 + serving must equal the non-speculative baseline (was 0 before the

4402 + fix: the all-groups fallback marked the mamba group as a drafter and

4403 + the volatile-tail pop consumed its only servable chunk)."""

4404 + baseline = self._roundtrip_served_tokens(self._make_scheduler(None))

4405 + assert baseline == 16

4406 + served = self._roundtrip_served_tokens(

4407 + self._make_scheduler(_shared_kv_mtp_config())

4408 + )

4409 + assert served == baseline

4410 +

4411 + def test_annotated_eagle_group_does_not_starve_coarse_sibling(self):

4412 + """Route 2 of #52735 (F1): a genuinely annotated eagle full-attention

4413 + group must not drag the confirmed boundary below a coarser sibling's

4414 + chunk granularity. The widened query makes the volatile-tail pop

4415 + land on an extra queried chunk, holding the boundary at 16 tokens

4416 + (was 0 before the fix: 4 hits -> pop -> 12 tokens < mamba chunk)."""

4417 + scheduler = self._make_scheduler(None, mamba_eagle=True)

4418 + assert [c.is_eagle_group for c in scheduler.config.kv_group_configs] == [

4419 + True,

4420 + False,

4421 + ]

4422 + assert self._roundtrip_served_tokens(scheduler) == 16

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -181,6 +181,7 @@ class RequestRunner:

181 extra_config_overrides: dict[str, Any] | None = None,

182 worker_count: int = 1,

183 retention_interval: int | None = None,

184 + speculative_config: Any | None = None,

185 ):

186 assert blocks_per_chunk == 1 or kv_cache_groups is None, (

187 "blocks_per_chunk > 1 requires all groups to have the same "

@@ -202,6 +203,8 @@ class RequestRunner:

203 vllm_config.scheduler_config.async_scheduling = async_scheduling

204 vllm_config.parallel_config.world_size = worker_count

205 vllm_config.cache_config.prefix_cache_retention_interval = retention_interval

206 + if speculative_config is not None:

207 + vllm_config.speculative_config = speculative_config

208

209 extra_config: dict[str, Any] = {

210 "spec_name": "MockOffloadingSpec",

@@ -679,6 +682,7 @@ def request_runner():

682 extra_config_overrides=None,

683 worker_count=1,

684 retention_interval=None,

685 + speculative_config=None,

686 ):

687 runner = RequestRunner(

688 block_size=block_size,

@@ -689,6 +693,7 @@ def request_runner():

693 extra_config_overrides=extra_config_overrides,

694 worker_count=worker_count,

695 retention_interval=retention_interval,

696 + speculative_config=speculative_config,

697 )

698 runners.append(runner)

699 return runner

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+31 -9

@@ -208,7 +208,21 @@ class SchedulerOffloadConfig(NamedTuple):

208 and vllm_config.speculative_config.use_eagle_block_drop()

209 )

210 if use_eagle_block_drop and not eagle_groups:

211 - eagle_groups = set(range(len(kv_cache_config.kv_cache_groups)))

211 + # No group is annotated as holding drafter layers. This happens

212 + # for shared-group MTP models (e.g. Qwen3.5-style), whose drafter

213 + # is a regular decoder layer merged into the target's

214 + # full-attention group, and for separate-draft models that miss

215 + # annotation. Treating every group as a drafter group here makes

216 + # the volatile-tail pop apply to ALL groups, which can zero the

217 + # whole request's offload hit whenever any group's servable

218 + # window is a single chunk (issue #52735). Drafter KV served one

219 + # chunk stale can only lower speculative acceptance rates — the

220 + # target model verifies every draft — so fail toward serving.

221 + logger.info_once(

222 + "KV offloading: speculative decoding is enabled but no "

223 + "KV-cache group is annotated as a drafter group; treating "

224 + "all groups as non-draft for offloading."

225 + )

226

227 if eagle_groups:

228 logger.info(

@@ -395,15 +409,19 @@ class RequestOffloadState:

409 accepted position may be rewritten after spec-token rejection. During

410 prefill the trailing chunk is stable (the draft input for a chunk's

411 last position is the next prompt token), so it is stored immediately.

398 - The exclusion must be applied consistently everywhere

399 - ``next_stored_chunk_idx`` is derived: otherwise the trailing chunk of

400 - each step is skipped on collection but jumped over by

401 - ``next_stored_chunk_idx``, so it is never re-considered and a

402 - permanent hole breaks prefix-reuse lookup.

412 + Once the request has finished, no further spec-token rejection can

413 + rewrite the tail, so the exclusion is lifted and the final chunk

414 + becomes storable (issue #52735). The exclusion must be applied

415 + consistently everywhere ``next_stored_chunk_idx`` is derived:

416 + otherwise the trailing chunk of each step is skipped on collection but

417 + jumped over by ``next_stored_chunk_idx``, so it is never re-considered

418 + and a permanent hole breaks prefix-reuse lookup. ``is_finished`` is

419 + monotonic, so the finish-time calls all see the lifted exclusion.

420 """

421 num_chunks = num_offloadable_tokens // group_config.tokens_per_chunk

422 is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens

406 - if group_config.is_eagle_group and is_decoding:

423 + # Finished requests have no pending speculation.

424 + if group_config.is_eagle_group and is_decoding and not self.req.is_finished():

425 num_chunks = max(0, num_chunks - 1)

426 num_allocated_chunks = (

427 len(group_state.block_ids) // self.config.blocks_per_chunk

@@ -761,9 +779,13 @@ class OffloadingConnectorScheduler:

779 )

780

781 # For eagle groups, query one extra chunk that will be popped.

764 - # We only need to increase the query size for sliding window groups.

782 + # Widening applies to every group type: without it, the pop

783 + # below shrinks max_hit_size_tokens past what was queried,

784 + # which can push the confirmed boundary under a coarser

785 + # sibling group's chunk granularity and zero the whole

786 + # request's hit (issue #52735).

787 query_max = max_hit_size_tokens

766 - if is_eagle_unverified and sliding_window_size_in_chunks is not None:

788 + if is_eagle_unverified:

789 query_max = min(

790 max_hit_size_tokens + tokens_per_chunk,

791 len(offload_keys) * tokens_per_chunk,