[KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots (#57169)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

tests/kernels/attention/test_deepgemm_attention.py

@@ -313,3 +313,49 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
             ref_logits = ref_logits.masked_fill(~mask, 0)
             diff = calc_diff(logits, ref_logits)
             assert diff < 1e-3, f"{diff=}"
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only")
+@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
+)
+def test_deepgemm_paged_mqa_packed_manager_block_stride():
+    torch.manual_seed(0)
+    num_blocks, block_size, packed_rows = 4, 256, 320
+    num_heads, head_dim, row_bytes = 64, 128, 132
+    page_size, pages_per_block = 64, block_size // 64
+    dense_stride, packed_stride = pages_per_block, packed_rows // page_size
+    compact_pages = kv_cache_cast_to_fp8(
+        torch.randn(
+            num_blocks, block_size, 1, head_dim, device="cuda", dtype=torch.bfloat16
+        ).view(num_blocks * pages_per_block, page_size, 1, head_dim)
+    )
+    packed_backing = torch.zeros(
+        num_blocks,
+        packed_rows,
+        row_bytes,
+        device="cuda",
+        dtype=torch.uint8,
+    )
+    packed_manager = packed_backing[:, :block_size]
+    packed_manager.copy_(compact_pages.squeeze(2).view(num_blocks, block_size, -1))
+    packed_pages = packed_backing.view(-1, page_size, row_bytes).unsqueeze(2)
+    offsets = torch.arange(pages_per_block, device="cuda", dtype=torch.int32)
+    tables = [
+        (stride + offsets).unsqueeze(0) for stride in (dense_stride, packed_stride)
+    ]
+
+    q = torch.randn(1, 1, num_heads, head_dim, device="cuda", dtype=torch.bfloat16).to(
+        torch.float8_e4m3fn
+    )
+    weights = torch.randn(1, num_heads, device="cuda", dtype=torch.float32)
+    context_lens = torch.full((1, 1), block_size, device="cuda", dtype=torch.int32)
+    schedule = get_paged_mqa_logits_metadata(context_lens, page_size, get_num_sms())
+    logits = [
+        fp8_fp4_paged_mqa_logits(
+            (q, None), cache, weights, context_lens, table, schedule, block_size, False
+        )
+        for cache, table in zip((compact_pages, packed_pages), tables)
+    ]
+    torch.testing.assert_close(*logits, rtol=0, atol=0)

tests/kernels/attention/test_rocm_aiter_mla_sink.py

@@ -463,7 +463,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):
     set_random_seed(412)
     num_tokens, num_heads, block_size = 2, 8, 16
     q = torch.randn(num_tokens, num_heads, Q_HEAD_DIM, device="cuda").to(dtype)
-    kv = torch.randn(2 * block_size, 1, Q_HEAD_DIM, device="cuda").to(dtype)
+    kv = torch.randn(2, 2, block_size, Q_HEAD_DIM, device="cuda").to(dtype)[:, 1]
     selected = torch.tensor([[1, 7, 18], [0, 3, 20]], dtype=torch.int32, device="cuda")
     sinks = torch.linspace(-3.0, 5.0, num_heads, device="cuda")
     impl = object.__new__(ROCMAiterMLASparseImpl)
@@ -504,7 +504,7 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):
         SimpleNamespace(_q_scale=None, _k_scale=None),
     )
     references = [
-        _sink_reference(q[i], kv[:, 0][selected[i].long()], sinks, SM_SCALE)
+        _sink_reference(q[i], kv.flatten(0, 1)[selected[i].long()], sinks, SM_SCALE)
         for i in range(num_tokens)
     ]
     expected, _, value_scale, score_error = (

tests/model_executor/test_jit_warmup_triton_launcher.py

@@ -553,12 +553,10 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
     owner = ComputeSlotMappingKernel()
     compile_key = owner.CompileKey(
         kv_cache_block_size=16,
-        blocks_per_kv_block=1,
         total_cp_world_size=2,
         total_cp_rank=1,
         cp_kv_cache_interleave_size=1,
         block_table_stride=128,
-        block_size=16,
     )
     launches: list[tuple[Any, ...]] = []
 
@@ -574,7 +572,6 @@ def test_compute_slot_mapping_uses_named_launcher_inputs(monkeypatch) -> None:
     assert inputs["block_table_stride"] == 128
     assert kwargs == {
         "KV_CACHE_BLOCK_SIZE": 16,
-        "BLOCKS_PER_KV_BLOCK": 1,
         "TOTAL_CP_WORLD_SIZE": 2,
         "TOTAL_CP_RANK": 1,
         "CP_KV_CACHE_INTERLEAVE_SIZE": 1,

tests/v1/attention/test_dspark_noncausal_sparse_mla.py

@@ -425,7 +425,7 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i
         and kv_cache_dtype != "fp8_ds_mla"
     ):
         pytest.skip("FlashMLA Sparse fp8 only supports fp8_ds_mla kv-cache dtype")
-    if block_size not in backend_cls.get_supported_kernel_block_sizes():
+    if not backend_cls.supports_block_size(block_size):
         pytest.skip(
             f"{backend_cls.get_name()} does not support block_size={block_size}"
         )

tests/v1/attention/test_flashinfer_mla_sparse_sm90.py

@@ -29,7 +29,7 @@ HEAD = 512
 TOPK = 128  # triton convert requires width % 128 == 0
 
 
-def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_):
+def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **kw):
     out = torch.full_like(token_indices, -1)
     counts = torch.zeros(token_indices.shape[0], dtype=torch.int32)
     for t in range(token_indices.shape[0]):
@@ -41,7 +41,7 @@ def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_):
             blk = int(block_table[int(req_id[t]), pos // BLOCK_SIZE])
             if blk < 0:
                 continue
-            vals.append(blk * BLOCK_SIZE + pos % BLOCK_SIZE)
+            vals.append(blk * kw["BLOCK_STRIDE_ROWS"] + pos % BLOCK_SIZE)
         out[t, : len(vals)] = torch.tensor(vals, dtype=out.dtype)
         counts[t] = len(vals)
     return out, counts
@@ -133,11 +133,8 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     meta.state = state
     q_nope = torch.randn(rows, impl.num_heads, HEAD)
     q_rope = torch.randn(rows, impl.num_heads, qk_rope)
-    cache = torch.zeros(
-        8 * BLOCK_SIZE,
-        impl.head_size,
-        dtype=torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16,
-    )
+    dtype = torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16
+    cache = torch.zeros(8, 2, BLOCK_SIZE, impl.head_size, dtype=dtype)[:, 0]
 
     out, lse = impl.forward_mqa(
         (q_nope, q_rope), cache, meta, SimpleNamespace(_k_scale_float=0.5)
@@ -147,7 +144,10 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     # Reserved buffers carry this step's slots; lengths are NOT refreshed
     # here (the builder plans them host-side before capture/replay).
     ref_slots, ref_counts = ref_convert(
-        meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer
+        meta.req_id_per_token,
+        meta.block_table,
+        impl.topk_indices_buffer,
+        BLOCK_STRIDE_ROWS=2 * BLOCK_SIZE,
     )
     offset = 0
     for t in range(rows):
@@ -161,7 +161,7 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
     assert state.wrapper.run_args is not None
     q_pe, ckv, kpe, kwargs = state.wrapper.run_args[1:]
     assert q_pe.shape == (rows, impl.num_heads, qk_rope)
-    assert ckv.shape == (8 * BLOCK_SIZE, 1, HEAD)
+    assert ckv.shape == ((8 - 1) * 2 * BLOCK_SIZE + BLOCK_SIZE, 1, HEAD)
     assert kpe.shape[-1] == qk_rope
     if impl.use_fp8_kv_cache:
         assert kwargs["ckv_scale"] == 0.5 and kwargs["kpe_scale"] == 1.0

tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py

@@ -24,6 +24,21 @@ from vllm.platforms.interface import DeviceCapability
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     QUANTIZED_DS_MLA_CACHE_FORMATS,
     FlashMLASparseBackend,
+    _fp8_kv_pages,
+)
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    KVCacheTensor,
+    KVQuantMode,
+    MLAAttentionSpec,
+)
+from vllm.v1.worker.utils import (
+    AttentionGroup,
+    allocate_kv_cache,
+    prepare_kernel_block_sizes,
+    select_common_block_size,
 )
 
 SM90 = DeviceCapability(major=9, minor=0)
@@ -83,6 +98,88 @@ def test_supported_head_sizes_include_512():
     assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
 
 
+def test_flashmla_bf16_nope_accepts_packed_manager_blocks():
+    spec = MLAAttentionSpec(
+        block_size=1152,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.bfloat16,
+        cache_dtype_str="bfloat16",
+        block_stride_alignment=1024,
+    )
+    num_blocks = 2
+    layers = ["layer.0", "layer.1"]
+    page_size = spec.page_size_bytes
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_blocks * len(layers) * page_size,
+                layers=layers,
+                layer_stride=page_size,
+                block_stride=len(layers) * page_size,
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(layers, spec)],
+    )
+    groups = [[AttentionGroup(FlashMLASparseBackend, layers, spec, 0)]]
+
+    kernel_block_sizes = prepare_kernel_block_sizes(config, groups)
+    caches = allocate_kv_cache(
+        config,
+        torch.device("cpu"),
+        KVCacheLayout.BLHNC,
+        kernel_block_sizes,
+    )
+
+    assert kernel_block_sizes == [1152]
+    assert caches["layer.0"].shape == (num_blocks, 1, 1152, 512)
+
+
+@pytest.mark.parametrize("sm100", [False, True])
+def test_flashmla_quantized_cache_kernel_pages(monkeypatch, sm100):
+    """Quantized blocks are read as 64-token pages; SM100 only row-aligns them
+    (TMA) and keeps 64-token blocks."""
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability_family",
+        lambda family, device_id=0: sm100 and family == 100,
+    )
+    spec = MLAAttentionSpec(
+        block_size=1152,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.float8_e4m3fn,
+        cache_dtype_str="fp8_ds_mla",
+        kv_quant_mode=KVQuantMode.FP8_PER_TENSOR,
+        state_content_bytes=656,
+    )
+
+    kernel_block_size = 64 if sm100 else 1152
+    assert (
+        select_common_block_size(1152, [FlashMLASparseBackend], [spec])
+        == kernel_block_size
+    )
+
+
+def test_fp8_kv_pages_address_packed_blocks():
+    """Token t of block b is row b * block_stride_rows + t of the 64-row pages;
+    64-token blocks are passed through."""
+    row, block_size, stride_rows = 656, 128, 192
+    backing = torch.arange(3 * stride_rows).repeat_interleave(row)
+    cache = backing.view(3, stride_rows, row)[:, :block_size]
+
+    pages, block_stride_rows = _fp8_kv_pages(cache, block_size)
+    assert block_stride_rows == stride_rows
+    assert pages.shape[1:] == (64, row)
+    slot = 2 * block_stride_rows + 100
+    assert torch.equal(pages[slot // 64, slot % 64], cache[2, 100])
+
+    pages, block_stride_rows = _fp8_kv_pages(cache[:, :64], 64)
+    assert block_stride_rows is None
+    assert torch.equal(pages, cache[:, :64])
+
+
 def test_quantized_ds_mla_formats_are_the_envelope_set():
     assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS
 

tests/v1/attention/test_glm5next_indexer_block_sizes.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""GLM-5.3-Flash indexer kernel block sizes are DeepGEMM pool pages."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.mla import indexer
+from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend
+
+KPOOL = 4
+
+
+@pytest.mark.parametrize(
+    "rocm,sm120,expected",
+    [
+        (False, False, [KPOOL * 32, KPOOL * 64]),
+        # SM120 DeepGEMM takes only 64-entry pages for the fp8 cache.
+        (False, True, [KPOOL * 64]),
+        # ROCm's AITER paged-MQA logits take the same pool pages.
+        (True, False, [KPOOL * 32, KPOOL * 64]),
+    ],
+)
+def test_glm5next_indexer_kernel_block_sizes(monkeypatch, rocm, sm120, expected):
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: rocm)
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability_family",
+        lambda family: sm120 and family == 120,
+    )
+    model_config = SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=KPOOL))
+    config = SimpleNamespace(model_config=model_config)
+    monkeypatch.setattr(indexer, "get_current_vllm_config", lambda: config)
+
+    spec = SimpleNamespace(tokens_per_state=KPOOL)
+    assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes(spec) == expected
+    # Platform block-size selection queries without a spec.
+    assert Glm5NextIndexerBackend.get_supported_kernel_block_sizes() == expected

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -122,7 +122,6 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe():
     builder = object.__new__(DeepseekV32IndexerMetadataBuilder)
     builder.dcp_world_size = 1
     builder.compress_ratio = 16
-    builder.kernel_block_size = None
     builder.num_sms = num_sms
     builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device)
     builder.kv_cache_spec = SimpleNamespace(num_states=64)
@@ -170,30 +169,24 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe():
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-@pytest.mark.parametrize(
-    "block_size,kernel_block_size,compress_ratio", [(64, 64, 1), (256, 64, 4)]
-)
-def test_indexer_draft_decode_refresh_follows_new_batch(
-    monkeypatch, block_size, kernel_block_size, compress_ratio
-):
+@pytest.mark.parametrize("compress_ratio", [1, 4])
+def test_indexer_draft_decode_refresh_follows_new_batch(monkeypatch, compress_ratio):
     """Without a prior build(), the refresh must not keep the last batch's state."""
     monkeypatch.setattr(indexer_module, "has_deep_gemm", lambda: False)
     device = torch.device("cuda")
     num_reqs, width = 3, 8
-    factor = block_size // kernel_block_size
 
     builder = object.__new__(DeepseekV32IndexerMetadataBuilder)
     builder.dcp_world_size = 1
     builder.compress_ratio = compress_ratio
-    builder.kernel_block_size = kernel_block_size
     builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device)
-    builder.kv_cache_spec = SimpleNamespace(block_size=block_size, num_states=64)
+    builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64)
 
     def stale(*shape: int) -> torch.Tensor:
         return torch.full(shape, -7, dtype=torch.int32, device=device)
 
     decode = DeepSeekV32IndexerDecodeMetadata(
-        block_table=stale(num_reqs, width // factor),
+        block_table=stale(num_reqs, width),
         seq_lens=stale(num_reqs, 1),
         decode_lens=stale(num_reqs),
         requires_padding=False,
@@ -219,8 +212,7 @@ def test_indexer_draft_decode_refresh_follows_new_batch(
 
     builder.update_draft_decode_metadata(metadata)
 
-    expected_block_table = request_block_table[:, ::factor] // factor
-    torch.testing.assert_close(decode.block_table, expected_block_table)
+    torch.testing.assert_close(decode.block_table, request_block_table)
     torch.testing.assert_close(decode.indices, builder.arange_buffer[:num_reqs])
     assert torch.all(decode.per_req_decode_lens == 1)
     assert torch.all(decode.decode_lens == 1)
@@ -471,7 +463,6 @@ def test_zero_token_pcp_rank_participates_in_compressed_mapping_gather(monkeypat
     builder.pcp_rank = 0
     builder.dcp_world_size = 1
     builder.cp_kv_cache_interleave_size = 1
-    builder.kernel_block_size = None
     builder.kv_cache_spec = SimpleNamespace(block_size=64, num_states=64)
     builder.compressed_slot_mapping_buffer = torch.zeros(8, dtype=torch.int64)
 

tests/v1/attention/test_kpool_indexer_block_sizes.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CPU tests for storage-block kernel-block selection in prepare_kernel_block_sizes."""
+"""CPU tests for the kpool indexer's kernel block selection."""
 
 from types import SimpleNamespace
 
@@ -11,20 +11,14 @@ from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.models.glm5next.common.attention import Glm5NextIndexerCache
 from vllm.v1.attention.backend import AttentionBackend, MultipleOf
 from vllm.v1.attention.backends.mla import indexer as indexer_mod
-from vllm.v1.attention.backends.mla.indexer import (
-    DeepseekV32IndexerBackend,
-    KpoolTailBackend,
-)
-from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
-    ROCMAiterMLASparseBackend,
-)
+from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
-from vllm.v1.worker.utils import prepare_kernel_block_sizes
+from vllm.v1.worker.utils import prepare_kernel_block_sizes, select_common_block_size
 
 pytestmark = pytest.mark.cpu_test
 
 # ``index_kpool`` of zai-org/GLM-5.3-Flash: one indexer state pools 4
-# compressed states (hence storage_block_size = page * 4).
+# compressed states (hence a kernel block is page * 4 tokens).
 INDEX_KPOOL = 4
 
 # The hybrid KDA/mamba manager floors at TP8..TP1; 640 is the geometry of the
@@ -59,14 +53,17 @@ def _prepare(
     spec: MLAAttentionSpec, backends: list[type[AttentionBackend]]
 ) -> list[int]:
     kv_cache_config = SimpleNamespace(
-        kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)]
+        kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec, layer_names=[])],
+        kv_cache_tensors=[],
     )
-    attn_groups = [[SimpleNamespace(backend=backend) for backend in backends]]
+    attn_groups = [
+        [SimpleNamespace(backend=backend, kv_cache_spec=spec) for backend in backends]
+    ]
     return prepare_kernel_block_sizes(kv_cache_config, attn_groups)
 
 
 @pytest.mark.parametrize("manager_block", TP_FLOOR_MANAGER_BLOCKS)
-def test_prepare_uses_storage_block_for_the_kpool_group(
+def test_kpool_group_maps_pool_pages(
     monkeypatch: pytest.MonkeyPatch, manager_block: int
 ):
     """The bug: on ROCm, kpool groups must get 128/256 - not the 640-4352
@@ -78,54 +75,43 @@ def test_prepare_uses_storage_block_for_the_kpool_group(
     """
     _mock_rocm_platform(monkeypatch)
     spec = _kpool_storage_spec(manager_block)
-    expected_storage = 256 if manager_block % 256 == 0 else 128
-    assert spec.storage_block_size == expected_storage
+    expected = 256 if manager_block % 256 == 0 else 128
 
-    selected = _prepare(
-        spec,
-        [ROCMAiterMLASparseBackend, DeepseekV32IndexerBackend, KpoolTailBackend],
-    )
-    assert selected == [expected_storage]
+    # The kpool group's AttentionGroup selects its kernel block this way.
+    selected = [
+        select_common_block_size(manager_block, [Glm5NextIndexerBackend], [spec])
+    ]
+    assert selected == [expected]
     assert selected != [manager_block]
     # The hybrid block-table split stays integral.
     assert manager_block % selected[0] == 0
 
 
-def test_prepare_falls_back_when_storage_block_is_not_supported():
-    """No storage block, or a storage block the backends reject (e.g. the
-    CUDA indexer's exact ``[64]`` vs a 128-token storage block): selection is
-    `select_common_block_size`'s answer, unchanged from no-fix behavior.
-    """
+def test_prepare_uses_the_backend_vote():
+    """Unpacked groups take `select_common_block_size`'s answer."""
 
     class Fixed64Backend:
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return [64]
 
         @staticmethod
         def get_name() -> str:
             return "FIXED64"
 
-    # Storage block the group's backend does not accept -> backend vote (64).
-    spec = _kpool_storage_spec(640)
-    assert spec.storage_block_size == 128
-    assert _prepare(spec, [Fixed64Backend]) == [64]
-
-    # No storage block at all (any non-kpool MLA model) -> backend vote.
     plain_spec = MLAAttentionSpec(
         block_size=640,
         num_kv_heads=1,
         head_size=128,
         dtype=torch.bfloat16,
     )
-    assert plain_spec.storage_block_size is None
     assert _prepare(plain_spec, [Fixed64Backend]) == [64]
 
     # ...and with backends that accept the manager block, the manager wins,
     # exactly as on main.
     class AcceptAllBackend:
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return [MultipleOf(1)]
 
         @staticmethod

tests/v1/attention/test_kpool_tail_slot_mapping.py

@@ -3,7 +3,7 @@
 """CPU tests for the kpool tail slot mapping (no GPU required).
 
 The kpool tail cache is a 1-block-per-request circular ring addressed by
-``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block
+``pos % kpool`` (``CircularBufferSpec`` / ``CircularBufferManager``: exactly one block
 allocated per request, never grown, so only column 0 of its block table is
 ever written; the rest stays zero-initialized).
 
@@ -32,27 +32,25 @@ from vllm.v1.attention.backends.mla.indexer import (
     KpoolTailMetadataBuilder,
     compute_kpool_tail_slot_mapping,
 )
-from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides
+from vllm.v1.kv_cache_interface import CircularBufferSpec, compute_layout_strides
 from vllm.v1.kv_cache_layout import KVCacheLayout
-from vllm.v1.worker.block_table import get_block_table_width
 
 KPOOL = 4
 
 
 def test_tail_backend_layout_matches_kernel_pointer_arithmetic():
     (layout,) = KpoolTailBackend.supported_kv_cache_layouts()
-    spec = KpoolTailSpec(
+    spec = CircularBufferSpec(
         block_size=KPOOL,
         num_kv_heads=2,
         head_size=128,
         head_size_v=0,
         dtype=torch.bfloat16,
-        sliding_window=KPOOL,
     )
     strides = compute_layout_strides(spec, num_blocks=8, num_layers=3, layout=layout)
     _, _, head_stride, state_stride, content_stride = strides
 
-    assert layout is KVCacheLayout.LBHNC
+    assert layout is KVCacheLayout.BLHNC
     assert head_stride == KPOOL * 128 * torch.bfloat16.itemsize
     assert state_stride == 128 * torch.bfloat16.itemsize
     assert content_stride == 1
@@ -77,37 +75,14 @@ def test_tail_ring_divides_the_attention_block(num_speculative_tokens, ring):
         SimpleNamespace(num_speculative_tokens=num_speculative_tokens)
     )
 
-    assert spec.block_size == spec.sliding_window == ring
+    assert isinstance(spec, CircularBufferSpec) and spec.block_size == ring
     assert ring >= KPOOL + num_speculative_tokens
     assert 640 % ring == 0
 
 
-def test_tail_spec_opts_out_of_generic_slot_mapping():
-    """The tail row is one block wide (padded to the block-table alignment), so
-    the generic kernel's ``pos // kpool`` column index runs off the end of the
-    allocation for long prompts. The spec must opt out of it entirely."""
-    spec = KpoolTailSpec(
-        block_size=KPOOL,
-        num_kv_heads=2,
-        head_size=128,
-        head_size_v=0,
-        dtype=torch.bfloat16,
-        sliding_window=KPOOL,
-    )
-    max_len = 1 << 20
-    width = get_block_table_width(
-        spec.max_num_blocks_per_req(None, max_len),
-        spec.block_size,
-        token_alignment=spec.block_table_token_alignment,
-    )
-
-    assert width * KPOOL < max_len
-    assert spec.uses_slot_mapping is False
-
-
 def make_tail_block_table(own_blocks, width=64):
     """Tail-group block table as BlockTables produces it: column 0 holds the
-    request's single KpoolTailManager block, the remaining columns are never
+    request's single CircularBufferManager block, the remaining columns are never
     written and stay zero."""
     bt = torch.zeros(len(own_blocks), width, dtype=torch.int32)
     bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
@@ -447,6 +422,7 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len)
     tokens between the last request boundary and num_actual_tokens (mapped to
     the last request), untouched padding beyond num_actual, and PAD for
     requests on the null block."""
+    ring_size = 2 * KPOOL
     positions, qsl, slot_mapping, _, num_reqs = make_batch(
         per_req, padded_len=padded_len
     )
@@ -455,14 +431,17 @@ def test_triton_mapping_matches_cpu(per_req, own_blocks, num_actual, padded_len)
     slot_mapping = torch.arange(padded_len, dtype=torch.int64) + 1000
     bt = make_tail_block_table(own_blocks)
 
-    ref = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
-    got = circular_tail_slots(
+    ref = compute_kpool_tail_slot_mapping(
+        slot_mapping, bt, qsl, positions, num_actual, num_reqs, ring_size
+    )
+    got = compute_kpool_tail_slot_mapping(
         slot_mapping.cuda(),
         bt.cuda(),
         qsl.cuda().to(torch.int32),
         positions.cuda(),
         num_actual,
         num_reqs,
+        ring_size,
     )
     torch.testing.assert_close(got.cpu(), ref)
     for req, blk in enumerate(own_blocks):

tests/v1/attention/test_mm_prefix.py

@@ -805,7 +805,6 @@ def test_composite_shared_cache_across_image_and_causal_steps(
             )
             layer.mm_prefix_clamp_sliding_window = True
         builder = backend.get_builder_cls()(spec, ["composite_test"], cfg, DEVICE)
-        builder.set_kernel_block_size(block_size)
         # Match production, where unwritten page tails are zero-initialized.
         cache = torch.zeros(
             num_blocks,

tests/v1/attention/test_rocm_glm5next_sparse.py

@@ -11,6 +11,7 @@ from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
+    ROCMAiterMLASparseBackend,
     _use_rocm_sparse_triton,
     fit_kpool_indices_to_aiter,
 )
@@ -19,6 +20,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     _validate_dsv4_sparse_dims,
     _validate_sparse_dims,
 )
+from vllm.v1.kv_cache_interface import KVCacheLayout
 
 
 @triton.jit
@@ -104,6 +106,10 @@ def test_fit_kpool_indices_rejects_narrow_input():
         )
 
 
