[KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models (#49612)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs)
         self._hand_shake_latency = hand_shake_latency
         self.kv_cache_layout = kv_cache_layout
-        # Mock register_kv_caches attribute needed for tests that do not call it.
+        # Mock register_kv_caches attributes needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
+        self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
         test_shape = self.attn_backends[0].get_kv_cache_shape(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
@@ -765,8 +766,9 @@ class TestNixlHandshake:
             assert remote_info.remote_tp_size == remote_tp_size
             assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
             # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
-            assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio
-            assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio
+            split_key = (-tp_ratio, worker.block_size)
+            assert split_key in worker.src_xfer_handles_by_tp_ratio
+            assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio
             assert remote_engine_id in worker.dst_xfer_side_handles
             assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set(
                 range(tp_ratio)
@@ -2091,7 +2093,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
         # Mock register_kv_cache which registers local handle
         worker.src_xfer_handles_by_block_size = {worker.block_size: 455}
         # P TP = 2 * D TP case, we should register 2 local handles
-        worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]}
+        worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]}
         worker.dst_xfer_side_handles = {"engine1": {0: 789}}
         worker._remote_agents = {"engine1": {(0, 0): "agent1"}}
         # _cleanup_remote_engine (called by shutdown) also clears these:

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch():
     assert list(result) == expected, f"Expected {expected}, got {list(result)}"
 
 
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_hetero_block_size_hybrid():
+    """With a block-size ratio, FA desc ids are ratio-expanded while SSM
+    desc ids keep the unexpanded logical stride (state blocks are never
+    sub-split)."""
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    worker = _make_mock_worker_for_desc_ids(
+        num_regions=2,
+        has_mamba=True,
+        group_spec_types=(FullAttentionSpec, MambaSpec),
+        block_len_per_layer=[100],
+    )
+
+    ratio = 4
+    # FA ids are already remote-granularity (expanded) sub-block ids.
+    fa_sub_blocks = [3, 5]
+    ssm_blocks = [1]
+    result = worker._compute_desc_ids(
+        block_ids=(fa_sub_blocks, ssm_blocks),
+        dst_num_blocks=100,
+        block_size_ratio=ratio,
+        physical_blocks_per_logical=1,
+    )
+
+    # FA regions have 100*4 entries each; SSM regions (4 per layer) start at
+    # 2*400 and stride by the unexpanded 100 logical blocks.
+    expected = [3, 5, 403, 405, 801, 901, 1001, 1101]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+def _bind_worker_method(worker, name):
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+
+    method = getattr(NixlConnectorWorker, name)
+    setattr(worker, name, method.__get__(worker, NixlConnectorWorker))
+
+
+@pytest.mark.cpu_test
+def test_map_block_ids_for_block_size_ratio_hybrid():
+    """Attention groups expand to remote granularity and clip to the remote
+    coverage; mamba state blocks pass through 1:1."""
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    worker = MagicMock(spec=NixlConnectorWorker)
+    worker._group_spec_types = (FullAttentionSpec, MambaSpec)
+    _bind_worker_method(worker, "get_mapped_blocks")
+    _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio")
+
+    local, remote = worker._map_block_ids_for_block_size_ratio(
+        [[1, 2, 3], [7]],
+        [list(range(30, 40)), [42]],
+        4,
+    )
+    # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks.
+    assert local == [list(range(4, 14)), [7]]
+    assert remote == [list(range(30, 40)), [42]]
+
+    # Attention-only full prefix hit: empty local list is preserved.
+    worker._group_spec_types = (FullAttentionSpec,)
+    local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4)
+    assert local == []
+
+
+@pytest.mark.cpu_test
+def test_post_process_zeroes_untransferred_tail():
+    """The untransferred sub-blocks of the last local block are zeroed on
+    receive; mamba state caches are untouched by the attention permute."""
+    from unittest.mock import MagicMock
+
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
+
+    ratio = 4
+    block_tokens = 8  # 2 tokens per remote sub-block
+
+    worker = MagicMock(spec=NixlConnectorWorker)
+    worker._group_spec_types = (FullAttentionSpec, MambaSpec)
+    worker.transfer_topo = MagicMock()
+    worker.device_type = "cpu"
+    worker.enable_permute_local_kv = False
+    attn_cache = torch.ones(6, block_tokens, 2, 4)
+    mamba_cache = torch.ones(6, 16)
+    worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
+    fa_group = MagicMock(layer_names=["attn.0"])
+    ssm_group = MagicMock(layer_names=["mamba.0"])
+    worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group])
+    # The cached property filters mamba layers out of the permuted caches.
+    attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
+    assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
+    worker._attention_kv_caches = attn_caches
+    _bind_worker_method(worker, "post_process_device_kv_on_receive")
+
+    # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received.
+    worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)])
+
+    # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens).
+    assert torch.all(attn_cache[2] == 1)
+    assert torch.all(attn_cache[3, :4] == 1)
+    assert torch.all(attn_cache[3, 4:] == 0)
+    # Untouched blocks and the mamba cache keep their content.
+    assert torch.all(attn_cache[4] == 1)
+    assert torch.all(mamba_cache == 1)
+
+
 @pytest.mark.cpu_test
 def test_nixl_metadata_hybrid_ssm_block_ids():
     """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -0,0 +1,619 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models
+under heterogeneous P/D block geometry (TP-sharded KDA-style state, so
+the mamba-aligned logical block size differs between P and D while the
+kernel-granularity pages stay equal).
+
+The invariant under test: every LOCAL byte range a request's READ transfers
+into must lie within that request's own blocks. A violation means an
+incoming transfer can overwrite a co-resident request's KV or mamba state
+mid-decode (silent corruption of an unrelated request).
+"""
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+
+from .utils import create_vllm_config
+
+
+class _RecordingNixl:
+    """Minimal NIXL wrapper stand-in that records descriptor lists and
+    prepared transfers so tests can resolve desc ids to byte ranges."""
+
+    def __init__(self, *args, **kwargs):
+        self.dlists: dict[int, np.ndarray] = {}
+        self.xfers: list[tuple] = []
+        self._next_handle = 1
+
+    def get_reg_descs(self, caches_data, mem_type):
+        return caches_data
+
+    def register_memory(self, descs, backends=None):
+        pass
+
+    def deregister_memory(self, descs):
+        pass
+
+    def get_agent_metadata(self):
+        return b"agent-meta"
+
+    def get_xfer_descs(self, blocks_data, mem_type):
+        return blocks_data
+
+    def prep_xfer_dlist(self, agent, descs):
+        handle = self._next_handle
+        self._next_handle += 1
+        self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3)
+        return handle
+
+    def add_remote_agent(self, metadata):
+        return "remote-agent"
+
+    def make_prepped_xfer(
+        self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None
+    ):
+        handle = self._next_handle
+        self._next_handle += 1
+        self.xfers.append(
+            (
+                op,
+                local_handle,
+                np.asarray(local_ids),
+                remote_handle,
+                np.asarray(remote_ids),
+            )
+        )
+        return handle
+
+    def transfer(self, handle):
+        pass
+
+    def check_xfer_state(self, handle):
+        return "DONE"
+
+    def get_xfer_telemetry(self, handle):
+        from types import SimpleNamespace
+
+        return SimpleNamespace(
+            xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1
+        )
+
+    def release_xfer_handle(self, handle):
+        pass
+
+    def release_dlist_handle(self, handle):
+        pass
+
+    def send_notif(self, agent, notif_msg=None):
+        pass
+
+    def get_new_notifs(self):
+        return {}
+
+    def remove_remote_agent(self, agent):
+        pass
+
+
+def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
+    """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
+        base_worker as bw,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+        NixlConnectorWorker,
+    )
+    from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+    from vllm.v1.kv_cache_interface import (
+        KVCacheConfig,
+        KVCacheGroupSpec,
+        KVCacheTensor,
+        MambaSpec,
+        MLAAttentionSpec,
+    )
+
+    mla_spec = MLAAttentionSpec(
+        block_size=local_block_size,
+        num_kv_heads=1,
+        head_size=6,
+        dtype=torch.float16,
+    )
+    unified_page = mla_spec.page_size_bytes
+    kda_spec = MambaSpec(
+        block_size=local_block_size,
+        shapes=((8, 3), (1, 4, 4)),
+        dtypes=(torch.float16, torch.float32),
+        page_size_padded=unified_page,
+        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=num_logical_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_logical_blocks * unified_page,
+                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+            )
+            for i in range(2)
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
+            KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
+            KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
+        ],
+    )
+
+    vllm_config = create_vllm_config(block_size=local_block_size)
+    vllm_config.cache_config.enable_prefix_caching = False
+    # kv_buffer_device defaults to the *real* platform's device type, which on
+    # a CPU-only test host would make this a host-buffer worker: host xfer
+    # buffers are per-layer, so the HMA shared-tensor regions this test builds
+    # would not be deduplicated. Pin it to the faked device type.
+    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
+
+    from unittest.mock import MagicMock
+
+    fake_backend = MagicMock()
+    fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
+    fake_backend.get_name.return_value = "FLASHMLA"
+    fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
+    fake_platform = MagicMock()
+    fake_platform.device_type = "cuda"
+    fake_platform.get_nixl_memory_type.return_value = "VRAM"
+
+    from vllm.config import set_current_vllm_config
+
+    with (
+        patch.object(bw, "NixlWrapper", _RecordingNixl),
+        patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
+        patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
+        patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
+        patch.object(bw, "current_platform", fake_platform),
+        patch(
+            "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
+            return_value="DS",
+        ),
+        set_current_vllm_config(vllm_config),
+    ):
+        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
+        worker.use_mla = True
+
+        # Attention caches are kernel-block granular on dim 0, as the
+        # receive post-process assumes.
+        ppl = local_block_size // kernel_block_size
+        tensors = [
+            torch.zeros(
+                num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8
+            )
+            for _ in range(2)
+        ]
+        worker.register_kv_caches(
+            {
+                "kda_a.0": tensors[0],
+                "mla.0": tensors[0],
+                "kda_b.0": tensors[0],
+                "kda_a.1": tensors[1],
+                "mla.1": tensors[1],
+                "kda_b.1": tensors[1],
+            }
+        )
+    # Keep tensors alive alongside the worker; flat views for byte checks.
+    worker._test_tensors = [t.view(-1) for t in tensors]
+    worker._test_tensors_2d = tensors
+    worker._test_unified_page = unified_page
+    return worker
+
+
+def _make_remote_meta(
+    worker,
+    remote_block_size,
+    remote_kernel_block_size,
+    remote_num_logical,
+    remote_ssm_sizes,
+):
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlAgentMetadata,
+    )
+
+    remote_ppl = remote_block_size // remote_kernel_block_size
+    # Kernel-granularity pages are TP-independent for MLA hybrids and must
+    # match the local ones for the handshake to pass, scaled down by the
+    # block-size ratio when the remote's kernel block is smaller.
+    block_size_ratio = worker.block_size // remote_kernel_block_size
+    kernel_page = worker.block_len_per_layer[0] // block_size_ratio
+    return NixlAgentMetadata(
+        engine_id="remote-engine",
+        agent_metadata=b"remote-agent-meta",
+        device_id=0,
+        kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
+        num_blocks=remote_num_logical * remote_ppl,
+        block_lens=[kernel_page, kernel_page],
+        kv_cache_layout=worker.kv_cache_layout,
+        block_size=remote_kernel_block_size,
+        ssm_sizes=remote_ssm_sizes,
+        attn_backend_name=worker.backend_name,
+        physical_blocks_per_logical_kv_block=remote_ppl,
+    )
+
+
+def _owned_byte_ranges(worker, group_logical_ids):
+    """Byte ranges owned by a request: for each HMA region tensor, every
+    logical block id of every group maps to one unified page."""
+    unified_page = worker._test_unified_page
+    bases = [t.data_ptr() for t in worker._test_tensors]
+    owned = []
+    for base in bases:
+        for ids in group_logical_ids:
+            for b in ids:
+                owned.append((base + b * unified_page, base + (b + 1) * unified_page))
+    return owned
+
+
+def _assert_local_writes_within(worker, owned_ranges):
+    nixl = worker.nixl_wrapper
+    assert nixl.xfers, "no transfers were posted"
+    violations = []
+    total_descs = 0
+    for op, local_handle, local_ids, _, remote_ids in nixl.xfers:
+        assert len(local_ids) == len(remote_ids)
+        desc_arr = nixl.dlists[local_handle]
+        for i in local_ids:
+            addr, length, _dev = desc_arr[int(i)]
+            addr, length = int(addr), int(length)
+            total_descs += 1
+            if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges):
+                violations.append((int(i), hex(addr), length))
+    assert not violations, (
+        f"{len(violations)}/{total_descs} local descriptors write outside "
+        f"the request's own blocks: {violations[:10]}"
+    )
+    return total_descs
+
+
+@pytest.mark.cpu_test
+def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
+    """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens
+    (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2),
+    equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and
+    TP-sharded KDA state. Every local descriptor of the request's reads must
+    stay within its own blocks."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+
+    worker = _make_mla_hybrid_worker(
+        local_block_size=12, kernel_block_size=4, num_logical_blocks=8
+    )
+    assert worker._physical_blocks_per_logical_kv_block == 3
+
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=8,
+        remote_kernel_block_size=4,
+        remote_num_logical=12,
+        remote_ssm_sizes=(24, 32),
+    )
+    for rank in (0, 1):
+        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
+
+    # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
+    # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
+    # Sparse, non-contiguous ids so neighbor blocks exist on all sides.
+    local_ids = ([2, 5], [1], [7])
+    remote_ids = [[1, 4], [5], [2]]
+
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="req-b",
+        local_block_ids=local_ids,
+        kv_transfer_params={
+            "remote_block_ids": remote_ids,
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-req-b",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": 2,
+        },
+    )
+    meta = metadata.reqs_to_recv["req-b"]
+    meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
+        meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
+    )
+    worker._recving_metadata["req-b"] = meta
+
+    worker._read_blocks_for_req("req-b", meta)
+
+    owned = _owned_byte_ranges(worker, local_ids)
+    total = _assert_local_writes_within(worker, owned)
+    # Multi-read: rank 0 carries the replicated MLA + its SSM shard,
+    # rank 1 carries only its SSM shard.
+    assert len(worker.nixl_wrapper.xfers) == 2
+    assert total > 0
+
+
+def _resolve(
+    desc_arr,
+    idx,
+    bases,
+    region_size,
+    unified_page,
+    desc_page,
+    logical_ids_attn,
+    block_tokens,
+):
+    """Resolve a desc id to (region, kind, token_start) where kind is 'attn'
+    (desc-page sized, sub-block-aligned, in the request's attention blocks)
+    or 'mamba'. token_start is the request-relative token offset, so local
+    and remote are comparable even when their kernel blocks differ in size."""
+    addr, length, _ = (int(x) for x in desc_arr[int(idx)])
+    for region, base in enumerate(bases):
+        off = addr - base
+        if 0 <= off < region_size:
+            b = off // unified_page
+            rem = off % unified_page
+            if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn:
+                pos = logical_ids_attn.index(b)
+                tokens_per_desc = block_tokens * desc_page // unified_page
+                sub = rem // desc_page
+                return (region, "attn", pos * block_tokens + sub * tokens_per_desc)
+            return (region, "mamba", None)
+    raise AssertionError(f"desc {idx} addr {addr:#x} not in any region")
+
+
+def _run_hetero_case(
+    local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None
+):
+    """Full pull-path run for one geometry; returns pairing records.
+
+    ``remote_kernel`` defaults to the local kernel block size; a smaller
+    value additionally exercises block_size_ratio > 1.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+        NixlConnectorMetadata,
+    )
+
+    remote_kernel = remote_kernel or kernel
+    block_size_ratio = kernel // remote_kernel
+    remote_ppl = remote_block // remote_kernel
+    matched = num_tokens - 1  # mamba N-1 rule
+    n_local = -(-num_tokens // local_block)
+    n_remote = -(-matched // remote_block)
+
+    worker = _make_mla_hybrid_worker(
+        local_block_size=local_block,
+        kernel_block_size=kernel,
+        num_logical_blocks=max(2 * n_local + 4, 8),
+    )
+    # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size
+    # shards of each.
+    meta_r = _make_remote_meta(
+        worker,
+        remote_block_size=remote_block,
+        remote_kernel_block_size=remote_kernel,
+        remote_num_logical=max(2 * n_remote + 4, 8),
+        remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
+    )
+    for rank in range(tp_size):
+        worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
+
+    # Sparse ids so neighbors exist between the request's blocks.
+    local_attn = [2 * i + 1 for i in range(n_local)]
+    remote_attn = [2 * i + 2 for i in range(n_remote)]
+    local_ids = (local_attn, [0], [2 * n_local + 2])
+    remote_ids = [remote_attn, [1], [0]]
+
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id="req-b",
+        local_block_ids=local_ids,
+        kv_transfer_params={
+            "remote_block_ids": remote_ids,
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-req-b",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": tp_size,
+        },
+    )
+    meta = metadata.reqs_to_recv["req-b"]
+    meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
+        meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
+    )
+    worker._recving_metadata["req-b"] = meta
+
+    # Sentinel-fill the local KV so untouched bytes are detectable.
+    for t in worker._test_tensors:
+        t.fill_(0xAA)
+
+    worker._read_blocks_for_req("req-b", meta)
+
+    # Invariant 1: all local writes within the request's own blocks.
+    owned = _owned_byte_ranges(worker, local_ids)
+    _assert_local_writes_within(worker, owned)
+
+    # Invariant 2: local<->remote attention pairs are token-aligned.
+    nixl = worker.nixl_wrapper
+    local_bases = [t.data_ptr() for t in worker._test_tensors]
+    remote_bases = [0x10_000_000, 0x20_000_000]
+    local_unified = worker._test_unified_page
+    remote_unified = (local_unified // local_block) * remote_block
+    # With block_size_ratio > 1 the local page is split into ratio sub-descs,
+    # each the size of a whole remote kernel page.
+    desc_page = worker.block_len_per_layer[0] // block_size_ratio
+    meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified
+    covered_tokens = set()
+    for op, lh, lids, rh, rids in nixl.xfers:
+        larr, rarr = nixl.dlists[lh], nixl.dlists[rh]
+        for li, ri in zip(lids, rids):
+            lreg, lkind, ltok = _resolve(
+                larr,
+                li,
+                local_bases,
+                len(worker._test_tensors[0]),
+                local_unified,
+                desc_page,
+                local_attn,
+                local_block,
+            )
+            rreg, rkind, rtok = _resolve(
+                rarr,
+                ri,
+                remote_bases,
+                meta_r_num_blocks_bytes,
+                remote_unified,
+                desc_page,
+                remote_attn,
+                remote_block,
+            )
+            assert lkind == rkind, (
+                f"pair kind mismatch: local {lkind} vs remote {rkind} "
+                f"(local desc {li}, remote desc {ri})"
+            )
+            assert lreg == rreg, (
+                f"region mismatch: local {lreg} vs remote {rreg} for "
+                f"tokens {ltok} vs {rtok}"
+            )
+            if lkind == "attn":
+                assert ltok == rtok, (
+                    f"TOKEN MISALIGNMENT: local sub-block holds tokens "
+                    f"[{ltok}..) but receives remote tokens [{rtok}..) "
+                    f"(geometry local_block={local_block}, "
+                    f"remote_block={remote_block}, N={num_tokens})"
+                )
+                covered_tokens.add(ltok)
+
+    # Invariant 3: full coverage of the matched tokens, at the finest
+    # transfer granularity (the remote kernel block).
+    needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)}
+    missing = needed - covered_tokens
+    assert not missing, (
+        f"tokens never transferred: {sorted(missing)[:8]} "
+        f"(geometry local_block={local_block}, remote_block={remote_block}, "
+        f"N={num_tokens}, matched={matched})"
+    )
+
+    # Invariant 4: no stale bytes after receive completion. The scheduler
+    # excludes the blocks covering the matched tokens from alloc-time KV
+    # zeroing (the zeroing would race the RDMA write), so every byte of
+    # those blocks must be either written by the transfer or zeroed by the
+    # receive post-process. Stale bytes surface as mid-response garbage
+    # once decode grows into the untransferred tail.
+    for op, lh, lids, rh, rids in nixl.xfers:
+        larr = nixl.dlists[lh]
+        for li in lids:
+            addr, length, _ = (int(x) for x in larr[int(li)])
+            for t in worker._test_tensors:
+                off = addr - t.data_ptr()
+                if 0 <= off < t.numel():
+                    t[off : off + length] = 0  # simulate the RDMA write
+                    break
+    done_sending, done_recving = worker.get_finished()
+    assert "req-b" in done_recving
+    n_excluded = -(-matched // local_block)
+    stale = []
+    for b in local_attn[:n_excluded]:
+        for region, t in enumerate(worker._test_tensors):
+            page = t[b * local_unified : (b + 1) * local_unified]
+            n_stale = int((page == 0xAA).sum())
+            if n_stale:
+                stale.append((region, b, n_stale))
+    assert not stale, (
+        f"stale (unzeroed, untransferred) bytes in matched-range attention "
+        f"blocks (region, block, bytes): {stale} "
+        f"(geometry local_block={local_block}, remote_block={remote_block}, "
+        f"N={num_tokens}, matched={matched})"
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "local_block,remote_block",
+    [
+        (12, 8),  # ppl 3 vs 2
+        (36, 8),  # ppl 9 vs 2 (large ppl asymmetry, scaled)
+        (24, 4),  # ppl 6 vs 1
+        (16, 24),  # remote larger than local (D_TP > P_TP direction)
+    ],
+)
+@pytest.mark.parametrize("num_tokens", list(range(2, 40)))
+def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens):
+    """Sweep prompt lengths across block-boundary residues for several
+    hetero-ppl geometries; assert neighbor-safety, token alignment, and
+    coverage of every transferred kernel block."""
+    _run_hetero_case(
+        local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "num_tokens",
+    # Residues around the remote kernel block (4), the local kernel block
+    # (8), the remote logical block (8) and the local logical block (24).
+    [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49],
+)
+def test_hetero_ppl_with_block_size_ratio(num_tokens):
+    """Both hetero regimes at once: kernel blocks differ (local 8 / remote
+    4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs
+    2). The transfer is clipped at remote sub-block granularity by the
+    pairing and front-trimmed by _apply_prefix_caching, so the
+    untransferred tail can span both a partial block and whole blocks —
+    the case each of the two former zeroing paths handled only half of."""
+    _run_hetero_case(
+        local_block=24,
+        kernel=8,
+        remote_block=8,
+        remote_kernel=4,
+        num_tokens=num_tokens,
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "num_tokens",
+    # Residues around every geometric boundary: kernel block (64), remote
+    # logical block (768), local logical block (5760), plus odd offsets.
+    [
+        2,
+        63,
+        64,
+        65,
+        127,
+        128,
+        300,
+        640,
+        767,
+        768,
+        769,
+        831,
+        832,
+        1000,
+        1535,
+        1536,
+        1537,
+        2303,
+        2304,
+        2305,
+        3001,
+        5759,
+        5760,
+        5761,
+        5824,
+        6528,
+        6529,
+    ],
+)
+def test_mla_hybrid_large_ppl_geometry(num_tokens):
+    """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode):
+    decode (local) logical block 5760 / kernel 64 (ppl=90), prefill
+    (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with
+    replicated MLA and 8-way TP-sharded KDA state."""
+    _run_hetero_case(
+        local_block=5760,
+        kernel=64,
+        remote_block=768,
+        num_tokens=num_tokens,
+        tp_size=8,
+    )

tests/v1/kv_connector/unit/test_tp_mapping.py

@@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles:
         # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0)
         # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0)
         assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)]
+
+    def test_hetero_block_size_splits(self):
+        """With a block-size ratio, single-source FA sub-block descs pass
+        through whole; SSM descs are unexpanded and split per source."""
+        plan = TPMapping(
+            source_ranks_per_group=((0,), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 0},
+            rank_offset_factor=0,
+        )
+
+        worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
+        # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded).
+        src_blocks_data = np.array(
+            [
+                (1000, 100, 0),
+                (1100, 100, 0),
+                (2000, 100, 0),
+                (2100, 100, 0),
+                (3000, 400, 0),
+            ],
+            dtype=np.uint64,
+        )
+
+        splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2))
+
+        assert len(splits) == 2
+        fa_passthrough = [
+            (1000, 100, 0),
+            (1100, 100, 0),
+            (2000, 100, 0),
+            (2100, 100, 0),
+        ]
+        assert splits[0] == fa_passthrough + [(3000, 200, 0)]
+        assert splits[1] == fa_passthrough + [(3200, 200, 0)]
+
+    def test_hetero_block_size_head_sharded_asserts(self):
+        """Head-sharded FA reads (multiple FA sources) are incompatible with
+        a block-size mismatch and must fail loudly."""
+        plan = TPMapping(
+            source_ranks_per_group=((0, 1), (0, 1)),
+            all_source_ranks=(0, 1),
+            rank_to_attention_slot={0: 0, 1: 1},
+            rank_offset_factor=0,
+        )
+
+        worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
+        src_blocks_data = np.array(
+            [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)],
+            dtype=np.uint64,
+        )
+
+        with pytest.raises(AssertionError, match="Head-sharded"):
+            list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2))

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -12,6 +12,7 @@ import uuid
 from collections import defaultdict
 from collections.abc import Iterator
 from concurrent.futures import Future, ThreadPoolExecutor
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
 import msgspec
@@ -67,6 +68,7 @@ from vllm.distributed.parallel_state import (
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.network_utils import make_zmq_path
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -122,9 +124,11 @@ class NixlBaseConnectorWorker:
             return (region_ids * num_blocks + block_arr).flatten()
 
         # Compute desc ids per group using the right stride: FA descs have
-        # num_blocks entries per region (kernel granularity), SSM descs have
-        # logical_blocks entries per region (no kernel splitting).
-        logical_blocks = num_blocks // physical_blocks_per_logical
+        # num_blocks entries per region (kernel granularity, expanded by
+        # block_size_ratio for heterogeneous block sizes), SSM descs have
+        # logical_blocks entries per region (no kernel splitting, and never
+        # ratio-expanded since state blocks are indivisible).
+        logical_blocks = dst_num_blocks // physical_blocks_per_logical
         all_descs: list[np.ndarray] = []
         for i, group in enumerate(block_ids):
             group_arr = np.asarray(group)
@@ -161,6 +165,7 @@ class NixlBaseConnectorWorker:
         plan: TPMapping,
         src_blocks_data: np.ndarray,
         num_fa_descs: int,
+        block_size_ratio: int = 1,
     ) -> Iterator[list[tuple[int, int, int]]]:
         """Build split handle data for P_TP > D_TP scenario.
 
@@ -187,6 +192,11 @@ class NixlBaseConnectorWorker:
 
         # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
         fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)
