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

[Bugfix][ROCm][Disagg] Allow MoRIIO HMA per-layer block_len (#58585)

Browse Source 
Signed-off-by: avininjamay8 <Avinash.Paul@amd.com>
Signed-off-by: avininjamay8 <avpaul@amd.com>
Co-authored-by: avininjamay8 <avpaul@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>
This commit is contained in:
Avinash Paul
2026-10-15 09:35:31 +08:00
committed by GitHub
co-authored by avininjamay8 Cursor Shanshan Shen
parent de2b869ee9
commit 429a3da665
2 changed files with 141 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
+111
View File
@@ -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
# --------------------------------------------------------------------------
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+30 -9
View File
@@ -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, "