A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 20:45 UTC
[KV Connector][Bugfix] MooncakeStore: don't double-apply Eagle prune in load_mask (#43516)
Signed-off-by: Dao Le <daole@inferact.ai> Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Dao007forever committed
May 25, 2026 at 22:11 UTC
7966fc7233ac5185fed9d638b7e8c8422bd9e118
2 files changed
+112
-7
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+82
@@ -300,3 +300,85 @@ def test_store_mask_fast_path_single_attention_group():
300
assert len(coord.attention_groups) == 1
301
masks = coord.store_mask(64)
302
assert masks == ([True] * 4, [True] * 4)
303
+
304
+
305
+# ----- Eagle / MTP interaction with load_mask -----
306
+
307
+
308
+def test_lookup_with_eagle_pops_last_full_attention_block():
309
+ """Sanity: with use_eagle, find_longest_cache_hit drops the last block.
310
+ Pairs with the load_mask test below to lock the round-trip contract."""
311
+ groups = [KVCacheGroupSpec(["L0"], _full(16))]
312
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
313
+ hs = _hashes(4)
314
+ cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
315
+ _masks, hit = coord.find_longest_cache_hit(
316
+ hs, max_length=64, cached_block_pool=cmap
317
+ )
318
+ # 4 blocks present, eagle pops 1 → 3 blocks = 48 tokens.
319
+ assert hit == 48
320
+
321
+
322
+def test_load_mask_with_eagle_does_not_double_prune_full_attention():
323
+ """Regression for silent KV corruption with MTP/EAGLE-3.
324
+
325
+ The recv side calls ``load_mask(block_hashes, token_len)`` where
326
+ ``token_len`` is already the eagle-pruned hit length from ``lookup``.
327
+ A second eagle pop here used to shorten the mask by one extra block;
328
+ ``process_tokens`` then yielded a chunk past the mask, which the worker
329
+ silently skipped — leaving the trailing block of the loaded prefix
330
+ uninitialized in local KV.
331
+ """
332
+ groups = [KVCacheGroupSpec(["L0"], _full(16))]
333
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
334
+ hs = _hashes(4)
335
+ cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
336
+ _masks, hit = coord.find_longest_cache_hit(
337
+ hs, max_length=64, cached_block_pool=cmap
338
+ )
339
+ assert hit == 48 # eagle popped 1 block
340
+
341
+ masks = coord.load_mask(hs, token_len=hit)
342
+ # Every chunk that process_tokens(token_len=48, ...) would yield must
343
+ # have a corresponding mask slot. process_tokens emits chunk_id 0..2
344
+ # (start=0, 16, 32), so the mask must be length 3, all True.
345
+ assert masks[0] == [True, True, True]
346
+
347
+
348
+def test_load_mask_with_eagle_hybrid_full_plus_swa():
349
+ """Hybrid (FullAttn + SWA) with eagle: load_mask must cover every chunk
350
+ in [0, token_len) for the FullAttn group; SWA group keeps its
351
+ tail-window mask."""
352
+ groups = [
353
+ KVCacheGroupSpec(["L0"], _full(16)),
354
+ KVCacheGroupSpec(["L1"], _swa(16, 32)),
355
+ ]
356
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
357
+ hs = _hashes(4)
358
+ exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
359
+ cmap = ExternalCachedBlockPool(exists)
360
+ _masks, hit = coord.find_longest_cache_hit(
361
+ hs, max_length=64, cached_block_pool=cmap
362
+ )
363
+ # FullAttn dictates the convergence; eagle pops one block off it.
364
+ assert hit == 48
365
+
366
+ masks = coord.load_mask(hs, token_len=hit)
367
+ # FullAttn: all chunks populated locally.
368
+ assert masks[0] == [True, True, True]
369
+ # SWA: tail-window only (ceil((32-1)/16) = 2 trailing blocks).
370
+ assert masks[1][-2:] == [True, True]
371
+
372
+
373
+def test_load_mask_without_eagle_unchanged():
374
+ """Sanity: when eagle is off, load_mask is identical to the pre-fix path."""
375
+ groups = [KVCacheGroupSpec(["L0"], _full(16))]
376
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=False)
377
+ hs = _hashes(4)
378
+ cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs})
379
+ _masks, hit = coord.find_longest_cache_hit(
380
+ hs, max_length=64, cached_block_pool=cmap
381
+ )
382
+ assert hit == 64
383
+ masks = coord.load_mask(hs, token_len=hit)
384
+ assert masks[0] == [True, True, True, True]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+30
-7
@@ -115,12 +115,21 @@ class MooncakeStoreCoordinator:
115
block_hashes: list[BlockHash],
116
max_length: int,
117
cached_block_pool: ExternalCachedBlockPool,
118
+ *,
119
+ apply_eagle: bool = True,
120
) -> tuple[tuple[list[bool], ...], int]:
121
"""Returns ``(load_mask_per_group, hit_length)``. ``mask[g][i]`` is True iff
122
group ``g`` populates chunk ``i`` locally (e.g. SWA and Mamba tail-only);
121
- recv-side callers skip False slots."""
123
+ recv-side callers skip False slots.
124
+
125
+ ``apply_eagle`` controls whether the per-spec ``use_eagle`` last-block
126
+ pop is applied. Lookup callers want it (the drafter requires recomputing
127
+ the last block); per-chunk mask callers must not, because ``token_len``
128
+ already reflects the eagle-pruned hit length and a second pop would
129
+ leave the trailing block unloaded.
130
+ """
131
blocks_per_group, hit_length = self._find_hit_blocks(
123
- block_hashes, max_length, cached_block_pool
132
+ block_hashes, max_length, cached_block_pool, apply_eagle=apply_eagle
133
)
134
masks = tuple(
135
[blk is not cached_block_pool.null_block for blk in blocks]
@@ -137,8 +146,17 @@ class MooncakeStoreCoordinator:
146
spec would populate chunk ``i`` locally at length ``token_len``
147
(e.g. SWA / Mamba tail-only).
148
"""
149
+ # ``apply_eagle=False`` because ``token_len`` is already the
150
+ # eagle-pruned hit length returned by ``client.lookup``. Re-applying
151
+ # the pop here would shorten the mask by one extra block; the recv
152
+ # thread would then silently skip the trailing chunk yielded by
153
+ # ``db.process_tokens`` and leave that block uninitialized in the
154
+ # local KV pool.
155
masks, _ = self.find_longest_cache_hit(
141
- block_hashes, token_len, ExternalCachedBlockPool()
156
+ block_hashes,
157
+ token_len,
158
+ ExternalCachedBlockPool(),
159
+ apply_eagle=False,
160
)
161
return masks
162
@@ -195,10 +213,17 @@ class MooncakeStoreCoordinator:
213
block_hashes: list[BlockHash],
214
max_length: int,
215
cached_block_pool: ExternalCachedBlockPool,
216
+ *,
217
+ apply_eagle: bool = True,
218
) -> tuple[tuple[list[KVCacheBlock], ...], int]:
219
"""Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
220
dispatches via spec_manager_map (we don't allocate managers).
221
+
222
+ When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` —
223
+ used by ``load_mask`` to avoid popping a second block on top of the
224
+ one already removed by the lookup.
225
"""
226
+ eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
227
if len(self.attention_groups) == 1:
228
spec, group_ids, manager_cls = self.attention_groups[0]
229
hashes = self.block_hashes_for_spec(block_hashes, spec)
@@ -208,7 +233,7 @@ class MooncakeStoreCoordinator:
233
kv_cache_group_ids=group_ids,
234
block_pool=cast(BlockPool, cached_block_pool),
235
kv_cache_spec=spec,
211
- use_eagle=(0 in self.eagle_attn_group_indices),
236
+ use_eagle=(0 in eagle_indices),
237
alignment_tokens=spec.block_size,
238
)
239
num_groups = len(self.kv_cache_groups)
@@ -237,9 +262,7 @@ class MooncakeStoreCoordinator:
262
)
263
continue
264
240
- use_eagle = (
241
- idx in self.eagle_attn_group_indices and idx not in eagle_verified
242
- )
265
+ use_eagle = idx in eagle_indices and idx not in eagle_verified
266
_max_length = curr_hit_length
267
if use_eagle:
268
_max_length = min(curr_hit_length + spec.block_size, max_length)