diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index c6df14d7a0..6d8dc24ea2 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -757,6 +757,117 @@ def test_worker_rejects_incompatible_mamba_specs_before_registration(): worker.register_kv_caches({"kda.0": object(), "kda.1": object()}) +def _make_register_worker(caches, geoms, *, disable_hma: bool = True): + worker = _FakeWorker( + mode=MoRIIOMode.READ, + _transfer_layer_names=set(caches), + block_size=16, + engine_id="e0", + dst_num_blocks={}, + block_lens={}, + layer_name_to_local_kv_cache_metadata={}, + kv_caches_base_addr={}, + kv_layer_mr_offset={}, + local_kv_cache_size=[], + block_window_per_layer=[], + moriio_wrapper=SimpleNamespace( + register_local_tensor=lambda _t: b"m", + get_agent_metadata=lambda: b"a", + async_wait_reqid=lambda: None, + ), + world_size=1, + layer_to_spec={}, + backend_name="TRITON_ATTN", + side_channel_port=6301, + tp_rank=0, + dp_rank=0, + is_producer=False, + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="x")), + scheduler_config=SimpleNamespace( + disable_hybrid_kv_cache_manager=disable_hma + ), + ), + ) + worker._is_mamba_layer = lambda _n: False + worker._is_mla_cache_layer = lambda _n: False + worker._get_layer_transfer_geometry = lambda n, remote_num_blocks=None: geoms[n] + worker._iter_layer_registration_regions = lambda n: [(caches[n], caches[n].nbytes)] + worker._moriio_handshake_listener = ( + lambda _m, ready_event, *_a, **_k: ready_event.set() + ) + return worker + + +def test_register_kv_caches_allows_hybrid_divergent_block_len(): + """GLM-5.2 / MiniMax HMA: divergent per-layer block_len is allowed.""" + Geom = moriio_layout.LayerTransferGeometry + caches = { + "layer0": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + "layer1": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + } + # Same tokens-per-block (16); different bytes-per-block (2112 vs 9216). + geoms = { + "layer0": Geom(4, 16, 2112, 132, 2112, None, None, 1, 1, False), + "layer1": Geom(4, 16, 9216, 576, 9216, None, None, 1, 1, False), + } + worker = _make_register_worker(caches, geoms) + worker.register_kv_caches(caches) + + assert worker.block_lens == {"layer0": 2112, "layer1": 9216} + assert worker.block_len == 9216 + + +def test_register_kv_caches_rejects_divergent_block_size(): + """Without HMA, block_size (tokens/page) must stay uniform.""" + Geom = moriio_layout.LayerTransferGeometry + caches = { + "layer0": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + "layer1": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + } + geoms = { + "layer0": Geom(4, 16, 2112, 132, 2112, None, None, 1, 1, False), + "layer1": Geom(4, 32, 2112, 66, 2112, None, None, 1, 1, False), + } + worker = _make_register_worker(caches, geoms, disable_hma=True) + with pytest.raises(ValueError, match="block size mismatch"): + worker.register_kv_caches(caches) + + +def test_register_kv_caches_rejects_divergent_block_size_even_if_block_len_differs(): + """Without HMA, size+len both differing must still raise.""" + Geom = moriio_layout.LayerTransferGeometry + caches = { + "layer0": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + "layer1": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + } + # Previously buggy: size+len both differ → registration continued. + geoms = { + "layer0": Geom(4, 16, 2112, 132, 2112, None, None, 1, 1, False), + "layer1": Geom(4, 32, 9216, 288, 9216, None, None, 1, 1, False), + } + worker = _make_register_worker(caches, geoms, disable_hma=True) + with pytest.raises(ValueError, match="block size mismatch"): + worker.register_kv_caches(caches) + + +def test_register_kv_caches_allows_hma_divergent_block_size(): + """Gemma4 HMA (#53721): SWA vs full-attn may differ in tokens/page.""" + Geom = moriio_layout.LayerTransferGeometry + caches = { + "layer0": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + "layer1": torch.zeros(2, 4, 16, 1, 8, dtype=torch.float16), + } + geoms = { + "layer0": Geom(4, 16, 2112, 132, 2112, None, None, 1, 1, False), + "layer1": Geom(4, 32, 4224, 132, 4224, None, None, 1, 1, False), + } + worker = _make_register_worker(caches, geoms, disable_hma=False) + worker.register_kv_caches(caches) + assert worker.block_lens == {"layer0": 2112, "layer1": 4224} + assert worker.block_len == 4224 + + # -------------------------------------------------------------------------- # _truncate_mamba_request_for_prefill # -------------------------------------------------------------------------- diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index b81ace6f9f..fa75eaee3d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2453,10 +2453,15 @@ class MoRIIOConnectorWorker: caches_data = [] self.layer_base_addr_index = {} base_addr_idx = 0 - + # HMA (#53721 Gemma4 SWA, GLM-5.2 DSA, MiniMax): transfers already + # size RDMA from per-layer geometry, so allow divergent block_len + # (bytes) always, and divergent block_size (tokens/page) only when + # the hybrid KV cache manager is enabled. Without HMA, block_size + # must stay uniform for the peer page handshake. hma_enabled = ( not self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager ) + for layer_name in kv_caches: self.layer_base_addr_index[layer_name] = base_addr_idx if self._is_mamba_layer(layer_name): @@ -2471,16 +2476,31 @@ class MoRIIOConnectorWorker: base_addr_idx += 1 continue geometry = self._get_layer_transfer_geometry(layer_name) - if hma_enabled: - if geometry.block_len != self.block_len: + if geometry.block_size != self.block_size: + if not hma_enabled: raise ValueError( - "MoRIIO KV cache block length mismatch for layer " - f"{layer_name}: {geometry.block_len} != {self.block_len}" + "MoRIIO KV cache block size mismatch for layer " + f"{layer_name}: {geometry.block_size} != " + f"{self.block_size}" ) - elif geometry.block_size != self.block_size: - raise ValueError( - "MoRIIO KV cache block size mismatch for layer " - f"{layer_name}: {geometry.block_size} != {self.block_size}" + logger.debug( + "MoRIIO HMA per-layer block_size for %s: %d " + "(page/default %d); continuing with per-layer sizes.", + layer_name, + geometry.block_size, + self.block_size, + ) + if geometry.block_len != self.block_len: + # Unlike block_size above, divergent block_len is allowed + # unconditionally (not gated on hma_enabled) -- it covers + # non-hybrid models too, e.g. GLM-5.2 DSA's per-layer + # geometry. Don't attribute it to HMA in the log. + logger.debug( + "MoRIIO per-layer block_len for %s: %d " + "(page/default %d); continuing with per-layer sizes.", + layer_name, + geometry.block_len, + self.block_len, ) # num_blocks is advertised as a single scalar to the peer, so must it # be uniform @@ -2510,6 +2530,7 @@ class MoRIIOConnectorWorker: # and a run that has three should not look identical in the log. distinct_block_lens = sorted(set(self.block_lens.values())) if len(distinct_block_lens) > 1: + self.block_len = max(distinct_block_lens) logger.info( "MoRIIO registered %d attention layers with %d distinct block " "lengths (bytes): %s; transfer sizes and strides are per layer, "