From 02e55d2af38adab51ec5ffb985063d355b20805b Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:48:37 +0800 Subject: [PATCH] [KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots (#57169) Signed-off-by: Yifan Qiao Signed-off-by: Dakai An Signed-off-by: Lucas Wilkinson Co-authored-by: Yifan Qiao Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Lucas Wilkinson --- .../attention/test_deepgemm_attention.py | 46 ++ .../attention/test_rocm_aiter_mla_sink.py | 4 +- .../test_jit_warmup_triton_launcher.py | 3 - .../test_dspark_noncausal_sparse_mla.py | 2 +- .../test_flashinfer_mla_sparse_sm90.py | 18 +- ...st_flashmla_nope_sm90_backend_selection.py | 97 +++ .../test_glm5next_indexer_block_sizes.py | 40 ++ .../test_indexer_deepseek_v4_slot_mapping.py | 19 +- .../test_kpool_indexer_block_sizes.py | 54 +- .../attention/test_kpool_tail_slot_mapping.py | 45 +- tests/v1/attention/test_mm_prefix.py | 1 - .../v1/attention/test_rocm_glm5next_sparse.py | 6 + .../v1/attention/test_sparse_mla_backends.py | 69 +- .../test_sparse_mla_kv_cache_layout.py | 21 +- .../test_partial_prefix_cache_hits.py | 8 +- tests/v1/core/test_contiguous_kv_packing.py | 42 +- tests/v1/core/test_kv_cache_utils.py | 602 ++---------------- .../core/test_single_type_kv_cache_manager.py | 14 +- tests/v1/kv_connector/unit/conftest.py | 2 +- .../test_mooncake_connector_hybrid_mamba.py | 9 +- .../unit/test_mooncake_store_hma_e2e.py | 4 +- .../unit/test_nixl_desc_geometry.py | 185 +----- tests/v1/simple_kv_offload/test_scheduler.py | 4 +- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- tests/v1/spec_decode/test_eagle.py | 6 + .../v1/spec_decode/test_llm_base_proposer.py | 24 +- .../test_gpu_model_runner_streaming.py | 1 - tests/v1/worker/test_attn_utils.py | 112 +++- tests/v1/worker/test_gpu_block_table.py | 31 +- tests/v1/worker/test_gpu_input_batch.py | 9 - .../test_gpu_kpool_tail_slot_mapping.py | 8 +- tests/v1/worker/test_gpu_model_runner.py | 108 +--- ...odel_runner_prefill_tail_classification.py | 3 - tests/v1/worker/test_gpu_model_runner_v2.py | 1 - tests/v1/worker/test_gpu_pcp_manager.py | 3 +- ...test_gpu_target_dependent_ar_speculator.py | 51 ++ tests/v1/worker/test_jit_warmup_migration.py | 18 +- .../worker/test_kv_cache_allocation_scope.py | 2 + tests/v1/worker/test_mamba_utils.py | 1 - tests/v1/worker/test_utils.py | 37 ++ .../kv_transfer/kv_connector/utils.py | 30 + .../v1/mooncake/mooncake_connector.py | 28 +- .../kv_connector/v1/nixl/base_worker.py | 124 +--- .../layers/attention/mla_attention.py | 8 - .../models/deepseek_v41/nvidia/model_state.py | 2 +- vllm/models/glm5next/common/attention.py | 62 +- .../glm5next/nvidia/ops/kpool_compress.py | 5 - vllm/platforms/cuda.py | 20 - vllm/platforms/interface.py | 39 +- vllm/utils/deep_gemm.py | 9 +- vllm/v1/attention/backend.py | 8 +- vllm/v1/attention/backends/composite.py | 5 - .../backends/mla/compressor_utils.py | 4 +- .../backends/mla/flashattn_mla_sparse.py | 7 +- .../backends/mla/flashinfer_mla_sparse.py | 22 +- .../mla/flashinfer_mla_sparse_sm120.py | 8 +- .../mla/flashinfer_mla_sparse_sm90.py | 25 +- .../attention/backends/mla/flashmla_sparse.py | 51 +- vllm/v1/attention/backends/mla/indexer.py | 66 +- .../backends/mla/rocm_aiter_mla_sparse.py | 25 +- .../v1/attention/backends/mla/sparse_utils.py | 9 +- .../attention/backends/mla/xpu_mla_sparse.py | 5 + vllm/v1/core/kv_cache_coordinator.py | 8 +- vllm/v1/core/kv_cache_utils.py | 354 ++-------- vllm/v1/core/single_type_kv_cache_manager.py | 13 - vllm/v1/engine/core.py | 4 +- vllm/v1/hisparse/binding.py | 8 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/kv_cache_interface.py | 56 +- vllm/v1/spec_decode/gemma4.py | 8 +- vllm/v1/spec_decode/llm_base_proposer.py | 17 +- vllm/v1/spec_decode/step3p5.py | 8 +- vllm/v1/worker/block_table.py | 129 +--- vllm/v1/worker/cp_utils.py | 6 +- vllm/v1/worker/gpu/attn_utils.py | 24 +- vllm/v1/worker/gpu/block_table.py | 26 +- vllm/v1/worker/gpu/input_batch.py | 17 +- vllm/v1/worker/gpu/model_runner.py | 8 +- .../gpu/spec_decode/dflash/speculator.py | 2 +- .../target_dependent_ar/speculator.py | 5 + vllm/v1/worker/gpu_input_batch.py | 2 - vllm/v1/worker/gpu_model_runner.py | 41 +- vllm/v1/worker/utils.py | 252 ++++++-- 83 files changed, 1280 insertions(+), 1986 deletions(-) create mode 100644 tests/v1/attention/test_glm5next_indexer_block_sizes.py diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index 5d11063c7a..f332fb0fea 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/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) diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index 0b007dbc87..f66d408c84 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/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 = ( diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index a92303dcf2..c19804eae4 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/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, diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 7cd7863b06..0942af4f45 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/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}" ) diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index 735748ae94..66113aca09 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/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 diff --git a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py b/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py index 6523588769..8532d2f43a 100644 --- a/tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py +++ b/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 diff --git a/tests/v1/attention/test_glm5next_indexer_block_sizes.py b/tests/v1/attention/test_glm5next_indexer_block_sizes.py new file mode 100644 index 0000000000..9f27d29946 --- /dev/null +++ b/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 diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 17dea5d03c..464fce66f2 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/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) diff --git a/tests/v1/attention/test_kpool_indexer_block_sizes.py b/tests/v1/attention/test_kpool_indexer_block_sizes.py index a737988dde..86e196d539 100644 --- a/tests/v1/attention/test_kpool_indexer_block_sizes.py +++ b/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 diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 183c63ec2f..062fd783e0 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/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): diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index f878ede486..b8085f0bc5 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/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, diff --git a/tests/v1/attention/test_rocm_glm5next_sparse.py b/tests/v1/attention/test_rocm_glm5next_sparse.py index e69686f500..6077b9eb2f 100644 --- a/tests/v1/attention/test_rocm_glm5next_sparse.py +++ b/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", diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 973a9e03e8..160b766da0 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/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, diff --git a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py index cd3c9b2181..d869141df8 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/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() ) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 51317deb8d..4bffe91a65 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/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, ), ), ], diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 638deeedc5..6a705b6e5e 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/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]), diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 53833ba837..b06479f1ca 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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. diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 415609aefb..f926f49651 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/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() == [] diff --git a/tests/v1/kv_connector/unit/conftest.py b/tests/v1/kv_connector/unit/conftest.py index 58c6e632d3..85862da685 100644 --- a/tests/v1/kv_connector/unit/conftest.py +++ b/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)]), ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index 23856b4541..969bf65f8c 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/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, ] diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6c2cff0a54..a8354390d7 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/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, diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 9f6bc56dee..4e7a7b8431 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/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( diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 08759bc1e3..6c0cc310c7 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/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 diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index 1d004c3667..3fc5caec2e 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/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, diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 2ddd0a9774..62a9a1635c 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/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] diff --git a/tests/v1/spec_decode/test_llm_base_proposer.py b/tests/v1/spec_decode/test_llm_base_proposer.py index fb1ecdd2df..086fd467ac 100644 --- a/tests/v1/spec_decode/test_llm_base_proposer.py +++ b/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 diff --git a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py b/tests/v1/streaming_input/test_gpu_model_runner_streaming.py index 9457c66050..7002176dcc 100644 --- a/tests/v1/streaming_input/test_gpu_model_runner_streaming.py +++ b/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, diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9c9c26c3c2..8d30403f30 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/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]) diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 5c2231a1f2..dac9f12ae6 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/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() diff --git a/tests/v1/worker/test_gpu_input_batch.py b/tests/v1/worker/test_gpu_input_batch.py index f37e824e6c..21157aae33 100644 --- a/tests/v1/worker/test_gpu_input_batch.py +++ b/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, ) diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py index 8a7445f7d4..d9d5c4de63 100644 --- a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py +++ b/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)) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 3851d5dad6..158befc150 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py b/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py index 92207aa38f..dd9046de61 100644 --- a/tests/v1/worker/test_gpu_model_runner_prefill_tail_classification.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index dec382725b..47173da081 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index d4b464e014..fcf171a89c 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/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, diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index b63bd7a2ff..a3ec5f4eb7 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/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, ): diff --git a/tests/v1/worker/test_jit_warmup_migration.py b/tests/v1/worker/test_jit_warmup_migration.py index ef887c5351..1875e605e5 100644 --- a/tests/v1/worker/test_jit_warmup_migration.py +++ b/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 diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 7f4bea52c2..52cf2a3741 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/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, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 4407d6fc54..7e588ccde2 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/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) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index f6f7227954..e17f214871 100644 --- a/tests/v1/worker/test_utils.py +++ b/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( diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index d314d6107a..d09e5da71b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/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]: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 9d5194f33e..41c1337ecd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/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: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 38a656954e..778eb0f756 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c..861c894c38 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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): diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index cc3ade80b8..7fc0bd25ff 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/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) diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index a264717f0d..b398a614c0 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/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, diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index db496d300b..eb289c65cf 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/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) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 68c3aa7851..401a98f66f 100644 --- a/vllm/platforms/cuda.py +++ b/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, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index dc3200759c..3df44a9275 100644 --- a/vllm/platforms/interface.py +++ b/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 diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 0eb979ce05..04c4508f0b 100644 --- a/vllm/utils/deep_gemm.py +++ b/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: diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 3387a31794..68bb0b3bdd 100644 --- a/vllm/v1/attention/backend.py +++ b/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( diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index 0ba40e99f2..f77e5e6d95 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/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( diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 41d973f296..bf939dc049 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/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) diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index aabd682240..0016eafedc 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/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: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 4ebda50178..21eb66aae4 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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, ) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 2c918c0e4a..2bfcee05c4 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 16d7c12e67..9e4e85761a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/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, ) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d441934c02..0ff34f8ac5 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/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, ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1005b8e424..e776057b05 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/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: diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index e3d982e882..e30fec715d 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/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 ) diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 2ec10218d9..2447df5ceb 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/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, diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py index 036dfe0173..87800ecd17 100644 --- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py +++ b/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 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c32173a635..019941a1bf 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/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. diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 53a01a5cab..919a7b4b3f 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/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: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 3e7ae1b706..e9f5f2f3c2 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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 diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index bfbaa38d94..d0dadae27e 100644 --- a/vllm/v1/engine/core.py +++ b/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 = [ diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index b8c0d5dff2..60c60a3ff5 100644 --- a/vllm/v1/hisparse/binding.py +++ b/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 diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 98636866f8..96d93650ad 100644 --- a/vllm/v1/hisparse/layout.py +++ b/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, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index ca3f485917..79e88ca91d 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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, ...], ...] diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index bea677dde8..af504b0a11 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/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[ diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index e4718bb3d7..79be2e8a4d 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/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( diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 7a73473adc..0853067f39 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/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[ diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index c63908dad0..2beec249bf 100644 --- a/vllm/v1/worker/block_table.py +++ b/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, diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index bcf56dc0bf..49dcccf443 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/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) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8aa9ba3cb8..d1feaef031 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index c7665c46b8..1ab78864c3 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/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) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index e2f68ad598..ea1c925dad 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 61d6197ceb..134b4bfb0e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index d59d87ae7e..ec46039533 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index 9a8ed7b8b4..4894ce77e9 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/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 diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 4707e4f2ca..73a8c0eb9a 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/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, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c02bdc74e2..22ad662b90 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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 diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index db0f81b81c..d10cdcec64 100644 --- a/vllm/v1/worker/utils.py +++ b/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 )