+def test_rocm_sparse_mla_supports_glm_packed_layout():
+    assert KVCacheLayout.BLHNC in ROCMAiterMLASparseBackend.supported_kv_cache_layouts()
+
+
 @pytest.mark.parametrize(
     (
         "kv_cache_dtype",

tests/v1/attention/test_sparse_mla_backends.py

@@ -140,7 +140,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
     metadata = SimpleNamespace(
         req_id_per_token=torch.zeros(1, dtype=torch.int32),
         block_table=torch.zeros((1, 1), dtype=torch.int32),
-        block_size=1,
+        block_size=32,
     )
     recorded_scale = None
 
@@ -178,7 +178,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
     )
     impl.forward_mqa(
         torch.zeros(1, 1, kv_lora_rank),
-        torch.zeros(1, 1, kv_lora_rank),
+        torch.zeros(1, 32, kv_lora_rank),
         metadata,
         SimpleNamespace(),
     )
@@ -437,6 +437,7 @@ def test_sparse_backend_decode_correctness(
     qk_rope_head_dim: int,
     v_head_dim: int,
     monkeypatch,
+    block_stride_rows: int | None = None,
 ):
     if qk_rope_head_dim == 0 and (
         backend_cls != FlashMLASparseBackend
@@ -478,8 +479,7 @@ def test_sparse_backend_decode_correctness(
         if device_capability is None or device_capability.major != 10:
             pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x")
 
-    supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()
-    if block_size not in supported_block_sizes:
+    if not backend_cls.supports_block_size(block_size):
         pytest.skip(
             f"{backend_cls.get_name()} does not support block_size={block_size}"
         )
@@ -719,6 +719,11 @@ def test_sparse_backend_decode_correctness(
         kv_cache_dtype=kv_cache_dtype,
         scale=kv_cache_scale,
     )
+    if block_stride_rows is not None:
+        pad = (0, 0, 0, block_stride_rows - block_size)
+        kv_cache = torch.nn.functional.pad(kv_cache, pad)[:, :, :block_size]
+        slot = int(common_attn_metadata.slot_mapping[-1])
+        query_row = kv_cache[slot // block_size, 0, slot % block_size].zero_()
 
     # The sparse builder clones the layer's dense-MHA prefill backend from
     # static_forward_context; register a mock layer carrying one.
@@ -814,6 +819,12 @@ def test_sparse_backend_decode_correctness(
             out_buffer,
         )
 
+    if block_stride_rows is not None and kv_cache_dtype == "auto":
+        expected_row = torch.cat((kv_c_vllm[-1], k_pe_vllm[-1, 0]))
+        torch.testing.assert_close(query_row, expected_row, rtol=0, atol=0)
+    elif block_stride_rows is not None:
+        assert query_row.any()
+
     assert backend_output.shape == sdpa_reference.shape
     assert backend_output.dtype == sdpa_reference.dtype
     assert torch.isfinite(backend_output).all()
@@ -828,6 +839,56 @@ def test_sparse_backend_decode_correctness(
         torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)
 
 
+@pytest.mark.parametrize(
+    ("backend_cls", "kv_cache_dtype", "block_size", "block_stride_rows", "head_dims"),
+    [
+        pytest.param(
+            FlashInferMLASparseTRTLLMBackend,
+            "auto",
+            256,
+            320,
+            (128, 128, 64, 128),
+            id="FlashInferTRTLLM",
+        ),
+        # GLM-5.3 on SM90: the fp8 kernel's 64-token pages in larger blocks.
+        pytest.param(
+            FlashMLASparseBackend,
+            "fp8_ds_mla",
+            128,
+            192,
+            (64, 256, 0, 256),
+            id="FlashMLA-fp8_ds_mla",
+        ),
+    ],
+)
+def test_sparse_mla_packed_stride(
+    default_vllm_config,
+    dist_init,
+    workspace_init,
+    monkeypatch,
+    backend_cls,
+    kv_cache_dtype,
+    block_size,
+    block_stride_rows,
+    head_dims,
+):
+    test_sparse_backend_decode_correctness(
+        default_vllm_config,
+        dist_init,
+        backend_cls,
+        "mixed_small",
+        kv_cache_dtype,
+        1,
+        block_size,
+        workspace_init,
+        1.0,
+        1.0,
+        *head_dims,
+        monkeypatch,
+        block_stride_rows=block_stride_rows,
+    )
+
+
 def _triton_convert_reference_impl(
     req_ids: torch.Tensor,
     block_table: torch.Tensor,

tests/v1/attention/test_sparse_mla_kv_cache_layout.py

@@ -12,6 +12,7 @@ from vllm.config import CacheConfig
 from vllm.model_executor.layers.attention.mla_attention import MLAAttention
 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
 from vllm.platforms import current_platform
+from vllm.v1.attention.backend import MultipleOf
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseTRTLLMBackend,
 )
@@ -25,7 +26,7 @@ from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
 from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
 from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs
 from vllm.v1.kv_cache_layout import KVCacheLayout
-from vllm.v1.worker.utils import allocate_kv_cache
+from vllm.v1.worker.utils import allocate_kv_cache, customize_attention_spec
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -48,10 +49,10 @@ pytestmark = pytest.mark.skip_global_cleanup
 @pytest.mark.parametrize(
     "backend,cache_dtype,sm100,stride_alignment",
     [
-        (FlashInferMLASparseTRTLLMBackend, "auto", True, 1152),
-        (FlashInferMLASparseTRTLLMBackend, "fp8", True, 576),
+        (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152),
+        (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576),
         (FlashMLASparseBackend, "auto", True, 1152),
-        (FlashMLASparseBackend, "fp8_ds_mla", False, None),
+        (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)),
         (FlashMLASparseBackend, "fp8_ds_mla", True, 656),
         (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352),
     ],
@@ -79,7 +80,9 @@ def test_allocation_and_warmup_follow_addressing_mode(
         non_causal_multi_token_decode=False,
     )
     layer._uses_flat_kv_cache = MethodType(MLAAttention._uses_flat_kv_cache, layer)
-    spec = MLAAttention.get_kv_cache_spec(layer, config)
+    spec = customize_attention_spec(
+        backend, MLAAttention.get_kv_cache_spec(layer, config)
+    )
     indexer = SimpleNamespace(
         cache_config=config.cache_config, head_dim=132, dtype=torch.uint8
     )
@@ -102,10 +105,8 @@ def test_allocation_and_warmup_follow_addressing_mode(
     MLAAttention.bind_kv_cache(layer, views["mla"])
     DeepseekV32IndexerCache.bind_kv_cache(indexer, views["indexer"])
     assert indexer.kv_cache.stride(0) == views["indexer"].stride(0)
-    if stride_alignment is not None:
-        assert (
-            views["mla"].stride(0) * views["mla"].element_size() % stride_alignment == 0
-        )
+    alignment = spec.get_block_stride_alignment()
+    assert views["mla"].stride(0) * views["mla"].element_size() % alignment == 0
     if not layer._uses_flat_kv_cache():
         register.assert_not_called()
     else:
@@ -113,7 +114,7 @@ def test_allocation_and_warmup_follow_addressing_mode(
         register.assert_called_once_with(config, block_stride_rows=stride)
         layer.kv_cache[1, 0].fill_(3)
         torch.testing.assert_close(rows[stride], layer.kv_cache[1, 0])
-    if stride_alignment is None:
+    if alignment == 1:
         assert cache.kv_cache_tensors[0].size == cache.num_blocks * sum(
             item.page_size_bytes for item in specs.values()
         )

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -26,8 +26,8 @@ from vllm.v1.core.kv_cache_utils import (
 )
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     MambaSpec,
@@ -2308,12 +2308,11 @@ def test_opted_out_scratch_group_keeps_partial_hash_hits():
             ),
             KVCacheGroupSpec(
                 ["tail"],
-                KpoolTailSpec(
+                CircularBufferSpec(
                     block_size=mamba_block_size,
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.float32,
-                    sliding_window=mamba_block_size,
                 ),
             ),
         ],
@@ -2366,12 +2365,11 @@ def test_kpool_tail_supports_128_token_partial_hash_hits():
             ),
             KVCacheGroupSpec(
                 ["tail"],
-                KpoolTailSpec(
+                CircularBufferSpec(
                     block_size=kpool_block_size,
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.float32,
-                    sliding_window=kpool_block_size,
                 ),
             ),
         ],

tests/v1/core/test_contiguous_kv_packing.py

@@ -88,7 +88,7 @@ def test_packed_alignment_preserves_hot_pages_and_generic_stride():
     hot = HiSparseHotSpec(block_size=16, page_size=5120, blocks_per_request=2)
     groups = [KVCacheGroupSpec(["attention"], spec), KVCacheGroupSpec(["hot"], hot)]
     alignment = lcm(1152, hot.page_size_bytes)
-    stride = _pool_bytes_per_block(groups)
+    stride = _pool_bytes_per_block(_mock_vllm_config("BLHNC"), groups)
     assert stride == alignment
     layout = KVCacheLayout.BLHNC
     tensors = _build_hisparse_kv_cache_tensors(groups, 3, 3 * stride, layout, stride)
@@ -254,7 +254,7 @@ class TestCSALinearGrouping:
     def test_every_group_fits_one_packed_block(self):
         config = _shared_layout_config()
         groups = get_kv_cache_groups(config, _make_csa_linear_specs())
-        bytes_per_block = _get_kv_cache_bytes_per_block(groups)
+        bytes_per_block = _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC)
 
         pages = _pages(groups)
         for group in groups:
@@ -418,7 +418,7 @@ class TestCSALinearGrouping:
         groups = _get_packed_kv_cache_groups(config, specs)
         gdn = [g for g in groups if g.layer_names[0].startswith("gdn.")]
 
-        assert _get_kv_cache_bytes_per_block(groups) == sum(
+        assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == sum(
             specs[name].page_size_bytes for name in specs if name.startswith("wide.")
         )
         # That block holds every GDN state at once, so the repeat pattern alone
@@ -465,7 +465,7 @@ class TestSlidingWindowBucketCap:
         pages = _pages(groups)
         main = next(g for g in groups if "layers.2.attn" in g.layer_names)
         main_bytes = sum(pages[n] for n in main.layer_names)
-        assert _get_kv_cache_bytes_per_block(groups) == main_bytes
+        assert _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC) == main_bytes
         swa_groups = [g for g in groups if g.layer_names[0].endswith(".swa")]
         per_group = main_bytes // pages["layers.0.swa"]
         assert len(swa_groups) == -(-43 // per_group)
@@ -479,9 +479,9 @@ class TestSlidingWindowBucketCap:
 class TestDensePacking:
     def test_bytes_per_block_is_largest_group(self):
         groups, g1, g2 = _mixed_page_groups()
-        assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block(
-            groups
-        )
+        assert _get_kv_cache_bytes_per_block(
+            groups, KVCacheLayout.BLHNC
+        ) == _expected_bytes_per_block(groups)
 
         config = get_kv_cache_config_from_groups(
             _mock_vllm_config("BLHNC"), groups, MEMORY
@@ -554,6 +554,20 @@ class TestDensePacking:
             assert idx_tensor.offset == 2 * _mla(512).page_size_bytes
             assert mla_tensor.block_stride == block_stride
 
+    def test_layer_outer_pads_only_unaligned_pages(self):
+        """Layer-outer block strides are pages, so only an unaligned page pads."""
+        mla_page = _mla(512).page_size_bytes
+        idx = replace(_mla(128), block_stride_alignment=12288)
+        groups = [_uniform_group({"mla.0": _mla(512), "idx.0": idx})]
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config("LBNHC"), groups, MEMORY
+        )
+        mla_tensor, idx_tensor = config.kv_cache_tensors
+        assert config.num_blocks == MEMORY // (mla_page + 12288)
+        assert (mla_tensor.block_stride, idx_tensor.block_stride) == (mla_page, 12288)
+        assert idx_tensor.offset == mla_page * config.num_blocks
+        _bind(config, "LBNHC")
+
     def test_overlaid_groups_alias_and_stay_isolated(self):
         groups, g1, g2 = _mixed_page_groups()
         # Overlay models resolve to a block-outer layout at backend selection (the
@@ -562,7 +576,9 @@ class TestDensePacking:
             _mock_vllm_config("BLNHC"), groups, MEMORY
         )
         assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups)
-        assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups)
+        assert _pool_bytes_per_block(
+            _mock_vllm_config("BLNHC"), groups
+        ) == _expected_bytes_per_block(groups)
 
         views = _bind(config, "BLNHC")
         assert set(views) == set(g1) | set(g2)
@@ -667,7 +683,10 @@ class TestCompressorRingGroup:
         assert len(ring_groups) == 1
         assert not ring_groups[0].kv_cache_spec.prefix_cacheable
         kv_cache_config = get_kv_cache_config_from_groups(
-            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+            config,
+            groups,
+            available_memory=64
+            * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC),
         )
         assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (128, 64)
         manager = KVCacheManager(
@@ -716,7 +735,10 @@ class TestSWABoundedReplayGrouping:
             for g in groups
         ) == [(False, 128)] * 4 + [(True, 0)]
         kv_cache_config = get_kv_cache_config_from_groups(
-            config, groups, available_memory=64 * _get_kv_cache_bytes_per_block(groups)
+            config,
+            groups,
+            available_memory=64
+            * _get_kv_cache_bytes_per_block(groups, KVCacheLayout.BLHNC),
         )
         manager = KVCacheManager(
             generate_scheduler_kv_cache_config([kv_cache_config]),

tests/v1/core/test_kv_cache_utils.py

@@ -10,7 +10,7 @@ from collections.abc import Callable
 from dataclasses import replace
 from pathlib import Path
 from types import SimpleNamespace
-from typing import Any, cast
+from typing import Any
 
 import pytest
 import torch
@@ -19,7 +19,6 @@ import vllm.v1.core.kv_cache_utils as kv_cache_utils
 import vllm.v1.hisparse.runtime as hisparse_runtime_module
 from vllm.config import (
     CacheConfig,
-    KVTransferConfig,
     ModelConfig,
     SchedulerConfig,
     VllmConfig,
@@ -69,7 +68,6 @@ from vllm.v1.kv_cache_interface import (
     HiddenStateCacheSpec,
     HiSparseHotSpec,
     HiSparseResidentSpec,
-    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheSpec,
@@ -1942,13 +1940,12 @@ def test_dcp_replicated_kpool_tail_keeps_block_interior_layout(monkeypatch):
     config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
     config.parallel_config.decode_context_parallel_size = 4
     config.cache_config.kv_cache_layout = None
-    tail = KpoolTailSpec(
+    tail = CircularBufferSpec(
         block_size=4,
         num_kv_heads=2,
         head_size=128,
         head_size_v=0,
         dtype=torch.bfloat16,
-        sliding_window=4,
         dcp_sharded=False,
     )
     layout = resolve_kv_cache_layout(config, [["LBHNC"]], [new_mla_spec(), tail])
@@ -2906,564 +2903,77 @@ def test_generate_scheduler_kv_cache_config():
     )
 
 
-def _glm5_like_kv_cache_spec(
-    mamba_spec_factory=new_mamba_spec,
-) -> tuple[dict[str, KVCacheSpec], list[str]]:
-    """(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba."""
+def test_packed_groups_glm5_like_hybrid():
+    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=8192))
+    vllm_config.cache_config.kv_cache_layout = "BLHNC"
+    kpool, block_size = 4, 1024
     kv_cache_spec: dict[str, KVCacheSpec] = {}
-    mamba_layers = []
-    for i in range(45):
-        if i % 4 == 3:
-            kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=576,
-                dtype=torch.bfloat16,
-            )
-            kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=132,
-                dtype=torch.uint8,
-                tokens_per_state=16,
-            )
-        else:
-            name = f"layers.{i}.linear_attn"
-            kv_cache_spec[name] = mamba_spec_factory()
-            mamba_layers.append(name)
-    return kv_cache_spec, mamba_layers
-
-
-def _glm5_like_kv_cache_spec_with_tail(
-    mamba_spec_factory=new_mamba_spec,
-) -> dict[str, KVCacheSpec]:
-    """Production kpool=4 proportions: (mamba*3, MLA + indexer + tail) * 11.
-
-    The tail's logical page (2 * kpool * 2*indexer_head_dim * bf16) must fit
-    inside the indexer page it parasitizes (block_size//kpool * 132 B), which
-    is why the fixture drops the tokens_per_state=16 shape used above.
-    """
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory)
-    for i in range(3, 45, 4):
-        kpool = 4
-        kv_cache_spec[f"layers.{i}.indexer"] = replace(
-            cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]),
-            tokens_per_state=kpool,
+    for i in range(11):
+        kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=576,
+            dtype=torch.bfloat16,
+            block_stride_alignment=32 * 576 * 2,
         )
-        kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec(
+        kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=132,
+            dtype=torch.uint8,
+            tokens_per_state=kpool,
+            block_stride_alignment=64 * 132,
+        )
+        kv_cache_spec[f"layers.{i}.tail"] = CircularBufferSpec(
             block_size=kpool,
             num_kv_heads=2,
             head_size=128,
             head_size_v=0,
             dtype=torch.bfloat16,
-            sliding_window=kpool,
         )
-    return kv_cache_spec
-
-
-def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]:
-    return {
-        layer_name: tensor
-        for tensor in kv_cache_config.kv_cache_tensors
-        for layer_name in tensor.layers
-    }
-
-
-def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int:
-    return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride
-
-
-def test_get_kv_cache_config_balanced_mamba_hybrid():
-    """Hybrid slot sharing: mamba layers co-own the MLA slot tensors."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
-
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    uniform_groups = [
-        group
-        for group in groups
-        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-    ]
-    mamba_groups = [
-        group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
-    ]
-
-    # Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer
-    # gets an MLA slot.
-    assert len(groups) == 5
-    assert len(uniform_groups) == 1
-    assert len(uniform_groups[0].layer_names) == 22
-    assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11
-    assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
-    for k, name in enumerate(mamba_layers):
-        assert name in mamba_groups[k % 4].layer_names
-    # Mamba pages are padded up to the MLA page.
-    for group in mamba_groups:
-        assert group.kv_cache_spec.page_size_padded == mla_page
-        assert group.kv_cache_spec.page_size_bytes == mla_page
-
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    assert bytes_per_block == 11 * mla_page + 11 * idx_page
-
-    # Every block id is charged the full per-block sum.
-    attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config)
-    mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
-    blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group
-    assert (
-        kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
-        == blocks_per_request * bytes_per_block
-    )
-
-    available_memory = bytes_per_block * 100 + 1
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, available_memory
-    )
-    assert kv_cache_config.num_blocks == 100
-
-    # Every logical layer has a view into one backing allocation. Mamba views
-    # alias their corresponding MLA slot by using the same byte offset.
-    assert len(kv_cache_config.kv_cache_tensors) == 56
-    assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
-    tensors = _tensor_by_layer(kv_cache_config)
-    mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)]
-    for i, mla_name in enumerate(mla_layer_names):
-        mla_tensor = tensors[mla_name]
-        assert mla_tensor.block_stride == mla_page
-        for group in mamba_groups:
-            if i < len(group.layer_names):
-                assert tensors[group.layer_names[i]].offset == mla_tensor.offset
-    for i in range(11):
-        assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page
-
-    total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors}))
-    assert total_allocated == bytes_per_block * 100
-    assert 0 <= available_memory - total_allocated < bytes_per_block
-
-    assert get_max_concurrency_for_kv_cache_config(
-        vllm_config, kv_cache_config
-    ) == pytest.approx(100 / blocks_per_request)
-
-
-def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor():
-    """The kpool tail parasitizes the indexer tensors instead of getting its
-    own: sibling idx/tail tensors paired by layer order, zero standalone tail
-    bytes, one shared block per request, and no prefix-caching leakage from
-    the kpool-sized scratch group."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    kv_cache_spec = _glm5_like_kv_cache_spec_with_tail()
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
-    tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes
-    assert tail_logical_page < idx_page
-
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    tail_group = next(
-        group
-        for group in groups
-        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-        and all(
-            isinstance(spec, KpoolTailSpec)
-            for spec in group.kv_cache_spec.kv_cache_specs.values()
+    for i in range(34):
+        kv_cache_spec[f"layers.{i}.kda"] = MambaSpec(
+            shapes=((3, 6144), (16, 128, 128)),
+            dtypes=(torch.bfloat16, torch.float32),
+            block_size=block_size,
         )
-    )
-    # The tail never prefix-caches; its page is padded up to the indexer page
-    # so the runner's strided view rides the indexer storage.
-    assert not tail_group.kv_cache_spec.prefix_cacheable
-    tail_inner = cast(
-        KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"]
-    )
-    assert tail_inner.page_size_padded == idx_page
-
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    assert bytes_per_block == 11 * mla_page + 11 * idx_page
-
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-    assert kv_cache_config.num_blocks == 100
-    # Tail layers get logical views but no additional storage: each view aliases
-    # its sibling indexer at the same offset in the shared allocation.
-    assert len(kv_cache_config.kv_cache_tensors) == 67
-    tensors = _tensor_by_layer(kv_cache_config)
-    for i in range(11):
-        idx_name = f"layers.{4 * i + 3}.indexer"
-        tail_name = f"layers.{4 * i + 3}.tail"
-        assert tensors[idx_name].offset == tensors[tail_name].offset
-    assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
-
-    # The layout detector surfaces the sibling names in layer order for the
-    # accounting and connector paths.
-    layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups)
-    assert layout is not None
-    assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)]
-    assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)]
-
-    # Accounting charges exactly one extra shared block per request for the
-    # tail, not a per-sequence allocation.
-    attn_group = next(
-        group
-        for group in groups
-        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-        and group is not tail_group
-    )
-    attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config)
-    mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
-    blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1
-    assert (
-        kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
-        == blocks_per_request * bytes_per_block
-    )
-
-
-def test_glm5_kpool_tail_does_not_drag_hash_block_size():
-    """The tail's kpool-sized scratch block (4 tokens) must not constrain the
-    prefix-cache hash granularity: participating groups alone decide it."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    def align_mamba():
-        return new_mamba_spec(mamba_cache_mode="align")
-
-    groups = kv_cache_utils.get_kv_cache_groups(
-        vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba)
-    )
-    kv_cache_config = KVCacheConfig(
-        num_blocks=1,
-        kv_cache_tensors=[],
-        kv_cache_groups=groups,
-    )
-    hash_vllm_config = SimpleNamespace(
-        cache_config=SimpleNamespace(
-            block_size=16,
-            enable_prefix_caching=True,
-            prefix_match_unit=None,
-        ),
-        parallel_config=SimpleNamespace(decode_context_parallel_size=1),
-        kv_transfer_config=object(),
-    )
-    # gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is
-    # the lcm including the tail (1024 % 4 == 0, so it coincides).
-    assert kv_cache_utils.resolve_kv_cache_block_sizes(
-        kv_cache_config, hash_vllm_config
-    ) == (1024, 16)
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible():
-    """Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    # (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page.
-    def big_mamba_spec():
-        return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,))
-
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec)
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    assert big_mamba_spec().page_size_bytes > mla_page
-
-    with pytest.raises(ValueError, match="does not fit the MLA page"):
-        kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer():
-    """Use the generic layout error when no kpool indexer is present."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-    vllm_config.cache_config.kv_cache_layout = "LBHNC"
-
-    kv_cache_spec: dict[str, KVCacheSpec] = {}
-    for i in range(27):
-        if i % 4 == 3 or i == 26:
-            kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=576,
-                dtype=torch.bfloat16,
-            )
-        else:
-            kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec(
-                shapes=((512, 1024),), dtypes=(torch.float32,)
-            )
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page
-
-    with pytest.raises(NotImplementedError, match="page size"):
-        kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba():
-    """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not
-    disable slot sharing; the layout re-pads them to the MLA page."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    def prepadded_mamba_spec():
-        return new_mamba_spec(page_size_padded=294_912)
-
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec)
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    assert prepadded_mamba_spec().page_size_bytes < mla_page
 
     groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    mamba_groups = [
-        group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
+    attn = [
+        g
+        for g in groups
+        if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs)
+        and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec)
     ]
