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 5, 03:30 UTC
[Mooncake] Fix mixed MLA+Eagle block-size validation (#39596)
Signed-off-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: Zhewen Li <zhewenli@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Zhewen Li committed
Apr 15, 2026 at 11:36 UTC
55e1a8e1035bddb0b5b63f9ddecc8b4e16fc3ef6
2 files changed
+65
-3
tests/v1/kv_connector/unit/test_mooncake_connector.py
+45
@@ -609,6 +609,51 @@ def test_register_kv_caches():
609
assert bl == tensor1[0].nbytes // tensor1.shape[1]
610
611
612
+def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
613
+ """Mixed MLA+Eagle caches should register by byte length, not shape."""
614
+
615
+ vllm_config = create_vllm_config(
616
+ kv_connector="MooncakeConnector", kv_role="kv_consumer"
617
+ )
618
+
619
+ with (
620
+ set_current_vllm_config(vllm_config),
621
+ patch_worker_dependencies(),
622
+ patch(
623
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
624
+ ),
625
+ patch(
626
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
627
+ ) as mock_thread,
628
+ ):
629
+ connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
630
+ worker = connector.connector_worker
631
+ mock_thread.return_value.is_alive.return_value = False
632
+
633
+ worker.use_mla = True
634
+ worker.kv_topo.is_mla = True
635
+
636
+ # MLA cache tensor: shape[-2] is the block size.
637
+ mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
638
+ # Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size.
639
+ eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
640
+ kv_caches = {"mla_layer": mla_cache, "eagle_layer": eagle_cache}
641
+
642
+ with patch.object(
643
+ worker.engine, "batch_register_memory", return_value=0
644
+ ) as mock_batch_register:
645
+ connector.register_kv_caches(kv_caches)
646
+
647
+ mock_batch_register.assert_called_once()
648
+ registered_ptrs, registered_lens = mock_batch_register.call_args[0]
649
+ assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()]
650
+ assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes]
651
+ assert worker.block_len_per_layer == [
652
+ mla_cache.nbytes // mla_cache.shape[0],
653
+ eagle_cache.nbytes // eagle_cache.shape[0],
654
+ ]
655
+
656
+
657
@pytest.mark.asyncio
658
@patch(
659
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+20
-3
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
23
EngineId,
24
TpKVTopology,
25
get_current_attn_backend,
26
+ get_current_attn_backends,
27
)
28
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
29
KVConnectorBase_V1,
@@ -47,6 +48,7 @@ from vllm.v1.attention.backend import AttentionMetadata
48
from vllm.v1.attention.backends.utils import get_kv_cache_layout
49
from vllm.v1.core.sched.output import SchedulerOutput
50
from vllm.v1.request import RequestStatus
51
+from vllm.v1.worker.utils import select_common_block_size
52
53
logger = init_logger(__name__)
54
@@ -751,6 +753,7 @@ class MooncakeConnectorWorker:
753
self.model_config = vllm_config.model_config
754
self.cache_config = vllm_config.cache_config
755
self.use_mla = self.model_config.use_mla
756
+ self._sync_block_size_with_kernel()
757
758
# Get the attention backend from the first layer
759
# NOTE (NickLucche) models with multiple backends are not supported yet
@@ -777,6 +780,23 @@ class MooncakeConnectorWorker:
780
self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata)
781
self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse)
782
783
+ def _sync_block_size_with_kernel(self) -> None:
784
+ # When speculative decoding (e.g. Eagle) is enabled, the main model
785
+ # and draft model may use different attention backends with different
786
+ # physical block sizes. Pick the common (smallest) block size so that
787
+ # KV-cache registration and transfer work correctly for both models.
788
+ backends = get_current_attn_backends(self.vllm_config)
789
+ kernel_block_size = select_common_block_size(self.block_size, backends)
790
+ if self.block_size != kernel_block_size:
791
+ logger.info_once(
792
+ "User-specified logical block size (%s) does not match"
793
+ " physical kernel block size (%s). Using the latter.",
794
+ self.block_size,
795
+ kernel_block_size,
796
+ )
797
+ assert self.block_size > kernel_block_size
798
+ self.block_size = kernel_block_size
799
+
800
def __del__(self):
801
self.shutdown()
802
@@ -1268,9 +1288,6 @@ class MooncakeConnectorWorker:
1288
self.block_len_per_layer.append(
1289
curr_tensor_size_bytes // self.num_blocks
1290
)
1271
-
1272
- kernel_block_size = cache.shape[-2 if self.use_mla else -3]
1273
- assert self.block_size == kernel_block_size
1291
kv_data_ptrs.append(base_addr)
1292
kv_data_lens.append(curr_tensor_size_bytes)
1293