+
+        assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), (
+            "Head-sharded attention reads with P_TP > D_TP and heterogeneous "
+            "block sizes are not supported"
+        )
         src_blocks_list = src_blocks_data.tolist()
 
         for p_idx, p_rank in enumerate(plan.all_source_ranks):
@@ -442,9 +452,12 @@ class NixlBaseConnectorWorker:
 
         # nixl_prepped_dlist_handle.
         self.src_xfer_handles_by_block_size: dict[int, int] = {}
+        # Local descriptor arrays per remote block size (block_size_ratio>1),
+        # kept for building per-tp-ratio splits at the same granularity.
+        self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {}
         # Populated dynamically during handshake based on remote configuration.
-        # Keep track of regions at different tp_ratio values. tp_ratio->handles
-        self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {}
+        # Per-source split handles, keyed by (tp_ratio, remote_block_size).
+        self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
         # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
         self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
 
@@ -1258,11 +1271,7 @@ class NixlBaseConnectorWorker:
             agent_metadata_bytes=encoder.encode(agent_metadata),
         )
 
-    def _build_mamba_local(
-        self,
-        base_addresses: list[int],
-        block_size_ratio: int,
-    ) -> np.ndarray:
+    def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray:
         """Build desc regions (conv sub-projections + ssm) per layer for
         local mamba blocks with DS conv layout, as an Nx3 uint64 array.
 
@@ -1289,16 +1298,17 @@ class NixlBaseConnectorWorker:
            | Key N-1 | Val N-1 |         |Conv N-1|   SSM N-1 |
            +-------------------+         +--------------------+
            |1st_split-2nd_split|         |1st_split-2nd_split |
+
+        Mamba state blocks are indivisible (not token-extent data), so the
+        descriptors always use the local page geometry regardless of any
+        attention block-size ratio; their desc ids are likewise never
+        ratio-expanded (see _compute_desc_ids).
         """