-    assert len(groups) == 5
-    assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
-    for group in mamba_groups:
-        assert group.kv_cache_spec.page_size_bytes == mla_page
-
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-    assert kv_cache_config.num_blocks == 100
-    assert len(kv_cache_config.kv_cache_tensors) == 56
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection():
-    """Round-robin mamba grouping keeps practical PP splits balanced: every
-    stage's largest projected mamba group slice fits its projected MLA
-    layers, so per-stage slot tensors all have an MLA owner."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec()
-    global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
-
-    # PP=2 split at transformer layer 22/23.
-    # Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer)
-    # layers, projected mamba groups [5, 5, 4, 4].
-    worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22}
-    groups = kv_cache_utils._project_kv_cache_groups_to_worker(
-        global_groups, worker_spec
-    )
-    mamba_groups = [
-        group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
+    rings = [
+        g
+        for g in groups
+        if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs)
+        and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec)
     ]
-    assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4]
-
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    assert bytes_per_block == 5 * mla_page + 5 * idx_page
-
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-    assert kv_cache_config.num_blocks == 100
-    tensors = _tensor_by_layer(kv_cache_config)
-    # Every projected Mamba slot aliases an MLA owner.
-    assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset
-    assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset
-    assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset
-    assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset
-    # The last slot only hosts the two groups with a 5th projected layer.
-    assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset
-    assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset
-    assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch):
-    """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA
-    layers but a round-robin slice of 3 under the minimum 4 mamba groups;
-    grouping bumps to 5 groups so every stage's projection keeps sharing
-    on instead of silently falling back on stage 0."""
-    from vllm.config import ParallelConfig
-    from vllm.distributed.utils import get_pp_indices
-
-    monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
-    vllm_config = VllmConfig(
-        model_config=ModelConfig(max_model_len=8192),
-        parallel_config=ParallelConfig(pipeline_parallel_size=4),
-    )
-
-    kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)]
-    assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6]
-    for k, name in enumerate(mamba_layers):
-        assert name in mamba_groups[k % 5].layer_names
-
-    for rank in range(4):
-        start, end = get_pp_indices(45, rank, 4)
-        worker_spec = {
-            n: s
-            for n, s in kv_cache_spec.items()
-            if start <= int(n.split(".")[1]) < end
-        }
-        projected = kv_cache_utils._project_kv_cache_groups_to_worker(
-            groups, worker_spec
+    mamba = [
+        g
+        for g in groups
+        if isinstance(
+            getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec
         )
-        assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch):
-    """Reject PP stages whose Mamba layers have no MLA slot to share."""
-    from vllm.config import ParallelConfig
-
-    monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
-    monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42")
-    vllm_config = VllmConfig(
-        model_config=ModelConfig(max_model_len=8192),
-        parallel_config=ParallelConfig(pipeline_parallel_size=2),
-    )
-
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec()
-    with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"):
-        kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag():
-    """Hybrid slot sharing must take precedence over the experimental
-    enable_cross_layers_blocks packed layout: generic packing would give the
-    MLA slots a strided view, breaking the contiguous virtual split."""
-    model_config = ModelConfig(max_model_len=8192)
-    kv_transfer_config = KVTransferConfig(
-        kv_connector="NixlConnector",
-        kv_role="kv_both",
-        kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
-    )
-    vllm_config = VllmConfig(
-        model_config=model_config, kv_transfer_config=kv_transfer_config
-    )
-
-    kv_cache_spec, _ = _glm5_like_kv_cache_spec()
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-    idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    assert bytes_per_block == 11 * mla_page + 11 * idx_page
-
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-    assert kv_cache_config.num_blocks == 100
-    assert len(kv_cache_config.kv_cache_tensors) == 56
-    assert all(
-        tensor.block_stride in (mla_page, idx_page)
-        for tensor in kv_cache_config.kv_cache_tensors
-    )
-    assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
-
-
-def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer():
-    """Kimi-Linear-like: MLA without indexer layers, idx_stride == 0."""
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-    vllm_config.cache_config.kv_cache_layout = "LBNHC"
-
-    # 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6].
-    kv_cache_spec: dict[str, KVCacheSpec] = {}
-    for i in range(27):
-        if i % 4 == 3 or i == 26:
-            kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=576,
-                dtype=torch.bfloat16,
-            )
-        else:
-            kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
-    mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
-
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    mamba_groups = [
-        group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
     ]
-    assert len(groups) == 4
-    assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6]
-    for group in mamba_groups:
-        assert group.kv_cache_spec.page_size_bytes == mla_page
-
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    assert bytes_per_block == 7 * mla_page
+    assert len(attn) == 1 and len(attn[0].layer_names) == 22
+    assert len(rings) == 1 and len(rings[0].layer_names) == 11
+    # 11 * (1.18 MB MLA + 34 KB indexer) per block fits 12 fp32-SSM KDA states.
+    assert len(mamba) == 3
+    assert sum(len(g.layer_names) for g in mamba) == 34
 
     kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-    assert kv_cache_config.num_blocks == 100
-    # Each group has one tensor descriptor and all groups alias the same slots.
-    assert len(kv_cache_config.kv_cache_tensors) == 4
-    tensors = _tensor_by_layer(kv_cache_config)
-    mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)]
-    for index, mla_name in enumerate(mla_names):
-        mla_offset = _layer_offset(tensors[mla_name], mla_name)
-        for group in mamba_groups:
-            if index < len(group.layer_names):
-                mamba_name = group.layer_names[index]
-                assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset
-    assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
-
-
-def test_get_kv_cache_capacity_after_scheduler_unwrap():
-    """max_concurrency must survive the scheduler-config unwrap.
-
-    Regression for the balanced-mamba hybrid layout:
-    ``generate_scheduler_kv_cache_config`` flattens the MLA
-    ``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the
-    scheduler config holds an MLA spec (page size A) next to ``MambaSpec``
-    groups (page size B). ``EngineCore._initialize_kv_caches`` calls
-    ``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size
-    path used to assert-fail (``assert len(page_sizes) == 1``) on this topology.
-    """
-    model_config = ModelConfig(max_model_len=8192)
-    vllm_config = VllmConfig(model_config=model_config)
-
-    kv_cache_spec: dict[str, KVCacheSpec] = {}
-    for i in range(45):
-        if i % 4 == 3:
-            kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=576,
-                dtype=torch.bfloat16,
-            )
-            kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
-                block_size=1024,
-                num_kv_heads=1,
-                head_size=132,
-                dtype=torch.uint8,
-                tokens_per_state=16,
-            )
-        else:
-            kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
-
-    groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
-    bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
-    kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
-        vllm_config, groups, bytes_per_block * 100 + 1
-    )
-
-    scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config])
-    # Confirm we are exercising the regression path: the MLA group is no longer
-    # UniformTypeKVCacheSpecs, so sibling specs have differing page sizes.
-    assert not any(
-        isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs)
-        for g in scheduler_config.kv_cache_groups
-    )
-
-    unwrapped_groups = scheduler_config.kv_cache_groups
-    expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups(
-        vllm_config, unwrapped_groups
-    )
-    expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups)
-    expected_blocks_per_request = (
-        expected_max_mem + expected_pool - 1
-    ) // expected_pool
-
-    _, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config)
-    assert max_concurrency > 0
-    assert max_concurrency == pytest.approx(
-        scheduler_config.num_blocks / expected_blocks_per_request
+        vllm_config, groups, available_memory=2 * 1024**3
     )
+    strides = {t.block_stride for t in kv_cache_config.kv_cache_tensors}
+    assert len(strides) == 1
+    (block_stride,) = strides
+    # Blocks sit a whole number of 32-row MLA pages and 64-state indexer pages
+    # apart so both caches can be re-paged for their kernels.
+    assert block_stride % (32 * 576 * 2) == 0
+    assert block_stride % (64 * 132) == 0
 
 
 def new_mla_spec(cache_dtype_str=None, block_size: int = 16):
@@ -4483,6 +3993,7 @@ def test_kv_cache_reserves_null_block_for_max_model_len(use_override):
     block_size = 16
     max_model_len = 512  # needs 512 / 16 = 32 blocks
     vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     spec = new_kv_cache_spec(block_size=block_size)
 
     # 32 blocks -> only 31 usable after the null block: one short -> reject.
@@ -4528,6 +4039,7 @@ def test_check_enough_kv_cache_memory_reserves_null_block():
     block_size = 16
     max_model_len = 512  # needs 512 / 16 = 32 blocks
     vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     spec = new_kv_cache_spec(block_size=block_size)
 
     # 32 blocks -> only 31 usable after the null block: one short -> reject.

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -20,7 +20,6 @@ from vllm.v1.core.single_type_kv_cache_manager import (
     ChunkedLocalAttentionManager,
     CircularBufferManager,
     FullAttentionManager,
-    KpoolTailManager,
     MambaManager,
     RSWAManager,
     SlidingWindowManager,
@@ -29,7 +28,6 @@ from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
     CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     MambaSpec,
     RSWASpec,
     SlidingWindowSpec,
@@ -406,15 +404,15 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime():
 
 @pytest.mark.parametrize("record_for_zeroing", [True, False])
 def test_external_kpool_tail_zeroing(record_for_zeroing):
-    """A transferred ring at index zero skips zeroing even with a local hit."""
-    spec = KpoolTailSpec(
+    """The kpool tail ring is never zeroed, transferred or not: a pool only
+    reads ring slots its own request wrote."""
+    spec = CircularBufferSpec(
         block_size=4,
         num_kv_heads=2,
         head_size=8,
         dtype=torch.bfloat16,
-        sliding_window=4,
     )
-    manager = KpoolTailManager(
+    manager = CircularBufferManager(
         spec,
         block_pool=BlockPool(10, enable_caching=True, hash_block_size=4),
         enable_caching=True,
@@ -428,9 +426,7 @@ def test_external_kpool_tail_zeroing(record_for_zeroing):
     )
     blocks = manager.req_to_blocks["request"]
     assert len(blocks) == 1
-    assert manager.take_new_block_ids() == (
-        [blocks[0].block_id] if record_for_zeroing else []
-    )
+    assert manager.take_new_block_ids() == []
     assert manager.allocate_new_blocks("request", 20, 20) == []
     assert manager.take_new_block_ids() == []
 

tests/v1/kv_connector/unit/conftest.py

@@ -23,5 +23,5 @@ def _mock_cpu_backend_block_sizes(monkeypatch):
     monkeypatch.setattr(
         CPUAttentionBackend,
         "get_supported_kernel_block_sizes",
-        staticmethod(lambda: [MultipleOf(16)]),
+        staticmethod(lambda kv_cache_spec=None: [MultipleOf(16)]),
     )

tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py

@@ -164,13 +164,12 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
         fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
         gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
         fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8)
-        gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8)
+        gdn_block_stride = gdn_spec.page_size_bytes + 64
+        gdn_raw = torch.empty(num_blocks, gdn_block_stride, dtype=torch.int8)
         (fa_cache,) = dense_kv_cache_views(
             fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC
         )
-        (gdn_cache,) = dense_kv_cache_views(
-            gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC
-        )
+        gdn_cache = gdn_raw[:, : gdn_spec.page_size_bytes]
 
         worker.register_kv_caches(
             {
@@ -184,7 +183,7 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
             "model.layers.0.self_attn",
             "model.layers.1.linear_attn",
         ]
-        assert worker.block_len_per_layer == [
+        assert worker.kv_block_len_per_layer == [
             fa_spec.page_size_bytes,
             gdn_spec.page_size_bytes,
         ]

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -31,7 +31,6 @@ from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
@@ -711,13 +710,12 @@ def test_worker_setup_tolerates_finer_scratch_group():
         mamba_cache_mode="align",
     )
     # Scratch block 4 is not divisible by the hash unit 8 below.
-    scratch = KpoolTailSpec(
+    scratch = CircularBufferSpec(
         block_size=4,
         num_kv_heads=2,
         head_size=64,
         head_size_v=0,
         dtype=torch.bfloat16,
-        sliding_window=4,
     )
     cfg = KVCacheConfig(
         num_blocks=4,

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -641,164 +641,43 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
 
 
 @pytest.mark.cpu_test
-@pytest.mark.parametrize("logical_block_size", [1152, 640])
-@pytest.mark.parametrize("tail_first", [False, True])
-def test_register_compressed_indexer_uses_virtual_transfer_pages(
-    logical_block_size, tail_first
-):
-    """Compressed indexer rows must split into contiguous NIXL transfer pages."""
-    from unittest.mock import MagicMock
+def test_nixl_keeps_packed_sparse_mla_manager_block():
+    """Like the worker, NIXL keeps packed sparse MLA + indexer manager blocks."""
+    from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+    from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
+    from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend
+    from vllm.v1.kv_cache_interface import UniformTypeKVCacheSpecs
 
-    from vllm.config import set_current_vllm_config
-    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
-        base_worker as bw,
+    specs = {
+        "mla": MLAAttentionSpec(
+            1152, num_kv_heads=1, head_size=512, dtype=torch.bfloat16
+        ),
+        "indexer": MLAAttentionSpec(
+            1152, num_kv_heads=1, head_size=128, dtype=torch.uint8, tokens_per_state=4
+        ),
+    }
+    backends = {"mla": FlashMLASparseBackend, "indexer": Glm5NextIndexerBackend}
+    layers = {
+        name: MagicMock(spec=AttentionLayerBase, **{"get_attn_backend.return_value": b})
+        for name, b in backends.items()
+    }
+    worker = object.__new__(bw.NixlBaseConnectorWorker)
+    worker.block_size, worker.num_blocks, worker.attn_backends = 1152, 2, []
+    worker._physical_blocks_per_logical_kv_block = 1
+    worker.vllm_config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
+        ),
+        compilation_config=SimpleNamespace(static_forward_context=layers),
     )
-    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
-        NixlConnectorWorker,
-    )
-    from vllm.v1.kv_cache_interface import (
-        KpoolTailSpec,
-        KVCacheConfig,
-        KVCacheGroupSpec,
-        KVCacheLayout,
-        KVCacheTensor,
-        MLAAttentionSpec,
-        UniformTypeKVCacheSpecs,
-        create_kv_cache_views,
+    worker.kv_cache_config = KVCacheConfig(
+        2, [], [KVCacheGroupSpec(list(specs), UniformTypeKVCacheSpecs(1152, specs))]
     )
 
-    num_logical_blocks = 3
-    transfer_block_size = 64
-    kernel_block_size = 128
-    tokens_per_state = 4
-    state_content_bytes = 132
+    worker._sync_block_size_with_kernel()
 
-    indexer_spec = MLAAttentionSpec(
-        block_size=logical_block_size,
-        num_kv_heads=1,
-        head_size=128,
-        head_size_v=0,
-        dtype=torch.uint8,
-        state_content_bytes=state_content_bytes,
-        tokens_per_state=tokens_per_state,
-    )
-    indexer_page_size = indexer_spec.page_size_bytes
-    tail_spec = KpoolTailSpec(
-        block_size=tokens_per_state,
-        num_kv_heads=2,
-        head_size=128,
-        head_size_v=0,
-        dtype=torch.bfloat16,
-        page_size_padded=indexer_page_size,
-        sliding_window=tokens_per_state,
-    )
-
-    allocation_size = num_logical_blocks * indexer_page_size
-    indexer_tensor = KVCacheTensor(
-        size=allocation_size,
-        layers=["indexer"],
-        layer_stride=allocation_size,
-        block_stride=indexer_page_size,
-    )
-    tail_tensor = KVCacheTensor(
-        size=allocation_size,
-        layers=["tail"],
-        layer_stride=allocation_size,
-        block_stride=indexer_page_size,
-    )
-    kv_cache_config = KVCacheConfig(
-        num_blocks=num_logical_blocks,
-        kv_cache_tensors=[indexer_tensor, tail_tensor],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["indexer"],
-                UniformTypeKVCacheSpecs(
-                    block_size=logical_block_size,
-                    kv_cache_specs={"indexer": indexer_spec},
-                ),
-            ),
-            KVCacheGroupSpec(
-                ["tail"],
-                UniformTypeKVCacheSpecs(
-                    block_size=tokens_per_state,
-                    kv_cache_specs={"tail": tail_spec},
-                ),
-            ),
-        ],
-    )
-
-    raw = torch.zeros(allocation_size, dtype=torch.int8)
-    (indexer_cache,) = create_kv_cache_views(
-        raw,
-        indexer_spec,
-        num_logical_blocks,
-        KVCacheLayout.LBHNC,
-        indexer_tensor,
-        kernel_block_size=kernel_block_size,
-    )
-    (tail_cache,) = create_kv_cache_views(
-        raw,
-        tail_spec,
-        num_logical_blocks,
-        KVCacheLayout.LBHNC,
-        tail_tensor,
-    )
-    assert indexer_cache.data_ptr() == tail_cache.data_ptr() == raw.data_ptr()
-
-    vllm_config = create_vllm_config(block_size=logical_block_size)
-    vllm_config.cache_config.kv_cache_layout = "LBHNC"
-    vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
-    fake_backend = MagicMock()
-    fake_backend.get_supported_kernel_block_sizes.return_value = [transfer_block_size]
-    fake_backend.get_name.return_value = "DEEPSEEK_V32_INDEXER"
-    fake_backend.full_cls_name.return_value = "fake.DEEPSEEK_V32_INDEXER"
-    fake_platform = MagicMock()
-    fake_platform.device_type = "cuda"
-    fake_platform.get_nixl_memory_type.return_value = "VRAM"
-
-    caches = [("indexer", indexer_cache), ("tail", tail_cache)]
-    if tail_first:
-        caches.reverse()
-
-    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),
-        set_current_vllm_config(vllm_config),
-    ):
-        worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
-        worker.use_mla = True
-        worker.register_kv_caches(dict(caches))
-
-    transfer_page_size = transfer_block_size // tokens_per_state * state_content_bytes
-    num_transfer_blocks = num_logical_blocks * (
-        logical_block_size // transfer_block_size
-    )
-    expected_descs = np.asarray(
-        [
-            [
-                raw.data_ptr() + block_idx * transfer_page_size,
-                transfer_page_size,
-                0,
-            ]
-            for block_idx in range(num_transfer_blocks)
-        ],
-        dtype=np.uint64,
-    )
-
-    assert worker.block_size == transfer_block_size
-    assert worker.num_regions == 1
-    assert worker.block_len_per_layer == [transfer_page_size]
-    assert worker.block_stride_per_layer == [transfer_page_size]
-    assert worker._region_is_mla == [True]
-    assert worker.kv_caches_base_addr[worker.engine_id][0] == [raw.data_ptr()]
-    assert worker._registered_descs[0] == [(raw.data_ptr(), raw.nbytes, 0, "")]
-    np.testing.assert_array_equal(worker.src_blocks_data, expected_descs)
-    assert expected_descs[-1, 0] + expected_descs[-1, 1] == (
-        raw.data_ptr() + raw.nbytes
-    )
+    assert worker.block_size == 1152
+    assert worker._physical_blocks_per_logical_kv_block == 1
 
 
 def _make_remote_meta(

tests/v1/simple_kv_offload/test_scheduler.py

@@ -47,7 +47,6 @@ from vllm.v1.core.single_type_kv_cache_manager import (
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
@@ -145,13 +144,12 @@ def _make_scratch_kv_cache_config(
     register_all_kvcache_specs(vllm_config=None)
     fa_config = _make_kv_cache_config(num_blocks, num_groups=1)
     scratch_layers = ["layer_scratch"]
-    scratch_spec = KpoolTailSpec(
+    scratch_spec = CircularBufferSpec(
         block_size=scratch_block_size,
         num_kv_heads=2,
         head_size=HEAD_SIZE,
         head_size_v=0,
         dtype=DTYPE,
-        sliding_window=scratch_block_size,
     )
     assert not scratch_spec.prefix_cacheable
     scratch_bytes = scratch_spec.page_size_bytes * num_blocks

tests/v1/spec_decode/test_dflash_prepare_inputs.py

@@ -102,7 +102,7 @@ def _run_prepare(
         block_tables=SimpleNamespace(
             slot_mappings=query_slot_mapping.unsqueeze(0),
             input_block_tables=[block_table],
-            kernel_block_sizes=[4],
+            block_sizes=[4],
             cp_rank=cp_rank,
             cp_size=cp_size,
             cp_interleave=cp_interleave,

tests/v1/spec_decode/test_eagle.py

@@ -963,6 +963,9 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
     # Mock draft_attn_groups for attention metadata building.
     mock_attn_group = mock.MagicMock()
     mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder
+    mock_attn_group.build_metadata_for_drafting = (
+        attn_metadata_builder.build_for_drafting
+    )
     mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)
     mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec
     proposer.draft_attn_groups = [mock_attn_group]
@@ -1066,6 +1069,9 @@ def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch):
     )
     mock_attn_group = mock.MagicMock()
     mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder
+    mock_attn_group.build_metadata_for_drafting = (
+        attn_metadata_builder.build_for_drafting
+    )
     mock_attn_group.layer_names = list(proposer._draft_attn_layer_names)
     mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec
     proposer.draft_attn_groups = [mock_attn_group]

tests/v1/spec_decode/test_llm_base_proposer.py

@@ -2,10 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for SpecDecodeBaseProposer.initialize_attn_backend.
 
-Block tables are stored at kernel-block granularity, so the proposer's
-``block_size`` (used for slot-mapping math) must be the kernel block size,
-not the KV cache manager's block size — the two differ when manager blocks
-are split for the attention kernel. The value must also be deterministic:
+Block tables are stored at manager-block granularity, so the proposer's
+``block_size`` (used for slot-mapping math) must be the KV cache manager's
+block size. Attention groups map to kernel blocks for their metadata builders.
+The value must also be deterministic:
 ``_draft_attn_layer_names`` is a set, whose iteration order varies across
 processes, so anything derived from iteration order must not leak into
 ``block_size``.
@@ -68,9 +68,9 @@ def _make_kv_cache_config(layer_names: set[str]) -> KVCacheConfig:
     return SimpleNamespace(kv_cache_groups=[group])  # type: ignore[return-value]
 
 
-def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):
-    """The proposer's slot-mapping math runs against the kernel-granularity
-    block table, so block_size must come from kernel_block_sizes."""
+def test_block_size_uses_manager_block_size(monkeypatch: pytest.MonkeyPatch):
+    """The proposer's slot-mapping math runs against the manager-granularity
+    block table; attention groups map kernel blocks for metadata builders."""
     layer_names = {"draft.0.self_attn.attn"}
     proposer = _make_proposer(monkeypatch, layer_names)
 
@@ -79,12 +79,8 @@ def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):
         kernel_block_sizes=[KERNEL_BLOCK_SIZE],
     )
 
-    assert proposer.block_size == KERNEL_BLOCK_SIZE
-    assert proposer.block_size != SCHEDULER_BLOCK_SIZE
-    # The metadata builder keeps receiving the kernel block size as well.
-    group = proposer.draft_attn_groups[0]
-    assert isinstance(group, _FakeAttentionGroup)
-    assert group.kernel_block_size == KERNEL_BLOCK_SIZE
+    assert proposer.block_size == SCHEDULER_BLOCK_SIZE
+    assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE
 
 
 def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch):
@@ -112,4 +108,4 @@ def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch)
         )
         assert len(proposer.draft_attn_groups) == 1
         assert proposer.draft_attn_groups[0].layer_names == expected_order
-        assert proposer.block_size == KERNEL_BLOCK_SIZE
+        assert proposer.block_size == SCHEDULER_BLOCK_SIZE

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

@@ -38,7 +38,6 @@ def mock_model_runner_with_input_batch():
         device="cpu",
         vocab_size=32000,
         block_sizes=[16],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[64],
         logitsprocs=None,
         is_pooling_model=False,

tests/v1/worker/test_attn_utils.py

@@ -15,7 +15,11 @@ import pytest
 import torch
 
 import vllm.v1.hisparse.binding as attn_utils_module
-from tests.v1.attention.utils import dense_kv_cache_views
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    dense_kv_cache_views,
+)
 from vllm.config.compilation import CompilationConfig, CUDAGraphMode
 from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
@@ -42,6 +46,8 @@ from vllm.v1.worker.utils import (
     AttentionGroup,
     allocate_kv_cache,
     copy_kv_cache_blocks_inplace,
+    customize_attention_spec,
+    map_kv_caches_to_kernel_blocks,
 )
 
 
@@ -79,7 +85,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size(
         backend = SimpleNamespace(
             get_name=lambda name=name: name,
             customize_spec=AttentionBackend.customize_spec,
-            get_supported_kernel_block_sizes=lambda sizes=sizes: sizes,
+            get_supported_kernel_block_sizes=lambda _=None, sizes=sizes: sizes,
         )
         layers[name] = SimpleNamespace(
             get_kv_cache_spec=lambda _, spec=specs[name]: spec,
@@ -99,7 +105,8 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size(
 
     resolved = attn_utils.get_kv_cache_spec(config)
     assert resolved["main"].block_size == resolved["indexer"].block_size == expected
-    assert resolved["dense"] is specs["dense"]
+    dense_backend = layers["dense"].get_attn_backend()
+    assert resolved["dense"] == customize_attention_spec(dense_backend, specs["dense"])
     assert all(spec.block_size == block_size for spec in specs.values())
 
 
@@ -495,20 +502,17 @@ def test_reshape_padded_kv_cache_strides_by_padded_page():
 @pytest.mark.parametrize(
     (
         "kernel_block_sizes",
-        "storage_block_size",
         "expected_num_blocks",
         "expected_num_states",
     ),
     [
-        (None, None, 4, 64),
-        ([256], None, 4, 64),
-        ([64], None, 16, 16),
-        ([64], 256, 4, 64),
+        (None, 4, 64),
+        ([256], 4, 64),
+        ([64], 16, 16),
     ],
 )
 def test_allocate_compressed_mla_cache(
     kernel_block_sizes: list[int] | None,
-    storage_block_size: int | None,
     expected_num_blocks: int,
     expected_num_states: int,
 ):
@@ -518,7 +522,6 @@ def test_allocate_compressed_mla_cache(
         head_size=128,
         dtype=torch.bfloat16,
         tokens_per_state=4,
-        storage_block_size=storage_block_size,
     )
     num_pages = 4
     config = KVCacheConfig(
@@ -744,3 +747,92 @@ def test_allocate_hisparse_kv_caches_host_pool_and_view_less_specs():
         backing.untyped_storage().data_ptr()
         == caches["indexer"].untyped_storage().data_ptr()
     )
+
+
+class _TableBuilder:
+    requires_block_table_width = False
+    supports_update_block_table = False
+
+    def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
+        self.kv_cache_spec = kv_cache_spec
+
+    def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
+        return common_attn_metadata.block_table_tensor
+
+
+_VLLM_CONFIG = SimpleNamespace(
+    model_config=SimpleNamespace(max_model_len=8192),
+    scheduler_config=SimpleNamespace(max_num_seqs=4),
+    parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+)
+
+
+def _packed_group(tokens_per_state, row_bytes, supported_sizes):
+    spec = MLAAttentionSpec(
+        block_size=1024,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.uint8,
+        state_content_bytes=row_bytes,
+        tokens_per_state=tokens_per_state,
+    )
+    backend = SimpleNamespace(
+        get_builder_cls=lambda: _TableBuilder,
+        get_supported_kernel_block_sizes=lambda _=None: supported_sizes,
+        get_name=lambda: "FAKE",
+        customize_spec=lambda spec: spec,
+    )
+    return AttentionGroup(backend, ["layer.0"], spec, 0)
+
+
+def test_packed_compressed_group_gets_kernel_block_table():
+    """A packed indexer cache whose kernel needs 64-state pages gets a table and
+    a bound view in kernel blocks that address the same packed-block rows."""
+    row, stride_rows = 132, 5 * 64
+    group = _packed_group(4, row, [128, 256])
+    group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024)
+    assert group.get_metadata_builder().kv_cache_spec.block_size == 256
+
+    backing = torch.arange(3 * stride_rows).repeat_interleave(row)
+    manager = backing.view(3, 1, stride_rows, row)[:, :, :256]
+    pages = map_kv_caches_to_kernel_blocks({"layer.0": manager}, [group])["layer.0"]
+    metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu")
+    metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32)
+    table = group.build_metadata(metadata)
+    assert table.tolist() == [[0, 1, 2, 3, 10, 11, 12, 13]]
+    assert torch.equal(pages[table[0, 4:], 0].flatten(0, 1), manager[2, 0])
+
+
+def test_split_group_maps_manager_block_table():
+    """Groups split in a dense layout keep manager block tables; builders get
+    b * blocks_per_kv_block + j against the view allocated in kernel blocks."""
+    group = _packed_group(4, 132, [128, 256])
+    group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=256)
+    kernel_view = torch.zeros(12, 1, 64, 132)
+    bound = map_kv_caches_to_kernel_blocks({"layer.0": kernel_view}, [group])
+    assert bound["layer.0"] is kernel_view
+
+    metadata = create_common_attn_metadata(BatchSpec([2048], [1]), 1024, "cpu")
+    metadata.block_table_tensor = torch.tensor([[0, 2]], dtype=torch.int32)
+    table = group.build_metadata(metadata)
+    assert table.tolist() == [[0, 1, 2, 3, 8, 9, 10, 11]]
+
+
+@pytest.mark.parametrize(
+    ("block_size", "page_states"), [(1024, 64), (384, 32), (256, 0)]
+)
+def test_split_blocks_align_to_kernel_blocks(block_size, page_states):
+    """Backends taking only fixed kernel blocks align packed blocks to the kernel
+    block the final block size is split into, and not at all when unsplit."""
+    group = _packed_group(4, 132, [128, 256])
+    spec = customize_attention_spec(group.backend, group.kv_cache_spec)
+    spec = spec.copy_with_new_block_size(block_size)
+    assert spec.get_block_stride_alignment() == max(page_states * 132, 1)
+
+
+def test_packed_token_cache_cannot_be_split():
+    group = _packed_group(1, 1152, [64])
+    group.create_metadata_builders(_VLLM_CONFIG, "cpu", kernel_block_size=1024)
+    packed = torch.zeros(3, 1, 2048, 1152, dtype=torch.uint8)[:, :, :1024]
+    with pytest.raises(ValueError, match="cannot be split"):
+        map_kv_caches_to_kernel_blocks({"layer.0": packed}, [group])

tests/v1/worker/test_gpu_block_table.py

@@ -22,7 +22,6 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
         max_num_batched_tokens=64,
         max_num_blocks_per_group=[8, 8, 8],
         device=device,
-        kernel_block_sizes=[16, 16, 8],
     )
 
     def fail_if_apply_write_called():
@@ -48,33 +47,31 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
         block_tables.block_tables[0].gpu[0, :2],
         torch.tensor([1, 2], dtype=torch.int32, device=device),
     )
-    # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two
-    # kernel block IDs.
     assert torch.equal(
-        block_tables.block_tables[1].gpu[0, :4],
-        torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device),
+        block_tables.block_tables[1].gpu[0, :2],
+        torch.tensor([10, 11], dtype=torch.int32, device=device),
     )
     assert torch.equal(
         block_tables.block_tables[0].gpu[1, :1],
         torch.tensor([3], dtype=torch.int32, device=device),
     )
     assert torch.equal(
-        block_tables.block_tables[1].gpu[1, :2],
-        torch.tensor([24, 25], dtype=torch.int32, device=device),
+        block_tables.block_tables[1].gpu[1, :1],
+        torch.tensor([12], dtype=torch.int32, device=device),
     )
     assert torch.equal(
         block_tables.block_tables[2].gpu[1, :2],
         torch.tensor([5, 6], dtype=torch.int32, device=device),
     )
     assert block_tables.num_blocks.np[0, 0] == 2
-    assert block_tables.num_blocks.np[1, 0] == 4
+    assert block_tables.num_blocks.np[1, 0] == 2
     assert block_tables.num_blocks.np[2, 0] == 0
     assert block_tables.num_blocks.np[0, 1] == 1
-    assert block_tables.num_blocks.np[1, 1] == 2
+    assert block_tables.num_blocks.np[1, 1] == 1
     assert block_tables.num_blocks.np[2, 1] == 2
     assert torch.equal(
         block_tables.num_blocks.gpu[:, :2],
-        torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device),
+        torch.tensor([[2, 1], [2, 1], [0, 2]], dtype=torch.int32, device=device),
     )
 
     for block_table in block_tables.block_tables:
@@ -96,15 +93,15 @@ def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch):
         torch.tensor([1, 2, 7], dtype=torch.int32, device=device),
     )
     assert torch.equal(
-        block_tables.block_tables[1].gpu[0, :6],
-        torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device),
+        block_tables.block_tables[1].gpu[0, :3],
+        torch.tensor([10, 11, 13], dtype=torch.int32, device=device),
     )
     assert torch.equal(
         block_tables.block_tables[2].gpu[0, :1],
         torch.tensor([8], dtype=torch.int32, device=device),
     )
     assert block_tables.num_blocks.np[0, 0] == 3
-    assert block_tables.num_blocks.np[1, 0] == 6
+    assert block_tables.num_blocks.np[1, 0] == 3
     assert block_tables.num_blocks.np[2, 0] == 1
 
 
@@ -116,7 +113,6 @@ def test_block_tables_apply_staged_writes_single_group():
         max_num_batched_tokens=16,
         max_num_blocks_per_group=[4],
         device=device,
-        kernel_block_sizes=[16],
     )
 
     block_tables.append_block_ids(
@@ -141,7 +137,6 @@ def test_block_tables_skip_custom_slot_mapping_groups():
         max_num_batched_tokens=4,
         max_num_blocks_per_group=[1, 1],
         device=device,
-        kernel_block_sizes=[8, 262144],
         slot_mapping_enabled=[False, True],
     )
     block_tables.append_block_ids(
@@ -170,7 +165,7 @@ def test_block_tables_skip_custom_slot_mapping_groups():
 
 
 @pytest.mark.parametrize("cp_rank", range(4))
-def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
+def test_dcp_slot_mapping_only_shards_sharded_groups(cp_rank: int):
     """Only sharded groups use DCP interleave in logical-block coordinates."""
     device = torch.device("cuda")
     block_tables = BlockTables(
@@ -179,7 +174,6 @@ def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
         max_num_batched_tokens=1024,
         max_num_blocks_per_group=[2, 8],
         device=device,
-        kernel_block_sizes=[64, 64],
         dcp_sharded=[True, False],
         cp_size=4,
         cp_rank=cp_rank,
@@ -230,7 +224,6 @@ def test_v1_block_table_move_row_clears_vacated_row():
         max_num_batched_tokens=64,
         pin_memory=False,
         device=torch.device("cuda"),
-        kernel_block_size=16,
         cp_kv_cache_interleave_size=1,
     )
     block_table.add_row([7, 8, 9], row_idx=0)
@@ -260,7 +253,6 @@ def test_get_dummy_block_tables_returns_zeroed_rows():
         max_num_batched_tokens=64,
         max_num_blocks_per_group=[8],
         device=device,
-        kernel_block_sizes=[16],
     )
     # Simulate a real step: stage a request's blocks and gather them into
     # the persistent input block tables.
@@ -292,7 +284,6 @@ def test_dummy_request_slot_mapping_is_pad():
         max_num_batched_tokens=16,
         max_num_blocks_per_group=[4],
         device=device,
-        kernel_block_sizes=[4],
     )
     block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True)
     block_tables.apply_staged_writes()

tests/v1/worker/test_gpu_input_batch.py

@@ -233,7 +233,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
         device=torch.device(device),
         vocab_size=1024,
         block_sizes=[1],
-        kernel_block_sizes=[1],
         max_num_blocks_per_req=[1024],
     )
     reqs: list[CachedRequestState] = []
@@ -327,7 +326,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         device=torch.device(device),
         vocab_size=1024,
         block_sizes=[1],
-        kernel_block_sizes=[1],
         max_num_blocks_per_req=[1024],
     )
     ref_input_batch: InputBatch = InputBatch(
@@ -337,7 +335,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         device=torch.device(device),
         vocab_size=1024,
         block_sizes=[1],
-        kernel_block_sizes=[1],
         max_num_blocks_per_req=[1024],
     )
 
@@ -388,7 +385,6 @@ def test_condense_clears_stale_allowed_token_ids_mask(device: str):
         device=torch.device(device),
         vocab_size=VOCAB_SIZE,
         block_sizes=[1],
-        kernel_block_sizes=[1],
         max_num_blocks_per_req=[1024],
     )
 
@@ -446,7 +442,6 @@ def test_swap_states_preserves_allowed_token_ids_mask(device: str):
         device=torch.device(device),
         vocab_size=VOCAB_SIZE,
         block_sizes=[1],
-        kernel_block_sizes=[1],
         max_num_blocks_per_req=[1024],
     )
 
@@ -516,7 +511,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
         device=torch.device(device),
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[64],
         is_pooling_model=True,
     )
@@ -552,7 +546,6 @@ def test_placeholder_spec_token_ids_written_verbatim():
         device=torch.device("cpu"),
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[1],
     )
     req = CachedRequestState(
@@ -599,7 +592,6 @@ def test_pooling_metadata_token_id_buffers(
         device=torch.device("cpu"),
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[64],
         is_pooling_model=True,
     )
@@ -630,7 +622,6 @@ def _make_input_batch(is_pooling_model: bool = False) -> InputBatch:
         device=torch.device("cpu"),
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[64],
         is_pooling_model=is_pooling_model,
     )

tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py

@@ -5,7 +5,7 @@
 The kpool tail cache holds one circular block per request, so its block table
 row is orders of magnitude shorter than the sequence it serves.
 `_compute_slot_mappings_kernel` indexes block tables by
-`position // kernel_block_size` with an effectively unmasked load, so
+`position // block_size` with an effectively unmasked load, so
 this leads to a read of neighboring requests' rows (and for the later requests, outside
 the buffer entirely).
 """