-        assert block_size_ratio == 1, (
-            "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
-            f"Got block_size_ratio={block_size_ratio}."
-        )
         assert base_addresses, "Local KV cache base addresses must not be empty."
         assert self._conv_decomp is not None
         conv_offsets = self._conv_decomp.local_conv_offsets
         conv_size, ssm_size = self._mamba_ssm_size
-        num_blocks = self._logical_num_blocks * block_size_ratio
+        num_blocks = self._logical_num_blocks
         physical_per_logical = self._physical_blocks_per_logical_kv_block
         device_id = self.device_id
         block_arange = np.arange(num_blocks, dtype=np.uint64)
@@ -1307,9 +1317,7 @@ class NixlBaseConnectorWorker:
         for i, base_addr in enumerate(base_addresses):
             # Jump one page_size, but ssm page_size may be bigger when kernel
             # locks block size to a specific value (physical_per_logical scale).
-            page_stride = (
-                self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
-            )
+            page_stride = self.block_len_per_layer[i] * physical_per_logical
             blk_addrs = base_addr + block_arange * page_stride
             for off, sz in conv_offsets:
                 parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
@@ -1459,7 +1467,7 @@ class NixlBaseConnectorWorker:
             self.device_id,
         )
         if self._has_mamba:
-            assert self.num_descs == len(blocks_data)
+            assert self.num_descs * block_size_ratio == len(blocks_data)
             # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
             # is unnecessary — a single conv desc per block suffices.  Consider
             # adding a fast path that falls back to the standard 2-region
@@ -1467,7 +1475,7 @@ class NixlBaseConnectorWorker:
             # remote has been seen.  Currently we always register 4 regions
             # because local descs are created before knowing the remote TP.
             logger.debug("Registering local Mamba descriptors (4 regions/layer)")
-            mamba = self._build_mamba_local(local_base_addresses, block_size_ratio)
+            mamba = self._build_mamba_local(local_base_addresses)
             blocks_data = np.concatenate([blocks_data, mamba])
 
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -1607,29 +1615,44 @@ class NixlBaseConnectorWorker:
 
         plan = self.tp_mappings[engine_id]
 
+        ### (Optional) Register a local handler at the remote engine's block
+        ### granularity (remote/prefill blocks smaller than local).
+        remote_block_size = nixl_agent_meta.block_size
+        src_blocks_data = self.src_blocks_data
+        if block_size_ratio > 1:
+            if remote_block_size not in self.src_xfer_handles_by_block_size:
+                handle, blocks_data = self.register_local_xfer_handler(
+                    remote_block_size
+                )
+                self.src_xfer_handles_by_block_size[remote_block_size] = handle
+                self.src_blocks_data_by_block_size[remote_block_size] = blocks_data
+            src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size]
+
         ### (Optional) Register local agent memory regions. MLA is not split.