@@ -15,7 +15,7 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
-from vllm.v1.kv_cache_interface import KpoolTailSpec
+from vllm.v1.kv_cache_interface import CircularBufferSpec
 from vllm.v1.worker.block_table import get_block_table_width
 from vllm.v1.worker.gpu.block_table import BlockTables
 
@@ -32,13 +32,12 @@ POISON_BLOCK_ID = 1 << 20
 
 
 def make_tail_spec():
-    return KpoolTailSpec(
+    return CircularBufferSpec(
         block_size=KPOOL,
         num_kv_heads=2,
         head_size=128,
         head_size_v=0,
         dtype=torch.bfloat16,
-        sliding_window=KPOOL,
     )
 
 
@@ -67,7 +66,6 @@ def test_kpool_tail_group_never_position_indexes_its_block_table():
         max_num_batched_tokens=len(positions),
         max_num_blocks_per_group=[mla_width, tail_width],
         device=device,
-        kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size],
         slot_mapping_enabled=[True, spec.uses_slot_mapping],
     )
     mla_blocks = list(range(1, mla_width + 1))

tests/v1/worker/test_gpu_model_runner.py

@@ -111,9 +111,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
         device=runner.device,
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
-        kernel_block_sizes=[
-            kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
-        ],
         max_num_blocks_per_req=[NUM_BLOCKS],
     )
     runner.initialize_attn_backend(kv_cache_config)
@@ -303,7 +300,7 @@ def _make_mock_backend_for_kernel_block_size(
             return "MOCK"
 
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return supported_sizes
 
     return _MockBackend()
@@ -441,7 +438,7 @@ def _mock_backend(supported: list, *, exact: bool = False):
             return "MOCK_EXACT" if exact else "MOCK"
 
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return list(supported)
 
         if exact:
@@ -1625,7 +1622,6 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch)
     runner.device = torch.device("cpu")
     runner.is_pooling_model = False
     runner._init_block_sizes = [16]
-    runner._init_kernel_block_sizes = [16]
     runner._init_max_num_blocks = [4]
     runner._init_slot_mapping_modes = [
         gpu_model_runner_module.SlotMappingMode.TOKEN_TO_KV_SLOT
@@ -1653,7 +1649,7 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch)
         lambda _: gpu_model_runner_module.KVCacheSpecKind.FULL_ATTENTION,
     )
 
-    runner.may_reinitialize_input_batch(kv_cache_config, [16])
+    runner.may_reinitialize_input_batch(kv_cache_config)
 
     assert input_batch_cls.call_count == 1
     assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16
@@ -1681,68 +1677,10 @@ def test_v2_runner_snapshots_late_interleave_adjustment(monkeypatch):
     assert runner.cp_interleave == 16
 
 
-def test_hybrid_block_table_initialization():
-    """Test hybrid block table with different kernel and kvcache_manager block
-    sizes."""
-    from vllm.v1.worker.block_table import BlockTable
-
-    # Test configuration: kvcache_manager block size = 32,
-    # kernel block size = 16
-    block_size = 32
-    kernel_block_sizes = [16]
-    max_num_reqs = 10
-    max_num_blocks_per_req = 20
-    max_num_batched_tokens = 512
-    cp_kv_cache_interleave_size = 8
-
-    block_table = BlockTable(
-        block_size=block_size,
-        max_num_reqs=max_num_reqs,
-        max_num_blocks_per_req=max_num_blocks_per_req,
-        max_num_batched_tokens=max_num_batched_tokens,
-        pin_memory=False,
-        device=torch.device(DEVICE_TYPE),
-        kernel_block_size=kernel_block_sizes[0],
-        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
-    )
-
-    # Verify hybrid block configuration
-    assert block_table.use_hybrid_blocks is True
-    assert block_table.block_size == kernel_block_sizes[0]
-    assert block_table.blocks_per_kv_block == (
-        block_size // kernel_block_sizes[0]
-    )  # Changed to use first element
-
-    # Test block table conversion logic
-    # One kvcache_manager block should map to multiple kernel blocks
-    kvcache_manager_blocks = [0, 1, 2]
-
-    # Verify that kvcache_manager blocks can be converted to kernel blocks
-    # and that block table operations work correctly.
-    req_index = 0
-    block_table.append_row(kvcache_manager_blocks, req_index)
-    # Get expected kernel blocks from the implementation for verification.
-    expected_kernel_blocks = block_table.map_to_kernel_blocks(
-        np.array(kvcache_manager_blocks),
-        block_table.blocks_per_kv_block,
-        block_table._kernel_block_arange,
-    )
-    # Verify block table state
-    assert block_table.num_blocks_per_row[req_index] == len(expected_kernel_blocks)
-    assert np.array_equal(
-        block_table.block_table.np[req_index, : len(expected_kernel_blocks)],
-        expected_kernel_blocks,
-    )
-
-
 def test_get_block_table_width_aligns_to_128_tokens():
     assert get_block_table_width(1875, 64) == 1876
 
 
-def test_get_block_table_width_splits_virtual_blocks():
-    assert get_block_table_width(235, 256, 64) == 940
-
-
 def test_mamba_state_table_width_is_not_aligned():
     block_tables = MultiGroupBlockTable(
         max_num_reqs=1,
@@ -1750,7 +1688,6 @@ def test_mamba_state_table_width_is_not_aligned():
         pin_memory=False,
         device=torch.device("cpu"),
         block_sizes=[39664],
-        kernel_block_sizes=[39664],
         max_num_blocks=[1],
         slot_mapping_modes=[SlotMappingMode.NONE],
     )
@@ -1758,42 +1695,6 @@ def test_mamba_state_table_width_is_not_aligned():
     assert block_tables[0].max_num_blocks_per_req == 1
 
 
-def test_input_batch_with_kernel_block_sizes():
-    """Test InputBatch initialization with kernel_block_sizes parameter."""
-    max_num_reqs = 10
-    max_model_len = 512
-    max_num_batched_tokens = 512
-    device = torch.device(DEVICE_TYPE)
-    vocab_size = 50272
-
-    # Test with different kernel block sizes
-    block_sizes = [32, 64]
-    kernel_block_sizes = [16, 32]
-
-    input_batch = InputBatch(
-        max_num_reqs=max_num_reqs,
-        max_model_len=max_model_len,
-        max_num_batched_tokens=max_num_batched_tokens,
-        device=device,
-        vocab_size=vocab_size,
-        block_sizes=block_sizes,
-        kernel_block_sizes=kernel_block_sizes,
-        max_num_blocks_per_req=[16, 8],
-    )
-
-    # Verify that block tables were created with kernel block sizes
-    assert len(input_batch.block_table.block_tables) == len(block_sizes)
-
-    for i, (kv_size, kernel_size) in enumerate(zip(block_sizes, kernel_block_sizes)):
-        block_table = input_batch.block_table.block_tables[i]
-        if kv_size != kernel_size:
-            assert block_table.use_hybrid_blocks is True
-            assert block_table.block_size == kernel_size
-        else:
-            assert block_table.use_hybrid_blocks is False
-            assert block_table.block_size == kernel_size
-
-
 def test_hybrid_cache_integration(default_vllm_config, dist_init):
     """Test hybrid cache architecture integration with GPUModelRunner."""
     # Create a new model runner with hybrid cache configuration
@@ -1843,9 +1744,8 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
         device=runner.device,
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
-        kernel_block_sizes=[16],
         max_num_blocks_per_req=[NUM_BLOCKS],
-    )  # Use kernel block size
+    )
 
     runner.initialize_attn_backend(kv_cache_config)
 

tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py

@@ -119,9 +119,6 @@ def _initialize_kv_cache(runner: GPUModelRunner) -> None:
         device=runner.device,
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
-        kernel_block_sizes=[
-            kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size
-        ],
         max_num_blocks_per_req=[NUM_BLOCKS],
     )
     runner.initialize_attn_backend(kv_cache_config)

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -237,7 +237,6 @@ def test_append_block_ids_rejects_write_past_row_capacity():
 
     block_tables = BlockTables.__new__(BlockTables)
     block_tables.num_kv_cache_groups = 1
-    block_tables.blocks_per_kv_block = [1]
     block_tables.block_tables = [_BlockTable()]
     block_tables.num_blocks = SimpleNamespace(
         np=torch.tensor([[0, 3]], dtype=torch.int32)

tests/v1/worker/test_gpu_pcp_manager.py

@@ -47,8 +47,7 @@ def _make_capture_manager(block_table: torch.Tensor):
     block_tables = SimpleNamespace(
         input_block_tables=(block_table,),
         num_kv_cache_groups=1,
-        kernel_block_sizes=(2,),
-        blocks_per_kv_block=(1,),
+        block_sizes=(2,),
     )
     manager = PCPManager(
         pcp_world_size=2,

tests/v1/worker/test_gpu_target_dependent_ar_speculator.py

@@ -20,6 +20,7 @@ from vllm.model_executor.models.mistral_large_3_eagle import (
 from vllm.v1.attention.backends import flash_attn as flash_attn_module
 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
 from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
@@ -30,6 +31,7 @@ from vllm.v1.worker.gpu.spec_decode.target_dependent_ar import speculator as spe
 from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import (
     TargetDependentARSpeculator,
 )
+from vllm.v1.worker.utils import AttentionGroup
 
 
 class _TestSpeculator(TargetDependentARSpeculator):
@@ -102,6 +104,7 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode):
                 get_metadata_builder=lambda _: SimpleNamespace(
                     build=build, supports_update_block_table=False
                 ),
+                build_metadata=lambda cm, ubatch_idx=0, **_: build(0, cm),
                 layer_names=["draft"],
             )
         ]
@@ -427,6 +430,8 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
         query_start_loc=torch.arange(3),
     )
     speculator.idx_mapping = torch.arange(2)
+    speculator.attn_groups = []
+    speculator.block_tables = SimpleNamespace(input_block_tables=[])
     generate_draft = Mock()
     speculator._generate_draft = generate_draft
     run_fullgraph = Mock()
@@ -450,6 +455,52 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
     assert run_fullgraph.call_count == expected_graph_replays
 
 
+def test_fused_full_graph_maps_split_group_block_tables():
+    """The fused FULL replay skips build(), so split groups' tables are mapped first."""
+    spec = MLAAttentionSpec(
+        block_size=1024, num_kv_heads=1, head_size=128, dtype=torch.uint8
+    )
+    group = AttentionGroup(
+        SimpleNamespace(),
+        ["layer.0"],
+        spec,
+        0,
+        kernel_block_size=256,
+        kernel_block_table=torch.zeros(1, 4, 32, dtype=torch.int32),
+        kernel_block_offsets=torch.arange(4, dtype=torch.int32),
+    )
+    speculator = object.__new__(_TestSpeculator)
+    speculator.input_buffers = SimpleNamespace(
+        positions=torch.arange(2), query_start_loc=torch.arange(3)
+    )
+    speculator.idx_mapping = torch.arange(2)
+    speculator.attn_groups = [[group]]
+    speculator.block_tables = SimpleNamespace(
+        input_block_tables=[torch.tensor([[1, 2], [0, 3]], dtype=torch.int32)]
+    )
+    replayed_tables = []
+    speculator.decode_cudagraph_manager = SimpleNamespace(
+        run_fullgraph=lambda _: replayed_tables.append(
+            group.kernel_block_table[0, :2, :8].tolist()
+        )
+    )
+
+    speculator._fused_multi_step_decode(
+        num_reqs=2,
+        skip_attn=False,
+        batch_desc=BatchExecutionDescriptor(
+            cg_mode=CUDAGraphMode.FULL, num_tokens=2, num_reqs=2
+        ),
+        num_tokens_across_dp=None,
+        seq_lens_cpu_upper_bound=None,
+        num_speculative_steps=4,
+    )
+
+    assert replayed_tables == [
+        [[4, 5, 6, 7, 8, 9, 10, 11], [0, 1, 2, 3, 12, 13, 14, 15]]
+    ]
+
+
 def test_update_draft_decode_metadata_updates_fa3_scheduler_metadata(
     monkeypatch,
 ):

tests/v1/worker/test_jit_warmup_migration.py

@@ -12,40 +12,24 @@ if not current_platform.is_cuda_alike():
 from vllm.v1.worker.block_table import ComputeSlotMappingKernel
 
 
-@pytest.mark.parametrize(
-    ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"),
-    [
-        (256, 1, 256, 16),
-        (256, 4, 64, 16),
-        (64, 1, 64, 16),
-        (8, 1, 8, 2),
-        (4, 1, 4, 2),
-    ],
-)
+@pytest.mark.parametrize("kv_cache_block_size", [256, 64, 8, 4])
 def test_compute_slot_mapping_warmup_matches_runtime_specializations(
     kv_cache_block_size: int,
-    blocks_per_kv_block: int,
-    block_size: int,
-    block_size_rep: int,
 ) -> None:
     kernel = ComputeSlotMappingKernel()
     kwargs = dict(
         kv_cache_block_size=kv_cache_block_size,
-        blocks_per_kv_block=blocks_per_kv_block,
         total_cp_world_size=1,
         total_cp_rank=0,
         cp_kv_cache_interleave_size=1,
         block_table_stride=32768,
-        block_size=block_size,
     )
     expected = kernel.CompileKey(
         kv_cache_block_size=kv_cache_block_size,
-        blocks_per_kv_block=blocks_per_kv_block,
         total_cp_world_size=1,
         total_cp_rank=0,
         cp_kv_cache_interleave_size=1,
         block_table_stride=16,
-        block_size=block_size_rep,
     )
 
     assert kernel.dispatch(**kwargs) == expected

tests/v1/worker/test_kv_cache_allocation_scope.py

@@ -112,6 +112,8 @@ def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
         model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="test")),
         compilation_config=SimpleNamespace(static_forward_context={}),
         kv_caches=[],
+        attn_groups=[],
+        speculative_config=None,
     )
     result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors(
         runner,

tests/v1/worker/test_mamba_utils.py

@@ -659,7 +659,6 @@ def test_v2_align_ctx_binds_request_slot_tables_for_state_copies():
         max_num_batched_tokens=64,
         max_num_blocks_per_group=[4],
         device=device,
-        kernel_block_sizes=[cfg.block_size],
     )
 
     state = object.__new__(MambaHybridModelState)

tests/v1/worker/test_utils.py

@@ -443,6 +443,43 @@ def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding():
     assert (backing[9] == -1).all()
 
 
+def test_bind_kv_cache_keeps_scheduler_block_views_for_runner():
+    """Block copies index scheduler blocks, so the runner keeps the unmapped
+    views while a packed layer is bound with kernel blocks."""
+
+    class _Layer:
+        def bind_kv_cache(self, kv_cache):
+            self.kv_cache = kv_cache
+
+    num_blocks, row = 4, 64
+    storage = torch.zeros(num_blocks * row)
+    # Two layers packed into one block row: an 8x4 page, then a 4x2 page
+    # split into two 2-row kernel blocks.
+    mla = storage.as_strided((num_blocks, 8, 4), (row, 4, 1))
+    indexer = storage.as_strided((num_blocks, 1, 4, 2), (row, 8, 2, 1), 40)
+    # The indexer in 2-row kernel blocks: block b's kernel block j is b * 16 + j.
+    mapped = indexer.as_strided((50, 1, 2, 2), (4, 8, 2, 1))
+    kv_caches = {"layers.0.attn": mla, "layers.1.indexer": indexer}
+    ctx = {name: _Layer() for name in kv_caches}
+    runner_kv_caches: list[torch.Tensor] = []
+
+    bind_kv_cache(
+        kv_caches,
+        ctx,
+        runner_kv_caches,
+        layer_kv_caches={**kv_caches, "layers.1.indexer": mapped},
+    )
+
+    assert ctx["layers.1.indexer"].kv_cache is mapped
+    assert runner_kv_caches == [mla, indexer]
+    copies = [KVCacheBlockCopy(1, 3)]
+    with pytest.raises(AssertionError, match="not divisible"):
+        copy_kv_cache_blocks_inplace([mapped], num_blocks, copies)
+    storage[row : 2 * row] = 7
+    copy_kv_cache_blocks_inplace(runner_kv_caches, num_blocks, copies)
+    assert (mla[3] == 7).all() and (indexer[3] == 7).all()
+
+
 @pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)])
 @pytest.mark.parametrize("has_copies", [False, True])
 def test_hisparse_host_copy_waits_for_writes_on_writer_only(

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 if TYPE_CHECKING:
@@ -351,6 +352,35 @@ def get_current_attn_backends(
         ]
 
 
+def get_current_attn_backends_and_specs(
+    vllm_config: VllmConfig,
+    kv_cache_config: KVCacheConfig,
+    fallback: list[type[AttentionBackend]],
+) -> tuple[list[type[AttentionBackend]], list[AttentionSpec | None]]:
+    """Distinct (backend, spec) pairs of the transfer layers, else ``fallback``.
+
+    Compressed caches are skipped: their kernel pages never split transferred
+    blocks.
+    """
+    pairs: dict[tuple[type[AttentionBackend], AttentionSpec | None], None] = {}
+    layer_type = cast(type[Any], AttentionLayerBase)
+    for group in kv_cache_config.transfer_groups:
+        specs = getattr(group.kv_cache_spec, "kv_cache_specs", {})
+        layers = get_layers_from_vllm_config(vllm_config, layer_type, group.layer_names)
+        for name, layer in layers.items():
+            spec = specs.get(name, group.kv_cache_spec)
+            if isinstance(spec, AttentionSpec) and spec.tokens_per_state > 1:
+                layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
+                assert layout.is_block_outermost, "Compressed caches must be packed"
+                continue
+            if not isinstance(spec, AttentionSpec):
+                spec = None
+            pairs[layer.get_attn_backend(), spec] = None
+    if not pairs:
+        return fallback, [None] * len(fallback)
+    return [backend for backend, _ in pairs], [spec for _, spec in pairs]
+
+
 def get_current_attn_backend(
     vllm_config: VllmConfig, layer_names: list[str] | None = None
 ) -> type[AttentionBackend]:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -26,6 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineId,
     TransferTopology,
     get_current_attn_backends,
+    get_current_attn_backends_and_specs,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
@@ -59,8 +60,8 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
+    CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     KVCacheSpec,
     MambaSpec,
     MLAAttentionSpec,
@@ -69,7 +70,6 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.metrics.cache_hit_source import CacheHitSource
 from vllm.v1.request import RequestStatus
-from vllm.v1.worker.block_table import BlockTable
 from vllm.v1.worker.utils import select_common_block_size
 
 logger = init_logger(__name__)
@@ -1350,7 +1350,10 @@ class MooncakeConnectorWorker:
         # physical block sizes. Pick the common (smallest) block size so that
         # KV-cache registration and transfer work correctly for both models.
         backends = get_current_attn_backends(self.vllm_config)
-        kernel_block_size = select_common_block_size(self.block_size, backends)
+        backends, specs = get_current_attn_backends_and_specs(
+            self.vllm_config, self.kv_cache_config, backends
+        )
+        kernel_block_size = select_common_block_size(self.block_size, backends, specs)
         if self.block_size != kernel_block_size:
             logger.info_once(
                 "User-specified logical block size (%s) does not match"
@@ -1727,11 +1730,12 @@ class MooncakeConnectorWorker:
         )
         group_specs = self.kv_cache_config.transfer_groups
         return [
-            BlockTable.map_to_kernel_blocks(
-                np.array(group),
-                self._physical_blocks_per_logical_kv_block,
-                block_arange,
-            ).tolist()
+            (
+                np.array(group)[:, None] * self._physical_blocks_per_logical_kv_block
+                + block_arange
+            )
+            .reshape(-1)
+            .tolist()
             if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
             else group
             for i, group in enumerate(block_ids)
@@ -2054,7 +2058,7 @@ class MooncakeConnectorWorker:
             # hetero PP can align by name (#56033). A non-contiguous row cannot
             # name a sub-span, so it is one region anchored at the allocation.
             use_packed_row = (
-                not isinstance(layer_spec, (MambaSpec, KpoolTailSpec))
+                not isinstance(layer_spec, (MambaSpec, CircularBufferSpec))
                 and storage_is_block_major
                 and block_stride > 0
                 and (is_mla_region or not hnc_contiguous)
@@ -2121,9 +2125,7 @@ class MooncakeConnectorWorker:
                 block_len = region_cache.stride(0) * region_cache.element_size()
                 region_base_addresses.append(base_addr)
 
-                if isinstance(layer_spec, KpoolTailSpec):
-                    kv_block_len = layer_spec.unpadded_page_size_bytes // 2
-                elif isinstance(layer_spec, AttentionSpec) and block_is_contiguous:
+                if isinstance(layer_spec, AttentionSpec) and block_is_contiguous:
                     assert (
                         layer_spec.page_size_bytes
                         % self._physical_blocks_per_logical_kv_block
@@ -2133,6 +2135,8 @@ class MooncakeConnectorWorker:
                         layer_spec.page_size_bytes
                         // self._physical_blocks_per_logical_kv_block
                     )
+                elif isinstance(layer_spec, MambaSpec) and block_is_contiguous:
+                    kv_block_len = layer_spec.page_size_bytes
                 else:
                     kv_block_len = block_len
                 if kv_block_len > block_len:

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

@@ -28,6 +28,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineTransferInfo,
     TransferTopology,
     get_current_attn_backends,
+    get_current_attn_backends_and_specs,
     kv_postprocess_blksize_and_layout_on_receive,
     kv_postprocess_blksize_on_receive,
 )
@@ -79,7 +80,6 @@ from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.kv_cache_interface import (
     CircularBufferSpec,
     FullAttentionSpec,
-    KpoolTailSpec,
     KVCacheLayout,
     KVCacheSpec,
     MambaSpec,
@@ -88,7 +88,6 @@ from vllm.v1.kv_cache_interface import (
     UniformTypeKVCacheSpecs,
     iter_layer_specs,
 )
-from vllm.v1.worker.block_table import BlockTable
 from vllm.v1.worker.utils import select_common_block_size
 
 if TYPE_CHECKING:
@@ -120,46 +119,6 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
     return len(block_strides) == len(storage_ptrs) == 1
 
 
-def _tensor_byte_span_end(cache: torch.Tensor) -> int:
-    """Return the exclusive end address touched by a nonnegative-stride view."""
-    if cache.numel() == 0:
-        return cache.data_ptr()
-    if any(stride < 0 for stride in cache.stride()):
-        raise ValueError("NIXL cache views must have nonnegative strides")
-    max_element_offset = sum(
-        (size - 1) * stride for size, stride in zip(cache.shape, cache.stride())
-    )
-    return cache.data_ptr() + (max_element_offset + 1) * cache.element_size()
-
-
-def _uses_dense_virtual_transfer_pages(
-    layer_spec: KVCacheSpec,
-    cache: torch.Tensor,
-    physical_page_size: int,
-    num_blocks: int,
-) -> bool:
-    """Return whether a compressed kernel view can be split into NIXL pages."""
-    if not (
-        isinstance(layer_spec, MLAAttentionSpec)
-        and layer_spec.tokens_per_state > 1
-        and cache.ndim == 4
-        and cache.shape[1] == 1
-        and cache.is_contiguous()
-        and physical_page_size > 0
-        and layer_spec.state_content_size_bytes > 0
-    ):
-        return False
-
-    block_stride = cache.stride(0) * cache.element_size()
-    return (
-        block_stride > physical_page_size
-        and block_stride % physical_page_size == 0
-        and physical_page_size % layer_spec.state_content_size_bytes == 0
-        and cache.shape[0] * (block_stride // physical_page_size) == num_blocks
-        and cache.nbytes == num_blocks * physical_page_size
-    )
-
-
 class NixlBaseConnectorWorker:
     """Base implementation of Worker side methods shared by pull and push."""
 
@@ -985,8 +944,10 @@ class NixlBaseConnectorWorker:
             )
 
     def _sync_block_size_with_kernel(self) -> None:
-        backends = get_current_attn_backends(self.vllm_config)
-        kernel_block_size = select_common_block_size(self.block_size, backends)
+        backends, specs = get_current_attn_backends_and_specs(
+            self.vllm_config, self.kv_cache_config, self.attn_backends
+        )
+        kernel_block_size = select_common_block_size(self.block_size, backends, specs)
         # Number of blocks not accounting for kernel block mismatches
         self._logical_num_blocks = self.num_blocks
         if self.block_size != kernel_block_size:
@@ -1425,7 +1386,7 @@ class NixlBaseConnectorWorker:
         if self._supports_pp_hma and self._has_packed_cache:
             # PP can reorder the same backends across stages of a packed model.
             backend_name = ",".join(
-                sorted(backend.get_name() for backend in self.attn_backends)
+                sorted({backend.get_name() for backend in self.attn_backends})
             )
         self.compat_hash = compute_nixl_compatibility_hash(
             self.vllm_config,
@@ -1478,31 +1439,6 @@ class NixlBaseConnectorWorker:
         ):
             raise NotImplementedError("PP push with packed KV caches requires MLA")
 
-        layer_specs: dict[str, KVCacheSpec] = {}
-        compressed_region_owners: dict[int, torch.Tensor] = {}
-        for layer_name, cache in xfer_buffers.items():
-            layer_spec = self._layer_specs.get(layer_name)
-            if isinstance(layer_spec, UniformTypeKVCacheSpecs):
-                layer_spec = layer_spec.kv_cache_specs[layer_name]
-            if layer_spec is None:
-                continue
-            layer_specs[layer_name] = layer_spec
-            physical_page_size = (
-                layer_spec.page_size_bytes
-                if isinstance(layer_spec, MambaSpec)
-                else layer_spec.page_size_bytes
-                // self._physical_blocks_per_logical_kv_block
-            )
-            num_blocks = (
-                self._logical_num_blocks
-                if isinstance(layer_spec, MambaSpec)
-                else self.num_blocks
-            )
-            if _uses_dense_virtual_transfer_pages(
-                layer_spec, cache, physical_page_size, num_blocks
-            ):
-                compressed_region_owners.setdefault(cache.data_ptr(), cache)
-
         track_region_layers = self._tracks_region_layers()
         region_layers: list[list[str]] = []
         packed_member_layouts: dict[str, tuple[int, int]] = {}
@@ -1527,7 +1463,7 @@ class NixlBaseConnectorWorker:
             # kernel requires a specific block size. This leads to SSM and FA layers
             # having different num_blocks.
             # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
-            layer_spec = layer_specs.get(layer_name)
+            layer_spec = self._layer_specs.get(layer_name)
             if layer_spec is None:
                 logger.debug(
                     "Skipping layer %s as no KVCache spec is present. "
@@ -1587,32 +1523,6 @@ class NixlBaseConnectorWorker:
                     region_device_id,
                 )
 
-            if isinstance(layer_spec, KpoolTailSpec):
-                compressed_owner = compressed_region_owners.get(cache.data_ptr())
-                if compressed_owner is not None:
-                    owner_storage = compressed_owner.untyped_storage()
-                    owner_end = compressed_owner.data_ptr() + compressed_owner.nbytes
-                    tail_is_covered = (
-                        compressed_owner.is_contiguous()
-                        and owner_storage.data_ptr() == storage_addr
-                        and _tensor_byte_span_end(cache) <= owner_end
-                    )
-                    if not tail_is_covered:
-                        raise AssertionError(
-                            "Kpool tail cache is not fully covered by its compressed "
-                            f"indexer region: layer={layer_name}, "
-                            f"tail_shape={tuple(cache.shape)}, "
-                            f"tail_stride={tuple(cache.stride())}, "
-                            f"owner_shape={tuple(compressed_owner.shape)}, "
-                            f"owner_stride={tuple(compressed_owner.stride())}"
-                        )
-                    logger.debug(
-                        "Skipping layer %s because its compressed indexer region "
-                        "covers the same storage",
-                        layer_name,
-                    )
-                    continue
-
             if isinstance(layer_spec, MambaSpec):
                 physical_ratio = self._physical_blocks_per_logical_kv_block
                 block_len = physical_page_size // physical_ratio
@@ -1647,9 +1557,6 @@ class NixlBaseConnectorWorker:
                     and cache[0].is_contiguous()
                     and cache[0].nbytes == physical_page_size
                 )
-                virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
-                    layer_spec, cache, physical_page_size, num_blocks
-                )
                 # Push workers address the pages of packed MLA rows by layer.
                 packed_mla_push = (
                     track_region_layers
@@ -1658,12 +1565,7 @@ class NixlBaseConnectorWorker:
                     and storage_is_block_major
                     and is_mla_region
                 )
-                if virtual_transfer_pages:
-                    # A compressed kernel row can contain multiple NIXL transfer pages.
-                    region_specs = [
-                        (cache.data_ptr(), physical_page_size, physical_page_size)
-                    ]
-                elif packed_mla_push and use_layer_name_routing:
+                if packed_mla_push and use_layer_name_routing:
                     # PP>1 prefill stage: register only this layer's page of each row.
                     region_specs = [
                         (cache.data_ptr(), physical_page_size, block_stride)
@@ -1671,7 +1573,7 @@ class NixlBaseConnectorWorker:
                 elif packed_mla_push:
                     # PP=1 decode or prefill: register whole rows, and advertise where
                     # this layer's page sits so a PP>1 prefill can write into it.
-                    storage_block_len = registration_len // num_blocks
+                    storage_block_len = block_stride
                     region_specs = [
                         (registration_base, storage_block_len, storage_block_len)
                     ]
@@ -3388,11 +3290,9 @@ class NixlBaseConnectorWorker:
                 physical_block_ids.append(group)
             else:
                 physical_block_ids.append(
-                    BlockTable.map_to_kernel_blocks(
-                        np.array(group),
-                        ratio,
-                        block_arange,
-                    ).tolist()
+                    (np.array(group)[:, None] * ratio + block_arange)
+                    .reshape(-1)
+                    .tolist()
                 )
         return physical_block_ids
 

vllm/model_executor/layers/attention/mla_attention.py

@@ -1376,14 +1376,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             is_index_group_leader=self.indexer is not None,
             non_causal_multi_token_decode=self.non_causal_multi_token_decode,
         )
-        # SM100 FlashMLA paged kernels also express TMA coordinates in token rows.
-        uses_tma_rows = (
-            self.attn_backend.get_name() == "FLASHMLA_SPARSE"
-            and self.kv_cache_dtype in ("fp8_ds_mla", "nvfp4_ds_mla")
-            and current_platform.is_device_capability_family(100)
-        )
-        if self._uses_flat_kv_cache() or uses_tma_rows:
-            spec = replace(spec, block_stride_alignment=spec.state_content_size_bytes)
         return spec
 
     def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):