+        split_key = (tp_ratio, remote_block_size)
         if (
             tp_ratio < 0
             and (not self.use_mla or len(plan.all_source_ranks) > 1)
-            and tp_ratio not in self.src_xfer_handles_by_tp_ratio
+            and split_key not in self.src_xfer_handles_by_tp_ratio
         ):
             # Remote tp_size > local tp_size: read from multiple remote ranks.
             # Logically "split" own regions into per-source chunks. Hybrid
             # MLA+SSM also needs this path: MLA is replicated and read once,
             # while the SSM state is sharded across every remote TP rank.
-            # We only do this once per remote tp_size (replica-friendly).
-            self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
+            # We only do this once per remote (tp_size, block_size).
+            self.src_xfer_handles_by_tp_ratio[split_key] = []
 
             for handle_data in self._build_local_splits_from_plan(
                 plan,
-                self.src_blocks_data,
-                self.num_descs,
+                src_blocks_data,
+                self.num_descs * block_size_ratio,
+                block_size_ratio,
             ):
                 descs = self.nixl_wrapper.get_xfer_descs(
                     handle_data, self.nixl_memory_type
                 )
                 handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
-                self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
+                self.src_xfer_handles_by_tp_ratio[split_key].append(handle)
 
         ### Register remote agent memory regions
         # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
@@ -1665,13 +1688,6 @@ class NixlBaseConnectorWorker:
             self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
         )
 
-        if block_size_ratio > 1:
-            # when prefill with smaller block_size, we need to init a
-            # new handler with same block_len to match
-            self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = (
-                self.register_local_xfer_handler(nixl_agent_meta.block_size)[0]
-            )
-
         return remote_agent_name
 
     def _validate_remote_agent_handshake(
@@ -1716,9 +1732,13 @@ class NixlBaseConnectorWorker:
                 "Disable prefix caching with --no-enable-prefix-caching."
             )
 
-        if self._is_hma_required:
-            assert block_size_ratio == 1, (
-                "HMA does not support different remote block size yet"
+        if block_size_ratio != 1:
+            # Heterogeneous block sizes transfer at remote-block granularity;
+            # the untransferred tail of the last local attention block is
+            # zeroed in the receive post-process, and mamba state pages
+            # transfer 1:1 (never sub-split).
+            assert not self.use_host_buffer, (
+                "Heterogeneous block sizes are not supported with host buffer"
             )
         kv_cache_layout = (
             self.kv_cache_layout
@@ -1875,27 +1895,57 @@ class NixlBaseConnectorWorker:
                     "d2h",
                 )
 
+    @cached_property
+    def _attention_kv_caches(self) -> list[torch.Tensor]:
+        """Device KV caches of attention layers (mamba states excluded),
+        as consumed by the receive post-process."""
+        assert self.device_kv_caches, (
+            "_attention_kv_caches accessed before register_kv_caches"
+        )
+        mamba_layers = {
+            name
+            for g, group in enumerate(self.kv_cache_config.kv_cache_groups)
+            if _is_ssm_spec(self._group_spec_types[g])
+            for name in group.layer_names
+        }
+        kv_caches = self.device_kv_caches
+        return [cache for name, cache in kv_caches.items() if name not in mamba_layers]
+
     def post_process_device_kv_on_receive(
         self,
         block_size_ratio: int,
-        block_ids_list: list[list[int]],
+        block_ids_list: list[tuple[list[int], int]],
+        convert: bool = True,
     ):
         """
         Post process device kv cache after receiving from remote.
 
-        3 types of post processing supported:
+        3 types of conversion supported (``convert``):
             * kv_cache_postprocess_layout => convert from HND to NHD
             * kv_cache_postprocess_blksize => convert from small block size
               to large block size
             * kv_cache_postprocess_blksize_and_layout => convert from small
               block size to large block size and convert from HND to NHD
 
+        The transfer only covers ``covered_sub_blocks`` remote-sized
+        sub-blocks of each request's local attention blocks; the rest was
+        clipped, either by remote-block pairing (block-size ratio) or by the
+        hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were
+        excluded from the scheduler's alloc-time KV zeroing (which would race
+        the RDMA write), so everything past the covered range is zeroed here.
+        Stale bytes would otherwise surface as garbage or NaNs once decode
+        grows into the untransferred tail.
         """
         if len(self.device_kv_caches) == 0:
             return
         assert block_size_ratio >= 1, "Only nP < nD supported currently."
         assert self.transfer_topo is not None
-        if self.enable_permute_local_kv and block_size_ratio > 1:
+        if not convert:
+            logger.debug(
+                "Post-processing device kv cache on receive by zeroing "
+                "untransferred blocks."
+            )
+        elif self.enable_permute_local_kv and block_size_ratio > 1:
             logger.debug(
                 "Post-processing device kv cache on receive by converting "
                 "block_size with %sx bigger and permuting layout from HND"
@@ -1914,18 +1964,45 @@ class NixlBaseConnectorWorker:
                 block_size_ratio,
             )
 
-        for block_ids in block_ids_list:
-            indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
+        attn_caches = self._attention_kv_caches
+        device = attn_caches[0].device
+        for block_ids, covered_sub_blocks in block_ids_list:
+            # Blocks the transfer didn't write: the token tail of the last
+            # partially covered block, then everything beyond it.
+            covered_blocks, sub_blocks_in_last = divmod(
+                covered_sub_blocks, block_size_ratio
+            )
+            first_stale = covered_blocks + (1 if sub_blocks_in_last else 0)
+            has_stale = first_stale < len(block_ids)
+            indices = None
+            if convert or has_stale:
+                indices = async_tensor_h2d(block_ids, device, torch.long)
 
-            for cache in self.device_kv_caches.values():
-                if self.enable_permute_local_kv and block_size_ratio > 1:
-                    kv_postprocess_blksize_and_layout_on_receive(
-                        cache, indices, block_size_ratio
-                    )
-                elif self.enable_permute_local_kv:
-                    kv_postprocess_layout_on_receive(cache, indices)
-                else:
-                    kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio)
+            if convert:
+                for cache in attn_caches:
+                    if self.enable_permute_local_kv and block_size_ratio > 1:
+                        kv_postprocess_blksize_and_layout_on_receive(
+                            cache, indices, block_size_ratio
+                        )
+                    elif self.enable_permute_local_kv:
+                        kv_postprocess_layout_on_receive(cache, indices)
+                    else:
+                        kv_postprocess_blksize_on_receive(
+                            cache, indices, block_size_ratio
+                        )
+
+            if sub_blocks_in_last:
+                last_block_id = block_ids[covered_blocks]
+                for cache in attn_caches:
+                    # Both post-processed layouts leave tokens on dim 1.
+                    sub_block_tokens = cache.shape[1] // block_size_ratio
+                    zero_from = sub_blocks_in_last * sub_block_tokens
+                    cache[last_block_id, zero_from:].zero_()
+            if has_stale:
+                assert indices is not None
+                stale_ids = indices[first_stale:]
+                for cache in attn_caches:
+                    cache.index_fill_(0, stale_ids, 0)
 
     def post_process_device_kv_on_receive_heterogeneous_attn(
         self, block_ids: list[int]
@@ -1995,18 +2072,33 @@ class NixlBaseConnectorWorker:
             if self.use_host_buffer:
                 self.sync_recved_kv_to_device(req_id, meta)
 
-            # post processing for heteroblocksize
+            # Post processing for heteroblocksize/layout, and for blocks the
+            # transfer clipped. The latter happens either at remote-block
+            # granularity (block_size_ratio > 1) or at kernel-block
+            # granularity, when equal kernel pages meet differing logical
+            # block sizes and _apply_prefix_caching front-trims to the
+            # minimum count (hybrid heterogeneous TP).
             remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
             block_size_ratio = self.transfer_topo.block_size_ratio(
                 remote_info.remote_block_size
             )
-            if not self.use_mla and (
-                block_size_ratio > 1 or self.enable_permute_local_kv
-            ):
-                assert not self._is_hma_required
-                block_ids_for_blocksize_post_process[block_size_ratio].append(
-                    meta.local_physical_block_ids[0]
-                )
+            hetero_ppl = (
+                remote_info.remote_physical_blocks_per_logical
+                != self._physical_blocks_per_logical_kv_block
+            )
+            if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl:
+                for g, local_group in enumerate(meta.local_physical_block_ids):
+                    if not local_group or _is_ssm_spec(self._group_spec_types[g]):
+                        continue
+                    # Number of remote-sized sub-blocks the transfer covered;
+                    # everything past this was clipped and must be zeroed.
+                    covered_sub_blocks = min(
+                        len(local_group) * block_size_ratio,
+                        len(meta.remote.block_ids[g]),
+                    )
+                    block_ids_for_blocksize_post_process[block_size_ratio].append(
+                        (local_group, covered_sub_blocks)
+                    )
             # post processing for heterogeneous attention
             if self.enable_heterogeneous_attn_post_process:
                 block_ids_for_heterogeneous_attn_post_process.append(
@@ -2016,7 +2108,14 @@ class NixlBaseConnectorWorker:
             block_size_ratio,
             block_ids_list,
         ) in block_ids_for_blocksize_post_process.items():
-            self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)
+            # MLA never needs the block-size/layout conversion, but its
+            # clipped blocks still need zeroing.
+            convert = not self.use_mla and (
+                block_size_ratio > 1 or self.enable_permute_local_kv
+            )
+            self.post_process_device_kv_on_receive(
+                block_size_ratio, block_ids_list, convert
+            )
 
         for block_ids in block_ids_for_heterogeneous_attn_post_process:
             self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