vllm/models/deepseek_v41/nvidia/model_state.py

@@ -669,7 +669,7 @@ class DeepseekV41ModelState(DefaultModelState):
                 copy.create_metadata_builders(
                     self.vllm_config,
                     self.device,
-                    group.metadata_builders[0].kernel_block_size,
+                    group.kernel_block_size,
                 )
                 copies.append(copy)
             own_groups.append(copies)

vllm/models/glm5next/common/attention.py

@@ -33,9 +33,9 @@ from vllm.model_executor.utils import maybe_disable_graph_partition
 from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
 from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool
 from vllm.platforms import current_platform
-from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
 from vllm.utils.math_utils import cdiv, next_power_of_2
-from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec
+from vllm.v1.attention.backends.mla.indexer import Glm5NextIndexerBackend
+from vllm.v1.kv_cache_interface import CircularBufferSpec, MLAAttentionSpec
 
 logger = init_logger(__name__)
 
@@ -89,15 +89,9 @@ class Glm5NextIndexerCache(DeepseekV32IndexerCache):
 
     The indexer shares one block with the co-located MLA (a single
     ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide
-    ``cache_config.block_size``. DeepGEMM's paged-MQA kernel
-    (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or
-    64, so the storage block is virtually split into pool pages of the
-    largest such size that tiles it (``storage_kernel_block_size``); this
-    needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for
-    ``index_kpool = 16``). A smaller block (e.g. the default 64) silently
-    collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at
-    the opaque C++ assert; ``get_kv_cache_spec`` guards this up front
-    instead.
+    ``cache_config.block_size``. DeepGEMM's paged-MQA kernel takes 32- or
+    64-state pages (only 64 on SM120), which ``Glm5NextIndexerBackend``
+    declares as its kernel block sizes.
     """
 
     def __init__(
@@ -113,12 +107,6 @@ class Glm5NextIndexerCache(DeepseekV32IndexerCache):
             head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
         )
         assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1"
-        # Keep chunked-prefill boundaries aligned to complete pools.
-        assert cache_config.block_size % index_kpool == 0, (
-            "Glm5NextIndexerCache: cache_config.block_size "
-            f"({cache_config.block_size}) must be a multiple of index_kpool "
-            f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned."
-        )
         self._index_kpool = index_kpool
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig):
@@ -128,33 +116,10 @@ class Glm5NextIndexerCache(DeepseekV32IndexerCache):
         # ``tokens_per_state`` is the KV-spec representation of kpool
         # compression in the current cache-layout API.
         assert isinstance(spec, MLAAttentionSpec)
-        spec = replace(spec, tokens_per_state=self._index_kpool)
+        return replace(spec, tokens_per_state=self._index_kpool)
 
-        # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block
-        # (= block_size // index_kpool) is virtually split into pool pages of
-        # the largest such size that tiles it, so it must be a multiple of 32.
-        storage_block_size = spec.block_size // self._index_kpool
-        assert (
-            spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0
-        ), (
-            "Glm5NextIndexerCache: kpool indexer requires cache block_size to "
-            f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so "
-            "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the "
-            f"storage block, got block_size={spec.block_size} -> "
-            f"storage_block_size={storage_block_size}."
-        )
-        max_page_size = max(PAGED_MQA_PAGE_SIZES)
-        min_page_size = min(PAGED_MQA_PAGE_SIZES)
-        if storage_block_size <= max_page_size:
-            page_size = storage_block_size
-        elif storage_block_size % max_page_size == 0:
-            page_size = max_page_size
-        else:
-            page_size = min_page_size
-        return replace(
-            spec,
-            storage_block_size=page_size * self._index_kpool,
-        )
+    def get_attn_backend(self):
+        return Glm5NextIndexerBackend
 
 
 class Glm5NextTailCache(DeepseekV32IndexerCache):
@@ -164,9 +129,9 @@ class Glm5NextTailCache(DeepseekV32IndexerCache):
     ``ring`` slots per request, overwritten in place by ``pos % ring``
     as decode/spec-decode advances. Prefill seeds it (instead of discarding the
     tail raw K+gate); the connector transfers it across PD; decode reads it to
-    compress the boundary pool correctly. ``KpoolTailSpec`` /
-    ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets
-    the in-progress pool survive across steps and across transfer.
+    compress the boundary pool correctly. ``CircularBufferSpec`` provides the
+    no-prune, 1-block/req allocation that lets the in-progress pool survive
+    across steps and across transfer.
 
     Stores raw bf16 K (``head_dim``) as the "K" half of each block and the
     bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed
@@ -202,13 +167,12 @@ class Glm5NextTailCache(DeepseekV32IndexerCache):
             f"({self.cache_config.block_size}) must be a multiple of the "
             f"tail ring ({ring})"
         )
-        return KpoolTailSpec(
+        return CircularBufferSpec(
             block_size=ring,
             num_kv_heads=2,
             head_size=self.head_dim,
             head_size_v=0,
             dtype=torch.bfloat16,
-            sliding_window=ring,
             # Tiny per-request ring; every DCP rank keeps the full tail so the
             # rank owning a pool can compress it.
             dcp_sharded=False,
@@ -303,7 +267,7 @@ class Indexer(nn.Module):
         # Paged tail cache (in-progress pool's raw K + gate score). Written by
         # prefill (seeds the boundary pool) and decode (per-step stash); read by
         # the decode kernel to compress the boundary pool. Transferred across PD
-        # so the decode side sees the prefill tail. See KpoolTailSpec/Manager.
+        # so the decode side sees the prefill tail. See CircularBufferManager.
         self.tail_cache = Glm5NextTailCache(
             head_dim=self.head_dim,
             dtype=torch.bfloat16,

vllm/models/glm5next/nvidia/ops/kpool_compress.py

@@ -396,11 +396,6 @@ def _kpool_tail_seed_kernel(
     ahead belongs to a different tail block (or is past the batch / padding,
     slot < 0). ``tslot = block * RING + pos % RING``; the destination is
     ``tail[block, {0:K, 1:score}, pos % RING, :]``.
-
-    The tail cache aliases the indexer cache with the indexer's (padded) block
-    stride, so blocks are addressed through ``TAIL_BLOCK_ELEMS`` /
-    ``KPOOL_HEAD`` (``tail.stride(0)`` / ``tail.stride(1)``), never as a dense
-    ``[num_blocks, 2, KPOOL, HEAD_DIM]`` array.
     """
     i = tl.program_id(0)
     t = tl.load(tslot_ptr + i).to(tl.int64)

vllm/platforms/cuda.py

@@ -414,26 +414,6 @@ class CudaPlatformBase(Platform):
 
         return valid_backends_priorities, invalid_reasons
 
-    @classmethod
-    def _get_indexer_block_alignment(cls, vllm_config: VllmConfig) -> int | None:
-        index_kpool = getattr(
-            vllm_config.model_config.hf_text_config, "index_kpool", None
-        )
-        if not index_kpool or index_kpool <= 1:
-            return None
-        from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
-
-        # kpool paged-MQA indexer: the storage block (block_size /
-        # index_kpool) is virtually split into pool pages, so block_size
-        # must be a multiple of index_kpool times a legal pool page.
-        page = min(PAGED_MQA_PAGE_SIZES)
-        if cls.is_device_capability_family(120):
-            # On sm120 the DeepGEMM paged-MQA kernel only accepts block_kv
-            # 64 for the fp8 indexer cache, so align to the largest pool
-            # page here to make the page split land on 64 not the min 32.
-            page = max(PAGED_MQA_PAGE_SIZES)
-        return index_kpool * page
-
     @classmethod
     def get_attn_backend_cls(
         cls,

vllm/platforms/interface.py

@@ -737,14 +737,14 @@ class Platform:
             for candidate in candidates:
                 if all(b.supports_block_size(candidate) for b in backend_classes):
                     return candidate
-        raise ValueError(
-            "The attention backends share no supported KV cache block size ("
-            + "; ".join(
-                f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}"
-                for b in backend_classes
+            raise ValueError(
+                "The attention backends share no supported KV cache block size ("
+                + "; ".join(
+                    f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}"
+                    for b in backend_classes
+                )
+                + ")."
             )
-            + ")."
-        )
 
     @classmethod
     def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -935,17 +935,6 @@ class Platform:
         if cache_config.mamba_page_size_padded is not None:
             cache_config.mamba_page_size_padded = shared_page
 
-    @classmethod
-    def _get_indexer_block_alignment(cls, vllm_config: "VllmConfig") -> int | None:
-        """Extra ``block_size`` multiple a sparse indexer needs, else ``None``.
-
-        The CUDA kpool paged-MQA indexer virtually splits each storage block
-        into pool pages, so ``block_size`` must be a multiple of
-        ``index_kpool * min(PAGED_MQA_PAGE_SIZES)`` — implemented in the CUDA
-        platform override. Other platforms impose no extra constraint.
-        """
-        return None
-
     @classmethod
     def _align_hybrid_block_size(
         cls,
@@ -1077,6 +1066,14 @@ class Platform:
                 ),
                 cache_config.block_size,
             )
+            # Also a multiple of every backend's smallest kernel block.
+            backend_block_alignment = lcm(
+                *(
+                    min(s.base if isinstance(s, MultipleOf) else s for s in sizes)
+                    for b in cls._find_non_ssm_backends(vllm_config)
+                    if (sizes := b.get_supported_kernel_block_sizes())
+                )
+            )
             if model_config.use_mla:
                 # TRTLLM/FlashInfer MLA decode kernels require the physical
                 # number of kernel blocks to be aligned to 128 / kernel_block_size.
@@ -1090,9 +1087,9 @@ class Platform:
             mamba_page_size,
             kernel_block_alignment_size * attn_page_size_1_token,
         )
-        indexer_align = cls._get_indexer_block_alignment(vllm_config)
-        if indexer_align:
-            attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align)
+        attn_block_size = backend_block_alignment * cdiv(
+            attn_block_size, backend_block_alignment
+        )
 
         if cache_config.block_size < attn_block_size:
             cache_config.block_size = attn_block_size

vllm/utils/deep_gemm.py

@@ -29,10 +29,11 @@ _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: set[str] = {
     "qwen3_5_moe_text",
 }
 
-# KV page sizes (in cache entries) supported by the paged-MQA logits kernels
-# (fp8_fp4_paged_mqa_logits / get_paged_mqa_logits_metadata). Larger storage
-# blocks must be virtually split into one of these page sizes.
-PAGED_MQA_PAGE_SIZES = (32, 64)
+
+def get_paged_mqa_page_sizes() -> tuple[int, ...]:
+    """KV page sizes (in cache entries) the paged-MQA logits kernels take (only
+    64 on SM120); larger storage blocks are split into one of these."""
+    return (64,) if current_platform.is_device_capability_family(120) else (32, 64)
 
 
 def should_auto_disable_deep_gemm(model_type: str | None) -> bool:

vllm/v1/attention/backend.py

@@ -48,12 +48,10 @@ class AttentionType(str, Enum):
     """Attention between dec. Q and enc. K/V for encoder-decoder."""
 
 
+@dataclass(frozen=True)
 class MultipleOf:
     base: int
 
-    def __init__(self, base: int):
-        self.base = base
-
     def __repr__(self) -> str:
         return f"MultipleOf({self.base})"
 
@@ -622,10 +620,6 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
         self.layer_names = layer_names
         self.vllm_config = vllm_config
         self.device = device