@@ -2206,6 +2305,45 @@ class NixlBaseConnectorWorker:
 
         return mapped_2d.flatten().astype(np.int64)
 
+    def _map_block_ids_for_block_size_ratio(
+        self,
+        local_block_ids: BlockIds,
+        remote_block_ids: BlockIds,
+        block_size_ratio: int,
+    ) -> tuple[BlockIds, BlockIds]:
+        """Map attention-group block ids to remote-block granularity.
+
+        Each local attention block is split into ``block_size_ratio``
+        sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the
+        remote list — the untransferred tail of the last local block — are
+        clipped here and zeroed in the receive post-process. Mamba state
+        blocks are indivisible and transfer 1:1, unexpanded.
+
+        ex: remote (prefill) block ids with block_size 4:
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        Local (decode) block ids with block_size 16: [1, 2, 3] expand to
+        [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote.
+        """
+        mapped_local: list[list[int]] = []
+        mapped_remote: list[list[int]] = []
+        for i, remote_group in enumerate(remote_block_ids):
+            local_group = local_block_ids[i] if local_block_ids else []
+            if _is_ssm_spec(self._group_spec_types[i]):
+                mapped_local.append(list(local_group))
+                mapped_remote.append(list(remote_group))
+                continue
+            mapped = self.get_mapped_blocks(
+                np.asarray(local_group), block_size_ratio
+            ).tolist()
+            if len(mapped) > len(remote_group):
+                mapped = mapped[: len(remote_group)]
+            mapped_local.append(mapped)
+            mapped_remote.append(list(remote_group))
+        if not any(mapped_local):
+            # Full prefix cache hit is indicated with an empty list.
+            return [], mapped_remote
+        return mapped_local, mapped_remote
+
     def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds:
         """
         Convert block ids to kernel physical block ids.
@@ -2300,13 +2438,18 @@ class NixlBaseConnectorWorker:
                     remote_block_ids[i] = remote_group[-num_local_blocks:]
                 else:
                     # TODO Handle prefix caching with different block_sizes
-                    max_padding = max(
-                        self._physical_blocks_per_logical_kv_block,
-                        remote_physical_per_logical,
+                    # Allocation rounding legitimately leaves up to
+                    # ppl - 1 trailing dead kernel blocks per side (plus one
+                    # extra local block for the recomputed final token), so
+                    # the counts may differ by up to the sum of the two
+                    # ratios; anything larger indicates mismatched lists.
+                    max_padding = (
+                        self._physical_blocks_per_logical_kv_block
+                        + remote_physical_per_logical
                     )
-                    assert abs(num_local_blocks - num_remote_blocks) < max_padding, (
+                    assert abs(num_local_blocks - num_remote_blocks) <= max_padding, (
                         f"Group {i}: |{num_local_blocks} - "
-                        f"{num_remote_blocks}| >= {max_padding}"
+                        f"{num_remote_blocks}| > {max_padding}"
                     )
                     num_blocks = min(num_local_blocks, num_remote_blocks)
                     local_block_ids[i] = local_block_ids[i][:num_blocks]

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

@@ -5,8 +5,6 @@
 import time
 from typing import TYPE_CHECKING
 
-import numpy as np
-
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
     NixlBaseConnectorWorker,
 )
@@ -178,10 +176,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
             )
             # Get side handles.
             if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1):
-                assert remote_block_size == self.block_size
                 # Remote tp_size > local tp_size: we must perform multiple
                 # reads. Get the memory chunk onto which we will write to.
-                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
+                split_key = (tp_ratio, remote_block_size)
+                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
             else:
                 # Single read from remote, we write to the whole memory region.
                 # Also handle remote block size different from local block size.
@@ -235,30 +233,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
             remote_info.remote_block_size
         )
         if block_size_ratio > 1:
-            # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
-            assert not self._is_hma_required
-            local_block_ids0 = local_block_ids[0] if local_block_ids else []
-            remote_block_ids0 = remote_block_ids[0]
-            local_block_ids_mapped = self.get_mapped_blocks(
-                np.asarray(local_block_ids0), block_size_ratio
-            ).tolist()
-            if len(local_block_ids_mapped) > len(remote_block_ids0):
-                # NOTE:
-                # get_mapped_blocks will always expand block_ids for n times.
-                # ex:
-                # prefill block_ids with block_size as 4:
-                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-                # Local decode block_ids with block_size as 16: [1, 2, 3]
-                # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
-                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
-                # Then we clip local to align with prefill
-                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
-                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
-                local_block_ids_mapped = local_block_ids_mapped[
-                    : len(remote_block_ids0)
-                ]
-            local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
-            remote_block_ids = [remote_block_ids0]
+            local_block_ids, remote_block_ids = (
+                self._map_block_ids_for_block_size_ratio(
+                    local_block_ids, remote_block_ids, block_size_ratio
+                )
+            )
         # NOTE(rob): having the staging blocks be on the READER side is
         # not going to work well (since we will have to call rearrange tensors).
         # after we detect the txn is complete (which means we cannot make the

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

@@ -39,7 +39,6 @@ from concurrent.futures import Future
 from typing import TYPE_CHECKING, Any
 
 import msgspec
-import numpy as np
 
 from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
@@ -553,8 +552,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
                 req_id,
             )
             if tp_ratio < 0 and not self.use_mla:
-                assert remote_block_size == self.block_size
-                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
+                split_key = (tp_ratio, remote_block_size)
+                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
             else:
                 local_xfer_side_handle = self.src_xfer_handles_by_block_size[
                     remote_block_size
@@ -606,18 +605,11 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
             remote_info.remote_block_size
         )
         if block_size_ratio > 1:
-            assert not self._is_hma_required
-            local_block_ids0 = local_block_ids[0] if local_block_ids else []
-            remote_block_ids0 = remote_block_ids[0]
-            local_block_ids_mapped = self.get_mapped_blocks(
-                np.asarray(local_block_ids0), block_size_ratio
-            ).tolist()
-            if len(local_block_ids_mapped) > len(remote_block_ids0):
-                local_block_ids_mapped = local_block_ids_mapped[
-                    : len(remote_block_ids0)
-                ]
-            local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
-            remote_block_ids = [remote_block_ids0]
+            local_block_ids, remote_block_ids = (
+                self._map_block_ids_for_block_size_ratio(
+                    local_block_ids, remote_block_ids, block_size_ratio
+                )
+            )
 
         notif_id = f"{remote_request_id}:{self.world_size}".encode()