-        self.kernel_block_size: int | None = None
-
-    def set_kernel_block_size(self, kernel_block_size: int) -> None:
-        self.kernel_block_size = kernel_block_size
 
     @classmethod
     def get_cudagraph_support(

vllm/v1/attention/backends/composite.py

@@ -297,11 +297,6 @@ def create_composite_attention_backend(
             )
             self.routing = routing_policy(layer_names, vllm_config)
 
-        def set_kernel_block_size(self, kernel_block_size):
-            super().set_kernel_block_size(kernel_block_size)
-            for builder in self._builders:
-                builder.set_kernel_block_size(kernel_block_size)
-
         @classmethod
         def get_cudagraph_support(cls, vllm_config, kv_cache_spec):
             return min(

vllm/v1/attention/backends/mla/compressor_utils.py

@@ -95,7 +95,9 @@ class CompressedSlotMappingKernel(
                 block_table_ptr + batch_idx * block_table_stride + block_ids,
                 mask=mask & is_valid,
             )
-            slot_ids = block_numbers * block_size + pos_after_compress % block_size
+            slot_ids = block_numbers.to(tl.int64) * block_size + (
+                pos_after_compress % block_size
+            )
 
             # NOTE
             slot_ids = tl.where(is_valid, slot_ids, PAD_ID)

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -24,6 +24,7 @@ from vllm.v1.attention.backend import (
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
     flat_kv_row_view,
     triton_convert_req_index_to_global_index,
 )
@@ -42,7 +43,11 @@ class FlashAttnMLASparseBackend(AttentionBackend):
 
     @staticmethod
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
-        return [64]
+        return [MultipleOf(64)]
+
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        return align_blocks_to_rows(spec)
 
     @staticmethod
     def get_name() -> str:

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """FlashInfer sparse MLA attention backend."""
 
+import math
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar
 
@@ -30,6 +31,7 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
     flat_kv_row_view,
     prepare_sparse_mla_safe_lengths,
     triton_convert_req_index_to_global_index,
@@ -69,6 +71,13 @@ class _FlashInferMLASparseBackendBase(AttentionBackend):
     def is_sparse(cls) -> bool:
         return True
 
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        # The kernels read the rows as pages of their kernel block size.
+        (page,) = cls.get_supported_kernel_block_sizes(spec)
+        assert isinstance(page, MultipleOf)
+        return align_blocks_to_rows(spec, page.base)
+
 
 class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
     """FlashInfer sparse MLA backend using the TRTLLM-gen launcher."""
@@ -84,7 +93,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
 
     @staticmethod
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
-        return [32, 64]
+        return [MultipleOf(32)]
 
     @staticmethod
     def get_impl_cls() -> type[MLAAttentionImpl]:
@@ -173,7 +182,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
 
     @staticmethod
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
-        return [64, 256]
+        return [MultipleOf(64)]
 
     @staticmethod
     def get_impl_cls() -> type[MLAAttentionImpl]:
@@ -575,7 +584,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
             assert prefill_lse is not None
             return output, torch.cat((decode_lse, prefill_lse))
 
-        _, block_stride_rows = flat_kv_row_view(
+        kv_rows, block_stride_rows = flat_kv_row_view(
             kv_c_and_k_pe_cache, attn_metadata.block_size
         )
 
@@ -602,7 +611,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
 
         return self._run_mqa_kernel(
             q,
-            kv_c_and_k_pe_cache,
+            # Block sizes and strides have been aligned to 32 rows.
+            kv_rows.view(
+                -1,
+                math.gcd(block_stride_rows, attn_metadata.block_size, 64),
+                kv_rows.shape[-1],
+            ),
             topk_indices_physical,
             seq_lens,
         )

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py

@@ -16,6 +16,7 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
 )
 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    flat_kv_row_view,
     triton_convert_req_index_to_global_index,
 )
 
@@ -184,6 +185,9 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada
             output = torch.cat(outputs) if len(outputs) > 1 else outputs[0]
             return output, None
 
+        kv_rows, block_stride_rows = flat_kv_row_view(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
         topk_indices_physical = cast(
             torch.Tensor,
             triton_convert_req_index_to_global_index(
@@ -191,13 +195,15 @@ class FlashInferMLASparseSM120Impl(SparseMLACommonImpl[FlashInferMLASparseMetada
                 attn_metadata.block_table,
                 topk_indices,
                 BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
             ),
         )
         return (
             self._run_mqa_kernel(
                 q,
-                kv_c_and_k_pe_cache,
+                # The SM120 kernels are built for 64-row pages only.
+                kv_rows.view(-1, 64, kv_rows.shape[-1]),
                 topk_indices_physical,
             ),
             None,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

@@ -26,7 +26,7 @@ needs no D2H sync. Per-step content (top-k slots) is written into the
 reserved buffers by kernels inside the captured forward, and captured runs
 read the refreshed plan buffers on replay.
 
-KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]``
+KV cache format: E4M3 ``[num_blocks, block_size, 512]`` (any block stride)
 (uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The
 per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the
 kernel but not wired yet (it needs a group-quantizing cache-write op).
@@ -57,9 +57,11 @@ from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseMetadataBuilder,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
+    flat_kv_row_view,
     triton_convert_req_index_to_global_index,
 )
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
+from vllm.v1.kv_cache_interface import AttentionSpec
 
 _FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
 _WORKSPACE_BYTES = 128 * 1024 * 1024
@@ -99,7 +101,12 @@ def _pack_topk_indices(
     end = tl.load(indptr + row + 1)
     mask = (cols < end - start) & (cols < WIDTH)
     values = tl.load(slots + row * WIDTH + cols, mask=mask, other=0)
-    tl.store(indices + start + cols, tl.maximum(values, 0), mask=mask)
+    first_slot = tl.load(slots + row * WIDTH)
+    tl.store(
+        indices + start + cols,
+        tl.where(values >= 0, values, tl.maximum(first_slot, 0)),
+        mask=mask,
+    )
 
 
 @triton.jit
@@ -149,6 +156,10 @@ class FlashInferMLASparseSM90Backend(AttentionBackend):
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(64)]
 
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        return align_blocks_to_rows(spec)
+
     @staticmethod
     def get_name() -> str:
         return "FLASHINFER_MLA_SPARSE_SM90"
@@ -222,10 +233,6 @@ class FlashInferMLASparseSM90Backend(AttentionBackend):
     ) -> tuple[int, ...]:
         return (num_blocks, block_size, head_size)
 
-    @classmethod
-    def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
-        return (KVCacheLayout.LBHNC,)
-
 
 class _SM90State:
     """Builder-owned wrapper, capture-stable buffers, and plan parameters.
@@ -602,6 +609,9 @@ class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Met
 
         assert self.topk_indices_buffer is not None
         topk_indices = self.topk_indices_buffer[:num_tokens]
+        kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
         # return_valid_counts=True keeps the compacted-prefix layout: valid
         # entries at [0, valid_count), -1 past it — exactly the prefix the
         # planned per-row lengths address.
@@ -610,6 +620,7 @@ class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Met
             attn_metadata.block_table,
             topk_indices,
             BLOCK_SIZE=attn_metadata.block_size,
+            BLOCK_STRIDE_ROWS=block_stride_rows,
             NUM_TOPK_TOKENS=topk_indices.shape[1],
             return_valid_counts=True,
         )

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, ClassVar
 
 import numpy as np
@@ -36,6 +36,7 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
     flat_kv_row_view,
     request_row_bounds,
     triton_convert_req_index_to_global_index,
@@ -53,7 +54,7 @@ from vllm.v1.attention.ops.flashmla import (
     flash_mla_with_kvcache,
     get_mla_metadata,
 )
-from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode
 from vllm.v1.worker.workspace import current_workspace_manager
 
 if TYPE_CHECKING:
@@ -132,6 +133,17 @@ QUANTIZED_DS_MLA_CACHE_FORMATS: frozenset[str] = frozenset(
 FP8_DS_MLA_ROPE_DIM = 64
 
 
+def _fp8_kv_pages(
+    kv_cache: torch.Tensor, block_size: int
+) -> tuple[torch.Tensor, int | None]:
+    """The cache as the fp8 kernel's 64-token pages, and the block stride in
+    rows when a block spans several pages."""
+    if block_size == 64:
+        return kv_cache, None
+    rows, block_stride_rows = flat_kv_row_view(kv_cache, block_size)
+    return rows.view(-1, 64, rows.shape[-1]), block_stride_rows
+
+
 class FlashMLASparseBackend(AttentionBackend):
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
@@ -144,7 +156,23 @@ class FlashMLASparseBackend(AttentionBackend):
 
     @staticmethod
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
-        return [64]
+        if (
+            isinstance(kv_cache_spec, AttentionSpec)
+            and kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
+            and current_platform.is_device_capability_family(100)
+        ):
+            return [64]
+        return [MultipleOf(64)]
+
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        # Off SM100 the quantized kernel reads 64-token pages, only needed
+        # when blocks span several.
+        if spec.kv_quant_mode != KVQuantMode.NONE and (
+            not current_platform.is_device_capability_family(100)
+        ):
+            return replace(spec, block_stride_alignment=MultipleOf(64))
+        return align_blocks_to_rows(spec)
 
     @staticmethod
     def get_name() -> str:
@@ -998,13 +1026,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
         # and the prefill rows per chunk, against the gathered workspace.
         lse: torch.Tensor | None = None
         topk_length: torch.Tensor | None = None
+        kv_pages, block_stride_rows = _fp8_kv_pages(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
         if self.pcp_dcp_kv_gather:
             pass
         elif num_prefill_tokens == 0 and not uses_host_cache:
             topk_indices, topk_length = self._convert_logical_to_physical_topk(
                 topk_indices,
                 attn_metadata,
-                block_stride_rows=None,
+                block_stride_rows=block_stride_rows,
                 return_valid_counts=True,
             )
         elif num_prefill_tokens > 0:
@@ -1013,6 +1044,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
                 attn_metadata.block_table,
                 topk_indices,
                 BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
                 HAS_PREFILL_WORKSPACE=has_prefill_workspace,
                 prefill_workspace_request_ids=prefill_request_ids,
@@ -1056,7 +1088,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
             topk_indices = topk_indices.view(num_decodes, seq_len, -1)
             attn_out, _ = self._fp8_flash_mla_kernel(
                 q=q,
-                kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
+                kv_c_and_k_pe_cache=kv_pages,
                 topk_indices=topk_indices,
                 kernel_metadata=fp8_metadata.decode.kernel_metadata,
             )
@@ -1173,6 +1205,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
         # req_id_per_token covers the whole batch; slice it to the MQA tokens
         # (q may exclude prefill tokens routed to dense MHA).
         req_id_per_token = attn_metadata.req_id_per_token[: topk_indices.shape[0]]
+        kv_pages, block_stride_rows = _fp8_kv_pages(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
         if self.dcp_world_size > 1:
             # The indexer emits global token ids; keep this rank's shard and
             # convert to local slots. compact_valid_to_front=False keeps the
@@ -1187,6 +1222,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
                 dcp_rank=self.dcp_rank,
                 cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size,
                 BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=topk_indices.shape[1],
                 compact_valid_to_front=False,
             )
@@ -1202,7 +1238,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
                 topk_indices = self._convert_logical_to_physical_topk(
                     topk_indices,
                     attn_metadata,
-                    block_stride_rows=None,
+                    block_stride_rows=block_stride_rows,
                     return_valid_counts=False,
                 )
             else:
@@ -1211,12 +1247,13 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
                     block_table,
                     topk_indices,
                     BLOCK_SIZE=attn_metadata.block_size,
+                    BLOCK_STRIDE_ROWS=block_stride_rows,
                     NUM_TOPK_TOKENS=topk_indices.shape[1],
                 )
 
         _attn_out, _lse = self._fp8_flash_mla_kernel(
             q=q.unsqueeze(0),  # unsqueeze to add batch_dim: (T, H, D) -> (1, T, H, D)
-            kv_c_and_k_pe_cache=kv_c_and_k_pe_cache,
+            kv_c_and_k_pe_cache=kv_pages,
             topk_indices=topk_indices.unsqueeze(0),  # (T, topk) -> (1, T, topk)
             kernel_metadata=fp8_metadata,
         )

vllm/v1/attention/backends/mla/indexer.py

@@ -7,7 +7,7 @@ import numpy as np
 import torch
 
 import vllm.envs as envs
-from vllm.config import VllmConfig
+from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
 from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs
@@ -22,6 +22,7 @@ from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.deep_gemm import (
     get_paged_mqa_logits_metadata,
+    get_paged_mqa_page_sizes,
     has_deep_gemm,
     native_next_n_supported,
 )
@@ -234,12 +235,25 @@ class DeepseekV32IndexerBackend(AttentionBackend):
         return DeepseekV32IndexerMetadataBuilder
 
 
+class Glm5NextIndexerBackend(DeepseekV32IndexerBackend):
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        if kv_cache_spec is not None:
+            # Worker get_kv_cache_spec: runs without the current config.
+            index_kpool = int(kv_cache_spec.tokens_per_state)
+        else:
+            # Platform block-size selection: no spec yet, under the current config.
+            model_config = get_current_vllm_config().model_config
+            index_kpool = model_config.hf_text_config.index_kpool
+        return [index_kpool * n for n in get_paged_mqa_page_sizes()]
+
+
 class KpoolTailBackend(DeepseekV32IndexerBackend):
     """Storage-only backend for the GLM-5.3-Flash kpool tail cache."""
 
     @classmethod
     def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
-        return (KVCacheLayout.LBHNC,)
+        return (KVCacheLayout.BLHNC,)
 
     @staticmethod
     def get_name() -> str:
@@ -1065,8 +1079,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                 dtype=torch.int32,
                 device=self.device,
             )
-        self.indexer_decode_block_table_buffer: torch.Tensor | None = None
-        self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens
         # DCP not supported yet
         self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
 
@@ -1353,16 +1365,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens
 
         compressed_slot_mapping = slot_mapping
-        indexer_block_table = block_table
         if self.compress_ratio > 1:
-            kernel_block_size = self.kernel_block_size
-            if (
-                kernel_block_size is not None
-                and self.kv_cache_spec.block_size != kernel_block_size
-                and self.kv_cache_spec.block_size % kernel_block_size == 0
-            ):
-                factor = self.kv_cache_spec.block_size // kernel_block_size
-                indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
             padded_num_tokens = num_tokens
             local_slot_mapping = slot_mapping
             if self.use_pcp:
@@ -1378,7 +1381,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                 local_slot_mapping,
                 query_start_loc,
                 seq_lens,
-                indexer_block_table,
+                block_table,
                 self.kv_cache_spec.num_states,
                 self.compress_ratio,
                 out=self.compressed_slot_mapping_buffer,
@@ -1497,7 +1500,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                     seq_lens,
                     compressed_seq_lens,
                     compressed_seq_lens_cpu,
-                    indexer_block_table,
+                    block_table,
                     self.compress_ratio,
                     query_slice=query_slice,
                     skip_kv_gather=query_slice.start > 0,
@@ -1632,27 +1635,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                     )
                 )
 
-            if self.compress_ratio > 1:
-                kernel_block_size = self.kernel_block_size
-                if (
-                    kernel_block_size is not None
-                    and self.kv_cache_spec.block_size != kernel_block_size
-                    and self.kv_cache_spec.block_size % kernel_block_size == 0
-                ):
-                    factor = self.kv_cache_spec.block_size // kernel_block_size
-                    compressed = block_table[:, ::factor] // factor
-                    rows, cols = compressed.shape
-                    if self.indexer_decode_block_table_buffer is None:
-                        self.indexer_decode_block_table_buffer = torch.zeros(
-                            (self._max_num_batched_tokens, cols),
-                            dtype=torch.int32,
-                            device=self.device,
-                        )
-                    self.indexer_decode_block_table_buffer[:rows, :cols].copy_(
-                        compressed
-                    )
-                    block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
-
             # Flattening always returns a buffer view, including single-token
             # batches. Keep its address stable across varlen graph replays.
             seq_lens_is_buffer_view = not use_native or next_n > 1
@@ -1741,17 +1723,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         # One query token per request, so the per-token decode block table is
         # the request block table (at the indexer kernel's page size). build()
         # may copy it into a builder buffer, so re-derive it for this batch.
-        block_table = metadata.block_table[: metadata.num_decode_tokens]
-        kernel_block_size = self.kernel_block_size
-        if (
-            self.compress_ratio > 1
-            and kernel_block_size is not None
-            and self.kv_cache_spec.block_size != kernel_block_size
-            and self.kv_cache_spec.block_size % kernel_block_size == 0
-        ):
-            factor = self.kv_cache_spec.block_size // kernel_block_size
-            block_table = block_table[:, ::factor] // factor
-        decode.block_table.copy_(block_table)
+        decode.block_table.copy_(metadata.block_table[: metadata.num_decode_tokens])
         if decode.indices is not None:
             decode.indices.copy_(self.arange_buffer[: metadata.num_decode_tokens])
         if decode.per_req_decode_lens is not None:

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -35,6 +35,10 @@ from vllm.v1.attention.backend import (
 from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
     AiterMLAHelper,
 )
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
+    flat_kv_row_view,
+)
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     rocm_sparse_attn_decode_bf16,
@@ -160,6 +164,7 @@ def _convert_req_index_to_global_index_kernel(
     # shapes (compile-time where possible)
     max_num_blocks_per_req: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
+    BLOCK_STRIDE_ROWS: tl.constexpr,
     BLOCK_N: tl.constexpr,  # tile width along columns
     # strides (in elements)
     bt_stride0,
@@ -201,9 +206,8 @@ def _convert_req_index_to_global_index_kernel(
     bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
     base = tl.load(bt_ptr, mask=valid_block, other=0)
 
-    # # If token == -1 OR block_id OOB, output 0; else base * BLOCK_SIZE + offset
     out_val = tl.where(
-        is_invalid_tok | (~valid_block), 0, base * BLOCK_SIZE + inblock_off
+        is_invalid_tok | (~valid_block), 0, base * BLOCK_STRIDE_ROWS + inblock_off
     )
     out_ptr_ij = out_ptr + seq_start + indice_id
     out_ptr_ij_mask = (seq_start + indice_id) < seq_end
@@ -218,13 +222,14 @@ def triton_convert_req_index_to_global_index(
     token_indices: torch.Tensor,  # int32 [num_tokens, NUM_TOPK_TOKENS]
     cu_seqlens: torch.Tensor,  # int32 [num_tokens + 1]
     paged_kv_indices: torch.Tensor,  # int32 [num_tokens * topk] out_buffer
+    BLOCK_STRIDE_ROWS: int,
     BLOCK_SIZE: int = 64,
     NUM_TOPK_TOKENS: int = 2048,
     BLOCK_N: int = 128,  # tile width along columns
 ):
     """out[token_id, indice_id] =
         block_table[req_id[token_id],
-            token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE
+            token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_STRIDE_ROWS
         + token_indices[token_id, indice_id] % BLOCK_SIZE
 
     Only when token_indices[token_id, indice_id] == -1 do we output -1.
@@ -264,6 +269,7 @@ def triton_convert_req_index_to_global_index(
         # shapes / constexprs
         max_num_blocks_per_req,
         BLOCK_SIZE,
+        BLOCK_STRIDE_ROWS,
         BLOCK_N,
         # strides
         bt_stride0,
@@ -345,6 +351,10 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [1, MultipleOf(16)]
 
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        return align_blocks_to_rows(spec)
+
     @staticmethod
     def get_name() -> str:
         return "ROCM_AITER_MLA_SPARSE"
@@ -371,8 +381,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
 
     @classmethod
     def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
-        # Global index conversion assumes contiguous pages within each layer.
-        return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC)
+        return (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC, KVCacheLayout.BLHNC)
 
     @classmethod
     def supports_sink(cls) -> bool:
@@ -1155,6 +1164,9 @@ class ROCMAiterMLASparseImpl(
             self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens
         )
 
+        kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view(
+            kv_c_and_k_pe_cache, attn_metadata.block_size
+        )
         triton_convert_req_index_to_global_index(
             attn_metadata.req_id_per_token,
             attn_metadata.block_table,
@@ -1162,6 +1174,7 @@ class ROCMAiterMLASparseImpl(
             attn_metadata.paged_kv_indptr,
             attn_metadata.paged_kv_indices,
             BLOCK_SIZE=attn_metadata.block_size,
+            BLOCK_STRIDE_ROWS=block_stride_rows,
             NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
         )
 
@@ -1179,5 +1192,5 @@ class ROCMAiterMLASparseImpl(
             return output, None
         mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
         return self._forward_mla(
-            layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata
+            layer, mla_padded_q, kv_c_and_k_pe_cache.unsqueeze(1), attn_metadata
         )

vllm/v1/attention/backends/mla/sparse_utils.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utility functions for sparse MLA backends."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 import numpy as np
@@ -22,6 +22,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
 )
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
+from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.block_table import get_block_table_width
 
 
@@ -40,6 +41,12 @@ def request_row_bounds(req_idx: np.ndarray) -> np.ndarray:
     return bounds
 
 
+def align_blocks_to_rows(spec: AttentionSpec, rows: int = 1) -> AttentionSpec:
+    """Sparse kernels read the cache as flat rows, viewed as ``rows``-row pages,
+    so blocks must be whole pages apart."""
+    return replace(spec, block_stride_alignment=rows * spec.state_content_size_bytes)
+
+
 def flat_kv_row_view(
     kv_cache: torch.Tensor,  # [num_blocks, block_size, head_dim]
     block_size: int,

vllm/v1/attention/backends/mla/xpu_mla_sparse.py

@@ -27,6 +27,7 @@ from vllm.v1.attention.backend import (
     MLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
+    align_blocks_to_rows,
     flat_kv_row_view,
     triton_convert_req_index_to_global_index,
 )
@@ -50,6 +51,10 @@ class XPUMLASparseBackend(AttentionBackend):
     def get_name() -> str:
         return "XPU_MLA_SPARSE"
 
+    @classmethod
+    def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
+        return align_blocks_to_rows(spec)
+
     @staticmethod
     def get_metadata_cls() -> type["XPUMLASparseMetadata"]:
         return XPUMLASparseMetadata

vllm/v1/core/kv_cache_coordinator.py

@@ -661,8 +661,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         self.hash_block_size = hash_block_size
         self.dcp_world_size = dcp_world_size
         # Only groups that participate in prefix caching must satisfy the
-        # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool
-        # tail, block_size=kpool) are scratch buffers and excluded.
+        # divisibility constraint; groups that opt out (e.g. circular buffers
+        # such as the GLM-5.3-Flash indexer tail) are scratch and excluded.
         cacheable_block_sizes = [
             manager.block_size
             for manager, group in zip(
@@ -724,8 +724,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         """
         self.attention_groups: list[SpecGroup] = []
         for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
-            # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
-            # kpool tail): their blocks are per-request scratch, never
+            # Skip groups that opt out of prefix caching (e.g. circular
+            # buffers): their blocks are per-request scratch, never
             # shareable, so they must not participate in hit lookup (their
             # manager-level hooks already no-op). Their slot in the per-group
             # hit tuple stays empty.

vllm/v1/core/kv_cache_utils.py

@@ -10,7 +10,7 @@ from collections import defaultdict
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field, replace
 from functools import partial, reduce
-from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
+from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, overload
 
 from vllm import envs
 from vllm.config import VllmConfig
@@ -32,7 +32,6 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     HiddenStateCacheSpec,
     HiSparseHotSpec,
-    KpoolTailSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheLayout,
@@ -1126,7 +1125,7 @@ def check_enough_kv_cache_memory(
         # of the specs since grouping may unify them in-place.
         groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec))
         check_memory = (
-            available_memory - _pool_bytes_per_block(groups)
+            available_memory - _pool_bytes_per_block(vllm_config, groups)
             if groups
             else available_memory
         )
@@ -1237,14 +1236,17 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
     return num_blocks
 
 
-def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
+def _pool_bytes_per_block(
+    vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec]
+) -> int:
     """Bytes consumed by one block in the worker's shared KV cache pool, mirroring
     the divisor used by `get_kv_cache_config_from_groups` to convert
     `available_memory` into `num_blocks`. Used to compute the effective KV cache
     capacity once `num_gpu_blocks_override` is applied.
     """
     return _get_kv_cache_bytes_per_block(
-        [group for group in kv_cache_groups if not group.host_resident]
+        [group for group in kv_cache_groups if not group.host_resident],
+        vllm_config.cache_config.get_resolved_kv_cache_layout(),
     )
 
 
@@ -1287,205 +1289,6 @@ def _get_kv_cache_groups_uniform_type(
     return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
 
 
-def _pp_balanced_mamba_group_count(
-    vllm_config: VllmConfig,
-    mamba_layer_names: list[str],
-    mla_layer_names: list[str],
-) -> int | None:
-    """Return a Mamba group count whose PP projections fit the MLA slots."""
-    num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names))
-    pp_size = vllm_config.parallel_config.pipeline_parallel_size
-    if pp_size == 1:
-        return num_groups
-
-    from vllm.distributed.utils import get_pp_indices
-    from vllm.model_executor.models.utils import extract_layer_index
-
-    total_layers = vllm_config.model_config.get_total_num_hidden_layers()
-    mamba_indices = [extract_layer_index(name) for name in mamba_layer_names]
-    mla_indices = [extract_layer_index(name) for name in mla_layer_names]
-    for rank in range(pp_size):
-        start, end = get_pp_indices(total_layers, rank, pp_size)
-        num_mamba = sum(start <= index < end for index in mamba_indices)
-        num_mla = sum(start <= index < end for index in mla_indices)
-        if not num_mamba:
-            continue
-        if not num_mla:
-            return None
-        num_groups = max(num_groups, cdiv(num_mamba, num_mla))
-    return num_groups
-
-
-def _get_kv_cache_groups_glm5_next(
-    vllm_config: VllmConfig,
-    kv_cache_spec: dict[str, KVCacheSpec],
-) -> list[KVCacheGroupSpec] | None:
-    """Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing."""
-    mamba_specs = {
-        name: spec
-        for name, spec in kv_cache_spec.items()
-        if isinstance(spec, MambaSpec)
-    }
-    tail_specs = {
-        name: spec
-        for name, spec in kv_cache_spec.items()
-        if isinstance(spec, KpoolTailSpec)
-    }
-    attn_specs = {
-        name: spec
-        for name, spec in kv_cache_spec.items()
-        if not isinstance(spec, (MambaSpec, KpoolTailSpec))
-    }
-    if not mamba_specs or not all(
-        type(spec) is MLAAttentionSpec for spec in attn_specs.values()
-    ):
-        return None
-
-    mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs)
-    idx_pages = {
-        spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1
-    }
-    if not idx_pages:
-        return None
-
-    assert all(spec.page_size_padded is None for spec in mla_specs.values())
-    assert len(idx_pages) == 1
-    mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1]
-    mla_pages = {mla_specs[name].page_size_bytes for name in mla_names}
-    assert len(mla_pages) == 1
-    mla_page = mla_pages.pop()
-    uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs)
-    assert uniform_spec is not None
-
-    tail_group: KVCacheGroupSpec | None = None
-    if tail_specs:
-        idx_page = next(iter(idx_pages))
-        padded_tail_specs: dict[str, KVCacheSpec] = {
-            name: replace(spec, page_size_padded=idx_page)
-            for name, spec in tail_specs.items()
-        }
-        tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs)
-        assert tail_uniform is not None
-        tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform)
-
-    any_mamba = next(iter(mamba_specs.values()))
-    assert all(spec == any_mamba for spec in mamba_specs.values())
-    if any_mamba.real_page_size_bytes > mla_page:
-        raise ValueError(
-            f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) "
-            f"does not fit the MLA page ({mla_page} bytes); increase tensor "
-            "parallelism or use a wider KV cache dtype"
-        )
-    padded_specs: dict[str, KVCacheSpec] = {
-        name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs
-    }
-    num_groups = _pp_balanced_mamba_group_count(
-        vllm_config, list(mamba_specs), mla_names
-    )
-    if num_groups is None:
-        raise ValueError(
-            "a pipeline stage has mamba layers but no MLA layer to share "
-            "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)"
-        )
-    mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)]
-    for index, name in enumerate(mamba_specs):
-        mamba_grouped_names[index % num_groups].append(name)
-
-    return (
-        [KVCacheGroupSpec(list(attn_specs), uniform_spec)]
-        + ([tail_group] if tail_group is not None else [])
-        + create_kv_cache_group_specs(padded_specs, mamba_grouped_names)
-    )
-
-
-def _glm5_next_tensor_layout(
-    kv_cache_groups: list[KVCacheGroupSpec],
-) -> (
-    tuple[
-        KVCacheGroupSpec,
-        list[KVCacheGroupSpec],
-        list[str],
-        list[str],
-        int,
-        int,
-        list[str],
-        int,
-    ]
-    | None
-):
-    """Recognize the GLM-5.3-Flash grouping after optional PP projection."""
-    uniform_groups = [
-        group
-        for group in kv_cache_groups
-        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-    ]
-    mamba_groups = [
-        group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec)
-    ]
-    attn_group: KVCacheGroupSpec | None = None
-    tail_group: KVCacheGroupSpec | None = None
-    for group in uniform_groups:
-        inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs
-        if all(type(spec) is MLAAttentionSpec for spec in inner.values()):
-            attn_group = group
-        elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()):
-            tail_group = group
-    if attn_group is None or not mamba_groups:
-        return None
-    if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups):
-        return None
-
-    attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
-    mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs)
-    if not all(
-        type(spec) is MLAAttentionSpec and spec.page_size_padded is None
-        for spec in mla_inner.values()
-    ):
-        return None
-    mla_names = [
-        name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1
-    ]
-    idx_names = [
-        name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1
-    ]
-    mla_pages = {mla_inner[name].page_size_bytes for name in mla_names}
-    idx_pages = {mla_inner[name].page_size_bytes for name in idx_names}
-    if len(mla_pages) != 1 or len(idx_pages) != 1:
-        return None
-    mla_page = mla_pages.pop()
-    idx_page = idx_pages.pop()
-    if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups):
-        return None
-
-    tail_names: list[str] = []
-    tail_page = 0
-    if tail_group is not None:
-        tail_names = list(tail_group.layer_names)
-        tail_inner = cast(
-            UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
-        ).kv_cache_specs
-        tail_pages = {
-            cast(KpoolTailSpec, spec).unpadded_page_size_bytes
-            for spec in tail_inner.values()
-        }
-        if len(tail_pages) != 1 or len(tail_names) != len(idx_names):
-            return None
-        tail_page = tail_pages.pop()
-        if tail_page > idx_page:
-            return None
-
-    return (
-        attn_group,
-        mamba_groups,
-        mla_names,
-        idx_names,
-        mla_page,
-        idx_page,
-        tail_names,
-        tail_page,
-    )
-
-
 def unify_kv_cache_spec_page_size(
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> dict[str, KVCacheSpec]:
@@ -1709,14 +1512,24 @@ def _get_per_layer_spec(
     return spec
 
 
+def _get_layer_block_stride(spec: KVCacheSpec) -> int:
+    """A layer-outer block stride: the page padded to its alignment."""
+    return round_up(spec.page_size_bytes, spec.get_block_stride_alignment())
+
+
 def _get_kv_cache_bytes_per_block(
-    kv_cache_groups: list[KVCacheGroupSpec],
+    kv_cache_groups: list[KVCacheGroupSpec], layout: KVCacheLayout
 ) -> int:
     """Return the largest cache group's bytes per block."""
-    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
-        _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout
-        return len(mla_names) * mla_page + len(idx_names) * idx_page
-
+    if not layout.is_block_outermost:
+        # Each layer's block stride is its own padded page.
+        return max(
+            sum(
+                _get_layer_block_stride(_get_per_layer_spec(group, layer_name))
+                for layer_name in group.layer_names
+            )
+            for group in kv_cache_groups
+        )
     bytes_per_block = max(
         sum(
             _get_per_layer_spec(group, layer_name).page_size_bytes
@@ -1731,7 +1544,7 @@ def _get_kv_cache_bytes_per_block(
         if isinstance(group.kv_cache_spec, HiSparseHotSpec)
     ]
     alignments.extend(
-        _get_per_layer_spec(group, layer_name).block_stride_alignment or 1
+        _get_per_layer_spec(group, layer_name).get_block_stride_alignment()
         for group in kv_cache_groups
         for layer_name in group.layer_names
     )
@@ -1806,72 +1619,9 @@ def get_kv_cache_config_from_groups(
             vllm_config, kv_cache_groups, available_memory, host_budget
         )
 
-    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
-        (
-            attn_group,
-            mamba_groups,
-            mla_names,
-            idx_names,
-            mla_page,
-            idx_page,
-            tail_names,
-            _,
-        ) = glm5_layout
-        bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page
-        num_blocks = may_override_num_blocks(
-            vllm_config, available_memory // bytes_per_block
-        )
-        size = bytes_per_block * num_blocks
-        attn_specs = cast(
-            UniformTypeKVCacheSpecs, attn_group.kv_cache_spec
-        ).kv_cache_specs
-
-        kv_cache_tensors: list[KVCacheTensor] = []
-
-        def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None:
-            kv_cache_tensors.append(
-                KVCacheTensor(
-                    size=size,
-                    layers=[layer_name],
-                    layer_stride=spec.page_size_bytes * num_blocks,
-                    block_stride=spec.page_size_bytes,
-                    offset=offset,
-                )
-            )
-
-        for index, mla_name in enumerate(mla_names):
-            offset = index * mla_page * num_blocks
-            add_tensor(mla_name, attn_specs[mla_name], offset)
-            for group in mamba_groups:
-                if index < len(group.layer_names):
-                    add_tensor(group.layer_names[index], group.kv_cache_spec, offset)
-
-        idx_base = len(mla_names) * mla_page * num_blocks
-        for index, idx_name in enumerate(idx_names):
-            offset = idx_base + index * idx_page * num_blocks
-            add_tensor(idx_name, attn_specs[idx_name], offset)
-            if tail_names:
-                tail_name = tail_names[index]
-                tail_group = next(
-                    group for group in kv_cache_groups if tail_name in group.layer_names
-                )
-                tail_specs = cast(
-                    UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
-                ).kv_cache_specs
-                add_tensor(tail_name, tail_specs[tail_name], offset)
-
-        return KVCacheConfig(
-            num_blocks=num_blocks,
-            kv_cache_tensors=kv_cache_tensors,
-            kv_cache_groups=kv_cache_groups,
-            prefix_cache_retention_interval=(
-                vllm_config.cache_config.prefix_cache_retention_interval
-            ),
-        )
-
     layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
     validate_kv_cache_layout(layout, kv_cache_groups)
-    bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups)
+    bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups, layout)
     interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None
 
     num_blocks = available_memory // bytes_per_block
@@ -1901,17 +1651,17 @@ def get_kv_cache_config_from_groups(
 
         byte_offset = 0
         for spec, layer_names in layers_by_spec.items():
+            block_stride = interleaved_block_stride
+            padded_page = _get_layer_block_stride(spec)
+            if block_stride is None and padded_page != spec.page_size_bytes:
+                assert layout.is_block_compact, f"{layout.name} cannot pad pages"
+                block_stride = padded_page
             layer_stride, block_stride, _, _, _ = compute_layout_strides(
                 spec,
                 num_blocks,
                 len(layer_names),
                 layout,
-                fixed_strides=(None, interleaved_block_stride, None, None, None),
-            )
-            offset = (
-                byte_offset
-                * max(layer_stride, spec.page_size_bytes)
-                // spec.page_size_bytes
+                fixed_strides=(None, block_stride, None, None, None),
             )
             kv_cache_tensors.append(
                 KVCacheTensor(
@@ -1919,10 +1669,13 @@ def get_kv_cache_config_from_groups(
                     layers=layer_names,
                     layer_stride=layer_stride,
                     block_stride=block_stride,
-                    offset=offset,
+                    offset=byte_offset,
                 )
             )
-            byte_offset += len(layer_names) * spec.page_size_bytes
+            # Block-outer layers share each block; layer-outer ones own regions.
+            byte_offset += len(layer_names) * (
+                spec.page_size_bytes if layout.is_block_outermost else layer_stride
+            )
 
     return KVCacheConfig(
         num_blocks=num_blocks,
@@ -2438,9 +2191,6 @@ def get_kv_cache_groups(
         # full attention, or all layers are sliding window attention with the
         # same window size). Put all layers into one group.
         return _get_kv_cache_groups_uniform_type(uniform_spec)
-    elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec):
-        return glm5_groups
-
     # Hidden-state layers use their own block table and must not be absorbed
     # into a compatible attention bucket.
     hidden_specs = {
@@ -2594,31 +2344,7 @@ def _max_memory_usage_bytes_from_groups(
     if not kv_cache_groups:
         return 0
 
-    if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
-        (
-            attn_group,
-            mamba_groups,
-            mla_names,
-            idx_names,
-            mla_page,
-            idx_page,
-            tail_names,
-            _,
-        ) = glm5_layout
-        uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
-        total_blocks = uniform_spec.max_memory_usage_pages(vllm_config)
-        total_blocks += sum(
-            cdiv(
-                group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
-                group.kv_cache_spec.page_size_bytes,
-            )
-            for group in mamba_groups
-        )
-        if tail_names:
-            total_blocks += 1
-        return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page)
-
-    bytes_per_block = _pool_bytes_per_block(kv_cache_groups)
+    bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups)
     total_blocks = 0
     for group in kv_cache_groups:
         spec = group.kv_cache_spec
@@ -2874,7 +2600,7 @@ def get_kv_cache_configs(
             if not groups:
                 adjusted_memory.append(avail_mem)
                 continue
-            bytes_per_block = _pool_bytes_per_block(groups)
+            bytes_per_block = _pool_bytes_per_block(vllm_config, groups)
             logger.info(
                 "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
                 avail_mem // bytes_per_block,
@@ -2887,7 +2613,7 @@ def get_kv_cache_configs(
     # the capacity check both plan against usable blocks. Allocation below
     # still uses the full memory.
     check_memory = [
-        avail_mem - _pool_bytes_per_block(groups) if groups else avail_mem
+        avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem
         for groups, avail_mem in zip(projected_groups_per_worker, available_memory)
     ]
 
@@ -2931,7 +2657,9 @@ def get_kv_cache_configs(
         # strides and offsets stay consistent with the shrunken allocation.
         groups = kv_cache_config.kv_cache_groups
         kv_cache_configs[i] = get_kv_cache_config_from_groups(
-            vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
+            vllm_config,
+            groups,
+            min_num_blocks * _pool_bytes_per_block(vllm_config, groups),
         )
 
     for kv_cache_config in kv_cache_configs:

vllm/v1/core/single_type_kv_cache_manager.py

@@ -7,7 +7,6 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, ClassVar
 
 from vllm.distributed.kv_events import MEDIUM_CPU
-from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.block_pool import BlockPool
 from vllm.v1.core.kv_cache_utils import (
@@ -27,7 +26,6 @@ from vllm.v1.kv_cache_interface import (
     HiddenStateCacheSpec,
     HiSparseHotSpec,
     HiSparseResidentSpec,
-    KpoolTailSpec,
     KVCacheGroupRole,
     KVCacheSpec,
     MambaSpec,
@@ -45,8 +43,6 @@ from vllm.v1.request import Request
 if TYPE_CHECKING:
     from vllm.v1.hisparse.coordinator import HiSparseCoordinator
 
-logger = init_logger(__name__)
-
 
 class SingleTypeKVCacheManager(ABC):
     """An abstract base class for a manager that handle the kv cache management
@@ -1299,10 +1295,6 @@ class CircularBufferManager(FullAttentionManager):
         return 0
 
 
-class KpoolTailManager(CircularBufferManager):
-    """One-block circular scratch manager for ``KpoolTailSpec``."""
-
-
 class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
     def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None:
         super().__init__(kv_cache_spec, **kwargs)
@@ -2687,11 +2679,6 @@ def register_all_kvcache_specs(vllm_config):
         SlidingWindowManager,
         uniform_type_base_spec=SlidingWindowMLASpec,
     )
-    KVCacheSpecRegistry.register(
-        KpoolTailSpec,
-        KpoolTailManager,
-        uniform_type_base_spec=KpoolTailSpec,
-    )
 
     KVCacheSpecRegistry.register(
         MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec

vllm/v1/engine/core.py

@@ -335,8 +335,8 @@ class EngineCore:
         vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
         kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
         if kv_cache_groups:
-            # Exclude groups that opt out of prefix caching (e.g. GLM-5.3-Flash
-            # kpool tail, a 1-block/req scratch buffer with block_size=kpool):
+            # Exclude groups that opt out of prefix caching (e.g. a circular
+            # buffer, a 1-block/req scratch buffer with a small ring block):
             # their small block_size would otherwise drag the global block_size
             # below the real allocator block size and desync it from mamba.
             participating = [

vllm/v1/hisparse/binding.py

@@ -55,8 +55,9 @@ def resolve_hisparse_specs(
     if len(block_sizes) != 1:
         raise ValueError("HiSparse requires one scheduler block size.")
     backends = [attn_layers[name].get_attn_backend() for name in mla_specs]
+    specs = list(mla_specs.values())
     try:
-        block_size = select_common_block_size(block_sizes.pop(), backends)
+        block_size = select_common_block_size(block_sizes.pop(), backends, specs)
     except ValueError as error:
         raise ValueError(
             "HiSparse requires a GPU block size supported by every sparse "
@@ -104,16 +105,13 @@ def allocate_hisparse_kv_caches(
             if isinstance(host_spec, UniformTypeKVCacheSpecs)
             else host_spec
         )
-        kernel_block_size = kernel_block_sizes[host_group_id]
-        if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
-            kernel_block_size = spec.storage_block_size
         views = create_kv_cache_views(
             backing,
             spec,
             kv_cache_config.num_blocks_of(tensor),
             layout,
             tensor,
-            kernel_block_size=kernel_block_size,
+            kernel_block_size=kernel_block_sizes[host_group_id],
         )
         kv_caches.update(zip(tensor.layers, views))
     return kv_caches

vllm/v1/hisparse/layout.py

@@ -333,7 +333,7 @@ def get_hisparse_kv_cache_config(
     device_groups = hisparse_layout.device_groups
     layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
     validate_kv_cache_layout(layout, device_groups)
-    bytes_per_block = _get_kv_cache_bytes_per_block(device_groups)
+    bytes_per_block = _get_kv_cache_bytes_per_block(device_groups, layout)
     num_blocks = may_override_num_blocks(
         vllm_config, available_memory // bytes_per_block
     )
@@ -345,7 +345,7 @@ def get_hisparse_kv_cache_config(
     host_groups = [hisparse_layout.source_group]
     host_layout = KVCacheLayout.LBNHC
     validate_kv_cache_layout(host_layout, host_groups)
-    host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups)
+    host_bytes_per_block = _get_kv_cache_bytes_per_block(host_groups, host_layout)
     host_size = host_bytes_per_block * hisparse_layout.host_num_blocks
     kv_cache_tensors[:0] = _build_hisparse_kv_cache_tensors(
         host_groups,

vllm/v1/kv_cache_interface.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import copy
+import math
 from collections import Counter
 from collections.abc import Collection, Sequence
 from dataclasses import dataclass, field, fields, replace
@@ -24,6 +25,7 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
+    from vllm.v1.attention.backend import MultipleOf
 
 logger = init_logger(__name__)
 
@@ -168,13 +170,28 @@ class KVCacheSpec:
     dcp_sharded: bool = field(default=False, kw_only=True)
     """Whether DCP shards this cache's token positions across ranks."""
 
-    block_stride_alignment: int | None = field(default=None, kw_only=True)
-    """Required byte alignment between physical blocks, including packed layers."""
+    block_stride_alignment: int | MultipleOf | None = field(default=None, kw_only=True)
+    """Required byte alignment between physical blocks, including packed layers.
+    ``MultipleOf(n)``: whole kernel blocks, when blocks are split into kernel
+    blocks of up to ``n`` tokens."""
 
     def __post_init__(self):
-        if self.block_stride_alignment is not None and self.block_stride_alignment <= 0:
+        alignment = self.block_stride_alignment
+        if isinstance(alignment, int) and alignment <= 0:
             raise ValueError("block_stride_alignment must be positive")
 
+    def get_block_stride_alignment(self) -> int:
+        """Bytes the stride between physical blocks must be a multiple of."""
+        from vllm.v1.attention.backend import MultipleOf
+
+        alignment = self.block_stride_alignment
+        if not isinstance(alignment, MultipleOf):
+            return alignment or 1
+        kernel_block_size = math.gcd(alignment.base, self.block_size)
+        if kernel_block_size == self.block_size:
+            return 1
+        return self.copy_with_new_block_size(kernel_block_size).page_size_bytes
+
     @property
     def prefix_cacheable(self) -> bool:
         """Whether this spec's group participates in prefix caching."""
@@ -674,8 +691,6 @@ class MLAAttentionSpec(FullAttentionSpec):
     model_version: str | None = None
     cache_role: SparseCacheRole = SparseCacheRole.SPARSE
     is_index_group_leader: bool = False
-    storage_block_size: int | None = None
-    """Token width used to view storage when it differs from the kernel block."""
     # Group capability enabled when any member flattens a non-causal query block
     # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
     non_causal_multi_token_decode: bool = False
@@ -696,7 +711,6 @@ class MLAAttentionSpec(FullAttentionSpec):
         model_version_set = set(spec.model_version for spec in specs)
         cache_role_set = {spec.cache_role for spec in specs}
         index_group_leader_set = {spec.is_index_group_leader for spec in specs}
-        storage_block_size_set = set(spec.storage_block_size for spec in specs)
         block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
         assert (
             len(cache_dtype_str_set) == 1
@@ -704,12 +718,11 @@ class MLAAttentionSpec(FullAttentionSpec):
             and len(model_version_set) == 1
             and len(cache_role_set) == 1
             and len(index_group_leader_set) == 1
-            and len(storage_block_size_set) == 1
             and len(block_stride_alignment_set) == 1
         ), (
             "All attention layers in the same KV cache group must use the same "
             "quantization method, tokens per state, model version, cache role, "
-            "index-sharing role, storage block size and block stride alignment."
+            "index-sharing role, and block stride alignment."
         )
         merged_spec = cls(
             block_size=specs[0].block_size,
@@ -727,7 +740,6 @@ class MLAAttentionSpec(FullAttentionSpec):
             model_version=model_version_set.pop(),
             cache_role=cache_role_set.pop(),
             is_index_group_leader=index_group_leader_set.pop(),
-            storage_block_size=storage_block_size_set.pop(),
             block_stride_alignment=block_stride_alignment_set.pop(),
             non_causal_multi_token_decode=any(
                 spec.non_causal_multi_token_decode for spec in specs
@@ -1019,32 +1031,6 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
         )
 
 
-@dataclass(frozen=True, kw_only=True)
-class KpoolTailSpec(SlidingWindowSpec):
-    """One-block circular scratch cache for a kpool indexer's raw tail."""
-
-    def max_admission_blocks_per_request(
-        self, max_in_flight_tokens: int, max_model_len: int
-    ) -> int:
-        return 1
-
-    def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        return 1
-
-    def is_uniform_with_collection(
-        self, kv_cache_specs: dict[str, KVCacheSpec]
-    ) -> bool:
-        return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values())
-
-    @property
-    def prefix_cacheable(self) -> bool:
-        return False
-
-    @property
-    def uses_slot_mapping(self) -> bool:
-        return False
-
-
 @dataclass(frozen=True)
 class MambaSpec(KVCacheSpec):
     shapes: tuple[tuple[int, ...], ...]

vllm/v1/spec_decode/gemma4.py

@@ -95,7 +95,7 @@ class Gemma4Proposer(SpecDecodeBaseProposer):
                 cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size]
             else:
                 cm = common_attn_metadata
-            attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
+            attn_metadata = attn_group.build_metadata_for_drafting(
                 common_attn_metadata=cm, draft_index=draft_index
             )
             per_group_attn_metadata.append(attn_metadata)
@@ -291,11 +291,7 @@ class Gemma4Proposer(SpecDecodeBaseProposer):
         self.draft_attn_groups = list(attention_groups.values())
         if self.draft_attn_groups:
             self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
-            self.block_size = (
-                self.draft_attn_groups[0]
-                .get_metadata_builder()
-                .kv_cache_spec.block_size
-            )
+            self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size
         else:
             self.kv_cache_gid = 0
             self.block_size = kv_cache_config.kv_cache_groups[

vllm/v1/spec_decode/llm_base_proposer.py

@@ -991,7 +991,7 @@ class SpecDecodeBaseProposer:
         per_group_attn_metadata: list[object] = []
         per_layer_attn_metadata: dict[str, object] = {}
         for attn_group in self.draft_attn_groups:
-            attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
+            attn_metadata = attn_group.build_metadata_for_drafting(
                 common_attn_metadata=common_attn_metadata, draft_index=draft_index
             )
             per_group_attn_metadata.append(attn_metadata)
@@ -1765,20 +1765,7 @@ class SpecDecodeBaseProposer:
                     attention_groups[backend_key].layer_names.append(layer_name)
 
         self.draft_attn_groups = list(attention_groups.values())
-        if kernel_block_sizes is not None and 0 <= self.kv_cache_gid < len(
-            kernel_block_sizes
-        ):
-            # Slot mappings are computed against the block table, which is
-            # stored at kernel-block granularity. Use the kernel block size
-            # rather than the KV cache manager's block size; the two differ
-            # when manager blocks are split for the attention kernel.
-            self.block_size = kernel_block_sizes[self.kv_cache_gid]
-        else:
-            self.block_size = (
-                self.draft_attn_groups[0]
-                .get_metadata_builder()
-                .kv_cache_spec.block_size
-            )
+        self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size
         logger.debug("Using block size %d for drafting layers", self.block_size)
 
     def _determine_batch_execution_and_padding(

vllm/v1/spec_decode/step3p5.py

@@ -144,7 +144,7 @@ class Step3p5MTPProposer(EagleProposer):
                     cm.slot_mapping = sm
             else:
                 cm = common_attn_metadata
-            attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
+            attn_metadata = attn_group.build_metadata_for_drafting(
                 common_attn_metadata=cm,
                 draft_index=draft_index,
             )
@@ -244,11 +244,7 @@ class Step3p5MTPProposer(EagleProposer):
         self.draft_attn_groups = list(attention_groups.values())
         if self.draft_attn_groups:
             self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
-            self.block_size = (
-                self.draft_attn_groups[0]
-                .get_metadata_builder()
-                .kv_cache_spec.block_size
-            )
+            self.block_size = self.draft_attn_groups[0].kv_cache_spec.block_size
         else:
             self.kv_cache_gid = 0
             self.block_size = kv_cache_config.kv_cache_groups[

vllm/v1/worker/block_table.py

@@ -29,24 +29,16 @@ logger = init_logger(__name__)
 def get_block_table_width(
     max_num_blocks: int,
     block_size: int,
-    kernel_block_size: int | None = None,
     *,
     token_alignment: int | None = 128,
 ) -> int:
-    """Return the width after optional alignment and virtual block splitting."""
-    if kernel_block_size is None:
-        kernel_block_size = block_size
-    if block_size % kernel_block_size != 0:
-        raise ValueError(
-            f"kernel_block_size {kernel_block_size} must divide "
-            f"block_size {block_size} evenly"
-        )
+    """Return the width after optional alignment."""
     if token_alignment is not None:
         if token_alignment <= 0:
             raise ValueError("token_alignment must be positive")
         block_alignment = token_alignment // math.gcd(token_alignment, block_size)
         max_num_blocks = cdiv(max_num_blocks, block_alignment) * block_alignment
-    return max_num_blocks * block_size // kernel_block_size
+    return max_num_blocks
 
 
 class SlotMappingMode(Enum):
@@ -63,7 +55,6 @@ class BlockTable:
         max_num_batched_tokens: int,
         pin_memory: bool,
         device: torch.device,
-        kernel_block_size: int,
         cp_kv_cache_interleave_size: int,
         slot_mapping_mode: SlotMappingMode = SlotMappingMode.TOKEN_TO_KV_SLOT,
     ):
@@ -74,9 +65,6 @@ class BlockTable:
         max_num_batched_tokens: Maximum number of tokens in a batch.
         pin_memory: Whether to pin memory for faster GPU transfers.
         device: Target device for the block table.
-        kernel_block_size: The block_size of underlying attention kernel.
-            Will be the same as `block_size` if `block_size` is supported
-            by the attention kernel.
         slot_mapping_mode: How this cache group maps scheduled tokens to
             cache slots. Mamba-like state caches do not use token slot
             mappings and should use SlotMappingMode.NONE.
@@ -86,30 +74,8 @@ class BlockTable:
         self.max_num_batched_tokens = max_num_batched_tokens
         self.pin_memory = pin_memory
         self.device = device
-        self.kv_cache_block_size = block_size
-
-        if kernel_block_size == block_size:
-            # Standard case: allocation and computation use same block size
-            # No block splitting needed, direct mapping
-            self.block_size = block_size
-            self.blocks_per_kv_block = 1
-            self.use_hybrid_blocks = False
-        else:
-            # Hybrid case: allocation block size differs from kernel block size
-            # Memory blocks are subdivided to match kernel requirements
-            # Example: 32-token memory blocks with 16-token kernel blocks
-            # → Each memory block corresponds to 2 kernel blocks
-            if block_size % kernel_block_size != 0:
-                raise ValueError(
-                    f"kernel_block_size {kernel_block_size} must divide "
-                    f"kv_manager_block_size size {block_size} evenly"
-                )
-
-            self.block_size = kernel_block_size
-            self.blocks_per_kv_block = block_size // kernel_block_size
-            self.use_hybrid_blocks = True
-
-        self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block
+        self.block_size = block_size
+        self.max_num_blocks_per_req = max_num_blocks_per_req
 
         self.block_table = self._make_buffer(
             self.max_num_reqs, self.max_num_blocks_per_req, dtype=torch.int32
@@ -120,13 +86,6 @@ class BlockTable:
             self.max_num_batched_tokens, dtype=torch.int64
         )
 
-        if self.use_hybrid_blocks:
-            self._kernel_block_arange = np.arange(0, self.blocks_per_kv_block).reshape(
-                1, -1
-            )
-        else:
-            self._kernel_block_arange = None
-
         try:
             self.pcp_world_size = get_pcp_group().world_size
             self.pcp_rank = get_pcp_group().rank_in_group
@@ -145,13 +104,11 @@ class BlockTable:
         self.slot_mapping_mode = slot_mapping_mode
         if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT:
             _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup(
-                kv_cache_block_size=self.kv_cache_block_size,
-                blocks_per_kv_block=self.blocks_per_kv_block,
+                kv_cache_block_size=self.block_size,
                 total_cp_world_size=self.dcp_world_size,
                 total_cp_rank=self.dcp_rank,
                 cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
                 block_table_stride=self.block_table.gpu.stride(0),
-                block_size=self.block_size,
             )
 
     def append_row(
@@ -162,11 +119,6 @@ class BlockTable:
         if not block_ids:
             return
 
-        if self.use_hybrid_blocks:
-            block_ids = self.map_to_kernel_blocks(
-                np.array(block_ids), self.blocks_per_kv_block, self._kernel_block_arange
-            )
-
         num_blocks = len(block_ids)
         start = self.num_blocks_per_row[row_idx]
         self.num_blocks_per_row[row_idx] += num_blocks
@@ -219,10 +171,8 @@ class BlockTable:
             positions,
             self.block_table.gpu,
             self.block_table.gpu.stride(0),
-            self.block_size,
             self.slot_mapping.gpu,
-            self.kv_cache_block_size,
-            self.blocks_per_kv_block,
+            self.block_size,
             self.dcp_world_size,
             self.dcp_rank,
             self.cp_kv_cache_interleave_size,
@@ -235,37 +185,6 @@ class BlockTable:
         self.block_table.gpu.fill_(0)
         self.block_table.cpu.fill_(0)
 
-    @staticmethod
-    def map_to_kernel_blocks(
-        kv_manager_block_ids: np.ndarray,
-        blocks_per_kv_block: int,
-        kernel_block_arange: np.ndarray,
-    ) -> np.ndarray:
-        """Convert kv_manager_block_id IDs to kernel block IDs.
-
-        Example:
-            # kv_manager_block_ids: 32 tokens,
-            # Kernel block size: 16 tokens
-            # blocks_per_kv_block = 2
-            >>> kv_manager_block_ids = np.array([0, 1, 2])
-            >>> Result: [0, 1, 2, 3, 4, 5]
-
-            # Each kv_manager_block_id maps to 2 kernel block id:
-            # kv_manager_block_id 0 → kernel block id [0, 1]
-            # kv_manager_block_id 1 → kernel block id [2, 3]
-            # kv_manager_block_id 2 → kernel block id [4, 5]
-
-        """
-        if blocks_per_kv_block == 1:
-            return kv_manager_block_ids
-
-        kernel_block_ids = (
-            kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block
-            + kernel_block_arange
-        )
-
-        return kernel_block_ids.reshape(-1)
-
     def get_device_tensor(self, num_reqs: int) -> torch.Tensor:
         """Returns the device tensor of the block table."""
         return self.block_table.gpu[:num_reqs]
@@ -296,16 +215,10 @@ class MultiGroupBlockTable:
         pin_memory: bool,
         device: torch.device,
         block_sizes: list[int],
-        kernel_block_sizes: list[int],
         max_num_blocks: list[int],
         cp_kv_cache_interleave_size: int = 1,
         slot_mapping_modes: list[SlotMappingMode] | None = None,
     ) -> None:
-        if len(kernel_block_sizes) != len(block_sizes):
-            raise ValueError(
-                f"kernel_block_sizes length ({len(kernel_block_sizes)}) "
-                f"must match block_sizes length ({len(block_sizes)})"
-            )
         if slot_mapping_modes is None:
             slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] * len(block_sizes)
         if len(slot_mapping_modes) != len(block_sizes):
@@ -339,17 +252,11 @@ class MultiGroupBlockTable:
                 max_num_batched_tokens,
                 pin_memory,
                 device,
-                kernel_block_size,
                 cp_kv_cache_interleave_size,
                 slot_mapping_mode=slot_mapping_mode,
             )
-            for (
-                block_size,
-                kernel_block_size,
-                max_num_blocks_per_req,
-                slot_mapping_mode,
-            ) in zip(
-                block_sizes, kernel_block_sizes, max_num_blocks, slot_mapping_modes
+            for block_size, max_num_blocks_per_req, slot_mapping_mode in zip(
+                block_sizes, max_num_blocks, slot_mapping_modes
             )
         ]
 
@@ -403,12 +310,10 @@ class ComputeSlotMappingKernel(
     @dataclass(frozen=True)
     class CompileKey:
         kv_cache_block_size: int
-        blocks_per_kv_block: int
         total_cp_world_size: int
         total_cp_rank: int
         cp_kv_cache_interleave_size: int
         block_table_stride: int
-        block_size: int
 
     @staticmethod
     @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"])
@@ -419,10 +324,8 @@ class ComputeSlotMappingKernel(
         positions_ptr,  # [num_tokens], int64
         block_table_ptr,  # [max_num_reqs, max_num_blocks_per_req], int32 (flat)
         block_table_stride,  # max_num_blocks_per_req
-        block_size,
         slot_mapping_ptr,  # [max_num_tokens], int64
         KV_CACHE_BLOCK_SIZE: tl.constexpr,
-        BLOCKS_PER_KV_BLOCK: tl.constexpr,
         TOTAL_CP_WORLD_SIZE: tl.constexpr,
         TOTAL_CP_RANK: tl.constexpr,
         CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
@@ -463,17 +366,12 @@ class ComputeSlotMappingKernel(
                 virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
             )
 
-            block_indices = (
-                virtual_block_indices * BLOCKS_PER_KV_BLOCK
-                + local_block_offsets // block_size
-            )
             block_numbers = tl.load(
-                block_table_ptr + row_offset + block_indices,
+                block_table_ptr + row_offset + virtual_block_indices,
                 mask=mask & is_local,
                 other=0,
             ).to(tl.int64)
-            slot_offsets = local_block_offsets % block_size
-            slot_ids = block_numbers * block_size + slot_offsets
+            slot_ids = block_numbers * KV_CACHE_BLOCK_SIZE + local_block_offsets
             slot_ids = tl.where(is_local, slot_ids, PAD_ID)
             tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)
 
@@ -481,13 +379,11 @@ class ComputeSlotMappingKernel(
         self,
         *,
         block_table_stride: int,
-        block_size: int,
         **compile_key_fields: int,
     ) -> CompileKey:
         return self.CompileKey(
             **compile_key_fields,
             block_table_stride=triton_scalar_specialization_rep(block_table_stride),
-            block_size=triton_scalar_specialization_rep(block_size),
         )
 
     def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]:
@@ -507,10 +403,8 @@ class ComputeSlotMappingKernel(
                 shape=(1, compile_key.block_table_stride),
             ),
             block_table_stride=compile_key.block_table_stride,
-            block_size=compile_key.block_size,
             slot_mapping=int64_ptr,
             kv_cache_block_size=compile_key.kv_cache_block_size,
-            blocks_per_kv_block=compile_key.blocks_per_kv_block,
             total_cp_world_size=compile_key.total_cp_world_size,
             total_cp_rank=compile_key.total_cp_rank,
             cp_kv_cache_interleave_size=compile_key.cp_kv_cache_interleave_size,
@@ -526,17 +420,14 @@ class ComputeSlotMappingKernel(
         positions: torch.Tensor,
         block_table: torch.Tensor,
         block_table_stride: int,
-        block_size: int,
         slot_mapping: torch.Tensor,
         kv_cache_block_size: int,
-        blocks_per_kv_block: int,
         total_cp_world_size: int,
         total_cp_rank: int,
         cp_kv_cache_interleave_size: int,
     ) -> LaunchSpec:
         return (num_reqs + 1,), dict(
             KV_CACHE_BLOCK_SIZE=kv_cache_block_size,
-            BLOCKS_PER_KV_BLOCK=blocks_per_kv_block,
             TOTAL_CP_WORLD_SIZE=total_cp_world_size,
             TOTAL_CP_RANK=total_cp_rank,
             CP_KV_CACHE_INTERLEAVE_SIZE=cp_kv_cache_interleave_size,

vllm/v1/worker/cp_utils.py

@@ -95,11 +95,9 @@ def prepare_dcp_dummy_context_metadata(
     max_valid_block_id = kv_cache_config.num_blocks - 1
     assert max_valid_block_id > 0
     for blk_table in input_batch.block_table.block_tables:
-        max_row_blocks = (
-            blk_table.max_num_blocks_per_req // blk_table.blocks_per_kv_block
-        )
         block_ids = [
-            (block_idx % max_valid_block_id) + 1 for block_idx in range(max_row_blocks)
+            (block_idx % max_valid_block_id) + 1
+            for block_idx in range(blk_table.max_num_blocks_per_req)
         ]
         for req_idx in range(num_reqs):
             blk_table.add_row(block_ids, req_idx)

vllm/v1/worker/gpu/attn_utils.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, cast
@@ -40,6 +40,8 @@ from vllm.v1.worker.utils import (
     add_kv_sharing_layers_to_kv_cache_groups,
     allocate_kv_cache,
     bind_kv_cache_to_layers,
+    customize_attention_spec,
+    map_kv_caches_to_kernel_blocks,
     prepare_kernel_block_sizes,
 )
 
@@ -141,7 +143,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
         # Skip modules that don't need KV cache (eg encoder-only attention)
         if spec := attn_module.get_kv_cache_spec(vllm_config):
             if isinstance(spec, AttentionSpec):
-                spec = attn_module.get_attn_backend().customize_spec(spec)
+                spec = customize_attention_spec(attn_module.get_attn_backend(), spec)
             kv_cache_spec[layer_name] = spec
     if vllm_config.attention_config.hisparse_config is not None:
         kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers)
@@ -382,6 +384,7 @@ def init_kv_cache(
     kv_cache_allocation_context: AbstractContextManager | None = None,
     *,
     block_tables: "BlockTables | None" = None,
+    attn_groups: Iterable[AttentionGroup] = (),
 ) -> dict[str, Any]:
     allocation_context = kv_cache_allocation_context or nullcontext()
     with allocation_context:
@@ -409,9 +412,10 @@ def init_kv_cache(
     num_attn_module = (
         2 if vllm_config.model_config.hf_config.model_type == "longcat_flash" else 1
     )
-    bindable_caches = {
-        name: cache for name, cache in kv_caches.items() if name in forward_context
-    }
+    bindable_caches = map_kv_caches_to_kernel_blocks(
+        {name: cache for name, cache in kv_caches.items() if name in forward_context},
+        attn_groups,
+    )
     bind_kv_cache_to_layers(
         bindable_caches,
         forward_context,
@@ -536,8 +540,8 @@ def build_attn_metadata(
                     cached_metadata[reuse_key], block_table, slot_mapping
                 )
             elif for_cudagraph_capture:
-                metadata = attn_metadata_builder.build_for_cudagraph_capture(
-                    common_attn_metadata
+                metadata = attn_group.build_metadata_for_cudagraph_capture(
+                    common_attn_metadata, ubatch_idx
                 )
             else:
                 attn_metadata_extra_kwargs = (
@@ -548,10 +552,8 @@ def build_attn_metadata(
                     if model_specific_attn_metadata is not None
                     else {}
                 )
-                metadata = attn_metadata_builder.build(
-                    common_prefix_len=0,
-                    common_attn_metadata=common_attn_metadata,
-                    **attn_metadata_extra_kwargs,
+                metadata = attn_group.build_metadata(
+                    common_attn_metadata, ubatch_idx, **attn_metadata_extra_kwargs
                 )
             if reuse_key is not None and reuse_key not in cached_metadata:
                 cached_metadata[reuse_key] = metadata

vllm/v1/worker/gpu/block_table.py

@@ -25,7 +25,6 @@ class BlockTables:
         max_num_batched_tokens: int,
         max_num_blocks_per_group: list[int],
         device: torch.device,
-        kernel_block_sizes: list[int],
         cp_size: int = 1,
         cp_rank: int = 0,
         cp_interleave: int = 1,
@@ -33,7 +32,6 @@ class BlockTables:
         dcp_sharded: list[bool] | None = None,
     ):
         self.block_sizes = block_sizes
-        self.kernel_block_sizes = kernel_block_sizes
         self.max_num_reqs = max_num_reqs
         self.max_num_batched_tokens = max_num_batched_tokens
         self.device = device
@@ -53,14 +51,9 @@ class BlockTables:
         assert len(dcp_sharded) == self.num_kv_cache_groups
         self.dcp_sharded = torch.tensor(dcp_sharded, dtype=torch.bool, device=device)
 
-        self.blocks_per_kv_block = [
-            bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
-        ]
-
         # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
         self.block_tables: list[StagedWriteTensor] = []
-        for i in range(self.num_kv_cache_groups):
-            max_num_blocks = max_num_blocks_per_group[i] * self.blocks_per_kv_block[i]
+        for max_num_blocks in max_num_blocks_per_group:
             block_table = StagedWriteTensor(
                 (self.max_num_reqs, max_num_blocks), dtype=torch.int32, device=device
             )
@@ -109,9 +102,6 @@ class BlockTables:
         self.block_sizes_tensor = torch.tensor(
             self.block_sizes, dtype=torch.int32, device=self.device
         )
-        self.kernel_block_sizes_tensor = torch.tensor(
-            self.kernel_block_sizes, dtype=torch.int32, device=self.device
-        )
         self.slot_mapping_enabled = torch.tensor(
             self._slot_mapping_enabled, dtype=torch.bool, device=self.device
         )
@@ -126,9 +116,6 @@ class BlockTables:
         for i in range(self.num_kv_cache_groups):
             start = self.num_blocks.np[i, req_index] if not overwrite else 0
             block_ids = new_block_ids[i]
-            bpk = self.blocks_per_kv_block[i]
-            if bpk > 1:
-                block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
             if self.redirect_writes_to_null_block:
                 block_ids = [0] * len(block_ids)
             end = start + len(block_ids)
@@ -218,7 +205,6 @@ class BlockTables:
             self.block_table_ptrs,
             self.block_table_strides,
             self.block_sizes_tensor,
-            self.kernel_block_sizes_tensor,
             self.slot_mapping_enabled,
             self.dcp_sharded,
             slot_mappings,
@@ -292,7 +278,6 @@ def _compute_slot_mappings_kernel(
     block_table_ptrs,  # [num_kv_cache_groups]
     block_table_strides,  # [num_kv_cache_groups]
     block_sizes,  # [num_kv_cache_groups]
-    kernel_block_sizes,  # [num_kv_cache_groups]
     slot_mapping_enabled,  # [num_kv_cache_groups]
     dcp_sharded,  # [num_kv_cache_groups]
     slot_mappings_ptr,  # [num_kv_cache_groups, max_num_tokens]
@@ -322,7 +307,6 @@ def _compute_slot_mappings_kernel(
     block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
     block_table_stride = tl.load(block_table_strides + group_id)
     kv_block_size = tl.load(block_sizes + group_id)
-    kernel_block_size = tl.load(kernel_block_sizes + group_id)
     mapping_enabled = tl.load(slot_mapping_enabled + group_id)
     if CP_SIZE != 1:
         sharded = tl.load(dcp_sharded + group_id)
@@ -354,16 +338,14 @@ def _compute_slot_mappings_kernel(
             local_positions = tl.where(sharded, local_positions, positions)
             is_local = ~sharded | is_local
 
-        block_indices = tl.where(
-            mapping_enabled, local_positions // kernel_block_size, 0
-        )
-        block_offsets = local_positions % kernel_block_size
+        block_indices = tl.where(mapping_enabled, local_positions // kv_block_size, 0)
+        block_offsets = local_positions % kv_block_size
         block_numbers = tl.load(
             block_table_ptr + req_state_idx * block_table_stride + block_indices,
             mask=is_local & is_real_req,
             other=0,
         )
-        slot_ids = block_numbers * kernel_block_size + block_offsets
+        slot_ids = block_numbers * kv_block_size + block_offsets
         if CP_SIZE != 1:
             slot_ids = tl.where(is_local, slot_ids, PAD_ID)
 

vllm/v1/worker/gpu/input_batch.py

@@ -252,17 +252,18 @@ def set_dummy_context(
     input_batch.positions.copy_(torch.from_numpy(local_pos + context_len))
 
     seq_len = context_len + query_len
-    for block_table, block_size, bpk in zip(
-        input_block_tables,
-        block_tables.kernel_block_sizes,
-        block_tables.blocks_per_kv_block,
-    ):
+    for block_table, block_size in zip(input_block_tables, block_tables.block_sizes):
         num_blocks = min(cdiv(seq_len, block_size), block_table.shape[1])
         # Spans are disjoint until the pool runs out, then they wrap and share
         # blocks: profiling only needs the reads to be realistic, not distinct.
-        block_ids = torch.arange(
-            num_reqs * num_blocks, dtype=block_table.dtype, device=block_table.device
-        ) % (num_kv_blocks * bpk)
+        block_ids = (
+            torch.arange(
+                num_reqs * num_blocks,
+                dtype=block_table.dtype,
+                device=block_table.device,
+            )
+            % num_kv_blocks
+        )
         block_table[:num_reqs, :num_blocks] = block_ids.view(num_reqs, num_blocks)
 
 

vllm/v1/worker/gpu/model_runner.py

@@ -688,7 +688,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             max_num_batched_tokens=self.max_num_tokens,
             max_num_blocks_per_group=max_num_blocks_per_group,
             device=self.device,
-            kernel_block_sizes=self.kernel_block_sizes,
             slot_mapping_enabled=slot_mapping_enabled,
             dcp_sharded=dcp_sharded,
             cp_size=self.dcp_size,
@@ -773,6 +772,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 self.vllm_config,
                 kv_cache_allocation_context=kv_cache_allocation_context,
                 block_tables=self.block_tables,
+                # Draft groups build their own metadata, so they map the views too.
+                attn_groups=(
+                    g
+                    for groups in self.attn_groups
+                    + getattr(self.speculator, "attn_groups", [])
+                    for g in groups
+                ),
             )
         self.kv_caches = [
             cache for cache in kv_caches_dict.values() if cache.device == self.device

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -432,7 +432,7 @@ class DFlashSpeculator(DraftModelSpeculator):
                 temperature,
                 seeds,
                 self.block_tables.input_block_tables[gid],
-                self.block_tables.kernel_block_sizes[gid],
+                self.block_tables.block_sizes[gid],
                 self.block_tables.cp_rank,
                 self.dcp_size,
                 self.block_tables.cp_interleave,

vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py

@@ -577,6 +577,11 @@ class TargetDependentARSpeculator(DraftModelSpeculator):
         idx_mapping = self.idx_mapping[:num_reqs]
 
         if batch_desc.cg_mode == CUDAGraphMode.FULL:
+            for groups in self.attn_groups:
+                for group in groups:
+                    group.map_to_kernel_block_table(
+                        self.block_tables.input_block_tables[group.kv_cache_group_id]
+                    )
             assert self.decode_cudagraph_manager is not None
             self.decode_cudagraph_manager.run_fullgraph(batch_desc)
             return

vllm/v1/worker/gpu_input_batch.py

@@ -96,7 +96,6 @@ class InputBatch:
         device: torch.device,
         vocab_size: int,
         block_sizes: list[int],  # The block_size of each kv cache group
-        kernel_block_sizes: list[int],
         max_num_blocks_per_req: list[int],
         logitsprocs: LogitsProcessors | None = None,
         logitsprocs_need_output_token_ids: bool = False,
@@ -187,7 +186,6 @@ class InputBatch:
             pin_memory=PIN_MEMORY,
             device=device,
             block_sizes=block_sizes,
-            kernel_block_sizes=kernel_block_sizes,
             max_num_blocks=max_num_blocks_per_req,
             cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
             slot_mapping_modes=slot_mapping_modes,

vllm/v1/worker/gpu_model_runner.py

@@ -233,6 +233,8 @@ from .utils import (
     allocate_kv_cache,
     bind_kv_cache,
     copy_kv_cache_blocks_inplace,
+    customize_attention_spec,
+    map_kv_caches_to_kernel_blocks,
     prepare_kernel_block_sizes,
     sanity_check_mm_encoder_outputs,
 )
@@ -718,7 +720,6 @@ class GPUModelRunner(
             max(self.max_model_len, self.max_encoder_len), placeholder_block_size
         )
         self._init_block_sizes = [placeholder_block_size]
-        self._init_kernel_block_sizes = [placeholder_block_size]
         self._init_max_num_blocks = [placeholder_max_num_blocks]
         self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT]
         self.cp_kv_cache_interleave_size = (
@@ -734,7 +735,6 @@ class GPUModelRunner(
             device=self.device,
             vocab_size=self.model_config.get_vocab_size(),
             block_sizes=[placeholder_block_size],
-            kernel_block_sizes=[placeholder_block_size],
             max_num_blocks_per_req=[placeholder_max_num_blocks],
             num_spec_tokens=self.num_spec_tokens,
             logitsprocs=build_logitsprocs(
@@ -2508,8 +2508,8 @@ class GPUModelRunner(
                 )
 
             if for_cudagraph_capture:
-                attn_metadata_i = builder.build_for_cudagraph_capture(
-                    common_attn_metadata
+                attn_metadata_i = attn_group.build_metadata_for_cudagraph_capture(
+                    common_attn_metadata, ubid or 0
                 )
             elif (
                 cache_key in cached_attn_metadata
@@ -2517,13 +2517,16 @@ class GPUModelRunner(
             ):
                 attn_metadata_i = builder.update_block_table(
                     cached_attn_metadata[cache_key],
-                    common_attn_metadata.block_table_tensor,
+                    attn_group.map_to_kernel_block_table(
+                        common_attn_metadata.block_table_tensor, ubid or 0
+                    ),
                     common_attn_metadata.slot_mapping,
                 )
             else:
-                attn_metadata_i = builder.build(
+                attn_metadata_i = attn_group.build_metadata(
+                    common_attn_metadata,
+                    ubid or 0,
                     common_prefix_len=cascade_attn_prefix_len,
-                    common_attn_metadata=common_attn_metadata,
                     **extra_attn_metadata_args,
                 )
                 if builder.supports_update_block_table:
@@ -7143,9 +7146,7 @@ class GPUModelRunner(
             return
         self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds)  # type: ignore[assignment]
 
-    def may_reinitialize_input_batch(
-        self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
-    ) -> None:
+    def may_reinitialize_input_batch(self, kv_cache_config: KVCacheConfig) -> None:
         """Re-initialize the input batch if the block sizes are different from
         what it was originally created with. This happens when the final
         block size (determined after model loading) differs from the
@@ -7154,7 +7155,6 @@ class GPUModelRunner(
 
         Args:
             kv_cache_config: The KV cache configuration.
-            kernel_block_sizes: The kernel block sizes for each KV cache group.
 
         """
         block_sizes = []
@@ -7179,14 +7179,12 @@ class GPUModelRunner(
 
         if (
             block_sizes != self._init_block_sizes
-            or kernel_block_sizes != self._init_kernel_block_sizes
             or max_num_blocks != self._init_max_num_blocks
             or slot_mapping_modes != self._init_slot_mapping_modes
             or self.cp_kv_cache_interleave_size
             != self.parallel_config.cp_kv_cache_interleave_size
         ):
             self._init_block_sizes = block_sizes
-            self._init_kernel_block_sizes = kernel_block_sizes
             self._init_max_num_blocks = max_num_blocks
             self._init_slot_mapping_modes = slot_mapping_modes
             self.cp_kv_cache_interleave_size = (
@@ -7201,7 +7199,6 @@ class GPUModelRunner(
                     device=self.device,
                     vocab_size=self.model_config.get_vocab_size(),
                     block_sizes=block_sizes,
-                    kernel_block_sizes=kernel_block_sizes,
                     max_num_blocks_per_req=max_num_blocks,
                     num_spec_tokens=self.num_spec_tokens,
                     logitsprocs=self.input_batch.logitsprocs,
@@ -7217,10 +7214,6 @@ class GPUModelRunner(
             f"InputBatch block_sizes {self._init_block_sizes} != "
             f"kv_cache block_sizes {block_sizes}"
         )
-        assert self._init_kernel_block_sizes == kernel_block_sizes, (
-            f"InputBatch kernel_block_sizes {self._init_kernel_block_sizes} "
-            f"!= kv_cache kernel_block_sizes {kernel_block_sizes}"
-        )
 
     def _attn_group_iterator(self) -> Iterator[AttentionGroup]:
         return itertools.chain.from_iterable(self.attn_groups)
@@ -7267,12 +7260,16 @@ class GPUModelRunner(
         num_attn_module = (
             2 if self.model_config.hf_config.model_type == "longcat_flash" else 1
         )
+        attn_groups = [g for groups in self.attn_groups for g in groups]
+        if self.speculative_config:
+            attn_groups += getattr(self.drafter, "draft_attn_groups", [])
         bind_kv_cache(
             kv_caches,
             self.compilation_config.static_forward_context,
             self.kv_caches,
             num_attn_module,
             kv_cache_groups=kv_cache_config.kv_cache_groups,
+            layer_kv_caches=map_kv_caches_to_kernel_blocks(kv_caches, attn_groups),
         )
         return kv_caches
 
@@ -7336,6 +7333,8 @@ class GPUModelRunner(
         # backends for that group only supports block_size 64, we will return
         # kernel_block_size 64 and split the 256-token-block to 4 blocks with 64
         # tokens each.
+        # Block-outermost packed groups are not split; their attention groups
+        # map kernel blocks themselves.
         kernel_block_sizes = prepare_kernel_block_sizes(
             kv_cache_config, self.attn_groups
         )
@@ -7345,7 +7344,7 @@ class GPUModelRunner(
         self.initialize_metadata_builders(kv_cache_config, kernel_block_sizes)
 
         # Reinitialize need to after initialize_attn_backend
-        self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
+        self.may_reinitialize_input_batch(kv_cache_config)
         # Capture warmup providers that depend on allocated KV-cache strides.
         with self.jit_warmup_registry.activate():  # type: ignore[attr-defined]
             kv_caches = self.initialize_kv_cache_tensors(
@@ -7422,7 +7421,9 @@ class GPUModelRunner(
             # Skip modules that don't need KV cache (eg encoder-only attention)
             if spec := attn_module.get_kv_cache_spec(self.vllm_config):
                 if isinstance(spec, AttentionSpec):
-                    spec = attn_module.get_attn_backend().customize_spec(spec)
+                    spec = customize_attention_spec(
+                        attn_module.get_attn_backend(), spec
+                    )
                 kv_cache_spec[layer_name] = spec
 
         return kv_cache_spec

vllm/v1/worker/utils.py

@@ -3,7 +3,7 @@
 import math
 from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from itertools import product as iprod
 from typing import Any
 
@@ -24,6 +24,7 @@ from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionMetadataBuilder,
+    CommonAttentionMetadata,
     MultipleOf,
 )
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
@@ -35,7 +36,6 @@ from vllm.v1.kv_cache_interface import (
     KVCacheLayout,
     KVCacheSpec,
     MambaSpec,
-    MLAAttentionSpec,
     UniformTypeKVCacheSpecs,
     create_kv_cache_views,
 )
@@ -166,6 +166,9 @@ class KVBlockZeroer:
                 kv = static_forward_context[layer_name].kv_cache
                 if not isinstance(kv, torch.Tensor):
                     continue
+                if group.kernel_block_stride:
+                    # Bound views address kernel blocks; zero by manager block.
+                    kv = group.map_kv_cache_to_manager_blocks(kv, num_blocks)
                 if kv.device.type != self.device.type:
                     continue
                 dp = kv.data_ptr()
@@ -255,12 +258,22 @@ class KVBlockZeroer:
             self.zero_block_ids([0])
 
 
-@dataclass
+@dataclass(eq=False)
 class AttentionGroup:
     backend: type[AttentionBackend]
     layer_names: list[str]
     kv_cache_spec: KVCacheSpec
     kv_cache_group_id: int
+    # Set when the backend needs blocks smaller than the KV cache manager's: its
+    # block tables and KV cache views then address kernel blocks, kernel block j
+    # of block b being b * kernel_block_stride + j. The stride is set when a
+    # packed KV cache is bound (other layers' pages sit between its blocks);
+    # otherwise it stays 0 and the kernel blocks per block are used.
+    kernel_block_size: int | None = None
+    kernel_block_stride: int = 0
+    # Persistent per ubatch: CUDA graphs capture the mapped block table.
+    kernel_block_table: torch.Tensor | None = None
+    kernel_block_offsets: torch.Tensor | None = None
     # When ubatching is enabled we will have a metadata builder for each ubatch
     # so that if they use internal persistent buffers for cudagraphs, and they
     # won't have to worry about conflicting with the other ubatches.
@@ -275,27 +288,44 @@ class AttentionGroup:
         kernel_block_size: int | None = None,
         num_metadata_builders: int = 1,
     ):
+        spec = self.kv_cache_spec
+        if kernel_block_size is not None and isinstance(spec, AttentionSpec):
+            # A backend in a packed KV cache group may need smaller blocks than
+            # the group is split into.
+            kernel_block_size = select_common_block_size(
+                kernel_block_size, [self.backend], [spec]
+            )
+            if kernel_block_size != spec.block_size:
+                self.kernel_block_size = kernel_block_size
         if kernel_block_size is None:
             kv_cache_spec_builder = self.kv_cache_spec
-        elif (
-            isinstance(self.kv_cache_spec, MLAAttentionSpec)
-            and self.kv_cache_spec.storage_block_size is not None
-        ):
-            kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size(
-                self.kv_cache_spec.storage_block_size
-            )
         else:
             kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size(
                 kernel_block_size
             )
         builder_cls = self.backend.get_builder_cls()
         builder_kwargs = {}
-        if builder_cls.requires_block_table_width:
+        if builder_cls.requires_block_table_width or self.kernel_block_size is not None:
             max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req(
                 vllm_config, vllm_config.model_config.max_model_len
             )
-            builder_kwargs["block_table_width"] = get_block_table_width(
-                max_num_blocks, self.kv_cache_spec.block_size, kernel_block_size
+            width = get_block_table_width(max_num_blocks, self.kv_cache_spec.block_size)
+            if kernel_block_size is not None:
+                width *= self.kv_cache_spec.block_size // kernel_block_size
+            if builder_cls.requires_block_table_width:
+                builder_kwargs["block_table_width"] = width
+        if self.kernel_block_size is not None:
+            self.kernel_block_table = torch.zeros(
+                num_metadata_builders,
+                vllm_config.scheduler_config.max_num_seqs,
+                width,
+                dtype=torch.int32,
+                device=device,
+            )
+            self.kernel_block_offsets = torch.arange(
+                spec.block_size // self.kernel_block_size,
+                dtype=torch.int32,
+                device=device,
             )
         self.metadata_builders = [
             builder_cls(
@@ -307,9 +337,109 @@ class AttentionGroup:
             )
             for _ in range(num_metadata_builders)
         ]
-        if kernel_block_size is not None:
-            for builder in self.metadata_builders:
-                builder.set_kernel_block_size(kernel_block_size)
+
+    def build_metadata(
+        self,
+        common_attn_metadata: CommonAttentionMetadata,
+        ubatch_idx: int = 0,
+        common_prefix_len: int = 0,
+        **kwargs,
+    ) -> Any:
+        return self.get_metadata_builder(ubatch_idx).build(
+            common_prefix_len,
+            self._with_kernel_block_table(common_attn_metadata, ubatch_idx),
+            **kwargs,
+        )
+
+    def build_metadata_for_cudagraph_capture(
+        self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0
+    ) -> Any:
+        return self.get_metadata_builder(ubatch_idx).build_for_cudagraph_capture(
+            self._with_kernel_block_table(common_attn_metadata, ubatch_idx)
+        )
+
+    def build_metadata_for_drafting(
+        self, common_attn_metadata: CommonAttentionMetadata, draft_index: int
+    ) -> Any:
+        return self.get_metadata_builder().build_for_drafting(
+            self._with_kernel_block_table(common_attn_metadata), draft_index
+        )
+
+    def _with_kernel_block_table(
+        self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0
+    ) -> CommonAttentionMetadata:
+        if self.kernel_block_size is None:
+            return common_attn_metadata
+        return replace(
+            common_attn_metadata,
+            block_table_tensor=self.map_to_kernel_block_table(
+                common_attn_metadata.block_table_tensor, ubatch_idx
+            ),
+        )
+
+    def map_to_kernel_block_table(
+        self, block_table: torch.Tensor, ubatch_idx: int = 0
+    ) -> torch.Tensor:
+        if self.kernel_block_size is None:
+            return block_table
+        assert self.kernel_block_table is not None
+        blocks_per_kv_block = self.kv_cache_spec.block_size // self.kernel_block_size
+        rows, cols = block_table.shape
+        out = self.kernel_block_table[ubatch_idx, :rows, : cols * blocks_per_kv_block]
+        torch.add(
+            self.kernel_block_offsets,
+            block_table.unsqueeze(-1),
+            alpha=self.kernel_block_stride or blocks_per_kv_block,
+            out=out.unflatten(1, (cols, -1)),
+        )
+        return out
+
+    def map_kv_cache_to_kernel_blocks(self, kv_cache: torch.Tensor) -> torch.Tensor:
+        """Views a packed cache in manager blocks as kernel blocks."""
+        assert self.kernel_block_size is not None
+        spec = self.kv_cache_spec
+        kernel_rows = spec.get_num_kernel_states(self.kernel_block_size)
+        if kv_cache.shape[-2] == kernel_rows:
+            return kv_cache
+        # Token and pooled caches are written through manager-block slot
+        # mappings, so only compressed caches (written from their block table)
+        # can be mapped.
+        if (
+            spec.tokens_per_state <= 1
+            or kv_cache.shape[1] != 1
+            or kv_cache.stride(0) % (kernel_rows * kv_cache.stride(-2))
+        ):
+            raise ValueError(
+                f"{self.backend.get_name()} needs {self.kernel_block_size}-token "
+                f"kernel blocks, which packed {spec.block_size}-token KV cache "
+                f"blocks cannot be split into. Reduce --block-size to "
+                f"{self.kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a "
+                "layer-compact layout (e.g. LBNHC)."
+            )
+        # Kernel block j of block b is b * kernel_block_stride + j. The view
+        # spans the other layers' pages between blocks (never addressing them),
+        # which ``view`` cannot express.
+        page_stride = kernel_rows * kv_cache.stride(-2)
+        self.kernel_block_stride = kv_cache.stride(0) // page_stride
+        num_pages = (kv_cache.shape[0] - 1) * self.kernel_block_stride + (
+            kv_cache.shape[-2] // kernel_rows
+        )
+        return kv_cache.as_strided(
+            (num_pages, *kv_cache.shape[1:-2], kernel_rows, kv_cache.shape[-1]),
+            (page_stride, *kv_cache.stride()[1:]),
+        )
+
+    def map_kv_cache_to_manager_blocks(
+        self, kv_cache: torch.Tensor, num_blocks: int
+    ) -> torch.Tensor:
+        """Views a packed cache in kernel blocks as manager blocks."""
+        assert self.kernel_block_stride
+        spec = self.kv_cache_spec
+        rows = spec.get_num_kernel_states(spec.block_size)
+        return kv_cache.as_strided(
+            (num_blocks, *kv_cache.shape[1:-2], rows, kv_cache.shape[-1]),
+            (kv_cache.stride(0) * self.kernel_block_stride, *kv_cache.stride()[1:]),
+        )
 
     def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder:
         assert len(self.metadata_builders) > ubatch_id
@@ -327,17 +457,48 @@ class AttentionGroup:
         self.get_metadata_builder().update_draft_decode_metadata(metadata)
 
 
+def map_kv_caches_to_kernel_blocks(
+    kv_caches: dict[str, torch.Tensor], attn_groups: Iterable[AttentionGroup]
+) -> dict[str, torch.Tensor]:
+    """The views to bind to layers. The manager views in ``kv_caches`` stay for
+    KV connectors."""
+    mapped = dict(kv_caches)
+    for group in attn_groups:
+        if group.kernel_block_size is not None:
+            mapped.update(
+                (name, group.map_kv_cache_to_kernel_blocks(kv_caches[name]))
+                for name in group.layer_names
+                if name in kv_caches
+            )
+    return mapped
+
+
+def customize_attention_spec(
+    backend: type[AttentionBackend], spec: AttentionSpec
+) -> AttentionSpec:
+    """The backend's spec. Blocks larger than every kernel block size a backend
+    takes are split into kernel blocks, which packed blocks must keep whole."""
+    spec = backend.customize_spec(spec)
+    sizes = backend.get_supported_kernel_block_sizes(spec)
+    fixed_sizes = [s for s in sizes if isinstance(s, int)]
+    if spec.block_stride_alignment is None and sizes and fixed_sizes == sizes:
+        spec = replace(spec, block_stride_alignment=MultipleOf(max(fixed_sizes)))
+    return spec
+
+
 def _block_size_is_supported(
-    backends: list[type[AttentionBackend]], block_size: int
+    backends: list[type[AttentionBackend]],
+    block_size: int,
+    specs: Sequence[KVCacheSpec | None],
 ) -> bool:
     """Check if the block size is supported by all backends.
 
     An exact ``int`` declaration must match exactly; a ``MultipleOf``
     declaration accepts any multiple of its base.
     """
-    for backend in backends:
+    for backend, spec in zip(backends, specs, strict=True):
         is_supported = False
-        for supported_size in backend.get_supported_kernel_block_sizes():
+        for supported_size in backend.get_supported_kernel_block_sizes(spec):
             if isinstance(supported_size, int):
                 if block_size == supported_size:
                     is_supported = True
@@ -354,6 +515,7 @@ def _block_size_is_supported(
 def select_common_block_size(
     kv_manager_block_size: int,
     backends: list[type[AttentionBackend]],
+    kv_cache_specs: Sequence[KVCacheSpec | None] | None = None,
 ) -> int:
     """Select a block size that is supported by all backends and is a factor of
     kv_manager_block_size.
@@ -364,6 +526,7 @@ def select_common_block_size(
     Args:
         kv_manager_block_size: Block size of KV cache.
         backends: List of attention backend classes.
+        kv_cache_specs: Per-backend cache specs for spec-specific constraints.
 
     Returns:
         The selected block size.
@@ -372,25 +535,27 @@ def select_common_block_size(
         ValueError: If no valid block size found.
 
     """
-    if _block_size_is_supported(backends, kv_manager_block_size):
+    specs = kv_cache_specs or [None] * len(backends)
+    if _block_size_is_supported(backends, kv_manager_block_size, specs):
         return kv_manager_block_size
 
     # MultipleOf constraints also accept the manager size if they accept a divisor.
     # Any remaining candidate must therefore be an explicit size from a backend.
     candidates = {
         size
-        for backend in backends
-        for size in backend.get_supported_kernel_block_sizes()
+        for backend, spec in zip(backends, specs, strict=True)
+        for size in backend.get_supported_kernel_block_sizes(spec)
         if isinstance(size, int) and kv_manager_block_size % size == 0
     }
 
     for size in sorted(candidates, reverse=True):
-        if _block_size_is_supported(backends, size):
+        if _block_size_is_supported(backends, size, specs):
             return size
     raise ValueError(
         f"No common block size for {kv_manager_block_size} ("
         + "; ".join(
-            f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}" for b in backends
+            f"{b.get_name()}: {b.get_supported_kernel_block_sizes(s)}"
+            for b, s in zip(backends, specs, strict=True)
         )
         + ")."
     )
@@ -450,9 +615,6 @@ def allocate_kv_cache(
         kernel_block_size = None
         if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
             kernel_block_size = kernel_block_sizes[group_id]
-        if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
-            kernel_block_size = spec.storage_block_size
-
         views = create_kv_cache_views(
             buf,
             spec,
@@ -497,24 +659,20 @@ def prepare_kernel_block_sizes(
             # This is an attention backend that supports virtual block splitting.
             kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
             group_backends = [g.backend for g in attn_groups[kv_cache_gid]]
-            storage_block_size = (
-                kv_cache_spec.storage_block_size
-                if isinstance(kv_cache_spec, MLAAttentionSpec)
-                else None
-            )
-            if storage_block_size is not None and _block_size_is_supported(
-                group_backends, storage_block_size
-            ):
-                # Storage-block specs (e.g. the GLM-5.3-Flash kpool indexer
-                # cache) address the cache in pool pages, and every other
-                # consumer (cache views, metadata builders, hisparse) already
-                # uses storage_block_size as the kernel block. Fall back to the
-                # backend vote when the group's backends do not accept it.
-                selected_kernel_size = storage_block_size
-            else:
-                selected_kernel_size = select_common_block_size(
-                    kv_manager_block_size, group_backends
+            group_specs = [g.kv_cache_spec for g in attn_groups[kv_cache_gid]]
+            # Block-outermost packing puts other layers' pages between the
+            # blocks, so they are never split group-wide.
+            selected_kernel_size = (
+                kv_manager_block_size
+                if any(
+                    t.block_stride > t.layer_stride
+                    for t in kv_cache_config.kv_cache_tensors
+                    if not set(t.layers).isdisjoint(kv_cache_group.layer_names)
                 )
+                else select_common_block_size(
+                    kv_manager_block_size, group_backends, group_specs
+                )
+            )
             kernel_block_sizes.append(selected_kernel_size)
         elif isinstance(kv_cache_spec, MambaSpec):
             # This is likely Mamba or other non-attention cache, no splitting.
@@ -656,6 +814,7 @@ def bind_kv_cache(
     runner_kv_caches: list[torch.Tensor],
     num_attn_module: int = 1,
     kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
+    layer_kv_caches: dict[str, torch.Tensor] | None = None,
 ) -> None:
     """Bind the allocated KV cache to both ModelRunner and forward context so
     that the KV cache can be used in the forward pass.
@@ -674,6 +833,9 @@ def bind_kv_cache(
         runner_kv_caches: The kv_cache declared by ModelRunner.
         kv_cache_groups: The KV cache groups of the model, used to resolve
             layers that share a KV cache.
+        layer_kv_caches: The views to bind to layers when they differ from
+            kv_caches (e.g. kernel-block views). The runner keeps kv_caches,
+            which are addressed by scheduler block IDs.
 
     """
     # Bind kv_caches to ModelRunner
@@ -699,7 +861,7 @@ def bind_kv_cache(
             runner_kv_caches.append(kv_caches[layer_name])
 
     bind_kv_cache_to_layers(
-        kv_caches, forward_context, num_attn_module, kv_cache_groups
+        layer_kv_caches or kv_caches, forward_context, num_attn_module, kv_cache_groups
     )