From df6c7ad03194a3051a197d57ca4a94e4a063014e Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:05:39 +0800 Subject: [PATCH 01/23] [Core][Refactor] Move KV cache planning logic to kv_cache_planning.py (#61183) Signed-off-by: Dakai An Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 5 + .buildkite/test_areas/lm_eval.yaml | 6 + tests/models/test_dspark_mla.py | 2 +- tests/models/test_initialization.py | 6 +- .../test_sparse_mla_kv_cache_layout.py | 2 +- tests/v1/core/test_contiguous_kv_packing.py | 6 +- tests/v1/core/test_kv_cache_planning.py | 2284 ++++++++++++++++ tests/v1/core/test_kv_cache_utils.py | 2348 +---------------- tests/v1/core/utils.py | 74 + tests/v1/engine/test_init_error_messaging.py | 4 +- .../unit/offloading_connector/test_config.py | 6 +- .../unit/test_hidden_states_connector.py | 2 +- tests/v1/simple_kv_offload/test_worker.py | 4 +- .../test_dsv4_packed_zeroer_geometry.py | 2 +- tests/v1/worker/test_gpu_model_runner.py | 5 +- vllm/v1/core/kv_cache_planning.py | 1637 ++++++++++++ vllm/v1/core/kv_cache_utils.py | 1615 +----------- vllm/v1/engine/core.py | 2 +- vllm/v1/hisparse/layout.py | 4 +- vllm/v1/worker/gpu/cudagraph_utils.py | 2 +- vllm/v1/worker/gpu_model_runner.py | 2 +- 21 files changed, 4063 insertions(+), 3955 deletions(-) create mode 100644 tests/v1/core/test_kv_cache_planning.py create mode 100644 vllm/v1/core/kv_cache_planning.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d2702e7d4f82..5145f5a272ff 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4321,6 +4321,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4354,6 +4355,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4387,6 +4389,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4420,6 +4423,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -4453,6 +4457,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f6a4a89f11cd..96e5d4fc5e9f 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -371,6 +371,7 @@ steps: - vllm/tokenizers/ - vllm/v1/attention/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/v1/spec_decode/metrics.py - vllm/v1/worker/ - vllm/_aiter_ops.py @@ -516,6 +517,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -559,6 +561,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -602,6 +605,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -645,6 +649,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py @@ -688,6 +693,7 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/core/kv_cache_utils.py + - vllm/v1/core/kv_cache_planning.py - vllm/model_executor/layers/layernorm.py - csrc/ - vllm/_aiter_ops.py diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index b2156af346c8..81019421ed13 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -394,7 +394,7 @@ def test_k3_dspark_mla_kv_cache_spec_groups_with_target_mla(): """The draft's MLA layers must share a KV cache group with the target's.""" from vllm.model_executor.layers.attention.mla_attention import MLAAttention from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention - from vllm.v1.core.kv_cache_utils import _get_kv_cache_groups_uniform_page_size + from vllm.v1.core.kv_cache_planning import _get_kv_cache_groups_uniform_page_size vllm_config = SimpleNamespace( model_config=None, cache_config=SimpleNamespace(block_size=64) diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 073f616af875..ee6dfee7ecbf 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -14,10 +14,8 @@ from vllm.platforms import current_platform from vllm.utils.mem_constants import GiB_bytes from vllm.v1.attention.backends.utils import resolve_kv_cache_layout -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - get_kv_cache_configs, -) +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.engine.core import EngineCore as V1EngineCore from ..utils import ( 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 d869141df81a..45ba1ac0dc55 100644 --- a/tests/v1/attention/test_sparse_mla_kv_cache_layout.py +++ b/tests/v1/attention/test_sparse_mla_kv_cache_layout.py @@ -23,7 +23,7 @@ flat_kv_row_view, ) 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.core.kv_cache_planning 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, customize_attention_spec diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 6a705b6e5e8e..64cd55054015 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -18,13 +18,15 @@ from vllm.config import CacheConfig from vllm.v1.core.kv_cache_manager import KVCacheManager -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, _get_packed_kv_cache_groups, _pool_bytes_per_block, - generate_scheduler_kv_cache_config, get_kv_cache_config_from_groups, get_kv_cache_groups, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, resolve_kv_cache_block_sizes, ) from vllm.v1.hisparse.layout import _build_hisparse_kv_cache_tensors diff --git a/tests/v1/core/test_kv_cache_planning.py b/tests/v1/core/test_kv_cache_planning.py new file mode 100644 index 000000000000..629286808148 --- /dev/null +++ b/tests/v1/core/test_kv_cache_planning.py @@ -0,0 +1,2284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import mmap +import subprocess +import sys +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.core.kv_cache_planning as kv_cache_planning +import vllm.v1.core.kv_cache_utils as kv_cache_utils +import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) +from vllm.config import ( + CacheConfig, + ModelConfig, + SchedulerConfig, + VllmConfig, +) +from vllm.config.attention import HiSparseConfig +from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.core.kv_cache_planning import ( + check_enough_kv_cache_memory, + estimate_max_model_len, + get_kv_cache_configs, + get_kv_cache_groups, + is_kv_cache_spec_uniform, + kv_cache_groups_tp_replicas, +) +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, +) +from vllm.v1.hisparse.layout import ( + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + HiSparseResidentSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + SparseCacheRole, + UniformTypeKVCacheSpecs, + iter_layer_specs, +) +from vllm.v1.kv_cache_layout import KVCacheLayout + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.parametrize("gpu_block_size", [32, 64]) +@pytest.mark.parametrize("shared_host_pool", [False, True]) +def test_hisparse_hma_uses_resolved_gpu_block_size( + monkeypatch, gpu_block_size, shared_host_pool +): + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + specs = { + "model.layers.0.self_attn": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ), + "model.layers.0.self_attn.indexer": MLAAttentionSpec( + block_size=gpu_block_size, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ), + } + config = SimpleNamespace( + attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), + model_config=SimpleNamespace( + hf_config=SimpleNamespace(index_topk=128), + max_model_len=gpu_block_size, + ), + parallel_config=SimpleNamespace( + tensor_parallel_size=2 if shared_host_pool else 1, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + world_size=2 if shared_host_pool else 1, + distributed_executor_backend="mp", + nnodes_within_dp=1, + ), + cache_config=SimpleNamespace( + num_gpu_blocks_override=7, + prefix_cache_retention_interval=None, + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, + ), + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 + ) + assert cache_config.num_blocks == 7 + assert cache_config.hisparse_host_num_blocks is not None + assert cache_config.hisparse_host_num_blocks > 7 + + host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups + assert cache_config.hisparse_shared_host_pool is shared_host_pool + host_page = host_group.kv_cache_spec.page_size_bytes + alignment = mmap.PAGESIZE if shared_host_pool else 1 + expected_host_stride = (host_page + alignment - 1) // alignment * alignment + assert cache_config.hisparse_host_block_stride == expected_host_stride + assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride + assert host_group.host_resident + assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) + host_layers = set(host_group.layer_names) + for tensor in cache_config.kv_cache_tensors: + assert all( + (name in host_layers) == tensor.host_resident for name in tensor.layers + ) + host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) + host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] + assert host_tensor.block_stride == host_spec.page_size_bytes + assert host_tensor.layer_stride == ( + host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks + ) + assert host_group.kv_cache_spec.block_size == gpu_block_size + assert indexer_group.kv_cache_spec.block_size == gpu_block_size + host_specs = host_group.kv_cache_spec.kv_cache_specs + gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs + assert set(host_specs) == {"model.layers.0.self_attn"} + assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} + assert indexer_group.kv_cache_spec.prefix_cacheable + assert host_group.enable_kv_transfer + auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] + assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) + assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) + assert all( + spec.block_size == gpu_block_size + for spec in auxiliary_specs + if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) + ) + scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( + cache_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(), + ), + ) + assert scheduler_block_size == hash_block_size == gpu_block_size + + +@pytest.mark.parametrize( + "max_model_len,num_gpu_blocks,ok", + [ + (8192, 160, False), + (8192, 1000, True), + (32768, 656, False), + (32768, 657, True), + (-1, 657, True), + ], +) +def test_hisparse_pool_must_fit_max_model_len( + monkeypatch, max_model_len, num_gpu_blocks, ok +): + """Each resident group takes its own blocks from HiSparse's shared GPU pool, + so a pool that fits only the indexer pages must be rejected at startup. + Resident pages are capped at the in-flight window: at 32768 tokens with 2048 + in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, + plus the null block. Auto-fit (-1) must pick a length that passes the same + check.""" + monkeypatch.setattr( + hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True + ) + monkeypatch.setattr( + kv_cache_planning, "get_hisparse_host_pool_bytes", lambda _: 2**30 + ) + specs: dict[str, KVCacheSpec] = {} + for i in range(4): + specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=576, + dtype=torch.bfloat16, + is_index_group_leader=True, + ) + specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + cache_role=SparseCacheRole.INDEXER, + ) + config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + config.scheduler_config.max_num_batched_tokens = 2048 + # Async scheduling doubles the in-flight window, and with it the bound. + config.scheduler_config.async_scheduling = False + config.attention_config.hisparse_config = HiSparseConfig() + config.model_config.hf_config.index_topk = 128 + config.cache_config.num_gpu_blocks_override = num_gpu_blocks + config.cache_config.kv_cache_layout = "BLHNC" + # Workers report the resident/hot caches HiSparse derives. + specs |= { + name: group.kv_cache_spec + for group in get_hisparse_kv_cache_groups(config, specs) or [] + if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) + for name in group.layer_names + } + + if ok: + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + else: + with pytest.raises(ValueError, match="max seq len"): + kv_cache_planning.get_kv_cache_configs(config, [specs], [2**34]) + + +def new_chunked_local_attention_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + attention_chunk_size=4, +): + return ChunkedLocalAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + attention_chunk_size=attention_chunk_size, + ) + + +def test_get_kv_cache_configs_multiple_workers(): + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + same_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + ] + + # Basic case. All things are the same. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different available memory. This is the case for TP. + # Use the smallest memory available. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + same_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 20, + ], + ) + assert kv_cache_configs == [expected, expected] + + # Different KV cache specs. This is the case for PP. + different_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer2": new_kv_cache_spec(), + "layer3": new_kv_cache_spec(), + }, + ] + + # Different workers have different layers. + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer2", "layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), + ], + ), + ] + + # Some layers are the same, some are different. This is the case for TP+PP + tp_pp_kv_cache_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + { + "layer3": new_kv_cache_spec(), + }, + ] + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + tp_pp_kv_cache_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + expected_12 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + ], + ) + expected_3 = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), + ], + ) + assert kv_cache_configs == [ + expected_12, + expected_12, + expected_3, + expected_3, + ] + + # Different workers have different types of layers. This is the case for + # hybrid models + PP. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_kv_cache_spec(), + }, + { + "layer3": new_sliding_window_spec(), + "layer4": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer1", "layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), + KVCacheGroupSpec([], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10 * 2, + layers=["layer3", "layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec([], ref_kv_cache_spec), + KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), + ], + ), + ] + + # When divided into multiple KVCacheGroups, need to ensure the number of + # layers per group is similar. + different_type_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + "layer2": new_sliding_window_spec(), + "layer3": new_sliding_window_spec(), + }, + { + "layer4": new_kv_cache_spec(), + "layer5": new_sliding_window_spec(), + "layer6": new_sliding_window_spec(), + }, + ] + kv_cache_configs = get_kv_cache_configs( + vllm_config, + different_type_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 10, + ref_kv_cache_spec.page_size_bytes * 10, + ], + ) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer3"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), + ], + ), + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer4"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer5"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * 10, + layers=["layer6"], + layer_stride=ref_kv_cache_spec.page_size_bytes * 10, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), + KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), + ], + ), + ] + + # Have conflicting layers. Need to raise an error. + conflicting_layer_specs = [ + { + "layer1": new_kv_cache_spec(), + }, + { + "layer1": new_sliding_window_spec(), + }, + ] + with pytest.raises(AssertionError): + get_kv_cache_configs( + vllm_config, + conflicting_layer_specs, + [ + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ref_kv_cache_spec.page_size_bytes * 2 * 10, + ], + ) + + +@pytest.mark.parametrize( + "asymmetric_memory", + [False, True], + ids=["symmetric", "asymmetric"], +) +def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): + model_config = ModelConfig(max_model_len=512) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + ref_kv_cache_spec = new_kv_cache_spec() + pp_kv_cache_specs = [ + {"layer1": ref_kv_cache_spec}, + {"layer2": ref_kv_cache_spec}, + ] + + expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 + avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks + + # With per-worker validation, each worker only needs memory for its own + # layers. Worker 2 having more memory shouldn't affect worker 1's config. + available_memory = ( + [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 + ) + + kv_cache_configs = get_kv_cache_configs( + vllm_config, + pp_kv_cache_specs, + available_memory, + ) + + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer1"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], + ), + KVCacheConfig( + num_blocks=expected_num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, + layers=["layer2"], + layer_stride=ref_kv_cache_spec.page_size_bytes + * expected_num_blocks, + block_stride=ref_kv_cache_spec.page_size_bytes, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], + ), + ] + + +def test_project_kv_cache_groups_to_worker(): + spec_a = new_kv_cache_spec() + spec_b = new_kv_cache_spec(num_kv_heads=4) + + global_groups = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), + ] + worker_spec = {"layer1": spec_a, "layer2": spec_a} + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, worker_spec + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer2"] + assert projected[0].kv_cache_spec is spec_a + + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups, {"layer4": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == [] + assert projected[0].kv_cache_spec is spec_a + + uniform_spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, + ) + global_groups_uniform = [ + KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), + ] + projected = kv_cache_planning._project_kv_cache_groups_to_worker( + global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} + ) + assert len(projected) == 1 + assert projected[0].layer_names == ["layer1", "layer3"] + proj_spec = projected[0].kv_cache_spec + assert isinstance(proj_spec, UniformTypeKVCacheSpecs) + assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} + + +@pytest.mark.parametrize("with_draft", [False, True]) +@pytest.mark.parametrize("indexer_alignment", [None, 512]) +def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): + from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view + from vllm.v1.worker.utils import allocate_kv_cache + + config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + config.cache_config.kv_cache_layout = "BLHNC" + config.parallel_config.decode_context_parallel_size = 4 + common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) + specs = { + "target": MLAAttentionSpec( + **common, + head_size=576, + state_content_bytes=656, + block_stride_alignment=656, + cache_dtype_str="fp8_ds_mla", + is_index_group_leader=True, + ), + "indexer": MLAAttentionSpec( + **common, + head_size=132, + cache_role=SparseCacheRole.INDEXER, + block_stride_alignment=indexer_alignment, + ), + "draft": SlidingWindowSpec( + block_size=64, + num_kv_heads=64, + head_size=64, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_sharded=False, + ), + } + if not with_draft: + del specs["draft"] + groups = get_kv_cache_groups(config, specs) + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( + config, groups, available_memory=8 * 1024 * 1024 + ) + caches = allocate_kv_cache( + cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) + ) + cache = caches["target"].squeeze(1) + if indexer_alignment is not None: + assert caches["indexer"].stride(0) % indexer_alignment == 0 + rows, stride_rows = flat_kv_row_view(cache, 64) + cache[1, 0].fill_(7) + torch.testing.assert_close(rows[stride_rows], cache[1, 0]) + assert torch.all(rows[stride_rows] == 7) + + +def test_is_kv_cache_spec_uniform(): + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_kv_cache_spec(num_kv_heads=32), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + } + assert is_kv_cache_spec_uniform(kv_cache_spec) + + kv_cache_spec = { + "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), + "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), + } + assert not is_kv_cache_spec_uniform(kv_cache_spec) + + script = """ +import sys + +from vllm.v1.core.kv_cache_planning import is_kv_cache_spec_uniform +from vllm.v1.kv_cache_interface import KVCacheSpec + +if sys.flags.optimize < 1: + raise RuntimeError("subprocess is not running with optimization enabled") +specs = { + "a": KVCacheSpec(block_size=1), + "b": KVCacheSpec(block_size=2), +} +if is_kv_cache_spec_uniform(specs): + raise RuntimeError("different specs were treated as uniform") +""" + subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[3], + check=True, + timeout=60, + ) + + +@pytest.mark.parametrize( + ("model_id", "max_model_len", "want_estimated_max_len"), + [ + ("Qwen/Qwen1.5-7B", 16385, 16384), + ("Qwen/Qwen1.5-7B", 16383, 16383), + ], +) +def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): + # Create a VllmConfig + model_config = ModelConfig( + model_id, + runner="generate", + dtype="float16", + max_model_len=max_model_len, + ) + scheduler_config = SchedulerConfig( + max_num_batched_tokens=32768, + max_model_len=model_config.max_model_len, + is_encoder_decoder=model_config.is_encoder_decoder, + ) + + vllm_config = VllmConfig( + model_config=model_config, + scheduler_config=scheduler_config, + ) + + # Create KV cache specs + kv_cache_spec = {} + for i in range(32): + layer_name = f"layer_{i}" + kv_cache_spec[layer_name] = FullAttentionSpec( + block_size=16, + num_kv_heads=32, + head_size=128, + dtype=torch.float16, + ) + # Estimate the maximum model length, 16384 model_len need 8GB + estimated_max_len = estimate_max_model_len( + vllm_config, kv_cache_spec, 8 * GiB_bytes + ) + assert estimated_max_len == want_estimated_max_len + + +def test_get_kv_cache_config_one_worker(): + # pass max_model_len to pass check_enough_kv_cache_memory + model_config = ModelConfig(max_model_len=16) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + vllm_config.cache_config.prefix_cache_retention_interval = None + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + # all layers are full attention -> single group + kv_cache_specs_full = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + kv_cache_config_full = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + print(kv_cache_config_full) + assert kv_cache_config_full == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + # all layers are sliding window -> single group + kv_cache_specs_sliding = { + "layer_1": new_sliding_window_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_sliding = get_kv_cache_configs( + vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_sliding == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) + ], + ) + + # full + sliding, but disable_hybrid_kv_cache_manager + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) + ), + ], + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + + # full + sliding, with hybrid_kv_cache_manager + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=64, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 64, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 64, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), + ], + ) + + # 2 full + 4 sliding, 2 layers per group + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_sliding_window_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_3", "layer_5"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 2, + layers=["layer_4", "layer_6"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), + KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), + ], + ) + + # 3 full + 7 sliding, pad to 3 full + 9 sliding + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_sliding_window_spec(), + "layer_5": new_sliding_window_spec(), + "layer_6": new_sliding_window_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1", "layer_2", "layer_3"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_4", "layer_7", "layer_10"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_5", "layer_8"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_6", "layer_9"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), + KVCacheGroupSpec( + ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() + ), + KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), + KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), + ], + ) + + # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss + # eagle where there is only one more full attention layer than sliding window layers + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + "layer_3": new_kv_cache_spec(), + "layer_4": new_kv_cache_spec(), + "layer_5": new_kv_cache_spec(), + "layer_6": new_kv_cache_spec(), + "layer_7": new_sliding_window_spec(), + "layer_8": new_sliding_window_spec(), + "layer_9": new_sliding_window_spec(), + "layer_10": new_sliding_window_spec(), + "layer_11": new_sliding_window_spec(), + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] + )[0] + print(kv_cache_config_hybrid) + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=[ + "layer_1", + "layer_2", + "layer_3", + "layer_4", + "layer_5", + "layer_6", + ], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 6, + layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], + new_kv_cache_spec(), + ), + KVCacheGroupSpec( + ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], + new_sliding_window_spec(), + ), + ], + ) + + # different hidden size but same type, use UniformTypeKVCacheSpecs + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=128), + "layer_2": new_kv_cache_spec(head_size=64), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] + )[0] + # Layers of different page sizes pack densely into one allocation: the + # 2x-larger layer_1 takes the first region, layer_2 the next. + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 2 * 32, + block_stride=mem_per_block_per_layer * 2, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32 * 3, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + offset=mem_per_block_per_layer * 2 * 32, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs_hybrid + ), + ) + ], + ) + + # Different hidden size and different type, align by different block size + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": new_sliding_window_spec(head_size=32), + } + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] + )[0] + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_1"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + KVCacheTensor( + size=mem_per_block_per_layer * 32, + layers=["layer_2"], + layer_stride=mem_per_block_per_layer * 32, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), + KVCacheGroupSpec( + ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) + ), + ], + ) + + # different hidden size that cannot be aligned by using different block size, + # but can be aligned by padding the smaller physical page. + swa_spec = new_sliding_window_spec(head_size=96) + kv_cache_specs_hybrid = { + "layer_1": new_kv_cache_spec(head_size=64), + "layer_2": swa_spec, + } + + kv_cache_config_hybrid = get_kv_cache_configs( + vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] + )[0] + padded_page_size = swa_spec.page_size_bytes + assert kv_cache_config_hybrid == KVCacheConfig( + num_blocks=42, + kv_cache_tensors=[ + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_1"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + KVCacheTensor( + size=padded_page_size * 42, + layers=["layer_2"], + layer_stride=padded_page_size * 42, + block_stride=padded_page_size, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1"], + new_kv_cache_spec( + head_size=64, + page_size_padded=padded_page_size, + ), + ), + KVCacheGroupSpec( + ["layer_2"], + new_sliding_window_spec(head_size=96), + ), + ], + ) + + # Test num_gpu_blocks_override + vllm_config.cache_config.num_gpu_blocks_override = 16 + kv_cache_config_override_blocks = get_kv_cache_configs( + vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] + )[0] + assert kv_cache_config_override_blocks == KVCacheConfig( + num_blocks=16, + kv_cache_tensors=[ + KVCacheTensor( + size=mem_per_block_per_layer * 16 * 2, + layers=["layer_1", "layer_2"], + layer_stride=mem_per_block_per_layer * 16, + block_stride=mem_per_block_per_layer, + ), + ], + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], + ) + + +def test_get_kv_cache_configs_attention_free(): + kv_cache_specs: dict[str, KVCacheSpec] = {} + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None + kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) + assert kv_cache_configs == [ + KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[], + ) + ] + + +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] = {} + 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}.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, + ) + 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, + ) + + groups = kv_cache_planning.get_kv_cache_groups(vllm_config, kv_cache_spec) + attn = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) + ] + rings = [ + g + for g in groups + if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) + and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) + ] + mamba = [ + g + for g in groups + if isinstance( + getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec + ) + ] + 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_planning.get_kv_cache_config_from_groups( + 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_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): + return SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version=model_version, + ) + + +def new_indexer_mla_spec(block_size=16): + # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller + # page size than the main MLA attention (uint8, small head), so their pages + # cannot be unified. + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + ) + + +def test_mixed_page_size_groups_use_spec_compatibility(): + specs = {} + for i in range(3): + specs[f"mla.{i}"] = new_mla_spec() + specs[f"indexer.{i}"] = new_indexer_mla_spec() + specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) + + config = _grouping_config() + config.cache_config = CacheConfig() + config.cache_config.kv_cache_layout = "BLNHC" + groups = get_kv_cache_groups(config, specs) + + assert len(groups) == 3 + assert {name for group in groups for name in group.layer_names} == set(specs) + assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] + + +def _grouping_config(): + cache_config = CacheConfig() + cache_config.kv_cache_layout = "LBNHC" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + cache_config=cache_config, + ) + + +def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): + block_size = 544 + full_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + # FullAttentionSpec stores both K and V, so this produces a + # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. + head_size=512, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + block_size=block_size, + shapes=((557056,),), + dtypes=(torch.bfloat16,), + mamba_cache_mode="align", + ) + hidden_spec = HiddenStateCacheSpec( + block_size=block_size, + num_kv_heads=3, + head_size=1024, + dtype=torch.bfloat16, + ) + assert full_spec.page_size_bytes == mamba_spec.page_size_bytes + + groups = get_kv_cache_groups( + _grouping_config(), + { + "model.full_attn": full_spec, + "model.mamba": mamba_spec, + "cache_only_layers.0": hidden_spec, + }, + ) + + hidden_group = next( + group + for group in groups + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ) + assert hidden_group.kv_cache_spec.block_size == 136 + + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + 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(), + ) + assert kv_cache_utils.resolve_kv_cache_block_sizes( + kv_cache_config, vllm_config + ) == (544, 136) + + +def test_multi_run_layer_compact_strides_place_hoisted_heads(): + """A layer-compact run region is its own dense allocation: under LHBNC the head + groups sit between the layers and the blocks, so a run's block stride is one head + group's slice, not the whole page (regression test for setStorage out-of-bounds + on ROCm hybrid models).""" + full = new_kv_cache_spec() + swa = new_sliding_window_spec(sliding_window=full.block_size * 2) + page = full.page_size_bytes + assert swa.page_size_bytes == page + + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.kv_cache_layout = "LHBNC" + config = kv_cache_planning.get_kv_cache_config_from_groups( + vllm_config, + [ + KVCacheGroupSpec(["full.0", "full.1"], full), + KVCacheGroupSpec(["swa.0", "swa.1"], swa), + ], + available_memory=4 * page * 8, + ) + + num_blocks = config.num_blocks + head_group = full.block_size * full.state_content_size_bytes + for tensor in config.kv_cache_tensors: + # One head group per block, head groups between the layers and blocks. + assert tensor.block_stride == head_group + assert tensor.layer_stride == full.num_heads * num_blocks * head_group + # The last block of the last layer's last head group must stay within the + # allocation (the old page-based stride overran it). + end = ( + tensor.offset + + tensor.layer_stride * (len(tensor.layers) - 1) + + (full.num_heads - 1) * num_blocks * head_group + + tensor.block_stride * (num_blocks - 1) + + head_group + ) + assert end <= tensor.size + + +def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): + specs = { + "target.0.attn": new_mla_spec(), + "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups(_grouping_config(), specs) + + assert len(groups) == 2 + assert all( + not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups + ) + + +def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): + # Sparse MLA pages cannot be padded safely. Keeping the draft's attention + # compute sliding-window while promoting only its allocation semantics lets + # every layer share the target's block table and remain contiguous. + draft = new_sliding_window_spec(block_size=16) + specs = { + "target.0.attn": new_mla_spec(block_size=64), + "target.0.indexer": new_indexer_mla_spec(block_size=64), + "draft.0": draft, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + assert len(groups) == 1 + assert set(groups[0].layer_names) == set(specs) + group_spec = groups[0].kv_cache_spec + assert isinstance(group_spec, UniformTypeKVCacheSpecs) + assert group_spec.block_size == 64 + promoted_draft = group_spec.kv_cache_specs["draft.0"] + assert isinstance(promoted_draft, FullAttentionSpec) + assert not isinstance(promoted_draft, SlidingWindowSpec) + assert promoted_draft.block_size == 64 + assert promoted_draft.sliding_window == draft.sliding_window + assert specs["draft.0"] is draft + assert "attention compute is unchanged" in caplog_vllm.text + + +def test_hidden_states_with_tp_scales_page_size(): + """When TP shrinks KV pages below the hidden-state per-token cost, + get_kv_cache_groups must scale up target block sizes so that the + common page accommodates the unsharded hidden states.""" + # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. + # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) + # = 16 * 2 * 64 * 2 = 4096 bytes. + kv_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + ) + # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) + # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. + hs_spec = HiddenStateCacheSpec( + block_size=16, + num_kv_heads=6, + head_size=512, + dtype=torch.bfloat16, + ) + specs = { + "target.0.attn": kv_spec, + "target.1.attn": kv_spec, + "cache_only_layers.48": hs_spec, + } + + groups = get_kv_cache_groups(_grouping_config(), specs) + + # The hidden-state layer should be present and no assertion should fire. + all_layers = {name for g in groups for name in g.layer_names} + assert "cache_only_layers.48" in all_layers + + # The target group block sizes must have been scaled up. + for g in groups: + if "cache_only_layers.48" not in g.layer_names: + assert g.kv_cache_spec.block_size > kv_spec.block_size + + +def test_auto_fit_max_model_len(): + """Test that max_model_len=-1 auto-fits to available GPU memory.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=1024) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # With enough memory, max_model_len stays at the derived max + large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [large_available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + # Reset for next test + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + # With limited memory, max_model_len should be reduced + # Need memory for at least max_model_len tokens + # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens + limited_memory = mem_per_block_per_layer * 2 * 32 + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [limited_memory] + ) + # Should be reduced to fit in memory + assert vllm_config.model_config.max_model_len < 1024 + assert vllm_config.model_config.max_model_len > 0 + + +def test_auto_fit_max_model_len_with_hybrid(): + """Test that auto-fit works with hybrid KV cache specs.""" + # Create config with original_max_model_len=-1 to trigger auto-fit + model_config = ModelConfig(max_model_len=8192) + # Simulate the user passing -1 by setting original_max_model_len + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer + gamma = 2 + kv_cache_specs = { + "layer_1": new_mamba_spec(num_speculative_blocks=gamma), + "layer_2": new_kv_cache_spec(), + } + + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [available_memory] + ) + assert vllm_config.model_config.max_model_len == 1024 + + +def test_auto_fit_max_model_len_not_triggered(): + """Test that auto-fit is not triggered when original_max_model_len is not -1.""" + model_config = ModelConfig(max_model_len=16) + # original_max_model_len should be None by default, not -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + + # This should work normally without auto-fit + _kv_cache_configs = get_kv_cache_configs( + vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] + ) + assert vllm_config.model_config.max_model_len == 16 + + +def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): + """Auto-fit must size max_model_len against the override-clamped pool, not + the raw `available_memory`. Without this, auto-fit could pick a + max_model_len that no longer fits once `num_gpu_blocks_override` is applied. + """ + model_config = ModelConfig(max_model_len=16384) + model_config.original_max_model_len = -1 # request auto-fit + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # Cap the cache to 32 blocks regardless of available memory. + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), # block_size=16 + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + # 32 blocks * block_size 16 = 512 token slots, so max_model_len must + # auto-fit at or below that. + assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 + + +def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): + """Admission check must use the override-clamped pool size, not raw + `available_memory`. Without this, startup could accept a max_model_len + that does not actually fit in `num_gpu_blocks_override` blocks. + """ + model_config = ModelConfig(max_model_len=16384) + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + # 32 blocks is far too small for max_model_len=16384 (would need 1024). + vllm_config.cache_config.num_gpu_blocks_override = 32 + + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 + kv_cache_specs = { + "layer_1": new_kv_cache_spec(), + "layer_2": new_kv_cache_spec(), + } + # Plenty of raw memory: a bytes-only check against this would pass. + large_available_memory = mem_per_block_per_layer * 2 * 1024 + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) + + +def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): + """DFlash drafters can have a smaller head size than the target model. + + For example, MiMo uses 192-dim target KV heads while its DFlash draft uses + 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer + block-size multiple, so the smaller page must be padded instead. + """ + # Attention layers read padded pages through a strided view, so padding is + # allowed. + target_spec = new_kv_cache_spec( + block_size=16, + num_kv_heads=1, + head_size=192, + dtype=torch.bfloat16, + ) + draft_spec = new_sliding_window_spec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=1024, + ) + + unified_specs = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "target_attn": target_spec, + "draft_attn": draft_spec, + } + ) + + assert unified_specs["target_attn"] == target_spec + unified_draft_spec = unified_specs["draft_attn"] + assert unified_draft_spec.block_size == draft_spec.block_size + assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes + assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes + assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes + + +def test_unify_hybrid_kv_cache_specs(): + # 1. has_full_attention and has_sliding_window + before_spec_1 = new_kv_cache_spec(block_size=64) + before_spec_2 = new_sliding_window_spec( + block_size=16, page_size_padded=32 * 1024, sliding_window=1024 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) + kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec(block_size=64) + expected_spec_2 = new_kv_cache_spec( + block_size=64, page_size_padded=64 * 1024, sliding_window=1024 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 + assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) + + # 2. has_full_attention and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + + # 3. has_full_attention, has_sliding_window and has_chunked_local_attention + before_spec_1 = new_kv_cache_spec() + before_spec_2 = new_sliding_window_spec( + page_size_padded=32 * 1024, sliding_window=1024 + ) + before_spec_3 = new_chunked_local_attention_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + kv_cache_spec = { + "layer_1": before_spec_1, + "layer_2": before_spec_2, + "layer_3": before_spec_3, + } + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + expected_spec_1 = new_kv_cache_spec() + expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) + expected_spec_3 = new_kv_cache_spec( + page_size_padded=32 * 1024, attention_chunk_size=512 + ) + assert kv_cache_spec["layer_1"] == expected_spec_1 + assert kv_cache_spec["layer_2"] == expected_spec_2 + assert kv_cache_spec["layer_3"] == expected_spec_3 + + # 4. No FullAttentionSpec, should not convert + kv_cache_spec = { + "layer_1": new_sliding_window_spec(sliding_window=1024), + "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), + } + + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + # Replicated Mamba state still requires the hybrid cache manager. + kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} + with pytest.raises(ValueError): + kv_cache_planning.unify_hybrid_kv_cache_specs(kv_cache_spec) + + +def test_unify_kv_cache_spec_page_size_mamba(): + """Regression test for https://github.com/vllm-project/vllm/issues/43626. + + MambaSpec's page_size_bytes is determined by its state shapes and does not + change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba + page instead of scaling its block_size. This situation arises when a layer + with a page larger than the (already platform-aligned) Mamba page joins the + specs, e.g. a dense draft model with more KV heads than the hybrid main + model. + """ + # 1. Hybrid main model (Mamba + full attention, pages already aligned at + # 16KB) plus a dense draft model layer with a 2x larger page (32KB). + # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so + # the old code scaled the Mamba block_size, which left page_size_bytes + # unchanged at 16384. + mamba_spec = new_mamba_spec() # page_size_bytes = 16384 + main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 + draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 + assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 + assert draft_attn_spec.page_size_bytes == 32768 + + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": mamba_spec, + "main_attn_layer": main_attn_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + # Mamba page is padded; block_size (caching granularity) is unchanged. + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + assert unified["mamba_layer"].block_size == mamba_spec.block_size + # Attention layer with smaller page still unifies by scaling block_size. + assert unified["main_attn_layer"].page_size_bytes == 32768 + assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size + assert unified["main_attn_layer"].page_size_padded is None + # Layer already at max page size is unchanged. + assert unified["draft_attn_layer"] == draft_attn_spec + + # 2. Mamba page already padded by the platform (state smaller than the + # padded page); the padding is re-applied at the new maximum. + padded_mamba_spec = new_mamba_spec( + shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 + ) + assert padded_mamba_spec.page_size_bytes == 16384 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": padded_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + assert unified["mamba_layer"].page_size_padded == 32768 + + # 3. Mamba page that does not evenly divide the maximum page size is + # padded as well (the divisibility constraint only applies to block_size + # scaling). + odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) + assert odd_mamba_spec.page_size_bytes == 24576 + assert 32768 % odd_mamba_spec.page_size_bytes != 0 + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "mamba_layer": odd_mamba_spec, + "draft_attn_layer": draft_attn_spec, + } + ) + assert unified["mamba_layer"].page_size_bytes == 32768 + + # 4. Attention layers with non-divisible page sizes are padded too: every + # backend reads a padded page through the view's block stride, so there is + # no longer a case that must raise. + unified = kv_cache_planning.unify_kv_cache_spec_page_size( + { + "attn_layer": new_kv_cache_spec(block_size=24), # 24576 + "draft_attn_layer": draft_attn_spec, # 32768 + } + ) + assert unified["attn_layer"].page_size_padded == 32768 + assert unified["attn_layer"].page_size_bytes == 32768 + assert unified["attn_layer"].block_size == 24 + + # 5. Uniform page sizes are returned unchanged. + specs = { + "mamba_layer": new_mamba_spec(), + "attn_layer": new_kv_cache_spec(), + } + assert kv_cache_planning.unify_kv_cache_spec_page_size(specs) == specs + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + 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. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + 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. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) + + +def _spec_decode_grouping_config(method="dspark", model_type=None): + """Grouping config with an EAGLE-family speculative method enabled.""" + return SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + get_resolved_kv_cache_layout=lambda: SimpleNamespace( + is_block_outermost=True + ) + ), + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), + speculative_config=SimpleNamespace( + method=method, + use_eagle=lambda: True, + use_eagle_block_drop=lambda: True, + ), + ) + + +def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): + """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a + DSpark-style draft MLA layer marked non_causal_multi_token_decode. + + The target's fp8 KV dtype is one difference from the draft on Kimi-K3 + (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass + draft_shares_target_spec to match that dtype while retaining the draft + marker as a group boundary. + """ + target_dtype = None if draft_shares_target_spec else "fp8_e4m3" + specs = { + "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), + "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), + } + if draft: + draft_spec = new_mla_spec(block_size=64) + specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) + return specs + + +def test_draft_group_annotated_on_hybrid_general_path(): + # A drafter's MLA layer carries non_causal_multi_token_decode, so its group + # is identifiable without keying off a model version. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "draft.attn.0" in flagged[0].layer_names + + +def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): + # Packed uniform-type groups can contain distinct target and draft layer + # specs; the combined group still holds volatile draft KV and must be + # flagged. What must never happen is a Mamba group being flagged: that + # widens its lookup window to two consecutive chunks, which align-mode + # checkpointing never produces, zeroing every lookup. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), + _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), + ) + + for group in groups: + if "draft.attn.0" in group.layer_names: + assert group.is_eagle_group + if isinstance(group.kv_cache_spec, MambaSpec): + assert not group.is_eagle_group + + +def test_draft_group_not_annotated_without_spec_decode(): + # The marker alone must not flag anything; the eagle semantics only apply + # when a speculative method is actually enabled. + config = _spec_decode_grouping_config() + config.speculative_config = None + groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) + + assert not any(g.is_eagle_group for g in groups) + + +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): + # No group carries the draft marker, so consumers fall back to + # conservative behavior that silently breaks reuse for Mamba groups. + # That must at least be visible. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) + ) + + assert not any(g.is_eagle_group for g in groups) + assert "no KV cache group could be identified as the draft model's" in ( + caplog_vllm.text + ) + + +def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): + # Pure-attention models degrade gracefully under the consumers' + # conservative fallback (a one-block hit drop at most), so the warning + # stays silent to avoid noise on every unannotated EAGLE deployment. + specs = { + "target.attn.0": new_mla_spec(block_size=64), + "target.attn.1": new_mla_spec(block_size=64), + } + groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def test_no_warning_when_draft_group_is_identified(caplog_vllm): + get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + assert "could be identified as the draft model's" not in caplog_vllm.text + + +def _deepseek_v4_specs(model_version="deepseek_v4"): + """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers + at differing page sizes, with the MTP draft layer registered last.""" + return { + "model.layers.0.self_attn.attn": new_mla_spec(), + "model.layers.1.self_attn.attn": new_mla_spec(), + "model.layers.2.self_attn.attn": new_swa_mla_spec( + head_size=1024, model_version=model_version + ), + # The MTP block registers last, and its sliding-window size differs, so + # it lands in a group of its own. + "model.layers.3.self_attn.attn": new_swa_mla_spec( + head_size=1024, sliding_window=256, model_version=model_version + ), + } + + +@pytest.mark.parametrize( + ("method", "model_type"), + [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], +) +def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): + # DeepseekV4's MTP block reuses the target's decoder layer, so its spec + # carries no draft marker and only the positional rule can find it. This + # pins the pre-existing behaviour that the unified annotator must preserve. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type=model_type), + _deepseek_v4_specs(model_version=None), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def test_trailing_layer_fallback_applies_to_any_mtp_model(): + # The positional rule is sound for every MTP drafter, not just DeepseekV4: + # MTP blocks reuse the target's decoder layer (no spec marker) and always + # register after every target layer. The model_type must not gate it. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="other"), + _deepseek_v4_specs(), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.3.self_attn.attn" in flagged[0].layer_names + + +def _qwen3_5_hybrid_specs(with_mtp_layer: bool): + """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with + the MTP drafter's full-attn layer (spec-identical to the target's) + registered last.""" + specs = {} + idx = 0 + for _ in range(2): + for _ in range(3): + specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( + mamba_cache_mode="align" + ) + idx += 1 + specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() + idx += 1 + if with_mtp_layer: + specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() + return specs + + +def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): + # A hybrid mamba + full-attention model with an MTP drafter that is + # spec-indistinguishable from the target reaches the general multi-group + # path. The trailing-layer rule must locate the draft group there so the + # Mamba groups are not swept up by the flag-all consumer fallback. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + flagged = [g for g in groups if g.is_eagle_group] + assert len(flagged) == 1 + assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names + for group in groups: + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + assert not group.is_eagle_group + assert "could be identified as the draft model's" not in caplog_vllm.text + + +@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) +def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): + # Other drafters are not covered by the trailing-layer rule, so an + # unidentifiable hybrid draft must still warn. + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method=method, model_type="qwen3_5"), + _qwen3_5_hybrid_specs(with_mtp_layer=True), + ) + + assert not any(g.is_eagle_group for g in groups) + assert "could be identified as the draft model's" in caplog_vllm.text + + +def test_trailing_layer_fallback_requires_exact_partition(): + # If the groups do not partition the layers exactly (e.g. a caller that + # dropped or duplicated layers), the positional rule is meaningless and + # must not fire. + from vllm.v1.core.kv_cache_planning import _annotate_eagle_groups + + specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) + config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") + groups = get_kv_cache_groups(config, specs) + for g in groups: + g.is_eagle_group = False + # Remove one layer from its group: no longer an exact partition. + trimmed = [ + KVCacheGroupSpec( + [n for n in g.layer_names if n != "model.layers.0.linear_attn"], + g.kv_cache_spec, + ) + for g in groups + ] + _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) + + assert not any(g.is_eagle_group for g in trimmed) + + +_GQA_SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 +) + + +_GQA_SWA_SPEC = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 +) + + +@pytest.mark.parametrize( + "specs,expected", + [ + pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": new_swa_mla_spec(), + "s.1": new_swa_mla_spec(), + }, + 4, + id="mla-and-swa-mla", + ), + pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), + pytest.param( + {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, + 1, + id="mla-and-gqa-uniform-group", + ), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": _GQA_SWA_SPEC, + "s.1": _GQA_SWA_SPEC, + }, + 1, + id="mla-and-gqa-swa", + ), + ], +) +def test_kv_tp_replicas(monkeypatch, specs, expected): + """Resolved per layer before scheduler flattening hides mixed groups.""" + from vllm.config import ParallelConfig + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "device_count", lambda: 4) + parallel_config = ParallelConfig(tensor_parallel_size=4) + vllm_config = VllmConfig( + model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config + ) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + mem = sum(s.page_size_bytes for s in specs.values()) * 10 + configs = get_kv_cache_configs(vllm_config, [specs], [mem]) + assert configs[0].kv_tp_replicas == expected + scheduler = generate_scheduler_kv_cache_config(configs) + assert scheduler.kv_tp_replicas == expected + + +@pytest.mark.parametrize( + "specs,tp_size,dcp_size,expected", + [ + pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), + pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), + pytest.param( + [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" + ), + pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), + pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), + pytest.param( + [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], + 8, + 1, + 2, + id="mla-and-gqa-partial", + ), + pytest.param( + [ + HiddenStateCacheSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 + ) + ], + 8, + 1, + 1, + id="hidden-state", + ), + ], +) +def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): + """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" + groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] + assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index b06479f1caf1..1489315ec255 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,20 +3,23 @@ import copy import hashlib import importlib -import mmap -import subprocess -import sys from collections.abc import Callable from dataclasses import replace -from pathlib import Path from types import SimpleNamespace from typing import Any import pytest import torch +import vllm.v1.core.kv_cache_planning as kv_cache_planning import vllm.v1.core.kv_cache_utils as kv_cache_utils import vllm.v1.hisparse.runtime as hisparse_runtime_module +from tests.v1.core.utils import ( + new_kv_cache_spec, + new_mamba_spec, + new_mla_spec, + new_sliding_window_spec, +) from vllm.config import ( CacheConfig, ModelConfig, @@ -33,25 +36,21 @@ ) from vllm.sampling_params import SamplingParams from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor -from vllm.utils.mem_constants import GiB_bytes from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager +from vllm.v1.core.kv_cache_planning import ( + get_kv_cache_groups, +) from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, - check_enough_kv_cache_memory, - estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, get_kv_cache_capacity, - get_kv_cache_configs, - get_kv_cache_groups, get_max_concurrency_for_kv_cache_config, get_request_block_hasher, hash_block_tokens, init_none_hash, - is_kv_cache_spec_uniform, - kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, @@ -62,10 +61,8 @@ get_hisparse_steady_state_concurrency, ) from vllm.v1.kv_cache_interface import ( - ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, - HiddenStateCacheSpec, HiSparseHotSpec, HiSparseResidentSpec, KVCacheConfig, @@ -93,110 +90,6 @@ pytestmark = pytest.mark.cpu_test -@pytest.mark.parametrize("gpu_block_size", [32, 64]) -@pytest.mark.parametrize("shared_host_pool", [False, True]) -def test_hisparse_hma_uses_resolved_gpu_block_size( - monkeypatch, gpu_block_size, shared_host_pool -): - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - specs = { - "model.layers.0.self_attn": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ), - "model.layers.0.self_attn.indexer": MLAAttentionSpec( - block_size=gpu_block_size, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ), - } - config = SimpleNamespace( - attention_config=SimpleNamespace(hisparse_config=HiSparseConfig()), - model_config=SimpleNamespace( - hf_config=SimpleNamespace(index_topk=128), - max_model_len=gpu_block_size, - ), - parallel_config=SimpleNamespace( - tensor_parallel_size=2 if shared_host_pool else 1, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, - world_size=2 if shared_host_pool else 1, - distributed_executor_backend="mp", - nnodes_within_dp=1, - ), - cache_config=SimpleNamespace( - num_gpu_blocks_override=7, - prefix_cache_retention_interval=None, - get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC, - ), - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, get_hisparse_kv_cache_groups(config, specs), available_memory=2**30 - ) - assert cache_config.num_blocks == 7 - assert cache_config.hisparse_host_num_blocks is not None - assert cache_config.hisparse_host_num_blocks > 7 - - host_group, indexer_group, *auxiliary_groups = cache_config.kv_cache_groups - assert cache_config.hisparse_shared_host_pool is shared_host_pool - host_page = host_group.kv_cache_spec.page_size_bytes - alignment = mmap.PAGESIZE if shared_host_pool else 1 - expected_host_stride = (host_page + alignment - 1) // alignment * alignment - assert cache_config.hisparse_host_block_stride == expected_host_stride - assert cache_config.hisparse_host_num_blocks == 2**30 // expected_host_stride - assert host_group.host_resident - assert not any(group.host_resident for group in [indexer_group, *auxiliary_groups]) - host_layers = set(host_group.layer_names) - for tensor in cache_config.kv_cache_tensors: - assert all( - (name in host_layers) == tensor.host_resident for name in tensor.layers - ) - host_tensor = next(t for t in cache_config.kv_cache_tensors if t.host_resident) - host_spec = host_group.kv_cache_spec.kv_cache_specs[host_tensor.layers[0]] - assert host_tensor.block_stride == host_spec.page_size_bytes - assert host_tensor.layer_stride == ( - host_spec.page_size_bytes * cache_config.hisparse_host_num_blocks - ) - assert host_group.kv_cache_spec.block_size == gpu_block_size - assert indexer_group.kv_cache_spec.block_size == gpu_block_size - host_specs = host_group.kv_cache_spec.kv_cache_specs - gpu_indexer_specs = indexer_group.kv_cache_spec.kv_cache_specs - assert set(host_specs) == {"model.layers.0.self_attn"} - assert set(gpu_indexer_specs) == {"model.layers.0.self_attn.indexer"} - assert indexer_group.kv_cache_spec.prefix_cacheable - assert host_group.enable_kv_transfer - auxiliary_specs = [group.kv_cache_spec for group in auxiliary_groups] - assert any(isinstance(spec, HiSparseResidentSpec) for spec in auxiliary_specs) - assert any(isinstance(spec, HiSparseHotSpec) for spec in auxiliary_specs) - assert all( - spec.block_size == gpu_block_size - for spec in auxiliary_specs - if isinstance(spec, (HiSparseResidentSpec, HiSparseHotSpec)) - ) - scheduler_block_size, hash_block_size = kv_cache_utils.resolve_kv_cache_block_sizes( - cache_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(), - ), - ) - assert scheduler_block_size == hash_block_size == gpu_block_size - - def test_hisparse_steady_state_concurrency_excludes_spilled_resident_pages(): """Running requests pin only their resident tail; the newest still needs its full in-flight window to be admitted.""" @@ -281,68 +174,6 @@ def test_hisparse_host_pool_must_fit_max_model_len(monkeypatch, extra_blocks, ok create_hisparse_layout(config, groups, host_budget=host_budget) -@pytest.mark.parametrize( - "max_model_len,num_gpu_blocks,ok", - [ - (8192, 160, False), - (8192, 1000, True), - (32768, 656, False), - (32768, 657, True), - (-1, 657, True), - ], -) -def test_hisparse_pool_must_fit_max_model_len( - monkeypatch, max_model_len, num_gpu_blocks, ok -): - """Each resident group takes its own blocks from HiSparse's shared GPU pool, - so a pool that fits only the indexer pages must be rejected at startup. - Resident pages are capped at the in-flight window: at 32768 tokens with 2048 - in flight, one request needs 512 indexer + 4 x 32 resident + 16 hot blocks, - plus the null block. Auto-fit (-1) must pick a length that passes the same - check.""" - monkeypatch.setattr( - hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True - ) - monkeypatch.setattr(kv_cache_utils, "get_hisparse_host_pool_bytes", lambda _: 2**30) - specs: dict[str, KVCacheSpec] = {} - for i in range(4): - specs[f"model.layers.{i}.self_attn"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=576, - dtype=torch.bfloat16, - is_index_group_leader=True, - ) - specs[f"model.layers.{i}.self_attn.indexer"] = MLAAttentionSpec( - block_size=64, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - cache_role=SparseCacheRole.INDEXER, - ) - config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) - config.scheduler_config.max_num_batched_tokens = 2048 - # Async scheduling doubles the in-flight window, and with it the bound. - config.scheduler_config.async_scheduling = False - config.attention_config.hisparse_config = HiSparseConfig() - config.model_config.hf_config.index_topk = 128 - config.cache_config.num_gpu_blocks_override = num_gpu_blocks - config.cache_config.kv_cache_layout = "BLHNC" - # Workers report the resident/hot caches HiSparse derives. - specs |= { - name: group.kv_cache_spec - for group in get_hisparse_kv_cache_groups(config, specs) or [] - if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec)) - for name in group.layer_names - } - - if ok: - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - else: - with pytest.raises(ValueError, match="max seq len"): - kv_cache_utils.get_kv_cache_configs(config, [specs], [2**34]) - - def test_hisparse_derived_specs_do_not_affect_kv_cache_layout(monkeypatch): """Workers report HiSparse's resident/hot caches, which lay out their own raw backing; KV cache layout resolution must skip them.""" @@ -496,28 +327,6 @@ def test_effective_attention_block_size_matches_events(dcp): ] == [block_size] -def new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=None, - attention_chunk_size=None, - kv_quant_mode=KVQuantMode.NONE, -): - return FullAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - attention_chunk_size=attention_chunk_size, - kv_quant_mode=kv_quant_mode, - ) - - def test_kv_cache_config_selects_only_transferable_groups(): """Connectors must see a stable projection of transfer-eligible groups.""" groups = [ @@ -581,60 +390,6 @@ def test_kv_cache_blocks_selects_requested_groups(): assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3]) -def new_sliding_window_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - sliding_window=1, -): - return SlidingWindowSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - sliding_window=sliding_window, - ) - - -def new_chunked_local_attention_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.float32, - page_size_padded=None, - attention_chunk_size=4, -): - return ChunkedLocalAttentionSpec( - block_size=block_size, - num_kv_heads=num_kv_heads, - head_size=head_size, - dtype=dtype, - page_size_padded=page_size_padded, - attention_chunk_size=attention_chunk_size, - ) - - -def new_mamba_spec( - block_size=16, - shapes=((2, 512), (3, 32, 32)), - dtypes=(torch.float32, torch.float32), - num_speculative_blocks=2, - mamba_cache_mode="none", - page_size_padded=None, -): - return MambaSpec( - block_size=block_size, - shapes=shapes, - dtypes=dtypes, - page_size_padded=page_size_padded, - mamba_cache_mode=mamba_cache_mode, - num_speculative_blocks=num_speculative_blocks, - ) - - @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor]) def test_none_hash(monkeypatch, hash_fn): import vllm.v1.core.kv_cache_utils @@ -1441,430 +1196,6 @@ def test_metrics_empty_stats(): assert metrics.hit_rate == 0.75 -def test_get_kv_cache_configs_multiple_workers(): - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - same_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - ] - - # Basic case. All things are the same. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different available memory. This is the case for TP. - # Use the smallest memory available. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - same_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 20, - ], - ) - assert kv_cache_configs == [expected, expected] - - # Different KV cache specs. This is the case for PP. - different_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer2": new_kv_cache_spec(), - "layer3": new_kv_cache_spec(), - }, - ] - - # Different workers have different layers. - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], new_kv_cache_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer2", "layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()), - ], - ), - ] - - # Some layers are the same, some are different. This is the case for TP+PP - tp_pp_kv_cache_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - { - "layer3": new_kv_cache_spec(), - }, - ] - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - tp_pp_kv_cache_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - expected_12 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - ], - ) - expected_3 = KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer3"], ref_kv_cache_spec), - ], - ) - assert kv_cache_configs == [ - expected_12, - expected_12, - expected_3, - expected_3, - ] - - # Different workers have different types of layers. This is the case for - # hybrid models + PP. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_kv_cache_spec(), - }, - { - "layer3": new_sliding_window_spec(), - "layer4": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer1", "layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec), - KVCacheGroupSpec([], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10 * 2, - layers=["layer3", "layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec([], ref_kv_cache_spec), - KVCacheGroupSpec(["layer3", "layer4"], new_sliding_window_spec()), - ], - ), - ] - - # When divided into multiple KVCacheGroups, need to ensure the number of - # layers per group is similar. - different_type_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - "layer2": new_sliding_window_spec(), - "layer3": new_sliding_window_spec(), - }, - { - "layer4": new_kv_cache_spec(), - "layer5": new_sliding_window_spec(), - "layer6": new_sliding_window_spec(), - }, - ] - kv_cache_configs = get_kv_cache_configs( - vllm_config, - different_type_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 10, - ref_kv_cache_spec.page_size_bytes * 10, - ], - ) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer3"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer1"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer2"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer3"], new_sliding_window_spec()), - ], - ), - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer4"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer5"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * 10, - layers=["layer6"], - layer_stride=ref_kv_cache_spec.page_size_bytes * 10, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer4"], ref_kv_cache_spec), - KVCacheGroupSpec(["layer5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer6"], new_sliding_window_spec()), - ], - ), - ] - - # Have conflicting layers. Need to raise an error. - conflicting_layer_specs = [ - { - "layer1": new_kv_cache_spec(), - }, - { - "layer1": new_sliding_window_spec(), - }, - ] - with pytest.raises(AssertionError): - get_kv_cache_configs( - vllm_config, - conflicting_layer_specs, - [ - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ref_kv_cache_spec.page_size_bytes * 2 * 10, - ], - ) - - -@pytest.mark.parametrize( - "asymmetric_memory", - [False, True], - ids=["symmetric", "asymmetric"], -) -def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): - model_config = ModelConfig(max_model_len=512) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - ref_kv_cache_spec = new_kv_cache_spec() - pp_kv_cache_specs = [ - {"layer1": ref_kv_cache_spec}, - {"layer2": ref_kv_cache_spec}, - ] - - expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1 - avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks - - # With per-worker validation, each worker only needs memory for its own - # layers. Worker 2 having more memory shouldn't affect worker 1's config. - available_memory = ( - [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2 - ) - - kv_cache_configs = get_kv_cache_configs( - vllm_config, - pp_kv_cache_specs, - available_memory, - ) - - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer1"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)], - ), - KVCacheConfig( - num_blocks=expected_num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks, - layers=["layer2"], - layer_stride=ref_kv_cache_spec.page_size_bytes - * expected_num_blocks, - block_stride=ref_kv_cache_spec.page_size_bytes, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)], - ), - ] - - -def test_project_kv_cache_groups_to_worker(): - spec_a = new_kv_cache_spec() - spec_b = new_kv_cache_spec(num_kv_heads=4) - - global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), - ] - worker_spec = {"layer1": spec_a, "layer2": spec_a} - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, worker_spec - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer2"] - assert projected[0].kv_cache_spec is spec_a - - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups, {"layer4": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == [] - assert projected[0].kv_cache_spec is spec_a - - uniform_spec = UniformTypeKVCacheSpecs( - block_size=16, - kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a}, - ) - global_groups_uniform = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec), - ] - projected = kv_cache_utils._project_kv_cache_groups_to_worker( - global_groups_uniform, {"layer1": spec_a, "layer3": spec_a} - ) - assert len(projected) == 1 - assert projected[0].layer_names == ["layer1", "layer3"] - proj_spec = projected[0].kv_cache_spec - assert isinstance(proj_spec, UniformTypeKVCacheSpecs) - assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} - - @pytest.mark.parametrize("sliding_window", [None, 256]) @pytest.mark.parametrize("disable_hybrid", [False, True]) @pytest.mark.parametrize("pcp_size", [1, 4]) @@ -1901,7 +1232,7 @@ def test_dcp_target_allocates_replicated_draft_independently( widths = [g.kv_cache_spec.max_num_blocks_per_req(config, 1024) for g in groups] assert widths == [16, 64] - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( + cache_config = kv_cache_planning.get_kv_cache_config_from_groups( config, groups, available_memory=16 * 1024 * 1024 ) scheduler_config = generate_scheduler_kv_cache_config([cache_config]) @@ -1971,58 +1302,6 @@ def test_full_attention_merge_preserves_replicated_cache_geometry(use_mla): type(draft).merge([draft, replace(draft, dcp_sharded=True)]) -@pytest.mark.parametrize("with_draft", [False, True]) -@pytest.mark.parametrize("indexer_alignment", [None, 512]) -def test_sparse_mla_preserves_physical_row_addressing(with_draft, indexer_alignment): - from vllm.v1.attention.backends.mla.sparse_utils import flat_kv_row_view - from vllm.v1.worker.utils import allocate_kv_cache - - config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - config.cache_config.kv_cache_layout = "BLHNC" - config.parallel_config.decode_context_parallel_size = 4 - common = dict(block_size=64, num_kv_heads=1, dtype=torch.uint8) - specs = { - "target": MLAAttentionSpec( - **common, - head_size=576, - state_content_bytes=656, - block_stride_alignment=656, - cache_dtype_str="fp8_ds_mla", - is_index_group_leader=True, - ), - "indexer": MLAAttentionSpec( - **common, - head_size=132, - cache_role=SparseCacheRole.INDEXER, - block_stride_alignment=indexer_alignment, - ), - "draft": SlidingWindowSpec( - block_size=64, - num_kv_heads=64, - head_size=64, - dtype=torch.bfloat16, - sliding_window=2048, - dcp_sharded=False, - ), - } - if not with_draft: - del specs["draft"] - groups = get_kv_cache_groups(config, specs) - cache_config = kv_cache_utils.get_kv_cache_config_from_groups( - config, groups, available_memory=8 * 1024 * 1024 - ) - caches = allocate_kv_cache( - cache_config, torch.device("cpu"), KVCacheLayout.BLHNC, [64] * len(groups) - ) - cache = caches["target"].squeeze(1) - if indexer_alignment is not None: - assert caches["indexer"].stride(0) % indexer_alignment == 0 - rows, stride_rows = flat_kv_row_view(cache, 64) - cache[1, 0].fill_(7) - torch.testing.assert_close(rows[stride_rows], cache[1, 0]) - assert torch.all(rows[stride_rows] == 7) - - @pytest.mark.parametrize( "layer_type,dcp_size,expected_width", [ @@ -2115,103 +1394,6 @@ def test_merge_kv_cache_spec(): assert merged_layer_spec.sliding_window == 1 -def test_is_kv_cache_spec_uniform(): - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_kv_cache_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_kv_cache_spec(num_kv_heads=32), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - } - assert is_kv_cache_spec_uniform(kv_cache_spec) - - kv_cache_spec = { - "layer_1": new_sliding_window_spec(num_kv_heads=32, sliding_window=1), - "layer_2": new_sliding_window_spec(num_kv_heads=32, sliding_window=2), - } - assert not is_kv_cache_spec_uniform(kv_cache_spec) - - script = """ -import sys - -from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform -from vllm.v1.kv_cache_interface import KVCacheSpec - -if sys.flags.optimize < 1: - raise RuntimeError("subprocess is not running with optimization enabled") -specs = { - "a": KVCacheSpec(block_size=1), - "b": KVCacheSpec(block_size=2), -} -if is_kv_cache_spec_uniform(specs): - raise RuntimeError("different specs were treated as uniform") -""" - subprocess.run( - [sys.executable, "-O", "-c", script], - cwd=Path(__file__).resolve().parents[3], - check=True, - timeout=60, - ) - - -@pytest.mark.parametrize( - ("model_id", "max_model_len", "want_estimated_max_len"), - [ - ("Qwen/Qwen1.5-7B", 16385, 16384), - ("Qwen/Qwen1.5-7B", 16383, 16383), - ], -) -def test_estimate_max_model_len(model_id, max_model_len, want_estimated_max_len): - # Create a VllmConfig - model_config = ModelConfig( - model_id, - runner="generate", - dtype="float16", - max_model_len=max_model_len, - ) - scheduler_config = SchedulerConfig( - max_num_batched_tokens=32768, - max_model_len=model_config.max_model_len, - is_encoder_decoder=model_config.is_encoder_decoder, - ) - - vllm_config = VllmConfig( - model_config=model_config, - scheduler_config=scheduler_config, - ) - - # Create KV cache specs - kv_cache_spec = {} - for i in range(32): - layer_name = f"layer_{i}" - kv_cache_spec[layer_name] = FullAttentionSpec( - block_size=16, - num_kv_heads=32, - head_size=128, - dtype=torch.float16, - ) - # Estimate the maximum model length, 16384 model_len need 8GB - estimated_max_len = estimate_max_model_len( - vllm_config, kv_cache_spec, 8 * GiB_bytes - ) - assert estimated_max_len == want_estimated_max_len - - def test_get_max_concurrency_for_kv_cache_config(): # Create a VllmConfig model_id = "Qwen/Qwen1.5-7B" @@ -2425,411 +1607,6 @@ def test_allocate_with_lookahead(): assert len(blocks.get_block_ids()[0]) == 2 -def test_get_kv_cache_config_one_worker(): - # pass max_model_len to pass check_enough_kv_cache_memory - model_config = ModelConfig(max_model_len=16) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - vllm_config.cache_config.prefix_cache_retention_interval = None - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - # all layers are full attention -> single group - kv_cache_specs_full = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - kv_cache_config_full = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - print(kv_cache_config_full) - assert kv_cache_config_full == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - # all layers are sliding window -> single group - kv_cache_specs_sliding = { - "layer_1": new_sliding_window_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_sliding = get_kv_cache_configs( - vllm_config, [kv_cache_specs_sliding], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_sliding == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec()) - ], - ) - - # full + sliding, but disable_hybrid_kv_cache_manager - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], new_kv_cache_spec(sliding_window=1) - ), - ], - ) - vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False - - # full + sliding, with hybrid_kv_cache_manager - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=64, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 64, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 64, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()), - ], - ) - - # 2 full + 4 sliding, 2 layers per group - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_sliding_window_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_3", "layer_5"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 2, - layers=["layer_4", "layer_6"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()), - KVCacheGroupSpec(["layer_3", "layer_5"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_4", "layer_6"], new_sliding_window_spec()), - ], - ) - - # 3 full + 7 sliding, pad to 3 full + 9 sliding - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_sliding_window_spec(), - "layer_5": new_sliding_window_spec(), - "layer_6": new_sliding_window_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1", "layer_2", "layer_3"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_4", "layer_7", "layer_10"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_5", "layer_8"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_6", "layer_9"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"], new_kv_cache_spec()), - KVCacheGroupSpec( - ["layer_4", "layer_7", "layer_10"], new_sliding_window_spec() - ), - KVCacheGroupSpec(["layer_5", "layer_8"], new_sliding_window_spec()), - KVCacheGroupSpec(["layer_6", "layer_9"], new_sliding_window_spec()), - ], - ) - - # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss - # eagle where there is only one more full attention layer than sliding window layers - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - "layer_3": new_kv_cache_spec(), - "layer_4": new_kv_cache_spec(), - "layer_5": new_kv_cache_spec(), - "layer_6": new_kv_cache_spec(), - "layer_7": new_sliding_window_spec(), - "layer_8": new_sliding_window_spec(), - "layer_9": new_sliding_window_spec(), - "layer_10": new_sliding_window_spec(), - "layer_11": new_sliding_window_spec(), - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32] - )[0] - print(kv_cache_config_hybrid) - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=[ - "layer_1", - "layer_2", - "layer_3", - "layer_4", - "layer_5", - "layer_6", - ], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 6, - layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"], - new_kv_cache_spec(), - ), - KVCacheGroupSpec( - ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"], - new_sliding_window_spec(), - ), - ], - ) - - # different hidden size but same type, use UniformTypeKVCacheSpecs - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=128), - "layer_2": new_kv_cache_spec(head_size=64), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32] - )[0] - # Layers of different page sizes pack densely into one allocation: the - # 2x-larger layer_1 takes the first region, layer_2 the next. - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 2 * 32, - block_stride=mem_per_block_per_layer * 2, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32 * 3, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - offset=mem_per_block_per_layer * 2 * 32, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs_hybrid - ), - ) - ], - ) - - # Different hidden size and different type, align by different block size - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": new_sliding_window_spec(head_size=32), - } - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 32] - )[0] - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=32, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_1"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - KVCacheTensor( - size=mem_per_block_per_layer * 32, - layers=["layer_2"], - layer_stride=mem_per_block_per_layer * 32, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec(["layer_1"], new_kv_cache_spec(head_size=64)), - KVCacheGroupSpec( - ["layer_2"], new_sliding_window_spec(head_size=32, block_size=32) - ), - ], - ) - - # different hidden size that cannot be aligned by using different block size, - # but can be aligned by padding the smaller physical page. - swa_spec = new_sliding_window_spec(head_size=96) - kv_cache_specs_hybrid = { - "layer_1": new_kv_cache_spec(head_size=64), - "layer_2": swa_spec, - } - - kv_cache_config_hybrid = get_kv_cache_configs( - vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32] - )[0] - padded_page_size = swa_spec.page_size_bytes - assert kv_cache_config_hybrid == KVCacheConfig( - num_blocks=42, - kv_cache_tensors=[ - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_1"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - KVCacheTensor( - size=padded_page_size * 42, - layers=["layer_2"], - layer_stride=padded_page_size * 42, - block_stride=padded_page_size, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1"], - new_kv_cache_spec( - head_size=64, - page_size_padded=padded_page_size, - ), - ), - KVCacheGroupSpec( - ["layer_2"], - new_sliding_window_spec(head_size=96), - ), - ], - ) - - # Test num_gpu_blocks_override - vllm_config.cache_config.num_gpu_blocks_override = 16 - kv_cache_config_override_blocks = get_kv_cache_configs( - vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32] - )[0] - assert kv_cache_config_override_blocks == KVCacheConfig( - num_blocks=16, - kv_cache_tensors=[ - KVCacheTensor( - size=mem_per_block_per_layer * 16 * 2, - layers=["layer_1", "layer_2"], - layer_stride=mem_per_block_per_layer * 16, - block_stride=mem_per_block_per_layer, - ), - ], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -def test_get_kv_cache_configs_attention_free(): - kv_cache_specs: dict[str, KVCacheSpec] = {} - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.prefix_cache_retention_interval = None - kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) - assert kv_cache_configs == [ - KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=[], - ) - ] - - def test_generate_uniform_type_kv_cache_specs(): # All layers are full attention, can be merged kv_cache_specs = { @@ -2879,223 +1656,28 @@ def test_generate_uniform_type_kv_cache_specs(): def test_generate_scheduler_kv_cache_config(): kv_cache_specs = { "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(head_size=128), - } - kv_cache_configs = [ - KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer_1", "layer_2"], - UniformTypeKVCacheSpecs( - block_size=16, kv_cache_specs=kv_cache_specs - ), - ), - ], - ) - ] - scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) - assert scheduler_kv_cache_config == KVCacheConfig( - num_blocks=10, - kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], - ) - - -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] = {} - 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}.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, - ) - 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, - ) - - groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec) - attn = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, MLAAttentionSpec) - ] - rings = [ - g - for g in groups - if isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) - and isinstance(g.kv_cache_spec.first_spec, CircularBufferSpec) - ] - mamba = [ - g - for g in groups - if isinstance( - getattr(g.kv_cache_spec, "first_spec", g.kv_cache_spec), MambaSpec - ) - ] - 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, 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): - # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - max_tp_shards=1, - head_size=576, - dtype=torch.float32, - cache_dtype_str=cache_dtype_str, - ) - - -def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): - return SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - max_tp_shards=1, - head_size=head_size, - dtype=torch.float32, - sliding_window=sliding_window, - model_version=model_version, - ) - - -def new_indexer_mla_spec(block_size=16): - # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller - # page size than the main MLA attention (uint8, small head), so their pages - # cannot be unified. - return MLAAttentionSpec( - block_size=block_size, - num_kv_heads=1, - head_size=132, - dtype=torch.uint8, - ) - - -def test_mixed_page_size_groups_use_spec_compatibility(): - specs = {} - for i in range(3): - specs[f"mla.{i}"] = new_mla_spec() - specs[f"indexer.{i}"] = new_indexer_mla_spec() - specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)}) - - config = _grouping_config() - config.cache_config = CacheConfig() - config.cache_config.kv_cache_layout = "BLNHC" - groups = get_kv_cache_groups(config, specs) - - assert len(groups) == 3 - assert {name for group in groups for name in group.layer_names} == set(specs) - assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6] - - -def _grouping_config(): - cache_config = CacheConfig() - cache_config.kv_cache_layout = "LBNHC" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - speculative_config=None, - cache_config=cache_config, - ) - - -def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity(): - block_size = 544 - full_spec = FullAttentionSpec( - block_size=block_size, - num_kv_heads=1, - # FullAttentionSpec stores both K and V, so this produces a - # 544 * 1 * (512 + 512) * 2 = 1,114,112-byte page. - head_size=512, - dtype=torch.bfloat16, - ) - mamba_spec = MambaSpec( - block_size=block_size, - shapes=((557056,),), - dtypes=(torch.bfloat16,), - mamba_cache_mode="align", - ) - hidden_spec = HiddenStateCacheSpec( - block_size=block_size, - num_kv_heads=3, - head_size=1024, - dtype=torch.bfloat16, - ) - assert full_spec.page_size_bytes == mamba_spec.page_size_bytes - - groups = get_kv_cache_groups( - _grouping_config(), - { - "model.full_attn": full_spec, - "model.mamba": mamba_spec, - "cache_only_layers.0": hidden_spec, - }, - ) - - hidden_group = next( - group - for group in groups - if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) - ) - assert hidden_group.kv_cache_spec.block_size == 136 - - kv_cache_config = KVCacheConfig( - num_blocks=1, + "layer_2": new_kv_cache_spec(head_size=128), + } + kv_cache_configs = [ + KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer_1", "layer_2"], + UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs + ), + ), + ], + ) + ] + scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs) + assert scheduler_kv_cache_config == KVCacheConfig( + num_blocks=10, kv_cache_tensors=[], - kv_cache_groups=groups, - ) - 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(), + kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())], ) - assert kv_cache_utils.resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) == (544, 136) def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): @@ -3133,125 +1715,6 @@ def test_resolve_dcp_kv_block_size_unwraps_uniform_type_specs(): assert kv_cache_utils.resolve_dcp_kv_cache_spec(wrapped_mamba, 4) is wrapped_mamba -def test_multi_run_layer_compact_strides_place_hoisted_heads(): - """A layer-compact run region is its own dense allocation: under LHBNC the head - groups sit between the layers and the blocks, so a run's block stride is one head - group's slice, not the whole page (regression test for setStorage out-of-bounds - on ROCm hybrid models).""" - full = new_kv_cache_spec() - swa = new_sliding_window_spec(sliding_window=full.block_size * 2) - page = full.page_size_bytes - assert swa.page_size_bytes == page - - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) - vllm_config.cache_config.kv_cache_layout = "LHBNC" - config = kv_cache_utils.get_kv_cache_config_from_groups( - vllm_config, - [ - KVCacheGroupSpec(["full.0", "full.1"], full), - KVCacheGroupSpec(["swa.0", "swa.1"], swa), - ], - available_memory=4 * page * 8, - ) - - num_blocks = config.num_blocks - head_group = full.block_size * full.state_content_size_bytes - for tensor in config.kv_cache_tensors: - # One head group per block, head groups between the layers and blocks. - assert tensor.block_stride == head_group - assert tensor.layer_stride == full.num_heads * num_blocks * head_group - # The last block of the last layer's last head group must stay within the - # allocation (the old page-based stride overran it). - end = ( - tensor.offset - + tensor.layer_stride * (len(tensor.layers) - 1) - + (full.num_heads - 1) * num_blocks * head_group - + tensor.block_stride * (num_blocks - 1) - + head_group - ) - assert end <= tensor.size - - -def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified(): - specs = { - "target.0.attn": new_mla_spec(), - "draft.0": new_sliding_window_spec(num_kv_heads=1, head_size=288), - } - assert len({spec.page_size_bytes for spec in specs.values()}) == 1 - - groups = get_kv_cache_groups(_grouping_config(), specs) - - assert len(groups) == 2 - assert all( - not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in groups - ) - - -def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm): - # Sparse MLA pages cannot be padded safely. Keeping the draft's attention - # compute sliding-window while promoting only its allocation semantics lets - # every layer share the target's block table and remain contiguous. - draft = new_sliding_window_spec(block_size=16) - specs = { - "target.0.attn": new_mla_spec(block_size=64), - "target.0.indexer": new_indexer_mla_spec(block_size=64), - "draft.0": draft, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - assert len(groups) == 1 - assert set(groups[0].layer_names) == set(specs) - group_spec = groups[0].kv_cache_spec - assert isinstance(group_spec, UniformTypeKVCacheSpecs) - assert group_spec.block_size == 64 - promoted_draft = group_spec.kv_cache_specs["draft.0"] - assert isinstance(promoted_draft, FullAttentionSpec) - assert not isinstance(promoted_draft, SlidingWindowSpec) - assert promoted_draft.block_size == 64 - assert promoted_draft.sliding_window == draft.sliding_window - assert specs["draft.0"] is draft - assert "attention compute is unchanged" in caplog_vllm.text - - -def test_hidden_states_with_tp_scales_page_size(): - """When TP shrinks KV pages below the hidden-state per-token cost, - get_kv_cache_groups must scale up target block sizes so that the - common page accommodates the unsharded hidden states.""" - # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank. - # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2) - # = 16 * 2 * 64 * 2 = 4096 bytes. - kv_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=2, - head_size=64, - dtype=torch.bfloat16, - ) - # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512) - # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page. - hs_spec = HiddenStateCacheSpec( - block_size=16, - num_kv_heads=6, - head_size=512, - dtype=torch.bfloat16, - ) - specs = { - "target.0.attn": kv_spec, - "target.1.attn": kv_spec, - "cache_only_layers.48": hs_spec, - } - - groups = get_kv_cache_groups(_grouping_config(), specs) - - # The hidden-state layer should be present and no assertion should fire. - all_layers = {name for g in groups for name in g.layer_names} - assert "cache_only_layers.48" in all_layers - - # The target group block sizes must have been scaled up. - for g in groups: - if "cache_only_layers.48" not in g.layer_names: - assert g.kv_cache_spec.block_size > kv_spec.block_size - - def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses(): assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION @@ -3525,179 +1988,6 @@ def test_request_with_prompt_embeds_and_mm_inputs(hash_fn: Callable[[Any], bytes assert block_hashes[1] == expected_hash2 -def test_auto_fit_max_model_len(): - """Test that max_model_len=-1 auto-fits to available GPU memory.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=1024) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # With enough memory, max_model_len stays at the derived max - large_available_memory = mem_per_block_per_layer * 2 * 1024 # plenty of memory - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [large_available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - # Reset for next test - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - # With limited memory, max_model_len should be reduced - # Need memory for at least max_model_len tokens - # 32 blocks worth of memory for 2 layers = can fit 32*16=512 tokens - limited_memory = mem_per_block_per_layer * 2 * 32 - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [limited_memory] - ) - # Should be reduced to fit in memory - assert vllm_config.model_config.max_model_len < 1024 - assert vllm_config.model_config.max_model_len > 0 - - -def test_auto_fit_max_model_len_with_hybrid(): - """Test that auto-fit works with hybrid KV cache specs.""" - # Create config with original_max_model_len=-1 to trigger auto-fit - model_config = ModelConfig(max_model_len=8192) - # Simulate the user passing -1 by setting original_max_model_len - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # 16KB per block per layer - gamma = 2 - kv_cache_specs = { - "layer_1": new_mamba_spec(num_speculative_blocks=gamma), - "layer_2": new_kv_cache_spec(), - } - - # One extra block on top of what a 1024-token request needs: the pool - # reserves one block as the null block. - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [available_memory] - ) - assert vllm_config.model_config.max_model_len == 1024 - - -def test_auto_fit_max_model_len_not_triggered(): - """Test that auto-fit is not triggered when original_max_model_len is not -1.""" - model_config = ModelConfig(max_model_len=16) - # original_max_model_len should be None by default, not -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - - # This should work normally without auto-fit - _kv_cache_configs = get_kv_cache_configs( - vllm_config, [kv_cache_specs], [mem_per_block_per_layer * 2 * 32] - ) - assert vllm_config.model_config.max_model_len == 16 - - -def test_auto_fit_max_model_len_respects_num_gpu_blocks_override(): - """Auto-fit must size max_model_len against the override-clamped pool, not - the raw `available_memory`. Without this, auto-fit could pick a - max_model_len that no longer fits once `num_gpu_blocks_override` is applied. - """ - model_config = ModelConfig(max_model_len=16384) - model_config.original_max_model_len = -1 # request auto-fit - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # Cap the cache to 32 blocks regardless of available memory. - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), # block_size=16 - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384). - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - # 32 blocks * block_size 16 = 512 token slots, so max_model_len must - # auto-fit at or below that. - assert 0 < vllm_config.model_config.max_model_len <= 32 * 16 - - -def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override(): - """Admission check must use the override-clamped pool size, not raw - `available_memory`. Without this, startup could accept a max_model_len - that does not actually fit in `num_gpu_blocks_override` blocks. - """ - model_config = ModelConfig(max_model_len=16384) - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - # 32 blocks is far too small for max_model_len=16384 (would need 1024). - vllm_config.cache_config.num_gpu_blocks_override = 32 - - mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 - kv_cache_specs = { - "layer_1": new_kv_cache_spec(), - "layer_2": new_kv_cache_spec(), - } - # Plenty of raw memory: a bytes-only check against this would pass. - large_available_memory = mem_per_block_per_layer * 2 * 1024 - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory]) - - -def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes(): - """DFlash drafters can have a smaller head size than the target model. - - For example, MiMo uses 192-dim target KV heads while its DFlash draft uses - 128-dim KV heads. The resulting page sizes are 3:2 rather than an integer - block-size multiple, so the smaller page must be padded instead. - """ - # Attention layers read padded pages through a strided view, so padding is - # allowed. - target_spec = new_kv_cache_spec( - block_size=16, - num_kv_heads=1, - head_size=192, - dtype=torch.bfloat16, - ) - draft_spec = new_sliding_window_spec( - block_size=16, - num_kv_heads=1, - head_size=128, - dtype=torch.bfloat16, - sliding_window=1024, - ) - - unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "target_attn": target_spec, - "draft_attn": draft_spec, - } - ) - - assert unified_specs["target_attn"] == target_spec - unified_draft_spec = unified_specs["draft_attn"] - assert unified_draft_spec.block_size == draft_spec.block_size - assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes - assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes - assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the # raw KV allocation, so they must be budgeted into the offload width. The @@ -3729,173 +2019,6 @@ def test_page_size_padded_wins(): assert spec.page_size_bytes == 65536 -def test_unify_hybrid_kv_cache_specs(): - # 1. has_full_attention and has_sliding_window - before_spec_1 = new_kv_cache_spec(block_size=64) - before_spec_2 = new_sliding_window_spec( - block_size=16, page_size_padded=32 * 1024, sliding_window=1024 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_spec["draft_layer_1"] = replace(before_spec_1, dcp_sharded=False) - kv_cache_spec["draft_layer_2"] = replace(before_spec_2, dcp_sharded=False) - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec(block_size=64) - expected_spec_2 = new_kv_cache_spec( - block_size=64, page_size_padded=64 * 1024, sliding_window=1024 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_2"].page_size_bytes == 64 * 1024 - assert kv_cache_spec["draft_layer_2"] == replace(expected_spec_2, dcp_sharded=False) - - # 2. has_full_attention and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - - # 3. has_full_attention, has_sliding_window and has_chunked_local_attention - before_spec_1 = new_kv_cache_spec() - before_spec_2 = new_sliding_window_spec( - page_size_padded=32 * 1024, sliding_window=1024 - ) - before_spec_3 = new_chunked_local_attention_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - kv_cache_spec = { - "layer_1": before_spec_1, - "layer_2": before_spec_2, - "layer_3": before_spec_3, - } - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - expected_spec_1 = new_kv_cache_spec() - expected_spec_2 = new_kv_cache_spec(page_size_padded=32 * 1024, sliding_window=1024) - expected_spec_3 = new_kv_cache_spec( - page_size_padded=32 * 1024, attention_chunk_size=512 - ) - assert kv_cache_spec["layer_1"] == expected_spec_1 - assert kv_cache_spec["layer_2"] == expected_spec_2 - assert kv_cache_spec["layer_3"] == expected_spec_3 - - # 4. No FullAttentionSpec, should not convert - kv_cache_spec = { - "layer_1": new_sliding_window_spec(sliding_window=1024), - "layer_2": new_chunked_local_attention_spec(attention_chunk_size=512), - } - - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - # Replicated Mamba state still requires the hybrid cache manager. - kv_cache_spec = {"attention": new_kv_cache_spec(), "mamba": new_mamba_spec()} - with pytest.raises(ValueError): - kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) - - -def test_unify_kv_cache_spec_page_size_mamba(): - """Regression test for https://github.com/vllm-project/vllm/issues/43626. - - MambaSpec's page_size_bytes is determined by its state shapes and does not - change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba - page instead of scaling its block_size. This situation arises when a layer - with a page larger than the (already platform-aligned) Mamba page joins the - specs, e.g. a dense draft model with more KV heads than the hybrid main - model. - """ - # 1. Hybrid main model (Mamba + full attention, pages already aligned at - # 16KB) plus a dense draft model layer with a 2x larger page (32KB). - # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so - # the old code scaled the Mamba block_size, which left page_size_bytes - # unchanged at 16384. - mamba_spec = new_mamba_spec() # page_size_bytes = 16384 - main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384 - draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768 - assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384 - assert draft_attn_spec.page_size_bytes == 32768 - - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": mamba_spec, - "main_attn_layer": main_attn_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - # Mamba page is padded; block_size (caching granularity) is unchanged. - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - assert unified["mamba_layer"].block_size == mamba_spec.block_size - # Attention layer with smaller page still unifies by scaling block_size. - assert unified["main_attn_layer"].page_size_bytes == 32768 - assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size - assert unified["main_attn_layer"].page_size_padded is None - # Layer already at max page size is unchanged. - assert unified["draft_attn_layer"] == draft_attn_spec - - # 2. Mamba page already padded by the platform (state smaller than the - # padded page); the padding is re-applied at the new maximum. - padded_mamba_spec = new_mamba_spec( - shapes=((2, 256), (3, 32, 32)), page_size_padded=16384 - ) - assert padded_mamba_spec.page_size_bytes == 16384 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": padded_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - assert unified["mamba_layer"].page_size_padded == 32768 - - # 3. Mamba page that does not evenly divide the maximum page size is - # padded as well (the divisibility constraint only applies to block_size - # scaling). - odd_mamba_spec = new_mamba_spec(shapes=((6144,),)) - assert odd_mamba_spec.page_size_bytes == 24576 - assert 32768 % odd_mamba_spec.page_size_bytes != 0 - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "mamba_layer": odd_mamba_spec, - "draft_attn_layer": draft_attn_spec, - } - ) - assert unified["mamba_layer"].page_size_bytes == 32768 - - # 4. Attention layers with non-divisible page sizes are padded too: every - # backend reads a padded page through the view's block stride, so there is - # no longer a case that must raise. - unified = kv_cache_utils.unify_kv_cache_spec_page_size( - { - "attn_layer": new_kv_cache_spec(block_size=24), # 24576 - "draft_attn_layer": draft_attn_spec, # 32768 - } - ) - assert unified["attn_layer"].page_size_padded == 32768 - assert unified["attn_layer"].page_size_bytes == 32768 - assert unified["attn_layer"].block_size == 24 - - # 5. Uniform page sizes are returned unchanged. - specs = { - "mamba_layer": new_mamba_spec(), - "attn_layer": new_kv_cache_spec(), - } - assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs - - def test_hma_not_disabled_when_kv_events_enabled(): """Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. @@ -3982,78 +2105,6 @@ def test_resolve_block_hashes_rejects_mismatched_view(): resolve_block_hashes(mismatched, 2, 4) -@pytest.mark.parametrize("use_override", [True, False]) -def test_kv_cache_reserves_null_block_for_max_model_len(use_override): - """A KV cache sized to exactly the blocks a max_model_len request needs must - be rejected. BlockPool keeps one block as the null block, so only - num_blocks - 1 are usable; accepting num_blocks == needed would leave the - request unschedulable and hang the engine. Covers both the - num_gpu_blocks_override path and the memory-derived block count. - """ - 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. - if use_override: - # Ample raw memory; the override alone constrains the block count. - vllm_config.cache_config.num_gpu_blocks_override = 32 - available_memory = [spec.page_size_bytes * 1024] - else: - available_memory = [spec.page_size_bytes * 32] - - with pytest.raises(ValueError, match="max seq len"): - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - -def test_auto_fit_max_model_len_reserves_null_block(): - """Auto-fit (max_model_len=-1) must size max_model_len against usable - blocks, not the total pool. With memory for exactly 64 blocks, one is the - null block, so auto-fit must settle on 63 * block_size; picking the full - pool would leave a max-length request unschedulable and hang the engine. - """ - block_size = 16 - model_config = ModelConfig(max_model_len=1024) - model_config.original_max_model_len = -1 - vllm_config = VllmConfig(model_config=model_config) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - spec = new_kv_cache_spec(block_size=block_size) - - # Exactly the 1024 / 16 = 64 blocks a full-length request would need. - available_memory = [spec.page_size_bytes * 64] - - get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) - - assert vllm_config.model_config.max_model_len == 63 * block_size - - -def test_check_enough_kv_cache_memory_reserves_null_block(): - """The public admission check must reject a KV cache sized to exactly the - blocks a max_model_len request needs. BlockPool keeps one block as the - null block, so only num_blocks - 1 are usable; accepting - num_blocks == needed would leave the request unschedulable and hang the - engine. - """ - 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. - with pytest.raises(ValueError, match="max seq len"): - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 - ) - - # 33 blocks -> 32 usable after the null block -> accept. - check_enough_kv_cache_memory( - vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 - ) - - def test_is_full_attention_spec_unwraps_uniform_type_specs(): """``UniformTypeKVCacheSpecs`` is not a ``FullAttentionSpec``, so callers scanning groups with a bare isinstance miss DeepSeek-V4-shaped configs.""" @@ -4105,340 +2156,3 @@ def test_wrapped_mamba_group_requires_block_zeroing(): ) assert config.needs_kv_cache_zeroing - - -def _spec_decode_grouping_config(method="dspark", model_type=None): - """Grouping config with an EAGLE-family speculative method enabled.""" - return SimpleNamespace( - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), - cache_config=SimpleNamespace( - get_resolved_kv_cache_layout=lambda: SimpleNamespace( - is_block_outermost=True - ) - ), - model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), - speculative_config=SimpleNamespace( - method=method, - use_eagle=lambda: True, - use_eagle_block_drop=lambda: True, - ), - ) - - -def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False): - """A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a - DSpark-style draft MLA layer marked non_causal_multi_token_decode. - - The target's fp8 KV dtype is one difference from the draft on Kimi-K3 - (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass - draft_shares_target_spec to match that dtype while retaining the draft - marker as a group boundary. - """ - target_dtype = None if draft_shares_target_spec else "fp8_e4m3" - specs = { - "target.attn.0": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.attn.1": new_mla_spec(block_size=64, cache_dtype_str=target_dtype), - "target.mamba.0": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - "target.mamba.1": new_mamba_spec(block_size=64, mamba_cache_mode="align"), - } - if draft: - draft_spec = new_mla_spec(block_size=64) - specs["draft.attn.0"] = replace(draft_spec, non_causal_multi_token_decode=True) - return specs - - -def test_draft_group_annotated_on_hybrid_general_path(): - # A drafter's MLA layer carries non_causal_multi_token_decode, so its group - # is identifiable without keying off a model version. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "draft.attn.0" in flagged[0].layer_names - - -def test_mamba_groups_never_flagged_even_when_draft_shares_a_group(): - # Packed uniform-type groups can contain distinct target and draft layer - # specs; the combined group still holds volatile draft KV and must be - # flagged. What must never happen is a Mamba group being flagged: that - # widens its lookup window to two consecutive chunks, which align-mode - # checkpointing never produces, zeroing every lookup. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), - _hybrid_specs_with_draft(draft=True, draft_shares_target_spec=True), - ) - - for group in groups: - if "draft.attn.0" in group.layer_names: - assert group.is_eagle_group - if isinstance(group.kv_cache_spec, MambaSpec): - assert not group.is_eagle_group - - -def test_draft_group_not_annotated_without_spec_decode(): - # The marker alone must not flag anything; the eagle semantics only apply - # when a speculative method is actually enabled. - config = _spec_decode_grouping_config() - config.speculative_config = None - groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True)) - - assert not any(g.is_eagle_group for g in groups) - - -def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so consumers fall back to - # conservative behavior that silently breaks reuse for Mamba groups. - # That must at least be visible. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) - ) - - assert not any(g.is_eagle_group for g in groups) - assert "no KV cache group could be identified as the draft model's" in ( - caplog_vllm.text - ) - - -def test_unidentifiable_draft_without_mamba_does_not_warn(caplog_vllm): - # Pure-attention models degrade gracefully under the consumers' - # conservative fallback (a one-block hit drop at most), so the warning - # stays silent to avoid noise on every unannotated EAGLE deployment. - specs = { - "target.attn.0": new_mla_spec(block_size=64), - "target.attn.1": new_mla_spec(block_size=64), - } - groups = get_kv_cache_groups(_spec_decode_grouping_config(), specs) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def test_no_warning_when_draft_group_is_identified(caplog_vllm): - get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - assert "could be identified as the draft model's" not in caplog_vllm.text - - -def _deepseek_v4_specs(model_version="deepseek_v4"): - """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers - at differing page sizes, with the MTP draft layer registered last.""" - return { - "model.layers.0.self_attn.attn": new_mla_spec(), - "model.layers.1.self_attn.attn": new_mla_spec(), - "model.layers.2.self_attn.attn": new_swa_mla_spec( - head_size=1024, model_version=model_version - ), - # The MTP block registers last, and its sliding-window size differs, so - # it lands in a group of its own. - "model.layers.3.self_attn.attn": new_swa_mla_spec( - head_size=1024, sliding_window=256, model_version=model_version - ), - } - - -@pytest.mark.parametrize( - ("method", "model_type"), - [("mtp", "deepseek_v4"), ("dspark", "deepseek_v4"), ("dspark", "deepseek_v41")], -) -def test_deepseek_v4_draft_group_annotated_on_packed_path(method, model_type): - # DeepseekV4's MTP block reuses the target's decoder layer, so its spec - # carries no draft marker and only the positional rule can find it. This - # pins the pre-existing behaviour that the unified annotator must preserve. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type=model_type), - _deepseek_v4_specs(model_version=None), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def test_trailing_layer_fallback_applies_to_any_mtp_model(): - # The positional rule is sound for every MTP drafter, not just DeepseekV4: - # MTP blocks reuse the target's decoder layer (no spec marker) and always - # register after every target layer. The model_type must not gate it. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="other"), - _deepseek_v4_specs(), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "model.layers.3.self_attn.attn" in flagged[0].layer_names - - -def _qwen3_5_hybrid_specs(with_mtp_layer: bool): - """Qwen3.5-shaped hybrid: repeating [GDN x3, full-attn x1] blocks, with - the MTP drafter's full-attn layer (spec-identical to the target's) - registered last.""" - specs = {} - idx = 0 - for _ in range(2): - for _ in range(3): - specs[f"model.layers.{idx}.linear_attn"] = new_mamba_spec( - mamba_cache_mode="align" - ) - idx += 1 - specs[f"model.layers.{idx}.self_attn.attn"] = new_kv_cache_spec() - idx += 1 - if with_mtp_layer: - specs["mtp.layers.0.self_attn.attn"] = new_kv_cache_spec() - return specs - - -def test_qwen3_5_mtp_draft_group_annotated_on_hybrid_path(caplog_vllm): - # A hybrid mamba + full-attention model with an MTP drafter that is - # spec-indistinguishable from the target reaches the general multi-group - # path. The trailing-layer rule must locate the draft group there so the - # Mamba groups are not swept up by the flag-all consumer fallback. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method="mtp", model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - flagged = [g for g in groups if g.is_eagle_group] - assert len(flagged) == 1 - assert "mtp.layers.0.self_attn.attn" in flagged[0].layer_names - for group in groups: - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - assert not group.is_eagle_group - assert "could be identified as the draft model's" not in caplog_vllm.text - - -@pytest.mark.parametrize("method", ["eagle", "eagle3", "dspark"]) -def test_non_mtp_eagle_hybrid_still_warns(caplog_vllm, method): - # Other drafters are not covered by the trailing-layer rule, so an - # unidentifiable hybrid draft must still warn. - groups = get_kv_cache_groups( - _spec_decode_grouping_config(method=method, model_type="qwen3_5"), - _qwen3_5_hybrid_specs(with_mtp_layer=True), - ) - - assert not any(g.is_eagle_group for g in groups) - assert "could be identified as the draft model's" in caplog_vllm.text - - -def test_trailing_layer_fallback_requires_exact_partition(): - # If the groups do not partition the layers exactly (e.g. a caller that - # dropped or duplicated layers), the positional rule is meaningless and - # must not fire. - from vllm.v1.core.kv_cache_utils import _annotate_eagle_groups - - specs = _qwen3_5_hybrid_specs(with_mtp_layer=True) - config = _spec_decode_grouping_config(method="mtp", model_type="qwen3_5") - groups = get_kv_cache_groups(config, specs) - for g in groups: - g.is_eagle_group = False - # Remove one layer from its group: no longer an exact partition. - trimmed = [ - KVCacheGroupSpec( - [n for n in g.layer_names if n != "model.layers.0.linear_attn"], - g.kv_cache_spec, - ) - for g in groups - ] - _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) - - assert not any(g.is_eagle_group for g in trimmed) - - -_GQA_SPEC = FullAttentionSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 -) -_GQA_SWA_SPEC = SlidingWindowSpec( - block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 -) - - -@pytest.mark.parametrize( - "specs,expected", - [ - pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": new_swa_mla_spec(), - "s.1": new_swa_mla_spec(), - }, - 4, - id="mla-and-swa-mla", - ), - pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), - pytest.param( - {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, - 1, - id="mla-and-gqa-uniform-group", - ), - pytest.param( - { - "l.0": new_mla_spec(), - "l.1": new_mla_spec(), - "s.0": _GQA_SWA_SPEC, - "s.1": _GQA_SWA_SPEC, - }, - 1, - id="mla-and-gqa-swa", - ), - ], -) -def test_kv_tp_replicas(monkeypatch, specs, expected): - """Resolved per layer before scheduler flattening hides mixed groups.""" - from vllm.config import ParallelConfig - from vllm.platforms import current_platform - - monkeypatch.setattr(current_platform, "device_count", lambda: 4) - parallel_config = ParallelConfig(tensor_parallel_size=4) - vllm_config = VllmConfig( - model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config - ) - vllm_config.cache_config.kv_cache_layout = "LBNHC" - mem = sum(s.page_size_bytes for s in specs.values()) * 10 - configs = get_kv_cache_configs(vllm_config, [specs], [mem]) - assert configs[0].kv_tp_replicas == expected - scheduler = generate_scheduler_kv_cache_config(configs) - assert scheduler.kv_tp_replicas == expected - - -@pytest.mark.parametrize( - "specs,tp_size,dcp_size,expected", - [ - pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), - pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), - pytest.param( - [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" - ), - pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), - pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), - pytest.param( - [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], - 8, - 1, - 2, - id="mla-and-gqa-partial", - ), - pytest.param( - [ - HiddenStateCacheSpec( - block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 - ) - ], - 8, - 1, - 1, - id="hidden-state", - ), - ], -) -def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): - """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" - groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] - assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index cf221fc22a34..541e55f5b60e 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -36,6 +36,10 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, + KVQuantMode, + MambaSpec, + MLAAttentionSpec, + SlidingWindowSpec, ) from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -333,3 +337,73 @@ def create_requests( ) requests.append(request) return requests + + +def new_kv_cache_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=None, + attention_chunk_size=None, + kv_quant_mode=KVQuantMode.NONE, +): + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + attention_chunk_size=attention_chunk_size, + kv_quant_mode=kv_quant_mode, + ) + + +def new_sliding_window_spec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.float32, + page_size_padded=None, + sliding_window=1, +): + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + page_size_padded=page_size_padded, + sliding_window=sliding_window, + ) + + +def new_mamba_spec( + block_size=16, + shapes=((2, 512), (3, 32, 32)), + dtypes=(torch.float32, torch.float32), + num_speculative_blocks=2, + mamba_cache_mode="none", + page_size_padded=None, +): + return MambaSpec( + block_size=block_size, + shapes=shapes, + dtypes=dtypes, + page_size_padded=page_size_padded, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + + +def new_mla_spec(cache_dtype_str=None, block_size: int = 16): + # head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576 + return MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=576, + dtype=torch.float32, + cache_dtype_str=cache_dtype_str, + ) diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index 5626dec17043..f9455389ef08 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory +from vllm.v1.core.kv_cache_planning import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -38,7 +38,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): config.parallel_config.decode_context_parallel_size = 1 monkeypatch.setattr( - "vllm.v1.core.kv_cache_utils.max_memory_usage_bytes", + "vllm.v1.core.kv_cache_planning.max_memory_usage_bytes", lambda c, s: 100 * 1024**3, # 100 GiB ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 07d530f97d1a..a6692b5cd5fa 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,10 +17,8 @@ SchedulerOffloadConfig, ) from vllm.platforms import current_platform -from vllm.v1.core.kv_cache_utils import ( - generate_scheduler_kv_cache_config, - kv_cache_groups_tp_replicas, -) +from vllm.v1.core.kv_cache_planning import kv_cache_groups_tp_replicas +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py index 735b042259b5..854632dce1d1 100644 --- a/tests/v1/kv_connector/unit/test_hidden_states_connector.py +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -10,7 +10,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 ExampleHiddenStatesConnector, ) -from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.core.kv_cache_planning import get_kv_cache_groups from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 16e142d7cb97..4f9a7c451ffc 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -29,11 +29,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import ( SimpleCPUOffloadConnector, ) -from vllm.v1.core.kv_cache_utils import ( +from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, is_kv_cache_spec_uniform, - resolve_kv_cache_block_sizes, ) +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 4cc8f9ffe2fd..abd0830a472b 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -104,7 +104,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span(): zeroed no matter which group owns it.""" from unittest.mock import MagicMock - from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups from vllm.v1.kv_cache_interface import ( KVCacheGroupSpec, KVCacheLayout, diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 1247bf193a76..534d81312436 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -43,7 +43,10 @@ ROCMAiterMLASparseBackend, ) from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs +from vllm.v1.core.kv_cache_planning import ( + estimate_max_model_len, + get_kv_cache_configs, +) from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py new file mode 100644 index 000000000000..6e258a68a9e7 --- /dev/null +++ b/vllm/v1/core/kv_cache_planning.py @@ -0,0 +1,1637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV cache planning: grouping, sizing, and memory-fit logic.""" + +import math +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from dataclasses import replace +from functools import partial, reduce +from typing import Any + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv, round_up +from vllm.utils.mem_utils import format_gib +from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.core.kv_cache_utils import _kv_cache_device_name +from vllm.v1.hisparse.layout import ( + get_hisparse_host_pool_bytes, + get_hisparse_kv_cache_config, + get_hisparse_kv_cache_groups, +) +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + ChunkedLocalAttentionSpec, + CircularBufferSpec, + FullAttentionSpec, + HiddenStateCacheSpec, + HiSparseHotSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheLayout, + KVCacheSpec, + KVCacheTensor, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, + compute_layout_strides, + iter_layer_specs, + replace_as, +) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry + +logger = init_logger(__name__) + + +def _check_enough_kv_cache_memory( + available_memory: int, + get_needed_memory: Callable[[], int], + max_model_len: int, + estimate_max_model_len: Callable[[int], int], +): + if available_memory <= 0: + raise ValueError( + "No available memory for the cache blocks. " + "Try increasing `gpu_memory_utilization` when initializing the engine " + "(this flag also controls CPU memory reservation on the CPU " + "backend, despite its name). " + "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + "for more details." + ) + + needed_memory = get_needed_memory() + + if needed_memory > available_memory: + estimated_max_len = estimate_max_model_len(available_memory) + estimated_msg = "" + if estimated_max_len > 0: + estimated_msg = ( + "Based on the available memory, " + f"the estimated maximum model length is {estimated_max_len}. " + ) + + raise ValueError( + f"To serve at least one request with the model's max seq len " + f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " + f"cache is needed, which is larger than the available KV cache " + f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" + f"Try increasing `gpu_memory_utilization` (which also controls " + f"CPU memory on the CPU backend) or decreasing `max_model_len` " + f"when initializing the engine. " + f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " + f"for more details." + ) + + +def max_memory_usage_bytes( + vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] +) -> int: + """Get the maximum memory usage in bytes for the given KV cache specs.""" + return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) + + +def estimate_max_model_len( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +) -> int: + """Estimates the maximum model length that can fit in the available memory + using binary search. + + This function temporarily modifies max_model_len during estimation but + restores the original value before returning, ensuring no side effects. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Returns: + The estimated maximum model length that can fit in the available memory. + + """ + # Save the original max_model_len to restore after estimation + original_max_model_len = vllm_config.model_config.max_model_len + + # Define a function to check if a given model length fits in memory + def fits_in_memory(model_len: int) -> bool: + # Temporarily modify the max_model_len for this calculation + vllm_config.model_config.max_model_len = model_len + # Calculate memory needed for the given model length + memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) + return memory_needed <= available_memory + + try: + # Binary search for the maximum model length + left, right = 1, original_max_model_len + + # If even the smallest model length doesn't fit, return 0 + if not fits_in_memory(left): + return 0 + + # Binary search for the maximum model length that fits + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits_in_memory(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + # Always restore the original max_model_len to avoid side effects + vllm_config.model_config.max_model_len = original_max_model_len + + +def check_enough_kv_cache_memory( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + available_memory: int, +): + """Checks whether `available_memory` is enough for the KV cache to hold at + least one request with the model's max_model_len. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + available_memory: Memory available for KV cache in bytes. + + Raises: + ValueError: If there is not enough memory available for the KV cache. + + """ + # No need to check for available memory if the kv_cache_spec is empty + if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # 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(vllm_config, groups) + if groups + else available_memory + ) + _check_enough_kv_cache_memory( + check_memory, + lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), + vllm_config.model_config.max_model_len, + lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), + ) + + +def create_kv_cache_group_specs( + kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] +) -> list[KVCacheGroupSpec]: + """Create KVCacheGroupSpec object for each kv cache group layer. + The layers in the same group should share the same + KVCacheSpec. + + Args: + kv_cache_spec: + A mapping from each layer name to its corresponding KVCacheSpec. + grouped_layer_names: + A list of kv cache groups, where each element is a list of layer + names that belong to the same group and should share the same + KVCacheSpec. + + Returns: + A list of KVCacheGroupSpec objects, one for each group. + + """ + kv_cache_groups = [] + for layer_names_one_group in grouped_layer_names: + layer_specs = [ + kv_cache_spec[layer_name] for layer_name in layer_names_one_group + ] + merged_layer_spec = layer_specs[0].merge(layer_specs) + kv_cache_groups.append( + KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) + ) + return kv_cache_groups + + +def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + """Whether all layers in the given KVCacheSpec have the same KV cache spec. + Note that we regard FullAttentionSpec with and without sliding window as + the same type. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + True if all layers have the same type, False otherwise. + + """ + if not kv_cache_spec: + # Encoder-only models do not have KV cache, kv_cache_type can be + # regarded as uniform. + return True + try: + kv_cache_spec_values = list(kv_cache_spec.values()) + _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) + except AssertionError: + return False + return True + + +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: + """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. + The override is logged once, at the call site in `get_kv_cache_configs`. + """ + if vllm_config.cache_config.num_gpu_blocks_override is not None: + num_blocks = vllm_config.cache_config.num_gpu_blocks_override + return num_blocks + + +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], + vllm_config.cache_config.get_resolved_kv_cache_layout(), + ) + + +def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: + """Get the page size of the KV cache.""" + page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} + assert len(page_sizes) == 1 + return page_sizes.pop() + + +def _get_kv_cache_groups_uniform_spec( + kv_cache_specs: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with the same KV cache + spec for all layers. + + Args: + kv_cache_specs: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) + + +def _get_kv_cache_groups_uniform_type( + spec: UniformTypeKVCacheSpecs, +) -> list[KVCacheGroupSpec]: + """Generates the KV cache configuration for a model with one type of KV cache + but different hidden sizes. All layers are merged into one group. + + Args: + spec: The UniformTypeKVCacheSpecs of the model + + Returns: + The generated KVCacheGroupSpecs + + """ + return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] + + +def unify_kv_cache_spec_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Unify the page size of the given KVCacheSpec. If the page size of all layers + are the same, return the original KVCacheSpec. If not same, unify the page + size by increasing the block size of layers with smaller page size. Two + cases cannot be unified by block size alone and pad their physical page to + the maximum instead: Mamba layers, whose page size comes from state shapes + and is independent of block size; and non-MLA attention layers whose page + does not evenly divide the maximum (the padded page is read through a + strided view). MLA is excluded because sparse MLA indexes the cache in + whole token rows (see ``flat_kv_row_view``), so its block stride can only + be padded by its own row-aligned ``alignment``, not to an arbitrary page + size. Raise NotImplementedError if failed to unify the page size; + ``get_kv_cache_groups`` catches it to try the full-allocation fallback + (e.g. MLA next to an incompatible sliding-window draft). + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + + Returns: + The updated KVCacheSpec with the same page_size_bytes. + + """ + page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} + if len(page_sizes) <= 1: + # All layers have the same page size, no need to unify. + return kv_cache_spec + + max_page_size = max(page_sizes) + new_kv_cache_spec = {} + for layer_name, layer_spec in kv_cache_spec.items(): + if layer_spec.page_size_bytes == max_page_size: + new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, MambaSpec): + # MambaSpec's page size is determined by its state shapes and does + # not scale with block_size, so pad the page instead. This is the + # same padding mechanism the platform uses to align Mamba pages + # with the main model's attention page size; it is needed here + # when another layer (e.g. from a draft model) has a larger page + # than the already-aligned Mamba page. + new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + else: + layer_page_size = layer_spec.page_size_bytes + if max_page_size % layer_page_size == 0: + ratio = max_page_size // layer_page_size + new_block_size = layer_spec.block_size * ratio + new_spec = replace(layer_spec, block_size=new_block_size) + elif isinstance(layer_spec, AttentionSpec) and not isinstance( + layer_spec, MLAAttentionSpec + ): + new_spec = replace(layer_spec, page_size_padded=max_page_size) + else: + raise NotImplementedError( + f"Layer {layer_name}: page size is not divisible by the " + "maximum page size and cannot be padded. Padding is only " + "supported for non-MLA attention layers." + ) + assert new_spec.page_size_bytes == max_page_size + new_kv_cache_spec[layer_name] = new_spec + return new_kv_cache_spec + + +def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: + # kv_cache_spec is an empty dict for attention free models + return not kv_cache_spec + + +def _get_kv_cache_groups_uniform_page_size( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Generates the KV cache groups for hybrid models with multiple + attention types but still with a uniform page size (physical memory per + block per layer) for all layers. + + Detailed explanation about kv cache management of hybrid models: + The layers in the models are repeated with some patterns, e.g., a model + with 10 full attention layers and 20 sliding window attention layers can be + regarded as repeating the pattern (1 * full, 2 * sw) 10 times. + The KVCacheManager allocates different block tables for each of the 3 layers + in the pattern, and repeats each of them 10 times to generate the + block_table for the 30 layers in the model. + Therefore, we can group the layers in the model into 3 kv_cache_groups, each + of which contains 10 layers in the model. + The KVCacheManager allocates the block_table for each group based on its + kv_cache spec, and the model runner applies the block table to each layer + in the group. + For example: + 1. A model only uses full attention. The pattern is + (num_hidden_layers * full), so there is only one group and the block table + is shared by all layers. It is already handled by + `_get_kv_cache_config_uniform_type`. + 2. A model with 10 full attention layers and 20 sliding window + attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so + there are 3 kv_cache_groups, each of which represents 10 layers. + + To simplify the implementation, we make the following assumptions: + 1. Physical memory per block: Must be the same across all KV cache groups. + Breaking this assumption is non-trivial due to memory fragmentation concerns + when allocating blocks of different sizes. + 2. Tokens per block (block_size): Currently, we directly use + `CacheConfig.block_size` for all layers. It can be extended to vary by KV + cache group, but within each KV cache group, all layers must share the same + block size. + 3. Physical memory per token per layer: This property is decided by model + config. Currently we only support models that have the same physical memory + per token per layer for all layers. Can be relaxed with a simple extension, + but still need to keep physical memory per block the same for all groups. + 4. Number of layers per group: Currently assumed the same for all layers. + Can be relaxed with a simple extension, but still need to keep physical + memory per block the same for all groups. + 5. Attention type within groups: All layers in a group must share the same + attention type. One exception is that, when + `--disable-hybrid-kv-cache-manager` is true, the single group for full + attention layers may also include attention layers using sliding window or + LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. + 6. Support for multiple attention types: The design for most components is + general to an arbitrary number of attention types. But + `find_longest_cache_hit` only supports one attention type or two + types of full-attention plus exactly one another type. The general + implementation of this function is feasible but we don't know how to + implement it cleanly yet. + + As we assume tokens per block, physical memory per token per layer, and + number of layers per group are the same now, we can ensure that physical + memory per block is the same for all groups. + + Args: + kv_cache_spec: The KVCacheSpec of each attention layer in the model + Returns: + The generated KVCacheGroupSpecs + + """ + # Group all layers by kv_cache_spec. + # E.g., 2 full attention layers and 3 sliding window attention layers, + # -> (full.0, full.1), (sw.0, sw.1, sw.2). + same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) + for layer_name, layer_spec in kv_cache_spec.items(): + same_type_layers[layer_spec].append(layer_name) + + # Attempt to further merge same-type layers based on whether their KV + # cache specs can be merged, to minimize the group count. This benefits + # situations where specs share a block layout and differ only in a + # property it can reconcile (e.g. full attention layers differing only in + # sliding window / attention chunk size). + layer_buckets: list[list[str]] = [] + spec_buckets: list[list[KVCacheSpec]] = [] + for layer_spec, layer_names in same_type_layers.items(): + for names, specs in zip(layer_buckets, spec_buckets): + try: + # A raise means that the specs are incompatible. + type(specs[0]).merge([*specs, layer_spec]) + except (AssertionError, ValueError): + continue + names.extend(layer_names) + specs.append(layer_spec) + break + else: + layer_buckets.append(list(layer_names)) + spec_buckets.append([layer_spec]) + + # Split each group into smaller groups, to make the number of layers in each + # group identical. Add padding to the last group of each type if necessary. + # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) + # split to 3 groups with 2 layers each: + # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). + # FIXME(Chen): At the moment of writing this code (2025-06-02), all + # open-source hybrid model follows a n:1 pattern between different attention + # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and + # full), so we can use the "1" in the n:1 pattern as the group size, which + # is the minimum number of layers among all attention types. Need a better + # strategy if we want to support more complex patterns (e.g., 20 full + 30 + # sw, where the group size should be 10). + min_num_layers = min([len(layers) for layers in layer_buckets]) + group_size = min_num_layers + max_num_layers = max([len(layers) for layers in layer_buckets]) + if max_num_layers < min_num_layers * 1.5: + # If the number of layers is not much larger than the minimum number of + # layers, use the maximum number of layers as the group size to avoid + # too many padding layers. A typical example is gpt-oss-20b + eagle, + # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of + # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding + # layers while accommodating speculative decoding drafters that add + # extra layers to one attention type. + group_size = max_num_layers + grouped_layers = [] + for layers in layer_buckets: + num_padding_layers = group_size - len(layers) % group_size + if num_padding_layers != group_size: + logger.warning( + "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa + num_padding_layers, + num_padding_layers / len(layers) * 100, + ) + num_groups = cdiv(len(layers), group_size) + # In PP case, say if we have + # - stage 0: full.0, sw.0, sw.1 + # - stage 1: full.1, sw.2, sw.3 + # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) + # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because + # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) + # and it will be padded to (full.0, padding), (sw.0, sw.1), + # (padding, padding) to ensure the number of layers in each group is + # the same and will cause memory waste. + # To avoid this, we assign layers[i::num_groups] to the i-th group + # instead of layers[i * group_size: (i + 1) * group_size] + for i in range(num_groups): + grouped_layers.append(layers[i::num_groups]) + return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) + + +def _get_per_layer_spec( + group: KVCacheGroupSpec, + layer_name: str, +) -> KVCacheSpec: + spec = group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + return spec.kv_cache_specs[layer_name] + 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], layout: KVCacheLayout +) -> int: + """Return the largest cache group's bytes per block.""" + 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 + for layer_name in group.layer_names + ) + for group in kv_cache_groups + ) + assert bytes_per_block > 0 + alignments = [ + group.kv_cache_spec.page_size_bytes + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiSparseHotSpec) + ] + alignments.extend( + _get_per_layer_spec(group, layer_name).get_block_stride_alignment() + for group in kv_cache_groups + for layer_name in group.layer_names + ) + return round_up(bytes_per_block, math.lcm(*alignments)) + + +def validate_kv_cache_layout( + layout: KVCacheLayout, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Validate that the resolved layout can express this model's packing. + + The layout was chosen once in the engine core from the backends' supported + sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 + indexer) declares block-outermost layouts there, so an inexpressible + layout reaching this point is an error. + """ + page_sizes = { + _get_per_layer_spec(group, layer_name).page_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + } + if len(page_sizes) == 1: + # A rectangular layer dim exists; every layout can express it. + return + + # Mixed page sizes pack pages side by side within a block, which needs each page + # to be one contiguous chunk inside its block (a block-compact layout) and, with + # multiple KV cache groups, the layer dim inside the block dim. + if not layout.is_block_compact or ( + len(kv_cache_groups) > 1 and layout.is_layer_compact + ): + raise ValueError( + f"KV cache layout {layout.name} cannot express this model's " + f"mixed page sizes ({sorted(page_sizes)}); a backend should " + "declare block-outermost supported layouts (e.g. BLHNC), or " + "set VLLM_KV_CACHE_LAYOUT=BLHNC." + ) + + +def get_kv_cache_config_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> KVCacheConfig: + """Generate the KV cache configuration from the KV cache groups and spec + of each layer. + + Args: + vllm_config: The global VllmConfig + kv_cache_groups: The KV cache groups + available_memory: Memory available for KV cache in bytes + Returns: + The generated KVCacheConfig + + """ + if len(kv_cache_groups) == 0: + # Attention free models do not have KV cache. + # Return num_blocks=1 as BlockPool always needs a null_block. + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + if vllm_config.attention_config.hisparse_config is not None: + host_budget = get_hisparse_host_pool_bytes(vllm_config) + return get_hisparse_kv_cache_config( + vllm_config, kv_cache_groups, available_memory, host_budget + ) + + 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, layout) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + + num_blocks = available_memory // bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + size = bytes_per_block * num_blocks + + # Groups alias from byte 0. Spec regions are laid out differently: + # + # block-outer (the same packing repeats for every block): + # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... + # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... + # |<--- bytes_per_block -->| + # + # layer-outer (only supported for uniform page sizes or single-group models): + # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | + # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | + + kv_cache_tensors = [] + for group in kv_cache_groups: + group_spec = group.kv_cache_spec + layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) + if isinstance(group_spec, UniformTypeKVCacheSpecs): + for layer_name, spec in group_spec.kv_cache_specs.items(): + layers_by_spec[spec].append(layer_name) + elif group.layer_names: + layers_by_spec[group_spec].extend(group.layer_names) + + 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, block_stride, None, None, None), + ) + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=byte_offset, + ) + ) + # 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, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + + +def _promote_local_kv_cache_specs( + kv_cache_spec: dict[str, KVCacheSpec], +) -> dict[str, KVCacheSpec]: + """Use full-attention allocation for local-attention cache specs. + + The returned specs affect KV cache management only. Attention modules keep + their original sliding-window or chunked-local compute behavior. + """ + promoted_specs = kv_cache_spec.copy() + + if is_kv_cache_spec_uniform( + promoted_specs + ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): + return promoted_specs + + has_full_attention = any( + isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() + ) + has_sliding_window = any( + isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() + ) + has_chunked_local_attention = any( + isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() + ) + full_block_sizes = { + spec.block_size + for spec in promoted_specs.values() + if isinstance(spec, FullAttentionSpec) + } + full_attention_block_size = ( + next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None + ) + + def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: + if spec.page_size_padded is None: + return None + unpadded_page_size = ( + spec.unpadded_page_size_bytes * block_size // spec.block_size + ) + return max(spec.page_size_padded, unpadded_page_size) + + promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { + SlidingWindowMLASpec: MLAAttentionSpec, + SlidingWindowSpec: FullAttentionSpec, + ChunkedLocalAttentionSpec: FullAttentionSpec, + } + + if has_full_attention and (has_sliding_window or has_chunked_local_attention): + for layer_name, spec in kv_cache_spec.items(): + target_cls = next( + (promotions[c] for c in type(spec).__mro__ if c in promotions), None + ) + if target_cls is None: + continue + assert isinstance(spec, AttentionSpec) + block_size = full_attention_block_size or spec.block_size + promoted_specs[layer_name] = replace_as( + spec, + target_cls, + # Promoted specs allocate blocks for all tokens and never free + # below the window, so the trailing-edge extension is moot. + drop=("extra_retained_tokens",), + block_size=block_size, + page_size_padded=promoted_page_size_padded(spec, block_size), + ) + + if not ( + is_kv_cache_spec_uniform(promoted_specs) + or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) + ): + raise ValueError("Failed to promote local KV cache specs to one unified type.") + + return promoted_specs + + +def _try_get_full_allocation_fallback_groups( + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Try a supported full-allocation fallback for local-attention layers.""" + if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): + return None + if any( + isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) + for spec in kv_cache_spec.values() + ): + return None + + has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) + has_regular_swa = any( + isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() + ) + if not (has_mla and has_regular_swa): + return None + + try: + promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) + except ValueError: + return None + uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) + if uniform_spec is None: + return None + logger.warning( + "KV cache page sizes cannot be unified; treating sliding-window " + "layers as full attention for cache allocation. Sliding-window " + "attention compute is unchanged." + ) + return _get_kv_cache_groups_uniform_type(uniform_spec) + + +def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): + """This function tries to convert the KV cache specs to one type if the model + is a hybrid model with multiple type of KV cache. It will convert all + SlidingWindowSpec to FullAttentionSpec if both types are present. + + Args: + kv_cache_spec: The kv cache spec of each attention layer in the model + + """ + groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) + for name, spec in kv_cache_spec.items(): + replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded + groups[replicated][name] = spec + for specs in groups.values(): + promoted_specs = _promote_local_kv_cache_specs(specs) + if promoted_specs == specs: + continue + logger.warning( + "Hybrid KV cache manager is disabled for this hybrid model, " + "This means we do not enable any optimizations for saving KV cache " + "memory (e.g., dropping the KV cache outside the sliding window). " + "The compute of layers like sliding window is still saved." + ) + kv_cache_spec.update(promoted_specs) + + +def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: + """Pick a chunk size that minimizes total upward padding. + + Each x is rounded up to a multiple of d: + + x -> ceil(x / d) * d + + Total padding is: + + pad(d) = sum_i (ceil(x_i / d) * d - x_i) + + We brute-force d in [lower_bound, max(values)] (fine for small lists / small + maxima) and return the d with minimum padding. Ties prefer larger d. + """ + if not values: + raise ValueError("values must be non-empty") + if any(x <= 0 for x in values): + raise ValueError(f"values must be positive, got: {list(values)!r}") + + min_d = max(1, lower_bound if lower_bound is not None else 1) + max_d = max(values) + if min_d > max_d: + return min_d + + best_d = min_d + best_pad: int | None = None + for d in range(min_d, max_d + 1): + pad = sum((d - (x % d)) % d for x in values) + if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): + best_pad = pad + best_d = d + + return best_d + + +def _get_packed_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Group mixed-page-size layers for contiguous block-outermost packing. + + Greedily buckets layers into uniform-type specs. Buckets with equal layer + counts per page size are treated as a repeating layer pattern (one layer + per page size) and split into groups covering the same number of pattern + repeats (picked by ``_approximate_gcd`` to minimize padding), so all + groups pack into the same per-block layout. Mamba buckets are additionally + split to fit the block the attention buckets already need. + Returns None when the layout is not block-outermost or all layers already + share one page size. + """ + layout = vllm_config.cache_config.get_resolved_kv_cache_layout() + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} + if not layout.is_block_outermost or len(page_sizes) <= 1: + return None + + buckets: list[dict[str, KVCacheSpec]] = [] + for name, spec in kv_cache_spec.items(): + for bucket in buckets: + candidate = {**bucket, name: spec} + if UniformTypeKVCacheSpecs.is_uniform_type(candidate): + bucket[name] = spec + break + else: + buckets.append({name: spec}) + + bucketed = [] + for bucket in buckets: + uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) + assert uniform_spec is not None + page_size_layers: dict[int, list[str]] = defaultdict(list) + for layer_name, layer_spec in bucket.items(): + page_size_layers[layer_spec.page_size_bytes].append(layer_name) + # Only 1:1 patterns (one layer of each page size per repeat) are + # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle + # repeat too, but such buckets are emitted whole instead. + balanced = len(set(map(len, page_size_layers.values()))) == 1 + bucketed.append((uniform_spec, page_size_layers, balanced)) + + # Balanced buckets that mix page sizes must stay whole, so the largest one + # sets a floor on the repeats per group; larger single-size buckets are + # split down toward it. No such bucket means nothing needs packing. + min_repeats_per_group = max( + ( + spec.get_max_layers_per_page_size() + for spec, page_size_layers, balanced in bucketed + if balanced and len(page_size_layers) > 1 + ), + default=0, + ) + repeats_per_group = ( + _approximate_gcd( + [ + spec.get_max_layers_per_page_size() + for spec, _, balanced in bucketed + if balanced + ], + lower_bound=min_repeats_per_group, + ) + if min_repeats_per_group + else None + ) + + def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: + if balanced and repeats_per_group is not None: + return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) + return 1 + + def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: + """Page bytes of the largest of the n groups a bucket splits into.""" + return sum( + cdiv(len(names), n) * page for page, names in page_size_layers.items() + ) + + # Bytes a block must hold however the state buckets end up split: mamba, + # circular-buffer and unsplit sliding-window buckets can go down to one + # state per group, every other bucket's split is fixed by the repeat pattern. + def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: + if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): + return True + return repeats_per_group is None and isinstance( + spec.first_spec, SlidingWindowSpec + ) + + anchor_bytes = max( + ( + widest_group_bytes( + page_size_layers, + len(spec.kv_cache_specs) + if is_state_bucket(spec) + else num_groups_for(spec, balanced), + ) + for spec, page_size_layers, balanced in bucketed + ), + default=0, + ) + + groups = [] + for spec, page_size_layers, balanced in bucketed: + num_groups = num_groups_for(spec, balanced) + # Cap a state group at the states a block already fits rather than let + # it widen the block. + if anchor_bytes and is_state_bucket(spec): + states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) + num_groups = max( + num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) + ) + if num_groups == 1: + groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) + continue + + pattern_repeats = list(zip(*page_size_layers.values())) + for i in range(num_groups): + group_layer_names = [ + name for repeat in pattern_repeats[i::num_groups] for name in repeat + ] + group_layer_specs = { + name: spec.kv_cache_specs[name] for name in group_layer_names + } + group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) + assert group_spec is not None + groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: + """Whether the drafter's KV layers can be located positionally. + + MTP drafters register their KV layers after the target layers, without a + distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. + The annotator separately checks that the groups partition the layers exactly. + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return False + if spec_config.method == "mtp": + return True + model_config = vllm_config.model_config + return model_config is not None and model_config.hf_config.model_type in ( + "deepseek_v4", + "deepseek_v41", + ) + + +def _groups_partition_layers_exactly( + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], +) -> bool: + """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. + + The trailing-layer draft fallback is only meaningful when the last + registered layer is guaranteed to appear in exactly one group. + """ + names = [name for group in kv_cache_groups for name in group.layer_names] + return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) + + +def _annotate_eagle_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + kv_cache_groups: list[KVCacheGroupSpec], + use_trailing_layer_fallback: bool = False, +) -> None: + """Flag the KV cache groups that hold drafter attention layers. + + Two detection rules, in order of preference: + + 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on + MLAAttentionSpec and set by drafter attention layers that run a + non-causal multi-token decode (today only Kimi-K3 DSpark). It survives + MLAAttentionSpec.merge, so it still identifies a group after per-group + spec merging, wherever grouping happens to land. It is sufficient but + not necessary: a drafter whose spec is indistinguishable from the + target's cannot be found this way. + 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose + MTP block reuses the target's own decoder layer and so carries no spec + marker. Their draft attention layers always register after every + target layer, so flag whichever group holds the last registered layer. + This rule is only valid where the groups partition exactly the layers + of ``kv_cache_spec``, which is re-checked here before applying it. + When a drafter's trailing caches span several groups, this rule flags + only the group holding the very last layer and must be generalized. + FIXME(yifan): avoid/generalize this hacky check. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_spec: The kv cache spec of each attention layer, in layer + registration order. Only read by rule 2. + kv_cache_groups: Groups to annotate in place. + use_trailing_layer_fallback: Enable rule 2. Callers gate this on + ``_uses_trailing_mtp_layers``. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle_block_drop(): + return + + for group in kv_cache_groups: + if any( + getattr(spec, "non_causal_multi_token_decode", False) + for spec in iter_layer_specs(group.kv_cache_spec) + ): + group.is_eagle_group = True + + if not use_trailing_layer_fallback: + return + if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): + return + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: + group.is_eagle_group = True + break + + +def _warn_if_unannotated_eagle_mamba( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Warn when no KV cache group could be identified as the draft model's. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_groups: Groups as they will be handed to consumers. + + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return + if any(group.is_eagle_group for group in kv_cache_groups): + return + mamba_groups = [ + idx + for idx, group in enumerate(kv_cache_groups) + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if not mamba_groups: + return + logger.warning( + "Speculative decoding (method=%s) is enabled but no KV cache group " + "could be identified as the draft model's.", + spec_config.method, + ) + + +def _largest_divisor_at_most(value: int, limit: int) -> int: + for candidate in range(min(value, limit), 0, -1): + if value % candidate == 0: + return candidate + return 1 + + +def _ensure_min_page_size( + groups: list[KVCacheGroupSpec], + common_page: int, + hidden_specs: dict[str, HiddenStateCacheSpec], +) -> tuple[list[KVCacheGroupSpec], int]: + """Scale up group block sizes so the common page is at least as large as + the biggest hidden-state per-token cost. + + This protects the hidden state extraction feature, where we store hidden states in + the KV cache, and return them to clients using a custom hidden state connector. + + Returns the (possibly scaled) groups and updated common page size. + """ + max_per_token = max( + spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + for spec in hidden_specs.values() + ) + if max_per_token <= common_page: + return groups, common_page + + scale = cdiv(max_per_token, common_page) + common_page *= scale + scaled: list[KVCacheGroupSpec] = [] + for g in groups: + s = g.kv_cache_spec + kw: dict[str, Any] = {"block_size": s.block_size * scale} + if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: + kw["page_size_padded"] = s.page_size_padded * scale + scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) + return scaled, common_page + + +def get_kv_cache_groups( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Split the layers in the model into groups with the same KV cache spec. + + Args: + vllm_config: The global VllmConfig + kv_cache_spec: The kv cache spec of each attention layer in the model + + Returns: + The generated KVCacheGroups + + """ + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + unify_hybrid_kv_cache_specs(kv_cache_spec) + + if is_kv_cache_type_attention_free(kv_cache_spec): + # This returns an empty list to allow for the KVCacheManager to handle + # attention free models. + return [] + + if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): + return hisparse_groups + + if is_kv_cache_spec_uniform(kv_cache_spec): + # KV cache of all layers are the same, which is true for + # most models. Allocate the same amount of memory for + # each layer. + return _get_kv_cache_groups_uniform_spec(kv_cache_spec) + elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): + # All layers need the same number of token slots (e.g., all layers are + # 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) + # Hidden-state layers use their own block table and must not be absorbed + # into a compatible attention bucket. + hidden_specs = { + k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) + } + filtered_spec = { + k: v + for k, v in kv_cache_spec.items() + if not isinstance(v, HiddenStateCacheSpec) + } + + if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): + # Block-outermost blocks are strided by the widest group, so hidden + # groups need no page alignment. + packed_groups += [ + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() + ] + return packed_groups + + # Prefer preserving each layer's cache semantics. If physical pages cannot + # be unified, try a supported allocation-only fallback before failing. + try: + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) + except NotImplementedError: + fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) + if fallback_groups is None: + raise + return fallback_groups + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) + + # Add hidden-state layers back with page aligned to the common page. + if hidden_specs: + common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) + # TP may shrink the common page below hidden-state per-token cost. + groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) + group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) + for name, spec in hidden_specs.items(): + per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) + max_block_size = max(common_page // per_token, 1) + new_bs = _largest_divisor_at_most(group_block_size, max_block_size) + wasted_bytes = common_page - new_bs * per_token + logger.info( + "Using block size %d for hidden-state cache layer %s; " + "page alignment wastes %d bytes (%.2f%%) per block", + new_bs, + name, + wasted_bytes, + wasted_bytes / common_page * 100, + ) + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) + groups.append(KVCacheGroupSpec([name], aligned)) + + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), + ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) + return groups + + +def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: + if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: + return 1 + if spec.dcp_sharded and dcp_size > 1: + return 1 + return max(1, tp_size // spec.max_tp_shards) + + +def kv_cache_groups_tp_replicas( + groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 +) -> int: + """Consecutive TP ranks holding identical KV for every layer.""" + specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] + if not specs: + return 1 + return reduce( + math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size + ) + + +def _max_memory_usage_bytes_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + """Calculate maximum memory usage in bytes from KV cache groups. + + This correctly accounts for padding in hybrid models. For example, if a + model has 8 full attention layers and 9 sliding window layers, they will + be padded to 9 full + 9 sliding window for uniform group sizes. + + Each group independently claims blocks from the shared pool, so a request consumes + the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. + """ + kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] + if not kv_cache_groups: + return 0 + + 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 + if isinstance(spec, UniformTypeKVCacheSpecs): + total_blocks += spec.max_memory_usage_pages(vllm_config) + else: + total_blocks += cdiv( + spec.max_memory_usage_bytes(vllm_config), + spec.page_size_bytes, + ) + + return bytes_per_block * total_blocks + + +def _estimate_max_model_len_from_groups( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> int: + """Binary search for the maximum model length that fits in available memory. + Returns 0 if even 1 token doesn't fit. + """ + original_max = vllm_config.model_config.max_model_len + hisparse_enabled = ( + vllm_config.attention_config.hisparse_config is not None + and bool(kv_cache_groups) + ) + + def fits(model_len: int) -> bool: + vllm_config.model_config.max_model_len = model_len + if hisparse_enabled: + try: + get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, available_memory + ) + except ValueError: + return False + return ( + _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) + <= available_memory + ) + + try: + left, right = 1, original_max + if not fits(left): + return 0 + result = 1 + while left <= right: + mid = (left + right) // 2 + if fits(mid): + result = mid + left = mid + 1 + else: + right = mid - 1 + return result + finally: + vllm_config.model_config.max_model_len = original_max + + +def _auto_fit_max_model_len( + vllm_config: VllmConfig, + projected_groups_per_worker: list[list[KVCacheGroupSpec]], + available_memory: list[int], +) -> None: + """When max_model_len is set to -1, this function estimates the largest + context length that can be supported with the available GPU memory. + It uses binary search to find the maximum length that fits across all + workers. + + Args: + vllm_config: The global VllmConfig (will be modified in-place) + projected_groups_per_worker: KV cache groups projected to each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + """ + original_max = vllm_config.model_config.max_model_len + + if all(not groups for groups in projected_groups_per_worker): + # All workers have empty specs (attention-free model) + logger.info_once( + "Auto-fit max_model_len: attention-free model, " + "using derived max_model_len=%d", + original_max, + ) + return + + # Find the max_model_len that fits across all workers. + auto_fit_max = original_max + limiting_worker_mem = available_memory[0] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + continue + worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) + if worker_max < auto_fit_max: + auto_fit_max = worker_max + limiting_worker_mem = avail_mem + + if auto_fit_max <= 0: + raise ValueError( + "Cannot auto-fit max_model_len: not enough GPU memory available " + "to serve even a single token. Try increasing `gpu_memory_utilization`." + ) + + kv_cache_memory = ( + f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" + ) + if vllm_config.attention_config.hisparse_config is not None: + host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) + kv_cache_memory += f" + {host_pool_gib} GiB host" + + if auto_fit_max >= original_max: + # The model's full context length fits in memory + logger.info_once( + "Auto-fit max_model_len: full model context length %d fits in " + "available KV cache memory (%s)", + original_max, + kv_cache_memory, + ) + else: + # Need to reduce max_model_len to fit in memory + vllm_config.model_config.max_model_len = auto_fit_max + logger.info_once( + "Auto-fit max_model_len: reduced from %d to %d to fit in " + "available KV cache memory (%s)", + original_max, + auto_fit_max, + kv_cache_memory, + ) + + +def _project_kv_cache_groups_to_worker( + global_kv_cache_groups: list[KVCacheGroupSpec], + worker_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Projects global KV cache groups onto a single worker's assigned layers. + + In pipeline parallelism, each worker only owns a subset of layers. This + function filters the global groups to include only layers present on the + given worker, adjusting UniformTypeKVCacheSpecs accordingly. + + Args: + global_kv_cache_groups: The global KV cache groups for the whole model. + worker_spec: The KV cache spec of each layer on this worker. + + Returns: + The projected KV cache groups containing only this worker's layers. + + """ + projected_groups: list[KVCacheGroupSpec] = [] + for group in global_kv_cache_groups: + worker_layer_names = [ + layer_name for layer_name in group.layer_names if layer_name in worker_spec + ] + if len(worker_layer_names) == len(group.layer_names): + projected_groups.append(group) + continue + group_spec = group.kv_cache_spec + if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): + group_spec = UniformTypeKVCacheSpecs( + block_size=group_spec.block_size, + kv_cache_specs={ + layer_name: group_spec.kv_cache_specs[layer_name] + for layer_name in worker_layer_names + }, + ) + projected_groups.append( + KVCacheGroupSpec( + worker_layer_names, + group_spec, + is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + ) + ) + return projected_groups + + +def get_kv_cache_configs( + vllm_config: VllmConfig, + kv_cache_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> list[KVCacheConfig]: + """Generates the KV cache configurations for a model. + Since we use a shared centralized controller for all workers, we need the + `kv_cache_config` to be consistent across all workers to make sure + the KV cache allocation can be applied to all workers. However, different + workers may have different memory available, and different type of layers + (when pipeline parallel is enabled). To handle the difference between + workers, the current implementation is: + 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for + the whole model. + 2. Generate the KV cache groups based on the layer ratio of the whole model. + This also handles spec unification for hybrid models. + 3. Handle auto-fit max_model_len and memory checks using per-worker + projected groups to account for PP sharding. + 4. Generate the KV cache configs for each worker based on the KV cache + grouping strategy. (This is reasonable because the layer ratio of + different PP stages are similar.) + 5. Change the num_blocks of each worker to the smallest among all workers + and shrink tensor sizes proportionally to avoid allocating unused memory. + + Args: + vllm_config: The global VllmConfig + kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. + available_memory: Memory available for KV cache in bytes for each + worker. + + Returns: + The generated KVCacheConfigs for each worker. + + """ + # Merge the KV cache specs of all workers. Different PP stages may have + # different layer names, and different TP ranks of the same PP stage should + # have the same KV cache spec. + merged_kv_cache_specs: dict[str, KVCacheSpec] = {} + for kv_cache_spec_one_worker in kv_cache_specs: + for layer_name, layer_spec in kv_cache_spec_one_worker.items(): + if layer_name not in merged_kv_cache_specs: + merged_kv_cache_specs[layer_name] = layer_spec + else: + assert merged_kv_cache_specs[layer_name] == layer_spec, ( + "The KV cache specs for the same layer are different " + "across workers. This is not supported yet." + ) + + # Check if the KV cache specs are registered correctly. + # This is to prevent that some layers are initialized with unregistered specs. + KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + + # When speculating with more than 1 speculative module (e.g. multi-layered MTP) + # tag every SlidingWindowSpec with how many extra tokens to retain in the window. + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + + # Get global KV cache groups. This also handles spec unification for + # hybrid models when disable_hybrid_kv_cache_manager is enabled. + # After this call, merged_kv_cache_specs may be modified in-place. + global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) + + # If original_max_model_len was -1, automatically + # determine the maximum model length that fits in available GPU memory. + # We use per-worker projected groups to account for PP sharding. + projected_groups_per_worker = [ + _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) + for worker_spec in kv_cache_specs + ] + + # If `num_gpu_blocks_override` is set, the cache size that will actually + # be allocated is decoupled from the profiled `available_memory`: + # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps + # `num_blocks` to the override. Reflect that in `available_memory` here so + # auto-fit, the admission check, and the per-worker config builder all + # plan against the same effective capacity. + override = vllm_config.cache_config.num_gpu_blocks_override + if override is not None: + adjusted_memory: list[int] = [] + for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + if not groups: + adjusted_memory.append(avail_mem) + continue + 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, + override, + ) + adjusted_memory.append(override * bytes_per_block) + available_memory = adjusted_memory + + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + 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) + ] + + if vllm_config.model_config.original_max_model_len == -1: + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) + + # Check if the available memory is enough per worker. + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): + if not groups: + continue + _check_enough_kv_cache_memory( + avail_mem, + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + vllm_config.model_config.max_model_len, + partial(_estimate_max_model_len_from_groups, vllm_config, groups), + ) + + kv_cache_configs: list[KVCacheConfig] = [] + for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( + projected_groups_per_worker, kv_cache_specs, available_memory + ): + assert sum(len(group.layer_names) for group in projected_groups) == len( + kv_cache_spec_one_worker + ), "Some layers are not assigned to any group." + kv_cache_configs.append( + get_kv_cache_config_from_groups( + vllm_config, projected_groups, available_memory_one_worker + ) + ) + + # Change the num_blocks of each rank to the smallest among all ranks. + # We also need to shrink the tensor size proportionally to avoid + # allocating unused memory. + min_num_blocks = min( + kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs + ) + for i, kv_cache_config in enumerate(kv_cache_configs): + if kv_cache_config.num_blocks == min_num_blocks: + continue + # Re-plan with exactly the memory the smallest rank can afford, so + # 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(vllm_config, groups), + ) + + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.decode_context_parallel_size, + ) + + return kv_cache_configs diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 03c6c5ff62ac..5c51a8d131ff 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -6,10 +6,8 @@ import hashlib import math import os -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, overload from vllm import envs @@ -17,34 +15,14 @@ from vllm.logger import init_logger from vllm.multimodal.utils import get_mm_features_in_window from vllm.utils.hashing import xxhash, xxhash_cbor -from vllm.utils.math_utils import cdiv, round_up -from vllm.utils.mem_utils import format_gib -from vllm.utils.torch_utils import get_dtype_size -from vllm.v1.hisparse.layout import ( - get_hisparse_host_pool_bytes, - get_hisparse_kv_cache_config, - get_hisparse_kv_cache_groups, -) +from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import ( - AttentionSpec, - ChunkedLocalAttentionSpec, - CircularBufferSpec, - FullAttentionSpec, - HiddenStateCacheSpec, - HiSparseHotSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheLayout, KVCacheSpec, - KVCacheTensor, MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, - compute_layout_strides, iter_layer_specs, - replace_as, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -1000,198 +978,6 @@ def request_block_hasher(request: Request) -> list[BlockHash]: return request_block_hasher -def _check_enough_kv_cache_memory( - available_memory: int, - get_needed_memory: Callable[[], int], - max_model_len: int, - estimate_max_model_len: Callable[[int], int], -): - if available_memory <= 0: - raise ValueError( - "No available memory for the cache blocks. " - "Try increasing `gpu_memory_utilization` when initializing the engine " - "(this flag also controls CPU memory reservation on the CPU " - "backend, despite its name). " - "See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - "for more details." - ) - - needed_memory = get_needed_memory() - - if needed_memory > available_memory: - estimated_max_len = estimate_max_model_len(available_memory) - estimated_msg = "" - if estimated_max_len > 0: - estimated_msg = ( - "Based on the available memory, " - f"the estimated maximum model length is {estimated_max_len}. " - ) - - raise ValueError( - f"To serve at least one request with the model's max seq len " - f"({max_model_len}), ({format_gib(needed_memory)} GiB KV " - f"cache is needed, which is larger than the available KV cache " - f"memory ({format_gib(available_memory)} GiB). {estimated_msg}" - f"Try increasing `gpu_memory_utilization` (which also controls " - f"CPU memory on the CPU backend) or decreasing `max_model_len` " - f"when initializing the engine. " - f"See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ " - f"for more details." - ) - - -def max_memory_usage_bytes( - vllm_config: VllmConfig, kv_cache_specs: Iterable[KVCacheSpec] -) -> int: - """Get the maximum memory usage in bytes for the given KV cache specs.""" - return sum(spec.max_memory_usage_bytes(vllm_config) for spec in kv_cache_specs) - - -def estimate_max_model_len( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -) -> int: - """Estimates the maximum model length that can fit in the available memory - using binary search. - - This function temporarily modifies max_model_len during estimation but - restores the original value before returning, ensuring no side effects. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Returns: - The estimated maximum model length that can fit in the available memory. - - """ - # Save the original max_model_len to restore after estimation - original_max_model_len = vllm_config.model_config.max_model_len - - # Define a function to check if a given model length fits in memory - def fits_in_memory(model_len: int) -> bool: - # Temporarily modify the max_model_len for this calculation - vllm_config.model_config.max_model_len = model_len - # Calculate memory needed for the given model length - memory_needed = max_memory_usage_bytes(vllm_config, kv_cache_spec.values()) - return memory_needed <= available_memory - - try: - # Binary search for the maximum model length - left, right = 1, original_max_model_len - - # If even the smallest model length doesn't fit, return 0 - if not fits_in_memory(left): - return 0 - - # Binary search for the maximum model length that fits - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits_in_memory(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - # Always restore the original max_model_len to avoid side effects - vllm_config.model_config.max_model_len = original_max_model_len - - -def check_enough_kv_cache_memory( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - available_memory: int, -): - """Checks whether `available_memory` is enough for the KV cache to hold at - least one request with the model's max_model_len. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - available_memory: Memory available for KV cache in bytes. - - Raises: - ValueError: If there is not enough memory available for the KV cache. - - """ - # No need to check for available memory if the kv_cache_spec is empty - if kv_cache_spec: - # Reserve the null block BlockPool permanently holds back, so the check - # plans against usable blocks, as in get_kv_cache_configs. Group a copy - # 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(vllm_config, groups) - if groups - else available_memory - ) - _check_enough_kv_cache_memory( - check_memory, - lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), - vllm_config.model_config.max_model_len, - lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), - ) - - -def create_kv_cache_group_specs( - kv_cache_spec: dict[str, KVCacheSpec], grouped_layer_names: list[list[str]] -) -> list[KVCacheGroupSpec]: - """Create KVCacheGroupSpec object for each kv cache group layer. - The layers in the same group should share the same - KVCacheSpec. - - Args: - kv_cache_spec: - A mapping from each layer name to its corresponding KVCacheSpec. - grouped_layer_names: - A list of kv cache groups, where each element is a list of layer - names that belong to the same group and should share the same - KVCacheSpec. - - Returns: - A list of KVCacheGroupSpec objects, one for each group. - - """ - kv_cache_groups = [] - for layer_names_one_group in grouped_layer_names: - layer_specs = [ - kv_cache_spec[layer_name] for layer_name in layer_names_one_group - ] - merged_layer_spec = layer_specs[0].merge(layer_specs) - kv_cache_groups.append( - KVCacheGroupSpec(layer_names_one_group, merged_layer_spec) - ) - return kv_cache_groups - - -def is_kv_cache_spec_uniform(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - """Whether all layers in the given KVCacheSpec have the same KV cache spec. - Note that we regard FullAttentionSpec with and without sliding window as - the same type. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - True if all layers have the same type, False otherwise. - - """ - if not kv_cache_spec: - # Encoder-only models do not have KV cache, kv_cache_type can be - # regarded as uniform. - return True - try: - kv_cache_spec_values = list(kv_cache_spec.values()) - _ = kv_cache_spec_values[0].merge(kv_cache_spec_values) - except AssertionError: - return False - return True - - def get_max_concurrency_for_kv_cache_config( vllm_config: VllmConfig, kv_cache_config: KVCacheConfig ) -> float: @@ -1227,1052 +1013,6 @@ def get_max_concurrency_for_kv_cache_config( return min(limits) -def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: - """Override the number of kv cache blocks if `num_gpu_blocks_override` is set. - The override is logged once, at the call site in `get_kv_cache_configs`. - """ - if vllm_config.cache_config.num_gpu_blocks_override is not None: - num_blocks = vllm_config.cache_config.num_gpu_blocks_override - return num_blocks - - -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], - vllm_config.cache_config.get_resolved_kv_cache_layout(), - ) - - -def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: - """Get the page size of the KV cache.""" - page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} - assert len(page_sizes) == 1 - return page_sizes.pop() - - -def _get_kv_cache_groups_uniform_spec( - kv_cache_specs: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with the same KV cache - spec for all layers. - - Args: - kv_cache_specs: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return create_kv_cache_group_specs(kv_cache_specs, [list(kv_cache_specs.keys())]) - - -def _get_kv_cache_groups_uniform_type( - spec: UniformTypeKVCacheSpecs, -) -> list[KVCacheGroupSpec]: - """Generates the KV cache configuration for a model with one type of KV cache - but different hidden sizes. All layers are merged into one group. - - Args: - spec: The UniformTypeKVCacheSpecs of the model - - Returns: - The generated KVCacheGroupSpecs - - """ - return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)] - - -def unify_kv_cache_spec_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Unify the page size of the given KVCacheSpec. If the page size of all layers - are the same, return the original KVCacheSpec. If not same, unify the page - size by increasing the block size of layers with smaller page size. Two - cases cannot be unified by block size alone and pad their physical page to - the maximum instead: Mamba layers, whose page size comes from state shapes - and is independent of block size; and non-MLA attention layers whose page - does not evenly divide the maximum (the padded page is read through a - strided view). MLA is excluded because sparse MLA indexes the cache in - whole token rows (see ``flat_kv_row_view``), so its block stride can only - be padded by its own row-aligned ``alignment``, not to an arbitrary page - size. Raise NotImplementedError if failed to unify the page size; - ``get_kv_cache_groups`` catches it to try the full-allocation fallback - (e.g. MLA next to an incompatible sliding-window draft). - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - - Returns: - The updated KVCacheSpec with the same page_size_bytes. - - """ - page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()} - if len(page_sizes) <= 1: - # All layers have the same page size, no need to unify. - return kv_cache_spec - - max_page_size = max(page_sizes) - new_kv_cache_spec = {} - for layer_name, layer_spec in kv_cache_spec.items(): - if layer_spec.page_size_bytes == max_page_size: - new_kv_cache_spec[layer_name] = layer_spec - elif isinstance(layer_spec, MambaSpec): - # MambaSpec's page size is determined by its state shapes and does - # not scale with block_size, so pad the page instead. This is the - # same padding mechanism the platform uses to align Mamba pages - # with the main model's attention page size; it is needed here - # when another layer (e.g. from a draft model) has a larger page - # than the already-aligned Mamba page. - new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - else: - layer_page_size = layer_spec.page_size_bytes - if max_page_size % layer_page_size == 0: - ratio = max_page_size // layer_page_size - new_block_size = layer_spec.block_size * ratio - new_spec = replace(layer_spec, block_size=new_block_size) - elif isinstance(layer_spec, AttentionSpec) and not isinstance( - layer_spec, MLAAttentionSpec - ): - new_spec = replace(layer_spec, page_size_padded=max_page_size) - else: - raise NotImplementedError( - f"Layer {layer_name}: page size is not divisible by the " - "maximum page size and cannot be padded. Padding is only " - "supported for non-MLA attention layers." - ) - assert new_spec.page_size_bytes == max_page_size - new_kv_cache_spec[layer_name] = new_spec - return new_kv_cache_spec - - -def is_kv_cache_type_attention_free(kv_cache_spec: dict[str, KVCacheSpec]) -> bool: - # kv_cache_spec is an empty dict for attention free models - return not kv_cache_spec - - -def _get_kv_cache_groups_uniform_page_size( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Generates the KV cache groups for hybrid models with multiple - attention types but still with a uniform page size (physical memory per - block per layer) for all layers. - - Detailed explanation about kv cache management of hybrid models: - The layers in the models are repeated with some patterns, e.g., a model - with 10 full attention layers and 20 sliding window attention layers can be - regarded as repeating the pattern (1 * full, 2 * sw) 10 times. - The KVCacheManager allocates different block tables for each of the 3 layers - in the pattern, and repeats each of them 10 times to generate the - block_table for the 30 layers in the model. - Therefore, we can group the layers in the model into 3 kv_cache_groups, each - of which contains 10 layers in the model. - The KVCacheManager allocates the block_table for each group based on its - kv_cache spec, and the model runner applies the block table to each layer - in the group. - For example: - 1. A model only uses full attention. The pattern is - (num_hidden_layers * full), so there is only one group and the block table - is shared by all layers. It is already handled by - `_get_kv_cache_config_uniform_type`. - 2. A model with 10 full attention layers and 20 sliding window - attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so - there are 3 kv_cache_groups, each of which represents 10 layers. - - To simplify the implementation, we make the following assumptions: - 1. Physical memory per block: Must be the same across all KV cache groups. - Breaking this assumption is non-trivial due to memory fragmentation concerns - when allocating blocks of different sizes. - 2. Tokens per block (block_size): Currently, we directly use - `CacheConfig.block_size` for all layers. It can be extended to vary by KV - cache group, but within each KV cache group, all layers must share the same - block size. - 3. Physical memory per token per layer: This property is decided by model - config. Currently we only support models that have the same physical memory - per token per layer for all layers. Can be relaxed with a simple extension, - but still need to keep physical memory per block the same for all groups. - 4. Number of layers per group: Currently assumed the same for all layers. - Can be relaxed with a simple extension, but still need to keep physical - memory per block the same for all groups. - 5. Attention type within groups: All layers in a group must share the same - attention type. One exception is that, when - `--disable-hybrid-kv-cache-manager` is true, the single group for full - attention layers may also include attention layers using sliding window or - LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details. - 6. Support for multiple attention types: The design for most components is - general to an arbitrary number of attention types. But - `find_longest_cache_hit` only supports one attention type or two - types of full-attention plus exactly one another type. The general - implementation of this function is feasible but we don't know how to - implement it cleanly yet. - - As we assume tokens per block, physical memory per token per layer, and - number of layers per group are the same now, we can ensure that physical - memory per block is the same for all groups. - - Args: - kv_cache_spec: The KVCacheSpec of each attention layer in the model - Returns: - The generated KVCacheGroupSpecs - - """ - # Group all layers by kv_cache_spec. - # E.g., 2 full attention layers and 3 sliding window attention layers, - # -> (full.0, full.1), (sw.0, sw.1, sw.2). - same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list) - for layer_name, layer_spec in kv_cache_spec.items(): - same_type_layers[layer_spec].append(layer_name) - - # Attempt to further merge same-type layers based on whether their KV - # cache specs can be merged, to minimize the group count. This benefits - # situations where specs share a block layout and differ only in a - # property it can reconcile (e.g. full attention layers differing only in - # sliding window / attention chunk size). - layer_buckets: list[list[str]] = [] - spec_buckets: list[list[KVCacheSpec]] = [] - for layer_spec, layer_names in same_type_layers.items(): - for names, specs in zip(layer_buckets, spec_buckets): - try: - # A raise means that the specs are incompatible. - type(specs[0]).merge([*specs, layer_spec]) - except (AssertionError, ValueError): - continue - names.extend(layer_names) - specs.append(layer_spec) - break - else: - layer_buckets.append(list(layer_names)) - spec_buckets.append([layer_spec]) - - # Split each group into smaller groups, to make the number of layers in each - # group identical. Add padding to the last group of each type if necessary. - # E.g., (full.0, full.1), (sw.0, sw.1, sw.2) - # split to 3 groups with 2 layers each: - # (full.0, full.1), (sw.0, sw.2), (sw.1, padding). - # FIXME(Chen): At the moment of writing this code (2025-06-02), all - # open-source hybrid model follows a n:1 pattern between different attention - # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and - # full), so we can use the "1" in the n:1 pattern as the group size, which - # is the minimum number of layers among all attention types. Need a better - # strategy if we want to support more complex patterns (e.g., 20 full + 30 - # sw, where the group size should be 10). - min_num_layers = min([len(layers) for layers in layer_buckets]) - group_size = min_num_layers - max_num_layers = max([len(layers) for layers in layer_buckets]) - if max_num_layers < min_num_layers * 1.5: - # If the number of layers is not much larger than the minimum number of - # layers, use the maximum number of layers as the group size to avoid - # too many padding layers. A typical example is gpt-oss-20b + eagle, - # with 12 sw + 13 full. We pad it to (13 sw, 13 full) instead of - # (12 sw, 24 full). 1.5 is a heuristic to avoid too many padding - # layers while accommodating speculative decoding drafters that add - # extra layers to one attention type. - group_size = max_num_layers - grouped_layers = [] - for layers in layer_buckets: - num_padding_layers = group_size - len(layers) % group_size - if num_padding_layers != group_size: - logger.warning( - "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa - num_padding_layers, - num_padding_layers / len(layers) * 100, - ) - num_groups = cdiv(len(layers), group_size) - # In PP case, say if we have - # - stage 0: full.0, sw.0, sw.1 - # - stage 1: full.1, sw.2, sw.3 - # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3) - # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because - # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group) - # and it will be padded to (full.0, padding), (sw.0, sw.1), - # (padding, padding) to ensure the number of layers in each group is - # the same and will cause memory waste. - # To avoid this, we assign layers[i::num_groups] to the i-th group - # instead of layers[i * group_size: (i + 1) * group_size] - for i in range(num_groups): - grouped_layers.append(layers[i::num_groups]) - return create_kv_cache_group_specs(kv_cache_spec, grouped_layers) - - -def _get_per_layer_spec( - group: KVCacheGroupSpec, - layer_name: str, -) -> KVCacheSpec: - spec = group.kv_cache_spec - if isinstance(spec, UniformTypeKVCacheSpecs): - return spec.kv_cache_specs[layer_name] - 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], layout: KVCacheLayout -) -> int: - """Return the largest cache group's bytes per block.""" - 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 - for layer_name in group.layer_names - ) - for group in kv_cache_groups - ) - assert bytes_per_block > 0 - alignments = [ - group.kv_cache_spec.page_size_bytes - for group in kv_cache_groups - if isinstance(group.kv_cache_spec, HiSparseHotSpec) - ] - alignments.extend( - _get_per_layer_spec(group, layer_name).get_block_stride_alignment() - for group in kv_cache_groups - for layer_name in group.layer_names - ) - return round_up(bytes_per_block, math.lcm(*alignments)) - - -def validate_kv_cache_layout( - layout: KVCacheLayout, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Validate that the resolved layout can express this model's packing. - - The layout was chosen once in the engine core from the backends' supported - sets; a backend whose model packs pages side by side (e.g. the DeepSeek-V4 - indexer) declares block-outermost layouts there, so an inexpressible - layout reaching this point is an error. - """ - page_sizes = { - _get_per_layer_spec(group, layer_name).page_size_bytes - for group in kv_cache_groups - for layer_name in group.layer_names - } - if len(page_sizes) == 1: - # A rectangular layer dim exists; every layout can express it. - return - - # Mixed page sizes pack pages side by side within a block, which needs each page - # to be one contiguous chunk inside its block (a block-compact layout) and, with - # multiple KV cache groups, the layer dim inside the block dim. - if not layout.is_block_compact or ( - len(kv_cache_groups) > 1 and layout.is_layer_compact - ): - raise ValueError( - f"KV cache layout {layout.name} cannot express this model's " - f"mixed page sizes ({sorted(page_sizes)}); a backend should " - "declare block-outermost supported layouts (e.g. BLHNC), or " - "set VLLM_KV_CACHE_LAYOUT=BLHNC." - ) - - -def get_kv_cache_config_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> KVCacheConfig: - """Generate the KV cache configuration from the KV cache groups and spec - of each layer. - - Args: - vllm_config: The global VllmConfig - kv_cache_groups: The KV cache groups - available_memory: Memory available for KV cache in bytes - Returns: - The generated KVCacheConfig - - """ - if len(kv_cache_groups) == 0: - # Attention free models do not have KV cache. - # Return num_blocks=1 as BlockPool always needs a null_block. - return KVCacheConfig( - num_blocks=1, - kv_cache_tensors=[], - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - if vllm_config.attention_config.hisparse_config is not None: - host_budget = get_hisparse_host_pool_bytes(vllm_config) - return get_hisparse_kv_cache_config( - vllm_config, kv_cache_groups, available_memory, host_budget - ) - - 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, layout) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None - - num_blocks = available_memory // bytes_per_block - num_blocks = may_override_num_blocks(vllm_config, num_blocks) - size = bytes_per_block * num_blocks - - # Groups alias from byte 0. Spec regions are laid out differently: - # - # block-outer (the same packing repeats for every block): - # group 0: | blk 0 [ A | B | pad ] | blk 1 [ A | B | pad ] | ... - # group 1: | blk 0 [ C | D ] | blk 1 [ C | D ] | ... - # |<--- bytes_per_block -->| - # - # layer-outer (only supported for uniform page sizes or single-group models): - # group 0: | A [ blk 0 | blk 1 | ... ] | B [ blk 0 | blk 1 | ... ] | - # group 1: | C [ blk 0 | blk 1 | ... ] | D [ blk 0 | blk 1 | ... ] | - - kv_cache_tensors = [] - for group in kv_cache_groups: - group_spec = group.kv_cache_spec - layers_by_spec: defaultdict[KVCacheSpec, list[str]] = defaultdict(list) - if isinstance(group_spec, UniformTypeKVCacheSpecs): - for layer_name, spec in group_spec.kv_cache_specs.items(): - layers_by_spec[spec].append(layer_name) - elif group.layer_names: - layers_by_spec[group_spec].extend(group.layer_names) - - 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, block_stride, None, None, None), - ) - kv_cache_tensors.append( - KVCacheTensor( - size=size, - layers=layer_names, - layer_stride=layer_stride, - block_stride=block_stride, - offset=byte_offset, - ) - ) - # 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, - kv_cache_tensors=kv_cache_tensors, - kv_cache_groups=kv_cache_groups, - prefix_cache_retention_interval=( - vllm_config.cache_config.prefix_cache_retention_interval - ), - ) - - -def _promote_local_kv_cache_specs( - kv_cache_spec: dict[str, KVCacheSpec], -) -> dict[str, KVCacheSpec]: - """Use full-attention allocation for local-attention cache specs. - - The returned specs affect KV cache management only. Attention modules keep - their original sliding-window or chunked-local compute behavior. - """ - promoted_specs = kv_cache_spec.copy() - - if is_kv_cache_spec_uniform( - promoted_specs - ) or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs): - return promoted_specs - - has_full_attention = any( - isinstance(spec, FullAttentionSpec) for spec in promoted_specs.values() - ) - has_sliding_window = any( - isinstance(spec, SlidingWindowSpec) for spec in promoted_specs.values() - ) - has_chunked_local_attention = any( - isinstance(spec, ChunkedLocalAttentionSpec) for spec in promoted_specs.values() - ) - full_block_sizes = { - spec.block_size - for spec in promoted_specs.values() - if isinstance(spec, FullAttentionSpec) - } - full_attention_block_size = ( - next(iter(full_block_sizes)) if len(full_block_sizes) == 1 else None - ) - - def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | None: - if spec.page_size_padded is None: - return None - unpadded_page_size = ( - spec.unpadded_page_size_bytes * block_size // spec.block_size - ) - return max(spec.page_size_padded, unpadded_page_size) - - promotions: dict[type[AttentionSpec], type[AttentionSpec]] = { - SlidingWindowMLASpec: MLAAttentionSpec, - SlidingWindowSpec: FullAttentionSpec, - ChunkedLocalAttentionSpec: FullAttentionSpec, - } - - if has_full_attention and (has_sliding_window or has_chunked_local_attention): - for layer_name, spec in kv_cache_spec.items(): - target_cls = next( - (promotions[c] for c in type(spec).__mro__ if c in promotions), None - ) - if target_cls is None: - continue - assert isinstance(spec, AttentionSpec) - block_size = full_attention_block_size or spec.block_size - promoted_specs[layer_name] = replace_as( - spec, - target_cls, - # Promoted specs allocate blocks for all tokens and never free - # below the window, so the trailing-edge extension is moot. - drop=("extra_retained_tokens",), - block_size=block_size, - page_size_padded=promoted_page_size_padded(spec, block_size), - ) - - if not ( - is_kv_cache_spec_uniform(promoted_specs) - or UniformTypeKVCacheSpecs.is_uniform_type(promoted_specs) - ): - raise ValueError("Failed to promote local KV cache specs to one unified type.") - - return promoted_specs - - -def _try_get_full_allocation_fallback_groups( - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Try a supported full-allocation fallback for local-attention layers.""" - if any(isinstance(spec, HiddenStateCacheSpec) for spec in kv_cache_spec.values()): - return None - if any( - isinstance(spec, (SlidingWindowMLASpec, ChunkedLocalAttentionSpec)) - for spec in kv_cache_spec.values() - ): - return None - - has_mla = any(isinstance(spec, MLAAttentionSpec) for spec in kv_cache_spec.values()) - has_regular_swa = any( - isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values() - ) - if not (has_mla and has_regular_swa): - return None - - try: - promoted_specs = _promote_local_kv_cache_specs(kv_cache_spec) - except ValueError: - return None - uniform_spec = UniformTypeKVCacheSpecs.from_specs(promoted_specs) - if uniform_spec is None: - return None - logger.warning( - "KV cache page sizes cannot be unified; treating sliding-window " - "layers as full attention for cache allocation. Sliding-window " - "attention compute is unchanged." - ) - return _get_kv_cache_groups_uniform_type(uniform_spec) - - -def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]): - """This function tries to convert the KV cache specs to one type if the model - is a hybrid model with multiple type of KV cache. It will convert all - SlidingWindowSpec to FullAttentionSpec if both types are present. - - Args: - kv_cache_spec: The kv cache spec of each attention layer in the model - - """ - groups: defaultdict[bool, dict[str, KVCacheSpec]] = defaultdict(dict) - for name, spec in kv_cache_spec.items(): - replicated = isinstance(spec, AttentionSpec) and not spec.dcp_sharded - groups[replicated][name] = spec - for specs in groups.values(): - promoted_specs = _promote_local_kv_cache_specs(specs) - if promoted_specs == specs: - continue - logger.warning( - "Hybrid KV cache manager is disabled for this hybrid model, " - "This means we do not enable any optimizations for saving KV cache " - "memory (e.g., dropping the KV cache outside the sliding window). " - "The compute of layers like sliding window is still saved." - ) - kv_cache_spec.update(promoted_specs) - - -def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int: - """Pick a chunk size that minimizes total upward padding. - - Each x is rounded up to a multiple of d: - - x -> ceil(x / d) * d - - Total padding is: - - pad(d) = sum_i (ceil(x_i / d) * d - x_i) - - We brute-force d in [lower_bound, max(values)] (fine for small lists / small - maxima) and return the d with minimum padding. Ties prefer larger d. - """ - if not values: - raise ValueError("values must be non-empty") - if any(x <= 0 for x in values): - raise ValueError(f"values must be positive, got: {list(values)!r}") - - min_d = max(1, lower_bound if lower_bound is not None else 1) - max_d = max(values) - if min_d > max_d: - return min_d - - best_d = min_d - best_pad: int | None = None - for d in range(min_d, max_d + 1): - pad = sum((d - (x % d)) % d for x in values) - if best_pad is None or pad < best_pad or (pad == best_pad and d > best_d): - best_pad = pad - best_d = d - - return best_d - - -def _get_packed_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec] | None: - """Group mixed-page-size layers for contiguous block-outermost packing. - - Greedily buckets layers into uniform-type specs. Buckets with equal layer - counts per page size are treated as a repeating layer pattern (one layer - per page size) and split into groups covering the same number of pattern - repeats (picked by ``_approximate_gcd`` to minimize padding), so all - groups pack into the same per-block layout. Mamba buckets are additionally - split to fit the block the attention buckets already need. - Returns None when the layout is not block-outermost or all layers already - share one page size. - """ - layout = vllm_config.cache_config.get_resolved_kv_cache_layout() - page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} - if not layout.is_block_outermost or len(page_sizes) <= 1: - return None - - buckets: list[dict[str, KVCacheSpec]] = [] - for name, spec in kv_cache_spec.items(): - for bucket in buckets: - candidate = {**bucket, name: spec} - if UniformTypeKVCacheSpecs.is_uniform_type(candidate): - bucket[name] = spec - break - else: - buckets.append({name: spec}) - - bucketed = [] - for bucket in buckets: - uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket) - assert uniform_spec is not None - page_size_layers: dict[int, list[str]] = defaultdict(list) - for layer_name, layer_spec in bucket.items(): - page_size_layers[layer_spec.page_size_bytes].append(layer_name) - # Only 1:1 patterns (one layer of each page size per repeat) are - # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle - # repeat too, but such buckets are emitted whole instead. - balanced = len(set(map(len, page_size_layers.values()))) == 1 - bucketed.append((uniform_spec, page_size_layers, balanced)) - - # Balanced buckets that mix page sizes must stay whole, so the largest one - # sets a floor on the repeats per group; larger single-size buckets are - # split down toward it. No such bucket means nothing needs packing. - min_repeats_per_group = max( - ( - spec.get_max_layers_per_page_size() - for spec, page_size_layers, balanced in bucketed - if balanced and len(page_size_layers) > 1 - ), - default=0, - ) - repeats_per_group = ( - _approximate_gcd( - [ - spec.get_max_layers_per_page_size() - for spec, _, balanced in bucketed - if balanced - ], - lower_bound=min_repeats_per_group, - ) - if min_repeats_per_group - else None - ) - - def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int: - if balanced and repeats_per_group is not None: - return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group) - return 1 - - def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: - """Page bytes of the largest of the n groups a bucket splits into.""" - return sum( - cdiv(len(names), n) * page for page, names in page_size_layers.items() - ) - - # Bytes a block must hold however the state buckets end up split: mamba, - # circular-buffer and unsplit sliding-window buckets can go down to one - # state per group, every other bucket's split is fixed by the repeat pattern. - def is_state_bucket(spec: UniformTypeKVCacheSpecs) -> bool: - if isinstance(spec.first_spec, (MambaSpec, CircularBufferSpec)): - return True - return repeats_per_group is None and isinstance( - spec.first_spec, SlidingWindowSpec - ) - - anchor_bytes = max( - ( - widest_group_bytes( - page_size_layers, - len(spec.kv_cache_specs) - if is_state_bucket(spec) - else num_groups_for(spec, balanced), - ) - for spec, page_size_layers, balanced in bucketed - ), - default=0, - ) - - groups = [] - for spec, page_size_layers, balanced in bucketed: - num_groups = num_groups_for(spec, balanced) - # Cap a state group at the states a block already fits rather than let - # it widen the block. - if anchor_bytes and is_state_bucket(spec): - states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1) - num_groups = max( - num_groups, cdiv(len(spec.kv_cache_specs), states_per_block) - ) - if num_groups == 1: - groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec)) - continue - - pattern_repeats = list(zip(*page_size_layers.values())) - for i in range(num_groups): - group_layer_names = [ - name for repeat in pattern_repeats[i::num_groups] for name in repeat - ] - group_layer_specs = { - name: spec.kv_cache_specs[name] for name in group_layer_names - } - group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs) - assert group_spec is not None - groups.append(KVCacheGroupSpec(group_layer_names, group_spec)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _uses_trailing_mtp_layers(vllm_config: VllmConfig) -> bool: - """Whether the drafter's KV layers can be located positionally. - - MTP drafters register their KV layers after the target layers, without a - distinguishing spec marker. DeepseekV4/V4.1 do this with ``dspark`` too. - The annotator separately checks that the groups partition the layers exactly. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return False - if spec_config.method == "mtp": - return True - model_config = vllm_config.model_config - return model_config is not None and model_config.hf_config.model_type in ( - "deepseek_v4", - "deepseek_v41", - ) - - -def _groups_partition_layers_exactly( - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], -) -> bool: - """Whether the groups cover every layer of ``kv_cache_spec`` exactly once. - - The trailing-layer draft fallback is only meaningful when the last - registered layer is guaranteed to appear in exactly one group. - """ - names = [name for group in kv_cache_groups for name in group.layer_names] - return len(names) == len(kv_cache_spec) and set(names) == set(kv_cache_spec) - - -def _annotate_eagle_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], - kv_cache_groups: list[KVCacheGroupSpec], - use_trailing_layer_fallback: bool = False, -) -> None: - """Flag the KV cache groups that hold drafter attention layers. - - Two detection rules, in order of preference: - - 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on - MLAAttentionSpec and set by drafter attention layers that run a - non-causal multi-token decode (today only Kimi-K3 DSpark). It survives - MLAAttentionSpec.merge, so it still identifies a group after per-group - spec merging, wherever grouping happens to land. It is sufficient but - not necessary: a drafter whose spec is indistinguishable from the - target's cannot be found this way. - 2. Positional fallback for MTP drafters (including DeepseekV4/V4.1 DSpark), whose - MTP block reuses the target's own decoder layer and so carries no spec - marker. Their draft attention layers always register after every - target layer, so flag whichever group holds the last registered layer. - This rule is only valid where the groups partition exactly the layers - of ``kv_cache_spec``, which is re-checked here before applying it. - When a drafter's trailing caches span several groups, this rule flags - only the group holding the very last layer and must be generalized. - FIXME(yifan): avoid/generalize this hacky check. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_spec: The kv cache spec of each attention layer, in layer - registration order. Only read by rule 2. - kv_cache_groups: Groups to annotate in place. - use_trailing_layer_fallback: Enable rule 2. Callers gate this on - ``_uses_trailing_mtp_layers``. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle_block_drop(): - return - - for group in kv_cache_groups: - if any( - getattr(spec, "non_causal_multi_token_decode", False) - for spec in iter_layer_specs(group.kv_cache_spec) - ): - group.is_eagle_group = True - - if not use_trailing_layer_fallback: - return - if not _groups_partition_layers_exactly(kv_cache_spec, kv_cache_groups): - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - -def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Warn when no KV cache group could be identified as the draft model's. - - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_groups: Groups as they will be handed to consumers. - - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return - if any(group.is_eagle_group for group in kv_cache_groups): - return - mamba_groups = [ - idx - for idx, group in enumerate(kv_cache_groups) - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if not mamba_groups: - return - logger.warning( - "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's.", - spec_config.method, - ) - - -def _largest_divisor_at_most(value: int, limit: int) -> int: - for candidate in range(min(value, limit), 0, -1): - if value % candidate == 0: - return candidate - return 1 - - -def _ensure_min_page_size( - groups: list[KVCacheGroupSpec], - common_page: int, - hidden_specs: dict[str, HiddenStateCacheSpec], -) -> tuple[list[KVCacheGroupSpec], int]: - """Scale up group block sizes so the common page is at least as large as - the biggest hidden-state per-token cost. - - This protects the hidden state extraction feature, where we store hidden states in - the KV cache, and return them to clients using a custom hidden state connector. - - Returns the (possibly scaled) groups and updated common page size. - """ - max_per_token = max( - spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - for spec in hidden_specs.values() - ) - if max_per_token <= common_page: - return groups, common_page - - scale = cdiv(max_per_token, common_page) - common_page *= scale - scaled: list[KVCacheGroupSpec] = [] - for g in groups: - s = g.kv_cache_spec - kw: dict[str, Any] = {"block_size": s.block_size * scale} - if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None: - kw["page_size_padded"] = s.page_size_padded * scale - scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw))) - return scaled, common_page - - -def get_kv_cache_groups( - vllm_config: VllmConfig, - kv_cache_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Split the layers in the model into groups with the same KV cache spec. - - Args: - vllm_config: The global VllmConfig - kv_cache_spec: The kv cache spec of each attention layer in the model - - Returns: - The generated KVCacheGroups - - """ - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - unify_hybrid_kv_cache_specs(kv_cache_spec) - - if is_kv_cache_type_attention_free(kv_cache_spec): - # This returns an empty list to allow for the KVCacheManager to handle - # attention free models. - return [] - - if hisparse_groups := get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec): - return hisparse_groups - - if is_kv_cache_spec_uniform(kv_cache_spec): - # KV cache of all layers are the same, which is true for - # most models. Allocate the same amount of memory for - # each layer. - return _get_kv_cache_groups_uniform_spec(kv_cache_spec) - elif uniform_spec := UniformTypeKVCacheSpecs.from_specs(kv_cache_spec): - # All layers need the same number of token slots (e.g., all layers are - # 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) - # Hidden-state layers use their own block table and must not be absorbed - # into a compatible attention bucket. - hidden_specs = { - k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec) - } - filtered_spec = { - k: v - for k, v in kv_cache_spec.items() - if not isinstance(v, HiddenStateCacheSpec) - } - - if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec): - # Block-outermost blocks are strided by the widest group, so hidden - # groups need no page alignment. - packed_groups += [ - KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() - ] - return packed_groups - - # Prefer preserving each layer's cache semantics. If physical pages cannot - # be unified, try a supported allocation-only fallback before failing. - try: - filtered_spec = unify_kv_cache_spec_page_size(filtered_spec) - except NotImplementedError: - fallback_groups = _try_get_full_allocation_fallback_groups(kv_cache_spec) - if fallback_groups is None: - raise - return fallback_groups - groups = _get_kv_cache_groups_uniform_page_size(filtered_spec) - - # Add hidden-state layers back with page aligned to the common page. - if hidden_specs: - common_page = get_uniform_page_size([g.kv_cache_spec for g in groups]) - # TP may shrink the common page below hidden-state per-token cost. - groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs) - group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups)) - for name, spec in hidden_specs.items(): - per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype) - max_block_size = max(common_page // per_token, 1) - new_bs = _largest_divisor_at_most(group_block_size, max_block_size) - wasted_bytes = common_page - new_bs * per_token - logger.info( - "Using block size %d for hidden-state cache layer %s; " - "page alignment wastes %d bytes (%.2f%%) per block", - new_bs, - name, - wasted_bytes, - wasted_bytes / common_page * 100, - ) - aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) - groups.append(KVCacheGroupSpec([name], aligned)) - - _annotate_eagle_groups( - vllm_config, - kv_cache_spec, - groups, - use_trailing_layer_fallback=_uses_trailing_mtp_layers(vllm_config), - ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) - return groups - - -def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: - if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: - return 1 - if spec.dcp_sharded and dcp_size > 1: - return 1 - return max(1, tp_size // spec.max_tp_shards) - - -def kv_cache_groups_tp_replicas( - groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 -) -> int: - """Consecutive TP ranks holding identical KV for every layer.""" - specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] - if not specs: - return 1 - return reduce( - math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size - ) - - def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: @@ -2334,359 +1074,6 @@ def update_kv_cache_capacity( ) -def _max_memory_usage_bytes_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> int: - """Calculate maximum memory usage in bytes from KV cache groups. - - This correctly accounts for padding in hybrid models. For example, if a - model has 8 full attention layers and 9 sliding window layers, they will - be padded to 9 full + 9 sliding window for uniform group sizes. - - Each group independently claims blocks from the shared pool, so a request consumes - the sum of the per-group block counts, i.e. ``bytes_per_block * total_blocks``. - """ - kv_cache_groups = [group for group in kv_cache_groups if not group.host_resident] - if not kv_cache_groups: - return 0 - - 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 - if isinstance(spec, UniformTypeKVCacheSpecs): - total_blocks += spec.max_memory_usage_pages(vllm_config) - else: - total_blocks += cdiv( - spec.max_memory_usage_bytes(vllm_config), - spec.page_size_bytes, - ) - - return bytes_per_block * total_blocks - - -def _estimate_max_model_len_from_groups( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], - available_memory: int, -) -> int: - """Binary search for the maximum model length that fits in available memory. - Returns 0 if even 1 token doesn't fit. - """ - original_max = vllm_config.model_config.max_model_len - hisparse_enabled = ( - vllm_config.attention_config.hisparse_config is not None - and bool(kv_cache_groups) - ) - - def fits(model_len: int) -> bool: - vllm_config.model_config.max_model_len = model_len - if hisparse_enabled: - try: - get_kv_cache_config_from_groups( - vllm_config, kv_cache_groups, available_memory - ) - except ValueError: - return False - return ( - _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) - <= available_memory - ) - - try: - left, right = 1, original_max - if not fits(left): - return 0 - result = 1 - while left <= right: - mid = (left + right) // 2 - if fits(mid): - result = mid - left = mid + 1 - else: - right = mid - 1 - return result - finally: - vllm_config.model_config.max_model_len = original_max - - -def _auto_fit_max_model_len( - vllm_config: VllmConfig, - projected_groups_per_worker: list[list[KVCacheGroupSpec]], - available_memory: list[int], -) -> None: - """When max_model_len is set to -1, this function estimates the largest - context length that can be supported with the available GPU memory. - It uses binary search to find the maximum length that fits across all - workers. - - Args: - vllm_config: The global VllmConfig (will be modified in-place) - projected_groups_per_worker: KV cache groups projected to each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - """ - original_max = vllm_config.model_config.max_model_len - - if all(not groups for groups in projected_groups_per_worker): - # All workers have empty specs (attention-free model) - logger.info_once( - "Auto-fit max_model_len: attention-free model, " - "using derived max_model_len=%d", - original_max, - ) - return - - # Find the max_model_len that fits across all workers. - auto_fit_max = original_max - limiting_worker_mem = available_memory[0] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - continue - worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) - if worker_max < auto_fit_max: - auto_fit_max = worker_max - limiting_worker_mem = avail_mem - - if auto_fit_max <= 0: - raise ValueError( - "Cannot auto-fit max_model_len: not enough GPU memory available " - "to serve even a single token. Try increasing `gpu_memory_utilization`." - ) - - kv_cache_memory = ( - f"{format_gib(limiting_worker_mem)} GiB {_kv_cache_device_name(vllm_config)}" - ) - if vllm_config.attention_config.hisparse_config is not None: - host_pool_gib = format_gib(get_hisparse_host_pool_bytes(vllm_config)) - kv_cache_memory += f" + {host_pool_gib} GiB host" - - if auto_fit_max >= original_max: - # The model's full context length fits in memory - logger.info_once( - "Auto-fit max_model_len: full model context length %d fits in " - "available KV cache memory (%s)", - original_max, - kv_cache_memory, - ) - else: - # Need to reduce max_model_len to fit in memory - vllm_config.model_config.max_model_len = auto_fit_max - logger.info_once( - "Auto-fit max_model_len: reduced from %d to %d to fit in " - "available KV cache memory (%s)", - original_max, - auto_fit_max, - kv_cache_memory, - ) - - -def _project_kv_cache_groups_to_worker( - global_kv_cache_groups: list[KVCacheGroupSpec], - worker_spec: dict[str, KVCacheSpec], -) -> list[KVCacheGroupSpec]: - """Projects global KV cache groups onto a single worker's assigned layers. - - In pipeline parallelism, each worker only owns a subset of layers. This - function filters the global groups to include only layers present on the - given worker, adjusting UniformTypeKVCacheSpecs accordingly. - - Args: - global_kv_cache_groups: The global KV cache groups for the whole model. - worker_spec: The KV cache spec of each layer on this worker. - - Returns: - The projected KV cache groups containing only this worker's layers. - - """ - projected_groups: list[KVCacheGroupSpec] = [] - for group in global_kv_cache_groups: - worker_layer_names = [ - layer_name for layer_name in group.layer_names if layer_name in worker_spec - ] - if len(worker_layer_names) == len(group.layer_names): - projected_groups.append(group) - continue - group_spec = group.kv_cache_spec - if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs): - group_spec = UniformTypeKVCacheSpecs( - block_size=group_spec.block_size, - kv_cache_specs={ - layer_name: group_spec.kv_cache_specs[layer_name] - for layer_name in worker_layer_names - }, - ) - projected_groups.append( - KVCacheGroupSpec( - worker_layer_names, - group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), - ) - ) - return projected_groups - - -def get_kv_cache_configs( - vllm_config: VllmConfig, - kv_cache_specs: list[dict[str, KVCacheSpec]], - available_memory: list[int], -) -> list[KVCacheConfig]: - """Generates the KV cache configurations for a model. - Since we use a shared centralized controller for all workers, we need the - `kv_cache_config` to be consistent across all workers to make sure - the KV cache allocation can be applied to all workers. However, different - workers may have different memory available, and different type of layers - (when pipeline parallel is enabled). To handle the difference between - workers, the current implementation is: - 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for - the whole model. - 2. Generate the KV cache groups based on the layer ratio of the whole model. - This also handles spec unification for hybrid models. - 3. Handle auto-fit max_model_len and memory checks using per-worker - projected groups to account for PP sharding. - 4. Generate the KV cache configs for each worker based on the KV cache - grouping strategy. (This is reasonable because the layer ratio of - different PP stages are similar.) - 5. Change the num_blocks of each worker to the smallest among all workers - and shrink tensor sizes proportionally to avoid allocating unused memory. - - Args: - vllm_config: The global VllmConfig - kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker. - available_memory: Memory available for KV cache in bytes for each - worker. - - Returns: - The generated KVCacheConfigs for each worker. - - """ - # Merge the KV cache specs of all workers. Different PP stages may have - # different layer names, and different TP ranks of the same PP stage should - # have the same KV cache spec. - merged_kv_cache_specs: dict[str, KVCacheSpec] = {} - for kv_cache_spec_one_worker in kv_cache_specs: - for layer_name, layer_spec in kv_cache_spec_one_worker.items(): - if layer_name not in merged_kv_cache_specs: - merged_kv_cache_specs[layer_name] = layer_spec - else: - assert merged_kv_cache_specs[layer_name] == layer_spec, ( - "The KV cache specs for the same layer are different " - "across workers. This is not supported yet." - ) - - # Check if the KV cache specs are registered correctly. - # This is to prevent that some layers are initialized with unregistered specs. - KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) - - # When speculating with more than 1 speculative module (e.g. multi-layered MTP) - # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) - for layer_name, layer_spec in merged_kv_cache_specs.items(): - if isinstance(layer_spec, SlidingWindowSpec): - merged_kv_cache_specs[layer_name] = replace( - layer_spec, extra_retained_tokens=extra_retained_tokens - ) - - # Get global KV cache groups. This also handles spec unification for - # hybrid models when disable_hybrid_kv_cache_manager is enabled. - # After this call, merged_kv_cache_specs may be modified in-place. - global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) - - # If original_max_model_len was -1, automatically - # determine the maximum model length that fits in available GPU memory. - # We use per-worker projected groups to account for PP sharding. - projected_groups_per_worker = [ - _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec) - for worker_spec in kv_cache_specs - ] - - # If `num_gpu_blocks_override` is set, the cache size that will actually - # be allocated is decoupled from the profiled `available_memory`: - # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps - # `num_blocks` to the override. Reflect that in `available_memory` here so - # auto-fit, the admission check, and the per-worker config builder all - # plan against the same effective capacity. - override = vllm_config.cache_config.num_gpu_blocks_override - if override is not None: - adjusted_memory: list[int] = [] - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): - if not groups: - adjusted_memory.append(avail_mem) - continue - 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, - override, - ) - adjusted_memory.append(override * bytes_per_block) - available_memory = adjusted_memory - - # Reserve the null block BlockPool permanently holds back, so auto-fit and - # the capacity check both plan against usable blocks. Allocation below - # still uses the full memory. - check_memory = [ - 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) - ] - - if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) - - # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, check_memory): - if not groups: - continue - _check_enough_kv_cache_memory( - avail_mem, - partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), - vllm_config.model_config.max_model_len, - partial(_estimate_max_model_len_from_groups, vllm_config, groups), - ) - - kv_cache_configs: list[KVCacheConfig] = [] - for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip( - projected_groups_per_worker, kv_cache_specs, available_memory - ): - assert sum(len(group.layer_names) for group in projected_groups) == len( - kv_cache_spec_one_worker - ), "Some layers are not assigned to any group." - kv_cache_configs.append( - get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker - ) - ) - - # Change the num_blocks of each rank to the smallest among all ranks. - # We also need to shrink the tensor size proportionally to avoid - # allocating unused memory. - min_num_blocks = min( - kv_cache_config.num_blocks for kv_cache_config in kv_cache_configs - ) - for i, kv_cache_config in enumerate(kv_cache_configs): - if kv_cache_config.num_blocks == min_num_blocks: - continue - # Re-plan with exactly the memory the smallest rank can afford, so - # 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(vllm_config, groups), - ) - - for kv_cache_config in kv_cache_configs: - kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( - kv_cache_config.kv_cache_groups, - vllm_config.parallel_config.tensor_parallel_size, - vllm_config.parallel_config.decode_context_parallel_size, - ) - - return kv_cache_configs - - class BlockHashListWithBlockSize: """Convert block-hash granularity from `hash_block_size` to `target_block_size`. Used when KV cache groups have different block sizes: `hash_block_size` diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 35cc8cc1860c..d8c976ff693b 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -50,10 +50,10 @@ from vllm.utils.network_utils import make_zmq_socket from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.attention.backends.utils import resolve_kv_cache_layout +from vllm.v1.core.kv_cache_planning import get_kv_cache_configs from vllm.v1.core.kv_cache_utils import ( BlockHash, generate_scheduler_kv_cache_config, - get_kv_cache_configs, get_request_block_hasher, init_none_hash, resolve_cache_hit_alignment_tokens, diff --git a/vllm/v1/hisparse/layout.py b/vllm/v1/hisparse/layout.py index 96d93650ad29..0a72427d4233 100644 --- a/vllm/v1/hisparse/layout.py +++ b/vllm/v1/hisparse/layout.py @@ -71,7 +71,7 @@ def get_hisparse_kv_cache_groups( if not mla_specs: return None - from vllm.v1.core.kv_cache_utils import get_kv_cache_groups + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups mla_group_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) assert mla_group_spec is not None @@ -323,7 +323,7 @@ def get_hisparse_kv_cache_config( available_memory: int, host_budget: int, ) -> KVCacheConfig: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( _get_kv_cache_bytes_per_block, may_override_num_blocks, validate_kv_cache_layout, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 73be62319a64..3a6409db5546 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -989,7 +989,7 @@ def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: """Allocate the smallest KV cache that still lets every graph be captured.""" - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e19c1b84f457..90059fd39f5b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6206,7 +6206,7 @@ def profile_run(self, randomize_inputs: bool = False) -> None: gc.collect() def _init_minimal_kv_cache_for_profiling(self) -> None: - from vllm.v1.core.kv_cache_utils import ( + from vllm.v1.core.kv_cache_planning import ( get_kv_cache_config_from_groups, get_kv_cache_groups, ) From c5422f4c466afd765b32b662151d0f37dad7baa5 Mon Sep 17 00:00:00 2001 From: RAHUL KUMAR <55033230+pctablet505@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:28:32 +0530 Subject: [PATCH 02/23] [Test][Renderer] Cover the reasoning -> reasoning_content mirror in chat message parsing (#56129) Signed-off-by: pctablet505 <55033230+pctablet505@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 --- tests/renderers/test_chat_utils.py | 46 ++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/renderers/test_chat_utils.py b/tests/renderers/test_chat_utils.py index d0a9b9e6c7a8..47fe49a63e6e 100644 --- a/tests/renderers/test_chat_utils.py +++ b/tests/renderers/test_chat_utils.py @@ -27,6 +27,7 @@ ChatCompletionMessageParam, ConversationMessage, _load_embeds_dict, + _parse_chat_message_content, _parse_metadata_array, _postprocess_messages, parse_chat_messages, @@ -3196,3 +3197,48 @@ def test_validate_chat_template_rejects_invalid_type(): ) as exc_info: validate_chat_template(123) assert exc_info.value.parameter == "chat_template" + + +def _parse_one(message) -> ConversationMessage: + # Plain-text content only reads these two tracker attributes. + tracker = MagicMock() + parser = tracker.create_parser.return_value + parser.model_config.enable_prompt_embeds = False + parser.mm_placeholder_storage.return_value = {} + result = _parse_chat_message_content( + message, tracker, content_format="string", interleave_strings=False + ) + assert len(result) == 1 + return result[0] + + +def test_assistant_reasoning_is_mirrored_to_reasoning_content(): + """``reasoning`` is mirrored to ``reasoning_content`` for legacy templates.""" + conv_msg = _parse_one( + { + "role": "assistant", + "content": "4", + "reasoning": "2+2 equals 4", + } + ) + + assert conv_msg.get("reasoning") == "2+2 equals 4" + assert conv_msg.get("reasoning_content") == "2+2 equals 4" + + +def test_assistant_without_reasoning_sets_neither_field(): + """No reasoning in, no reasoning keys out: no empty thinking block in templates.""" + conv_msg = _parse_one({"role": "assistant", "content": "Hi"}) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg + + +def test_non_assistant_reasoning_is_not_mirrored(): + """Only assistant turns carry reasoning; a user turn must not.""" + conv_msg = _parse_one( + {"role": "user", "content": "Hi", "reasoning": "should be ignored"} + ) + + assert "reasoning" not in conv_msg + assert "reasoning_content" not in conv_msg From 657a2e94174fadcae7ae729b076d61001509d804 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:22:48 -0700 Subject: [PATCH 03/23] [Rust Frontend] Accept whole-number floats for integer tool arguments (#61299) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- rust/src/parser/src/tool/parameters.rs | 31 +++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/rust/src/parser/src/tool/parameters.rs b/rust/src/parser/src/tool/parameters.rs index 391ca90f105f..afb3d4a9f2cb 100644 --- a/rust/src/parser/src/tool/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -294,7 +294,9 @@ fn try_convert_value(param_type: &JsonParamType<'_>, input: &ParamInput) -> Opti fn try_convert_text_value(param_type: &JsonParamType<'_>, value: &str) -> Option { match param_type { JsonParamType::String => Some(Value::String(value.to_string())), - JsonParamType::Integer => convert_integer_text(value), + JsonParamType::Integer => { + convert_integer_text(value).or_else(|| convert_whole_number_text(value)) + } JsonParamType::Number => convert_number_text(value), JsonParamType::Boolean => try_convert_boolean(value), JsonParamType::Object(_) if value.trim().is_empty() => Some(Value::Object(Map::new())), @@ -401,6 +403,14 @@ pub(crate) fn convert_integer_text(value: &str) -> Option { .map(Value::Number) } +/// Convert a whole number written as a float (`3.0`) to a JSON integer, since JSON Schema +/// counts any number with a zero fractional part as an `integer`. +fn convert_whole_number_text(value: &str) -> Option { + let number = without_digit_separators(value.trim())?.parse::().ok()?; + let in_range = number >= i64::MIN as f64 && number < i64::MAX as f64; + (number.fract() == 0.0 && in_range).then(|| Value::Number(Number::from(number as i64))) +} + /// Convert raw text to a JSON number, as for a `number` parameter. pub(crate) fn convert_number_text(value: &str) -> Option { try_convert_number(&without_digit_separators(value.trim())?) @@ -591,6 +601,25 @@ mod tests { } } + #[test] + fn integers_accept_whole_number_floats() { + let params = ToolSchema::from_schema(json!({ + "type": "object", + "properties": { + "count": { "type": "integer" }, + "either": { "type": ["integer", "number"] } + } + })); + + assert_eq!(params.convert("count", text("3.0")), json!(3)); + assert_eq!(params.convert("count", text(" -2.0 ")), json!(-2)); + assert_eq!(params.convert("count", text("1_000.0")), json!(1000)); + assert_eq!(params.convert("either", text("3.5")), json!(3.5)); + for invalid in ["3.5", "inf", "NaN", "1e300"] { + assert_eq!(params.convert("count", text(invalid)), json!(invalid)); + } + } + #[test] fn nullable_schemas_accept_null() { let params = ToolSchema::from_schema(json!({ From 28d679e762a0b6ed560d767a76ebe02699bf5c46 Mon Sep 17 00:00:00 2001 From: dskwe Date: Tue, 13 Oct 2026 12:23:45 +0800 Subject: [PATCH 04/23] [Bugfix] Preserve independent tokenizer code_revision in snapshot resolution (#61272) Signed-off-by: Yuan Li Co-authored-by: Cyrus Leung --- tests/tokenizers_/test_basic.py | 85 +++++++++++++++++++++++++++++++++ vllm/tokenizers/registry.py | 13 ++--- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/tests/tokenizers_/test_basic.py b/tests/tokenizers_/test_basic.py index a46be988e5c8..a91124549457 100644 --- a/tests/tokenizers_/test_basic.py +++ b/tests/tokenizers_/test_basic.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from pathlib import Path from typing import _get_protocol_attrs # type: ignore from unittest.mock import patch import pytest from transformers import ( + AutoTokenizer, PreTrainedTokenizerBase, TokenizersBackend, tokenization_utils_base, @@ -13,6 +15,7 @@ from vllm.tokenizers import TokenizerLike, get_tokenizer from vllm.tokenizers.hf import HfTokenizer from vllm.tokenizers.mistral import MistralTokenizer +from vllm.tokenizers.registry import resolve_tokenizer_args def _get_missing_attrs(obj: object, target: type): @@ -58,6 +61,88 @@ def test_tokenizer_revision(tokenizer_name: str): get_tokenizer(tokenizer_name, revision="never") +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision, error_match", + [ + pytest.param("never", "not a valid git identifier", id="unknown-revision"), + pytest.param( + "73293ac504e43f4a895688b57c65cebe24f01a38", + r"tokenization_fast\.py", + id="revision-without-code", + ), + ], +) +def test_tokenizer_code_revision_rejects_missing_code( + code_revision: str, error_match: str +): + """A data snapshot must not shadow the requested remote-code revision.""" + with pytest.raises(OSError, match=error_match): + get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision="08a35780b7a78794436d8654aafa26ee0165c149", + code_revision=code_revision, + trust_remote_code=True, + ) + + +@pytest.mark.skip_global_cleanup +def test_tokenizer_code_revision_distinct_from_revision(): + """Distinct code/data revisions: the resolver must keep the Hub ID. + + Dynamic code is read from a local directory as-is, so resolving to a + data-revision snapshot would silently shadow an independent code + revision (the same contract already enforced for processors). + """ + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + repo = "hf-internal-testing/test_dynamic_tokenizer" + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision=revision, + trust_remote_code=True, + ) + assert Path(name).is_dir() + assert kwargs["code_revision"] == revision + + _, name, _, kwargs = resolve_tokenizer_args( + repo, + tokenizer_mode="hf", + revision=revision, + code_revision="73293ac504e43f4a895688b57c65cebe24f01a38", + trust_remote_code=True, + ) + assert name == repo + assert kwargs["code_revision"] == "73293ac504e43f4a895688b57c65cebe24f01a38" + + +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize( + "code_revision", [None, "08a35780b7a78794436d8654aafa26ee0165c149"] +) +def test_tokenizer_code_revision_default_or_matching(code_revision: str | None): + revision = "08a35780b7a78794436d8654aafa26ee0165c149" + tokenizer = get_tokenizer( + "hf-internal-testing/test_dynamic_tokenizer", + tokenizer_mode="hf", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + reference = AutoTokenizer.from_pretrained( + "hf-internal-testing/test_dynamic_tokenizer", + revision=revision, + code_revision=code_revision, + trust_remote_code=True, + ) + assert isinstance(tokenizer, PreTrainedTokenizerBase) + assert tokenizer.encode("hello world") == reference.encode("hello world") + assert tokenizer.name_or_path == "hf-internal-testing/test_dynamic_tokenizer" + _assert_tokenizer_like(tokenizer) + + def test_hub_tokenizer_loads_from_local_snapshot(): """Transformers makes a Hub request per file when given a repo ID.""" with patch.object( diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index a7adf4597f6a..8b1863023c99 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -165,12 +165,13 @@ def resolve_tokenizer_args( "mistral_format=True is not supported with tokenizer_mode='hf'" ) kwargs["mistral_format"] = False - tokenizer_name = get_non_weight_snapshot_path( - str(tokenizer_name), - revision=revision, - cache_dir=download_dir, - token=kwargs.get("token"), - ) + if kwargs.get("code_revision") in (None, revision): + tokenizer_name = get_non_weight_snapshot_path( + str(tokenizer_name), + revision=revision, + cache_dir=download_dir, + token=kwargs.get("token"), + ) return tokenizer_mode, tokenizer_name, args, kwargs From 69d00bc2645cac2144096be3ef1a12a73f95f366 Mon Sep 17 00:00:00 2001 From: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:55:20 -0700 Subject: [PATCH 05/23] [Bugfix] Accept whole-number floats for integer tool arguments in the response parser (#61297) Signed-off-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Signed-off-by: Bugen Zhao Co-authored-by: Bugen Zhao --- tests/parser/test_chat_parsing.py | 8 ++++++++ vllm/parser/chat_parsing/response_parser.py | 9 ++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/parser/test_chat_parsing.py b/tests/parser/test_chat_parsing.py index 679d97f15b20..e3b8dba17220 100644 --- a/tests/parser/test_chat_parsing.py +++ b/tests/parser/test_chat_parsing.py @@ -1555,6 +1555,14 @@ def test_coerce_tool_calls_casts_declared_types(self): def test_coerce_falls_back_to_raw_on_failure(self): self.assertEqual(_coerce("not-a-number", ("integer",)), "not-a-number") + def test_coerce_integer_accepts_whole_number_floats(self): + self.assertEqual(_coerce("3.0", ("integer",)), 3) + self.assertIsInstance(_coerce("3.0", ("integer",)), int) + self.assertEqual(_coerce("3.5", ("integer",)), "3.5") + self.assertEqual(_coerce("3.5", ("integer", "number")), 3.5) + self.assertEqual(_coerce("inf", ("integer",)), "inf") + self.assertEqual(_coerce("NaN", ("integer",)), "NaN") + def test_coerce_handles_any_of_and_null(self): self.assertEqual(_coerce("5", ("integer", "null")), 5) self.assertIsNone(_coerce("null", ("integer", "null"))) diff --git a/vllm/parser/chat_parsing/response_parser.py b/vllm/parser/chat_parsing/response_parser.py index 556df803a70a..0416fe0fe3ea 100644 --- a/vllm/parser/chat_parsing/response_parser.py +++ b/vllm/parser/chat_parsing/response_parser.py @@ -50,7 +50,14 @@ def _coerce(raw: str, types: tuple[str, ...]) -> Any: for type_name in types: try: if type_name == "integer": - return int(raw) + try: + return int(raw) + except ValueError: + # JSON Schema integers include whole numbers written as floats, e.g. "3.0" + number = float(raw) + if not number.is_integer(): + continue + return int(number) if type_name == "number": number = float(raw) if number != number or number in (float("inf"), float("-inf")): From 5555a2d8f7116b2b4877f66516f4128176cf722a Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:02:12 +0900 Subject: [PATCH 06/23] [Bugfix][Model] Fix EXAONE-4.5 image inputs: ViT received un-normalised pixels (#61104) Signed-off-by: godmook Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com> --- vllm/model_executor/models/exaone4_5.py | 6 ++++++ vllm/model_executor/models/exaone4_5_mtp.py | 1 + 2 files changed, 7 insertions(+) diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py index 5a53e60b7d5e..ad4049c8eeb0 100644 --- a/vllm/model_executor/models/exaone4_5.py +++ b/vllm/model_executor/models/exaone4_5.py @@ -36,6 +36,7 @@ from vllm.distributed import utils as dist_utils from vllm.logger import init_logger from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.quantization import QuantizationConfig @@ -272,6 +273,7 @@ def __init__( vision_config: Exaone4_5_VisionConfig, norm_eps: float = 1e-6, quant_config: QuantizationConfig | None = None, + input_norm: nn.Module | None = None, prefix: str = "", use_data_parallel: bool = False, ) -> None: @@ -279,6 +281,7 @@ def __init__( vision_config=vision_config, norm_eps=norm_eps, quant_config=quant_config, + input_norm=input_norm, prefix=prefix, ) depth = vision_config.depth @@ -335,7 +338,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.config = config + self.model_config = vllm_config.model_config self.multimodal_config = multimodal_config + self.video_pruning_rate = multimodal_config.video_pruning_rate self.is_multimodal_pruning_enabled = ( multimodal_config.is_multimodal_pruning_enabled() ) @@ -345,6 +350,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vision_config, norm_eps=getattr(config, "rms_norm_eps", 1e-6), quant_config=self.quant_config, + input_norm=build_mm_input_norm(self.model_config), prefix=maybe_prefix(prefix, "visual"), use_data_parallel=self.use_data_parallel, ) diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py index d9c7960010a5..41235929867a 100644 --- a/vllm/model_executor/models/exaone4_5_mtp.py +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -75,6 +75,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.layers = nn.ModuleList( Exaone4DecoderLayer( text_config, + cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=f"{prefix}.layers.{idx}", ) From b337523c118d1cdb211732199f0a7d5b1e942bb9 Mon Sep 17 00:00:00 2001 From: Chaojun Zhang Date: Tue, 13 Oct 2026 13:12:10 +0800 Subject: [PATCH 07/23] [XPU] [Test] Extend spawn-safe sitecustomize monkeypatch to XPU (#57911) Signed-off-by: Chaojun Zhang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: xiangdong <40376367+zxd1997066@users.noreply.github.com> --- .buildkite/intel_jobs/distributed_intel.yaml | 33 ++++++++++++++++++++ tests/v1/shutdown/conftest.py | 15 +++++++-- tests/v1/shutdown/test_forward_error.py | 8 ++--- tests/v1/shutdown/test_startup_error.py | 8 ++--- 4 files changed, 53 insertions(+), 11 deletions(-) create mode 100644 .buildkite/intel_jobs/distributed_intel.yaml diff --git a/.buildkite/intel_jobs/distributed_intel.yaml b/.buildkite/intel_jobs/distributed_intel.yaml new file mode 100644 index 000000000000..54800b4103d2 --- /dev/null +++ b/.buildkite/intel_jobs/distributed_intel.yaml @@ -0,0 +1,33 @@ +group: Distributed Intel +depends_on: + - image-build-xpu +steps: +- label: ":intel: (BMG) Distributed Torchrun + Shutdown" + key: distributed-torchrun-shutdown-tests-2-gpus + timeout_in_minutes: 50 + device: intel_gpu + agent_tags: + label: production + gpu: 2+ + mem: 24+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/distributed/ + - "!vllm/distributed/kv_transfer/" + - vllm/engine/ + - vllm/executor/ + - vllm/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && pytest -v -s v1/shutdown' diff --git a/tests/v1/shutdown/conftest.py b/tests/v1/shutdown/conftest.py index b276d0fab267..ef54b188fefa 100644 --- a/tests/v1/shutdown/conftest.py +++ b/tests/v1/shutdown/conftest.py @@ -10,9 +10,18 @@ @pytest.fixture -def rocm_sitecustomize_factory(monkeypatch, tmp_path: Path): - """Return a function that installs a given sitecustomize payload.""" - if not current_platform.is_rocm(): +def spawn_sitecustomize_factory(monkeypatch, tmp_path: Path): + """Return a function that installs a given sitecustomize payload. + + Both ROCm and XPU force the `spawn` multiprocessing start method (see + `vllm.platforms.rocm`/`vllm.platforms.xpu`), which re-imports modules in + child processes instead of inheriting them via `fork`'s copy-on-write + memory. A plain `monkeypatch.setattr` in the parent test process is + therefore invisible to spawned EngineCore/worker subprocesses on these + platforms, so the payload must be installed via `sitecustomize.py` on + `PYTHONPATH` instead. + """ + if not (current_platform.is_rocm() or current_platform.is_xpu()): return lambda _: None def install(lines: Iterable[str]) -> None: diff --git a/tests/v1/shutdown/test_forward_error.py b/tests/v1/shutdown/test_forward_error.py index 8bc09a64cac9..c40a41cc3b41 100644 --- a/tests/v1/shutdown/test_forward_error.py +++ b/tests/v1/shutdown/test_forward_error.py @@ -41,21 +41,21 @@ def evil_forward(self, *args, **kwargs): @pytest.fixture -def rocm_evil_forward(rocm_sitecustomize_factory): +def evil_forward_fixture(spawn_sitecustomize_factory): lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", "from vllm.model_executor.models.llama import LlamaForCausalLM", inspect.getsource(evil_forward), f"LlamaForCausalLM.forward = {evil_forward.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.asyncio @pytest.mark.parametrize("tensor_parallel_size", [2, 1]) @pytest.mark.parametrize("model", MODELS) async def test_async_llm_model_error( - monkeypatch, rocm_evil_forward, tensor_parallel_size: int, model: str + monkeypatch, evil_forward_fixture, tensor_parallel_size: int, model: str ) -> None: """Test that AsyncLLM propagates a forward pass error and frees memory. @@ -118,7 +118,7 @@ async def generate(request_id: str): @pytest.mark.parametrize("model", MODELS) def test_llm_model_error( monkeypatch, - rocm_evil_forward, + evil_forward_fixture, tensor_parallel_size: int, enable_multiprocessing: bool, model: str, diff --git a/tests/v1/shutdown/test_startup_error.py b/tests/v1/shutdown/test_startup_error.py index 6d6d0d91e77f..d70f69409ee5 100644 --- a/tests/v1/shutdown/test_startup_error.py +++ b/tests/v1/shutdown/test_startup_error.py @@ -30,7 +30,7 @@ def evil_method(self, *args, **kwargs): @pytest.fixture -def rocm_evil_method(rocm_sitecustomize_factory, request): +def evil_method_fixture(spawn_sitecustomize_factory, request): failing_method = request.getfixturevalue("failing_method") lines = [ "from vllm.distributed import get_tensor_model_parallel_rank", @@ -38,7 +38,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): inspect.getsource(evil_method), f"LlamaForCausalLM.{failing_method} = {evil_method.__name__}", ] - rocm_sitecustomize_factory(lines) + spawn_sitecustomize_factory(lines) @pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC) @@ -47,7 +47,7 @@ def rocm_evil_method(rocm_sitecustomize_factory, request): @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_async_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, failing_method: str, @@ -84,7 +84,7 @@ def test_async_llm_startup_error( @pytest.mark.parametrize("failing_method", ["forward", "load_weights"]) def test_llm_startup_error( monkeypatch, - rocm_evil_method, + evil_method_fixture, model: str, tensor_parallel_size: int, enable_multiprocessing: bool, From f7ba007b5e766128468fed38c5da9777b521b8ee Mon Sep 17 00:00:00 2001 From: Divakar Verma <137818590+divakar-amd@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:13:33 -0400 Subject: [PATCH 08/23] [ROCm][MoE] Split AiterExperts by weight-quantization key (#60322) Signed-off-by: Divakar Verma Co-authored-by: Claude Sonnet 5 --- .buildkite/test-amd.yaml | 27 +++++- docs/design/moe_kernel_features.md | 2 +- .../moe/modular_kernel_tools/common.py | 8 +- .../moe/modular_kernel_tools/mk_objects.py | 40 +++++--- .../moe/test_modular_kernel_combinations.py | 44 ++++++--- tests/kernels/moe/test_mori_moe.py | 8 +- tests/kernels/moe/test_rocm_aiter_moe.py | 4 +- tests/kernels/moe/test_topk_softplus_sqrt.py | 2 +- .../layers/fused_moe/__init__.py | 12 ++- .../fused_moe/experts/rocm_aiter_moe.py | 97 ++++++++++++++----- .../layers/fused_moe/oracle/fp8.py | 6 +- .../layers/fused_moe/oracle/mxfp4.py | 8 +- .../layers/fused_moe/oracle/unquantized.py | 4 +- 13 files changed, 198 insertions(+), 64 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 5145f5a272ff..d815ee9268b2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1528,7 +1528,32 @@ steps: commands: - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":amd: (MI300) MoE Kernels Multi-GPU" # TBD + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_2 + num_gpus: 2 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe/test_modular_kernel_combinations.py + - tests/kernels/moe/modular_kernel_tools/ + - tests/kernels/utils.py + - vllm/model_executor/layers/fused_moe/ + - vllm/distributed/device_communicators/ + - vllm/distributed/parallel_state.py + - vllm/envs.py + - vllm/config + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "multigpu" - label: ":amd: (MI300) vLLM IR" # TBD timeout_in_minutes: 50 diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 6c71aa1f46d4..2528bc9ea459 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.experts.marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | hpc | standard | fp8 | G(128),T | silu | Y | Y | [`HPCExperts`][vllm.model_executor.layers.fused_moe.hpc_moe.HPCExperts] | -| rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | +| rocm aiter moe | standard | none,
mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai,
situ | Y | Y | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.rocm_aiter_fused_experts],
[`AiterUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterUnquantizedExperts],
[`AiterFp8BlockExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8BlockExperts],
[`AiterFp8TensorExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8TensorExperts],
[`AiterFp8ChannelExperts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterFp8ChannelExperts],
[`AiterMxfp4Experts`][vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe.AiterMxfp4Experts] | | cpu_moe | standard | N/A | N/A | silu, gelu,
gelu_tanh,
swigluoai | Y | N | [`X86CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.X86CPUUnquantizedExperts],
[`ArmCPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.ArmCPUUnquantizedExperts],
[`CPUUnquantizedExperts`][vllm.model_executor.layers.fused_moe.experts.cpu_moe.CPUUnquantizedExperts] | | naive batched4 | batched | int8,
fp8 | G,A,T | silu, gelu | 6 | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] | diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 5f6349a32880..1f79b62bf8c8 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -167,6 +167,9 @@ def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: vllm_config.model_config = SimpleNamespace( enforce_eager=True, is_moe=True, + enable_sleep_mode=False, + sleep_mode_backend="cumem", + sleep_mode_offload_cudagraph=False, ) vllm_config.parallel_config.data_parallel_size = self.world_size vllm_config.parallel_config.enable_expert_parallel = True @@ -722,6 +725,9 @@ def _maybe_convert_weights_for_experts( rank_weights: WeightTensors, ) -> WeightTensors: """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK).""" + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( Fp8MoeBackend, convert_to_fp8_moe_kernel_format, @@ -734,7 +740,7 @@ def _maybe_convert_weights_for_experts( # intermediate sizes that are not a multiple of 256, but only for the fp8 per-tensor # and per-token schemes. Those keep JIT-compiling until # https://github.com/ROCm/aiter/issues/5766 is fixed. - if fe_name == "AiterExperts" and ( + if issubclass(fe_type, AiterExperts) and ( config.quant_dtype is None or config.quant_block_shape is not None ): return _shuffle_weights_for_aiter(rank_weights) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 6ca644c5fce1..56119ffddeb7 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -93,6 +93,7 @@ class ExpertInfo: ] common_float_and_int_types = common_float_types + [torch.int8] nvfp4_types = ["nvfp4"] +mxfp4_types = ["mxfp4"] fp8_types = [current_platform.fp8_dtype()] @@ -318,20 +319,37 @@ def expert_info(kind) -> ExpertInfo: if has_aiter(): from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) - register_experts( - AiterExperts, - standard_format, - # AiterExperts also supports the fully-unquantized (None, None) - # scheme (see SUPPORTED_W_A in rocm_aiter_moe.py), not just fp8. - common_float_types, - blocked_quantization_support=True, - needs_aiter=True, - ) + # Registration dtypes gate which quant configs the modular-kernel test + # harness will ever pair with a given class (see Config.is_valid()): + # they must match what each subclass's _supports_quant_scheme() actually + # declares, or the harness will generate a case the class rejects. + for aiter_experts_cls, aiter_supported_dtypes in ( + (AiterUnquantizedExperts, common_float_types), + (AiterFp8BlockExperts, fp8_types), + (AiterFp8TensorExperts, fp8_types), + (AiterFp8ChannelExperts, fp8_types), + (AiterMxfp4Experts, mxfp4_types), + ): + register_experts( + aiter_experts_cls, + standard_format, + aiter_supported_dtypes, + blocked_quantization_support=True, + needs_aiter=True, + ) else: - AiterExperts = None + AiterUnquantizedExperts = None + AiterFp8BlockExperts = None + AiterFp8TensorExperts = None + AiterFp8ChannelExperts = None + AiterMxfp4Experts = None if has_deep_gemm() and is_deep_gemm_supported(): register_experts( diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index f29c22085240..da643befecd9 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -16,7 +16,13 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts +from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterUnquantizedExperts, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTensorSym, kFp8StaticTensorSym, @@ -80,10 +86,10 @@ def assert_aiter_quant_scheme_case(config: Config) -> None: indirectly through the general quant-config sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return - if config.quant_config is None: + if config.quant_config is None or config.quant_dtype is None: w_key, a_key = None, None else: w_key, a_key = config.fp8_quant_key_pair() @@ -101,7 +107,7 @@ def assert_aiter_activation_case(config: Config) -> None: AiterExperts being reached only indirectly through the general activation sweep. See https://github.com/vllm-project/vllm/issues/54966.""" fe_cls = config.fused_experts_type - if fe_cls is not AiterExperts: + if not issubclass(fe_cls, AiterExperts): return assert fe_cls._supports_activation(config.activation), ( @@ -156,7 +162,7 @@ def rank_worker( if ( topk == 1 and config.supports_apply_weight_on_input() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_block_shape is not None ): print( @@ -172,7 +178,7 @@ def rank_worker( # https://github.com/vllm-project/vllm/issues/57029 if ( config.world_size > 1 - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and getattr(config.prepare_finalize_type, "__name__", "") in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize") ): @@ -210,7 +216,7 @@ def rank_worker( is_aiter_fp8 = ( _cp.is_rocm() - and config.fused_experts_type is AiterExperts + and issubclass(config.fused_experts_type, AiterExperts) and config.quant_config is not None ) if is_aiter_fp8: @@ -270,7 +276,7 @@ def is_nyi_config(config: Config) -> bool: return True if config.activation != MoEActivation.SILU: - if config.fused_experts_type is not AiterExperts: + if not issubclass(config.fused_experts_type, AiterExperts): return True # AITER-only for this axis, for now if config.quant_dtype is not None: return True # unquantized-only for this axis, for now @@ -591,7 +597,7 @@ def test_aiter_moe_sorting_backend_dispatch_env_matrix( dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -671,7 +677,7 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) dtype=torch.bfloat16, quant_config=None, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=AiterUnquantizedExperts, world_size=1, ) assert config.is_valid()[0] @@ -738,6 +744,20 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) ] +def _padding_matrix_fe_cls( + quant_config: TestMoEQuantConfig | None, +) -> type[AiterExperts]: + """Map a _PADDING_QUANT_CONFIGS entry to the AiterExperts leaf class that + supports it (mirrors Config.fp8_quant_key_pair's scheme derivation).""" + if quant_config is None: + return AiterUnquantizedExperts + if quant_config.block_shape is not None: + return AiterFp8BlockExperts + if quant_config.per_out_ch_quant: + return AiterFp8ChannelExperts + return AiterFp8TensorExperts + + def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor: """Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real gate/up halves of `real_per_half` size each.""" @@ -950,7 +970,7 @@ def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig | dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None, intermediate_size_per_partition_unpadded=( @@ -1113,7 +1133,7 @@ def test_aiter_moe_token_padding_garbage_rows( dtype=torch.bfloat16, quant_config=quant_config, prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, - fused_experts_type=AiterExperts, + fused_experts_type=_padding_matrix_fe_cls(quant_config), world_size=1, ) assert config.is_valid()[0] diff --git a/tests/kernels/moe/test_mori_moe.py b/tests/kernels/moe/test_mori_moe.py index 9869b6bc5467..f06afc8bd63e 100644 --- a/tests/kernels/moe/test_mori_moe.py +++ b/tests/kernels/moe/test_mori_moe.py @@ -42,7 +42,10 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): from vllm._aiter_ops import rocm_aiter_ops - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( + AiterFp8ChannelExperts, + AiterUnquantizedExperts, + ) from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import ( MoriPrepareAndFinalize, ) @@ -94,7 +97,8 @@ def _mori_worker(pgi, vllm_config, cpu_group, m, n, k, fp8, graph): assert isinstance(prepare, MoriPrepareAndFinalize) assert prepare.use_fp8_dispatch == fp8 assert prepare.num_dispatchers() == 2 - experts = AiterExperts(config, quant) + experts_cls = AiterFp8ChannelExperts if fp8 else AiterUnquantizedExperts + experts = experts_cls(config, quant) assert not experts.expects_unquantized_inputs kernel = FusedMoEKernel(prepare, experts) local_w1, local_w2 = rocm_aiter_ops.shuffle_weights( diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index abd46753515e..af2f7beb1bd5 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -848,13 +848,13 @@ def test_activation_method_enum_values(): def test_aiter_mxfp4_quant_scheme_support_matches_gfx950(): """AITER MXFP4 MoE support should stay gfx950-only.""" from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kMxfp4Static, ) - assert AiterExperts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() + assert AiterMxfp4Experts._supports_quant_scheme(kMxfp4Static, None) is on_gfx950() @pytest.mark.skipif(not on_gfx950(), reason="gfx950 ROCm only") diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index bc438c66931b..a6f7f3080583 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -87,7 +87,7 @@ def _torch_topk_softplus_sqrt( "backend,expected_experts", [ ("aiter_triton_mxfp4_bf16", "AiterW4A16ExpertsMonolithic"), - ("aiter", "AiterExperts"), + ("aiter", "AiterMxfp4Experts"), ("triton_unfused", "UnfusedOAITritonExperts"), ], ) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index 120fbd152668..e68310f98d0d 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -106,7 +106,11 @@ def get_config() -> dict[str, Any] | None: BatchedTritonExperts, ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, + AiterMxfp4Experts, + AiterUnquantizedExperts, ) from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -132,7 +136,11 @@ def get_config() -> dict[str, Any] | None: ) __all__ += [ - "AiterExperts", + "AiterUnquantizedExperts", + "AiterFp8BlockExperts", + "AiterFp8TensorExperts", + "AiterFp8ChannelExperts", + "AiterMxfp4Experts", "fused_topk", "fused_experts", "get_config_file_name", diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 07eddbf4a092..5c506eae123a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -482,29 +482,6 @@ def _supports_current_device() -> bool: def _supports_no_act_and_mul() -> bool: return False - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - SUPPORTED_W_A = [ - (None, None), - (kFp8Static128BlockSym, kFp8Dynamic128Sym), - (kFp8StaticTensorSym, kFp8StaticTensorSym), - (kFp8StaticTensorSym, kFp8DynamicTensorSym), - (kFp8StaticChannelSym, kFp8DynamicTokenSym), - (kMxfp4Static, None), - (kMxfp4Static, kMxfp4Dynamic), - ] - if (weight_key, activation_key) not in SUPPORTED_W_A: - return False - if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950, on_gfx1250 - - if not on_gfx950() or on_gfx1250(): - return False - return True - @staticmethod def _supports_activation(activation: MoEActivation) -> bool: return activation in [ @@ -597,3 +574,77 @@ def apply( output.set_(result) else: output.copy_(result) + + +class AiterUnquantizedExperts(AiterExperts): + """Unquantized (bf16/fp16) weights and activations.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(None, None)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8BlockExperts(AiterExperts): + """128x128 block-quantized weight scale, 128-grouped dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8TensorExperts(AiterExperts): + """Per-tensor static weight scale, either per-tensor static or + per-tensor dynamic activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kFp8StaticTensorSym, kFp8DynamicTensorSym), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterFp8ChannelExperts(AiterExperts): + """Per-output-channel static weight scale, per-token dynamic + activation scale.""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [(kFp8StaticChannelSym, kFp8DynamicTokenSym)] + return (weight_key, activation_key) in SUPPORTED_W_A + + +class AiterMxfp4Experts(AiterExperts): + """MXFP4-quantized weights, with either unquantized (bf16/fp16) or + MXFP4-quantized dynamic activations. gfx950-only (not gfx1250).""" + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + SUPPORTED_W_A = [ + (kMxfp4Static, None), + (kMxfp4Static, kMxfp4Dynamic), + ] + if (weight_key, activation_key) not in SUPPORTED_W_A: + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() and not on_gfx1250() diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 446825c59047..074acd2a0942 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -211,10 +211,12 @@ def backend_to_kernel_cls( elif backend == Fp8MoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterFp8BlockExperts, + AiterFp8ChannelExperts, + AiterFp8TensorExperts, ) - return [AiterExperts] + return [AiterFp8BlockExperts, AiterFp8TensorExperts, AiterFp8ChannelExperts] elif backend == Fp8MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.experts.triton_cutlass_moe import ( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index cbbc39b94310..88e80c755d1f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -284,10 +284,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( @@ -305,10 +305,10 @@ def backend_to_kernel_cls( elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterMxfp4Experts, ) - return [AiterExperts] + return [AiterMxfp4Experts] elif backend == Mxfp4MoeBackend.XPU: from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4 diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 25b55670d1f1..1af0450dfd13 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -126,10 +126,10 @@ def backend_to_kernel_cls( elif backend == UnquantizedMoeBackend.AITER: from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - AiterExperts, + AiterUnquantizedExperts, ) - return [AiterExperts] + return [AiterUnquantizedExperts] elif backend == UnquantizedMoeBackend.TRITON: from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( From 2ce8012fa0c7fbb37335545325e07dacd8141023 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:21:52 +0300 Subject: [PATCH 09/23] [ROCm][Perf] Skip cleaning sparse prefill MQA logits (#51314) Signed-off-by: Sriram Kumar Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sparse.py | 54 +++++++++++++++++++ .../v1/attention/ops/rocm_aiter_mla_sparse.py | 4 +- 2 files changed, 57 insertions(+), 1 deletion(-) create mode 100644 tests/v1/attention/test_rocm_aiter_mla_sparse.py diff --git a/tests/v1/attention/test_rocm_aiter_mla_sparse.py b/tests/v1/attention/test_rocm_aiter_mla_sparse.py new file mode 100644 index 000000000000..1d28b57bebec --- /dev/null +++ b/tests/v1/attention/test_rocm_aiter_mla_sparse.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.platforms import current_platform + +if not current_platform.is_rocm() or not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse indexer test requires ROCm with AITER.", + allow_module_level=True, + ) + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.v1.attention.ops import rocm_aiter_mla_sparse as sparse_mod + + +def test_rocm_fp8_mqa_logits_disables_clean_logits(monkeypatch): + captured = {} + expected = torch.empty((1, 1)) + + def fake_fp8_mqa_logits(*args, clean_logits): + captured["clean_logits"] = clean_logits + return expected + + monkeypatch.setattr(sparse_mod, "_ON_GFX942", False) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True) + monkeypatch.setattr( + sparse_mod, + "mqa_logits_module", + lambda: SimpleNamespace(fp8_mqa_logits=fake_fp8_mqa_logits), + ) + + q = torch.empty((1, 1, 1)) + k_fp8 = torch.empty((1, 1)) + scale = torch.empty((1,)) + weights = torch.empty((1, 1)) + cu_seqlen_ks = torch.zeros((1,), dtype=torch.int32) + cu_seqlen_ke = torch.ones((1,), dtype=torch.int32) + + result = sparse_mod.rocm_fp8_mqa_logits( + q, + (k_fp8, scale), + weights, + cu_seqlen_ks, + cu_seqlen_ke, + ) + + assert result is expected + assert captured["clean_logits"] is False diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 6a93d2151127..78ce85f6ae97 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -933,7 +933,9 @@ def rocm_fp8_mqa_logits( if aiter_mqa_logits_module is not None: fp8_mqa_logits = aiter_mqa_logits_module.fp8_mqa_logits - return fp8_mqa_logits(q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke) + return fp8_mqa_logits( + q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=False + ) else: return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke) From 0d2689e0dfd5e9f74f361cab7701df0feee33ef3 Mon Sep 17 00:00:00 2001 From: Sriram Kumar Date: Tue, 13 Oct 2026 08:25:56 +0300 Subject: [PATCH 10/23] [ROCm][Perf] Skip redundant sparse index remap on non-indexer layers (#51309) Signed-off-by: Sriram Kumar Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> --- .../attention/test_rocm_aiter_mla_sink.py | 2 + .../test_rocm_aiter_mla_sparse_index_remap.py | 295 ++++++++++++++++++ ...est_rocm_aiter_mla_sparse_metadata_sync.py | 4 + .../test_rocm_aiter_triton_sparse_mla.py | 2 + .../backends/mla/rocm_aiter_mla_sparse.py | 62 +++- 5 files changed, 349 insertions(+), 16 deletions(-) create mode 100644 tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sink.py b/tests/kernels/attention/test_rocm_aiter_mla_sink.py index f66d408c8495..399a1fd0e3d0 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sink.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sink.py @@ -476,6 +476,8 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype): (num_tokens, 128), -1, dtype=torch.int32, device="cuda" ) impl.topk_indices_buffer[:, : selected.shape[1]] = selected + impl.owns_indexer = True + impl.index_epoch = 0 impl.q_concat_buffer = torch.empty_like(q) metadata = SimpleNamespace( attn_out_dtype=dtype, diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py new file mode 100644 index 000000000000..2ad9d3a70a52 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_index_remap.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_rocm(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires ROCm.", + allow_module_level=True, + ) + +from vllm._aiter_ops import is_aiter_found_and_supported + +if not is_aiter_found_and_supported(): + pytest.skip( + "ROCm AITER sparse MLA index remap test requires a supported AITER " + "installation.", + allow_module_level=True, + ) + +from vllm.model_executor.layers.attention import sparse_mla_attention as sparse_mla_mod +from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod + +NUM_HEADS = 16 +KV_LORA_RANK = 512 +QK_NOPE_HEAD_DIM = 128 +QK_ROPE_HEAD_DIM = 64 +HEAD_SIZE = 576 +NUM_TOKENS = 2 +TOPK_TOKENS = 4 + + +@pytest.fixture(autouse=True) +def _reset_index_epochs(): + # The counter is keyed by id(), and a freed test buffer can hand its id to + # the next one, so a leftover count would leak across tests. + sparse_mod._INDEX_EPOCHS.clear() + yield + sparse_mod._INDEX_EPOCHS.clear() + + +def _make_metadata(): + return SimpleNamespace( + num_actual_tokens=NUM_TOKENS, + req_id_per_token=torch.zeros(NUM_TOKENS, dtype=torch.int32, device="cpu"), + block_table=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8), + paged_kv_indptr=torch.zeros(NUM_TOKENS + 1, dtype=torch.int32, device="cpu"), + paged_kv_indices=torch.zeros( + NUM_TOKENS * TOPK_TOKENS, dtype=torch.int32, device="cpu" + ), + block_size=1, + topk_tokens=TOPK_TOKENS, + remapped_buf=None, + remapped_epoch=-1, + ) + + +def _patch_init_deps(monkeypatch, *, speculative: bool = False): + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "check_num_heads_validity", + staticmethod(lambda num_heads: None), + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + model_config=SimpleNamespace(dtype=torch.bfloat16), + speculative_config=SimpleNamespace() if speculative else None, + attention_config=SimpleNamespace(hisparse_config=None), + ) + monkeypatch.setattr(sparse_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr(sparse_mla_mod, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr( + sparse_mla_mod, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + sparse_mod, + "current_workspace_manager", + lambda: SimpleNamespace( + get_simultaneous=lambda *specs: tuple( + torch.zeros(shape, dtype=dtype, device="cpu") for shape, dtype in specs + ) + ), + ) + + +def _build_impl(indexer, topk_indices_buffer): + return sparse_mod.ROCMAiterMLASparseImpl( + num_heads=NUM_HEADS, + head_size=HEAD_SIZE, + scale=1.0, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + topk_indices_buffer=topk_indices_buffer, + indexer=indexer, + q_lora_rank=None, + kv_lora_rank=KV_LORA_RANK, + qk_nope_head_dim=QK_NOPE_HEAD_DIM, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM, + v_head_dim=QK_NOPE_HEAD_DIM, + kv_b_proj=None, + ) + + +def _new_buffer(): + return torch.zeros(NUM_TOKENS, TOPK_TOKENS, dtype=torch.int32, device="cpu") + + +def _patch_forward_deps(monkeypatch, impl, remap_calls): + monkeypatch.setattr( + sparse_mod, + "triton_convert_req_index_to_global_index", + lambda *args, **kwargs: remap_calls.append(args), + ) + monkeypatch.setattr( + sparse_mod.AiterMLAHelper, + "get_mla_padded_q", + staticmethod(lambda num_heads, q: q), + ) + attn_out = torch.zeros(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK, device="cpu") + monkeypatch.setattr( + type(impl), + "_forward_mla", + lambda self, layer, q, kv_cache, attn_metadata: (attn_out, None), + ) + return attn_out + + +def _run_forward(impl, metadata): + return impl.forward_mqa( + torch.zeros(NUM_TOKENS, NUM_HEADS, HEAD_SIZE, device="cpu"), + torch.zeros(4, 1, HEAD_SIZE, device="cpu"), + metadata, + SimpleNamespace(_q_scale=None, _k_scale=None), + ) + + +@pytest.mark.parametrize("has_indexer", [True, False]) +def test_owns_indexer_is_structural(monkeypatch, has_indexer): + """owns_indexer reflects the layer, not the serving config.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer = SimpleNamespace(topk_indices_buffer=buffer) if has_indexer else None + + impl = _build_impl(indexer, buffer) + + assert impl.owns_indexer is has_indexer + assert impl.topk_indices_buffer is buffer + + +def test_forward_mqa_remaps_once_per_metadata(monkeypatch): + """Skip layers sharing a metadata share its remap; the first one pays. + + All layers read one top-k buffer, so the reuse has to key off the metadata + that owns ``paged_kv_indices``, not off the buffer being distinct. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first = _build_impl(None, buffer) + second = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + attn_out = _patch_forward_deps(monkeypatch, first, remap_calls) + metadata = _make_metadata() + + out, lse = _run_forward(first, metadata) + assert out is attn_out + assert lse is None + assert len(remap_calls) == 1 + assert metadata.remapped_buf is buffer + + _run_forward(second, metadata) + assert len(remap_calls) == 1 + + +def test_forward_mqa_remaps_for_each_metadata(monkeypatch): + """Skip layers on a second metadata must remap its own paged_kv_indices. + + Regression: gating on ``indexer is not None`` skipped every layer in the + non-indexer metadata group, leaving its ``paged_kv_indices`` unwritten. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + indexer_layer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_layer = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, indexer_layer, remap_calls) + + indexer_metadata = _make_metadata() + _run_forward(indexer_layer, indexer_metadata) + assert len(remap_calls) == 1 + assert indexer_metadata.remapped_buf is buffer + + # A distinct metadata instance carries its own, still-unwritten indices. + skip_metadata = _make_metadata() + _run_forward(skip_layer, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_buf is buffer + + +def test_forward_mqa_remaps_after_a_later_indexer_write(monkeypatch): + """A skip metadata must be refreshed once per indexer write, not once ever. + + Regression for cross-layer index sharing (``index_topk_freq``). With + freq=4 on a 78-layer model the 57 skip layers share one metadata while the + 21 indexer layers share another, so the skip metadata has to pick up every + new selection. Recording only the buffer left layers 4 onwards attending + to the first selection of the pass, which collapsed RULER niah_single_2 to + 0.026 on GLM-5.3 and took GPQA from 0.884 to 0.672. + """ + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + # Construction order mirrors the layer stack: indexer, three skips, indexer. + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + first_skip = _build_impl(None, buffer) + reuse_skip = _build_impl(None, buffer) + _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + next_skip = _build_impl(None, buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, first_skip, remap_calls) + skip_metadata = _make_metadata() + + _run_forward(first_skip, skip_metadata) + assert len(remap_calls) == 1 + + # Same selection, same metadata: reuse. + _run_forward(reuse_skip, skip_metadata) + assert len(remap_calls) == 1 + + # A later indexer layer has since rewritten the buffer in place. The object + # is unchanged, so only the epoch can reveal it. + _run_forward(next_skip, skip_metadata) + assert len(remap_calls) == 2 + assert skip_metadata.remapped_epoch == next_skip.index_epoch + + +def test_index_epoch_tracks_the_preceding_indexer_write(monkeypatch): + """Layers between two indexer layers share the earlier one's epoch.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + first_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_a = _build_impl(None, buffer) + skip_b = _build_impl(None, buffer) + second_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + skip_c = _build_impl(None, buffer) + + assert skip_a.index_epoch == first_indexer.index_epoch + assert skip_b.index_epoch == first_indexer.index_epoch + assert second_indexer.index_epoch == first_indexer.index_epoch + 1 + assert skip_c.index_epoch == second_indexer.index_epoch + + +def test_forward_mqa_remaps_when_buffer_differs(monkeypatch): + """A layer reading a different buffer cannot reuse the metadata's remap.""" + _patch_init_deps(monkeypatch) + remapped_buffer = _new_buffer() + other_buffer = _new_buffer() + layer = _build_impl(None, other_buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, layer, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = remapped_buffer + + _run_forward(layer, metadata) + assert len(remap_calls) == 1 + assert metadata.remapped_buf is other_buffer + + +def test_forward_mqa_indexer_layer_always_remaps(monkeypatch): + """The indexer rewrites its buffer in place, so its remap is never stale.""" + _patch_init_deps(monkeypatch) + buffer = _new_buffer() + impl = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer) + + remap_calls: list[tuple] = [] + _patch_forward_deps(monkeypatch, impl, remap_calls) + metadata = _make_metadata() + metadata.remapped_buf = buffer + metadata.remapped_epoch = impl.index_epoch + + _run_forward(impl, metadata) + assert len(remap_calls) == 1 diff --git a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py index dc1ff79eeabd..4289aac1e7cd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py @@ -523,6 +523,8 @@ def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch): impl.kv_cache_dtype = "auto" impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) @@ -570,6 +572,8 @@ def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache impl.kv_cache_dtype = kv_cache_dtype impl.num_heads = 4 impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4) + impl.owns_indexer = True + impl.index_epoch = 0 expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16) forward_mla = Mock(return_value=(expected_output, None)) monkeypatch.setattr(impl, "_forward_mla", forward_mla) diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 645eec29796d..6849c4e3ebde 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -96,6 +96,8 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs( impl.sinks = torch.randn(num_heads, generator=gen).to(device) impl.kv_cache_dtype = kv_cache_dtype impl.topk_indices_buffer = topk_indices.to(device) + impl.owns_indexer = True + impl.index_epoch = 0 impl.use_aiter_sparse_mla = True layer = SimpleNamespace( _q_scale=torch.tensor([0.02], device=device), 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 a0393e772c3e..037b5432ed03 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -401,6 +401,18 @@ def supports_sink(cls) -> bool: return on_mi3xx() +# Indexer writes per shared top-k buffer, in layer construction order. +_INDEX_EPOCHS: dict[int, int] = {} + + +def _index_epoch(buf: object, owns_indexer: bool) -> int: + """Number of indexer writes up to and including this layer.""" + # id(): tensors are unhashable, and the buffer outlives every reader. + epoch = _INDEX_EPOCHS.get(id(buf), 0) + (1 if owns_indexer else 0) + _INDEX_EPOCHS[id(buf)] = epoch + return epoch + + @dataclass(kw_only=True) class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): qo_indptr: torch.Tensor | None = None @@ -409,6 +421,10 @@ class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata): paged_kv_indptr: torch.Tensor | None = None attn_out_dtype: torch.dtype | None = None + # Source buffer and index_epoch of paged_kv_indices; -1 before any remap. + remapped_buf: object = None + remapped_epoch: int = -1 + # Persistent MLA metadata (only populated when persistent mode is enabled, # i.e. when the aiter sparse decode kernel supports work-stealing splits). work_meta_data: torch.Tensor | None = None @@ -818,6 +834,10 @@ def __init__( self.masked_mha_available = False vllm_config = get_current_vllm_config() + self.owns_indexer: bool = indexer is not None + self.index_epoch: int = _index_epoch( + self.topk_indices_buffer, self.owns_indexer + ) max_tokens = vllm_config.scheduler_config.max_num_batched_tokens q_concat_shape = (max_tokens, num_heads, head_size) (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( @@ -1162,25 +1182,35 @@ def forward_mqa( num_actual_toks = q.shape[0] - # Get topk indices - assert self.topk_indices_buffer is not None - topk_indices = fit_kpool_indices_to_aiter( - 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[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - 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, - ) + + # Indexer layers may rewrite their buffer, so always remap. Non-indexer + # layers have their own KV cache group: remap once per indexer write. + if ( + self.owns_indexer + or attn_metadata.remapped_epoch != self.index_epoch + or attn_metadata.remapped_buf is not self.topk_indices_buffer + ): + # Get topk indices + assert self.topk_indices_buffer is not None + topk_indices = fit_kpool_indices_to_aiter( + self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens + ) + + triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + 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, + ) + attn_metadata.remapped_buf = self.topk_indices_buffer + attn_metadata.remapped_epoch = self.index_epoch # write the latent and rope to kv cache if fp8_attention: From 380638afbb2c5658f22124f6d1d38fcd464ea7d9 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 13 Oct 2026 13:46:37 +0800 Subject: [PATCH 11/23] [Model] Add Clef-Flash support (#61308) Signed-off-by: Tianyu Guo Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 26 +- tests/config/test_model_arch_config.py | 34 ++ .../structured_decisions/test_strategies.py | 65 ++-- .../multimodal/pooling/test_clef_omni.py | 7 +- tests/models/registry.py | 4 + .../structured_decisions/strategies.py | 56 +++ vllm/model_executor/models/clef_flash.py | 44 +++ .../models/clef_joint_schema.py | 343 ++++++++++++++++++ vllm/model_executor/models/clef_omni.py | 339 +---------------- vllm/model_executor/models/registry.py | 1 + .../model_arch_config_convertor.py | 22 ++ 11 files changed, 580 insertions(+), 361 deletions(-) create mode 100644 vllm/model_executor/models/clef_flash.py create mode 100644 vllm/model_executor/models/clef_joint_schema.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index ce09d6f81fee..4a7ec00dd2ea 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,30 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Flash + +[Clef-Flash](https://huggingface.co/Cloudflare/clef-flash) uses a Qwen3.5-9B +backbone and a joint schema head to score all questions in one forward pass. +vLLM detects the head and serves the checkpoint as a pooling model. + +```bash +vllm serve Cloudflare/clef-flash --max-model-len 16384 +``` + +Send requests to `/v1/systemone` using `state` and `questions`, as in the +examples above. `choice`, `noul` and `score` questions are supported. + +- Text and JSON states use the checkpoint's fixed prompt. JSON is compact + with sorted keys. `instructions` returns 400; `chat_template_kwargs` + has no effect. +- The state is truncated to fit 16,384 tokens, or `max_model_len` if smaller. + A schema that exceeds this bound returns 400. +- `images` and `videos` are placed before the state. Videos are sampled at + 2 fps and processed by the Qwen3.5 vision processor. Audio inputs return + 400, and video soundtracks are ignored. +- Each request performs one pooling read. `usage.input_tokens` counts the + prompt once, `output_tokens` is 0, and `label_mass` is 1. + ## Clef-Omni [Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: @@ -188,7 +212,7 @@ vllm serve LiquidAI/d1-3B | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Flash](#clef-flash), images/video; [Clef-Omni](#clef-omni); [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 8ec16ee7f9fb..2ff992a20562 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -70,6 +70,40 @@ def _load_groundtruth(filename: str) -> dict: return json.load(f) +@pytest.mark.parametrize("has_joint_head", [False, True]) +@pytest.mark.parametrize( + "architecture", ["Qwen3_5ForConditionalGeneration", "ColQwen3_5"] +) +def test_qwen35_joint_head_selects_clef_flash(tmp_path, has_joint_head, architecture): + """Only a Qwen3.5 generation checkpoint with a joint head selects Clef.""" + from transformers import Qwen3_5Config + + head_config = { + "hidden_size": 4096, + "width": 1024, + "routing_layers": 2, + "layers": 4, + "heads": 16, + "feedforward": 4096, + } + if has_joint_head: + (tmp_path / "joint_head_config.json").write_text(json.dumps(head_config)) + config = Qwen3_5Config(architectures=[architecture]) + config.name_or_path = str(tmp_path) + model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config.text_config, + _supports_multimodal_for_mm_prefix=lambda: True, + ) + converted = ModelConfig.get_model_arch_config(model_config) + if has_joint_head and architecture == "Qwen3_5ForConditionalGeneration": + assert converted.architectures == ["ClefFlashForJointSchema"] + assert config.joint_head_config == head_config + else: + assert converted.architectures == [architecture] + assert not hasattr(config, "joint_head_config") + + def _assert_model_arch_config( model_config, expected: dict, check_head_size: bool = True ): diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index dae545d1f04c..505f0e35eef1 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -33,6 +33,7 @@ parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + ClefFlashStrategy, DiffusionGemmaCanvasStrategy, JointSchemaStrategy, LiquidStrategy, @@ -94,6 +95,8 @@ def test_strategy_selection(): # Clef-Omni scores options with its head, so it reads no logprobs. clef = model("ClefOmniForJointSchema", "raw_logits") assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy + clef_flash = model("ClefFlashForJointSchema", "raw_logits") + assert select_read_strategy(clef_flash, ("token_classify",)) is ClefFlashStrategy def test_route_is_registered_by_default(): @@ -366,8 +369,8 @@ def test_canvas_read(): ) -@pytest.fixture -def joint_schema(qwen, engine_and_renderer): +@pytest.fixture(params=[JointSchemaStrategy, ClefFlashStrategy]) +def joint_schema(request, qwen, engine_and_renderer): tokenizer, _ = qwen engine, online_renderer = engine_and_renderer # The probability the head gives each option, by option id. @@ -382,9 +385,7 @@ async def encode(prompt, params, request_id, **kwargs): yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) engine.encode.side_effect = encode - strategy = JointSchemaStrategy( - ReadContext(engine, online_renderer, None, "auto", {}) - ) + strategy = request.param(ReadContext(engine, online_renderer, None, "auto", {})) return serving(engine, strategy), engine, tokenizer @@ -468,7 +469,10 @@ async def test_media_needs_a_model_that_takes_it(serving_and_engine): @pytest.mark.asyncio -async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): +@pytest.mark.parametrize("num_images", [1, 2]) +async def test_joint_schema_reads_media_before_the_state( + joint_schema, monkeypatch, num_images +): serving, engine, tokenizer = joint_schema strategy = serving.strategy strategy.__dict__["_connector"] = MediaConnector() @@ -491,23 +495,30 @@ async def process_for_engine_async(prompt, arrival_time): Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() response = await serving.create_decision( - StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image] * num_images) ) assert isinstance(response, StructuredDecisionResponse) - placeholder = tokenizer.encode( - "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False - ) + if isinstance(strategy, ClefFlashStrategy): + media_text = "<|vision_start|><|image_pad|><|vision_end|>" * num_images + "\n" + else: + media_text = "<|vision_start|><|image_pad|><|vision_end|>\n" * num_images + placeholder = tokenizer.encode(media_text, add_special_tokens=False) assert rendered["prompt_token_ids"] == strategy.prefix + placeholder assert rendered["multi_modal_data"]["image"][0].size == (8, 8) - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE} - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } prompt = engine.encode.call_args.args[0] assert prompt["type"] == "multimodal" ids = prompt["prompt_token_ids"] - expanded = len(strategy.prefix) + len(placeholder) + 2 - assert ids[:expanded].count(image_pad) == 3 + expanded = len(strategy.prefix) + len(placeholder) + 2 * num_images + assert ids[:expanded].count(image_pad) == 3 * num_images state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] assert tokenizer.decode(ids[expanded:]) == state_and_questions assert response.answers["team"]["choice"] == "billing" @@ -515,7 +526,9 @@ async def process_for_engine_async(prompt, arrival_time): @pytest.mark.asyncio -async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): +async def test_joint_schema_video_audio_follows_model_modalities( + joint_schema, monkeypatch +): serving, engine, tokenizer = joint_schema strategy = serving.strategy video = MediaWithBytes( @@ -532,6 +545,7 @@ async def load_from_url_async(self, url, media_io, *, fetch_timeout=None): return video async def fetch_audio_async(self, url): + assert not isinstance(strategy, ClefFlashStrategy) if url == "https://example.org/silent.mp4": raise ValueError("No audio stream found.") return track @@ -552,11 +566,17 @@ async def process_for_engine_async(prompt, arrival_time): ) ) assert isinstance(response, StructuredDecisionResponse) - assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE}, - "use_audio_in_video": True, - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } # A video without a decodable soundtrack is watched silently. response = await serving.create_decision( @@ -566,7 +586,8 @@ async def process_for_engine_async(prompt, arrival_time): ) assert isinstance(response, StructuredDecisionResponse) assert rendered["multi_modal_data"] == {"video": [video]} - assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + if not isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False # Audio clips cannot be read with videos that have soundtracks. response = await serving.create_decision( diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index ccdc913b3132..381ac23d53cc 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +"""The shared Clef pooler against a port of the checkpoint's joint schema head, which loops over questions where the pooler batches them.""" import math @@ -11,7 +11,10 @@ import torch import torch.nn.functional as F -from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.model_executor.models.clef_joint_schema import ( + JointSchemaHead, + JointSchemaPooler, +) from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates diff --git a/tests/models/registry.py b/tests/models/registry.py index 0a7027461ef3..b08efd9fb0f1 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -802,6 +802,10 @@ def check_available_online( "Salesforce/blip2-opt-2.7b", extras={"6b": "Salesforce/blip2-opt-6.7b"}, ), + "ClefFlashForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-flash", + max_model_len=4096, + ), "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025", extras={"command-a-plus": "CohereLabs/command-a-plus-05-2026-bf16"}, diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 5567cb808ab5..66238cf7f103 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -809,7 +809,63 @@ async def read( return reads +class ClefFlashStrategy(JointSchemaStrategy): + """Clef-Flash's joint schema prompt with Qwen3.5 image/video processing.""" + + modalities = frozenset({"image", "video"}) + MAX_LENGTH = 16384 + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>", + } + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + connector = self._connector + video_io = VideoMediaIO(ImageMediaIO(), num_frames=-1, fps=self.VIDEO_FPS) + try: + images, videos = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["video"] * len(videos) + + "\n" + ) + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if videos: + mm_data["video"] = videos + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs={"videos_kwargs": {"do_sample_frames": False}}, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + try: + engine_input = ( + await self.context.online_renderer.renderer.process_for_engine_async( + prompt, time.time() + ) + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefFlashForJointSchema": ClefFlashStrategy, "ClefOmniForJointSchema": JointSchemaStrategy, "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, diff --git a/vllm/model_executor/models/clef_flash.py b/vllm/model_executor/models/clef_flash.py new file mode 100644 index 000000000000..32a9e0a59003 --- /dev/null +++ b/vllm/model_executor/models/clef_flash.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Flash: the Qwen3.5 backbone with a joint schema decision head.""" + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler +from .interfaces_base import default_pooling_type +from .qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5ProcessingInfo +from .qwen3_vl import Qwen3VLDummyInputsBuilder, Qwen3VLMultiModalProcessor + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class ClefFlashForJointSchema(Qwen3_5ForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**model_config.hf_config.joint_head_config), + self._embed_lexical, + ) + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/clef_joint_schema.py b/vllm/model_executor/models/clef_joint_schema.py new file mode 100644 index 000000000000..fad2beb09899 --- /dev/null +++ b/vllm/model_executor/models/clef_joint_schema.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared joint schema head for Clef decision models. + +The head reads the final hidden states of the whole prompt, so it runs as a +``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Set +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py index 7cd0d676873b..1533dde27378 100644 --- a/vllm/model_executor/models/clef_omni.py +++ b/vllm/model_executor/models/clef_omni.py @@ -1,47 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). - -The Qwen3-Omni thinker with a joint schema head that scores every option of -every question in one forward pass. The head reads the final hidden states of -the whole prompt, so it runs as a ``token_classify`` pooler. - -Each request carries its schema in ``PoolingParams.extra_kwargs`` under -``"joint_schema"``: - -- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. -- ``question_spans``: the ``[start, end)`` of each question's instructions. -- ``option_spans``: the ``[start, end)`` of each option, question by question. -- ``option_counts``: the number of options of each question. - -Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), -so they hold after multimodal placeholders are expanded. The output has one -value per option, in ``option_spans`` order: its probability within its -question, or its logit with ``use_activation=False``. -""" +"""The Qwen3-Omni thinker with a Clef joint schema head.""" import math -from collections.abc import Callable, Mapping, Set -from dataclasses import dataclass +from collections.abc import Mapping from typing import Any import numpy as np import torch import torch.nn.functional as F -from torch import nn from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig -from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate -from vllm.model_executor.layers.pooler.tokwise.methods import AllPool from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.tasks import PoolingTask -from vllm.utils.torch_utils import async_tensor_h2d -from vllm.v1.outputs import PoolerOutput -from vllm.v1.pool.metadata import PoolingMetadata +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler from .interfaces_base import default_pooling_type from .qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, @@ -50,314 +25,6 @@ Qwen3OmniMoeThinkerProcessingInfo, ) -JOINT_HEAD_WEIGHTS = "joint_head.safetensors" -NUM_QUESTION_TYPES = 3 -MAX_LOGIT_SCALE = math.log(100.0) - - -def segment_mean( - values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor -) -> torch.Tensor: - sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) - sums.index_add_(0, segment_ids, values.float()) - return (sums / lengths[:, None]).to(values.dtype) - - -@dataclass -class SchemaIndices: - """Where one prompt's questions and options are, on the model's device.""" - - positions: torch.Tensor - """Every token of every span, questions first.""" - span_ids: torch.Tensor - """The span of each of ``positions``.""" - span_lengths: torch.Tensor - """``[Q + O]``.""" - question_types: torch.Tensor - """``[Q]``.""" - option_question: torch.Tensor - """The question of each option, ``[O]``.""" - option_slots: torch.Tensor - """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" - option_mask: torch.Tensor - """``[Q, max_options]``, true at the slots of options.""" - - -class EvidenceRoutingLayer(nn.Module): - def __init__(self, width: int, heads: int, feedforward: int) -> None: - super().__init__() - self.query_norm = nn.LayerNorm(width) - self.memory_norm = nn.LayerNorm(width) - self.attention = nn.MultiheadAttention(width, heads, batch_first=True) - self.feedforward_norm = nn.LayerNorm(width) - # The Identity takes the place of the checkpoint's dropout. - self.feedforward = nn.Sequential( - nn.Linear(width, feedforward), - nn.GELU(), - nn.Identity(), - nn.Linear(feedforward, width), - ) - - def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: - # Separate key and value tensors, like the reference, keep attention's - # unpacked input projections. - routed, _ = self.attention( - self.query_norm(queries), - self.memory_norm(memory), - self.memory_norm(memory), - need_weights=False, - ) - queries = queries + routed - return queries + self.feedforward(self.feedforward_norm(queries)) - - -class JointSchemaHead(nn.Module): - def __init__( - self, - hidden_size: int, - width: int, - routing_layers: int, - layers: int, - heads: int, - feedforward: int, - ) -> None: - super().__init__() - self.hidden_norm = nn.LayerNorm(hidden_size) - self.memory_projection = nn.Linear(hidden_size, width, bias=False) - self.question_projection = nn.Linear(hidden_size, width, bias=False) - self.option_question_projection = nn.Linear(hidden_size, width, bias=False) - self.global_projection = nn.Linear(hidden_size, width, bias=False) - self.option_context_projection = nn.Linear(hidden_size, width, bias=False) - self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) - self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) - self.evidence_layers = nn.ModuleList( - EvidenceRoutingLayer(width, heads, feedforward) - for _ in range(routing_layers) - ) - self.option_summary_norm = nn.LayerNorm(width) - self.layers = nn.ModuleList( - nn.TransformerDecoderLayer( - d_model=width, - nhead=heads, - dim_feedforward=feedforward, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=True, - ) - for _ in range(layers) - ) - self.field_norm = nn.LayerNorm(width) - self.option_norm = nn.LayerNorm(width) - self.residual_scorer = nn.Sequential( - nn.Linear(width * 4, width), - nn.GELU(), - nn.Identity(), - nn.Linear(width, 1), - ) - self.prior_logit_scale = nn.Parameter(torch.zeros(())) - self.joint_logit_scale = nn.Parameter(torch.zeros(())) - self.residual_gate = nn.Parameter(torch.zeros(())) - - def forward( - self, - hidden_states: torch.Tensor, - lexical: torch.Tensor, - schema: SchemaIndices, - ) -> torch.Tensor: - """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` - final hidden states and each option's ``[O, hidden]`` mean output - embedding.""" - num_questions = schema.question_types.shape[0] - option_question = schema.option_question - hidden_states = self.hidden_norm(hidden_states) - memory = self.memory_projection(hidden_states).unsqueeze(0) - global_vector = hidden_states[-1] - means = segment_mean( - hidden_states[schema.positions], schema.span_ids, schema.span_lengths - ) - question_vectors, contexts = means[:num_questions], means[num_questions:] - - options = ( - self.option_context_projection(contexts) - + self.option_lexical_projection(lexical) - + self.option_question_projection(question_vectors)[option_question] - ).unsqueeze(0) - for layer in self.evidence_layers: - options = layer(options, memory) - options = options[0] - - # Each field summarizes its options, weighted by their match to it. - base_fields = self.question_projection(question_vectors) - width = options.shape[-1] - grid = options.new_zeros(schema.option_mask.numel(), width) - grid[schema.option_slots] = options - grid = grid.view(*schema.option_mask.shape, width) - routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) - routing = routing.masked_fill(~schema.option_mask, float("-inf")) - summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) - fields = ( - base_fields - + self.option_summary_norm(summaries) - + self.global_projection(global_vector).unsqueeze(0) - + self.type_embedding(schema.question_types) - ).unsqueeze(0) - for layer in self.layers: - fields = layer(fields, memory) - fields = self.field_norm(fields[0])[option_question] - - anchors = F.normalize(question_vectors + global_vector, dim=-1) - prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - prior = prior_scale * ( - F.normalize(lexical, dim=-1) * anchors[option_question] - ).sum(dim=-1) - options = self.option_norm(options) - cosine = F.cosine_similarity(fields, options, dim=-1) - features = torch.cat( - [fields, options, fields * options, torch.abs(fields - options)], dim=-1 - ) - residual = self.residual_scorer(features).squeeze(-1) - joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - joint = joint_scale * cosine + residual - return prior + torch.sigmoid(self.residual_gate) * joint - - -def parse_joint_schema( - extra_kwargs: dict[str, Any] | None, prompt_len: int -) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: - """The question types, the absolute spans of the questions then of the - options, and the option counts. None if the schema is missing or invalid.""" - schema = (extra_kwargs or {}).get("joint_schema") - if not isinstance(schema, dict): - return None - try: - types = [int(t) for t in schema["question_types"]] - counts = [int(c) for c in schema["option_counts"]] - question_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["question_spans"] - ] - option_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["option_spans"] - ] - except (KeyError, TypeError, ValueError): - return None - spans = question_spans + option_spans - if ( - not types - or len(counts) != len(types) - or len(question_spans) != len(types) - or len(option_spans) != sum(counts) - or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) - or any(c < 1 for c in counts) - or any(not 0 <= start < end <= prompt_len for start, end in spans) - ): - return None - return types, spans, counts - - -class JointSchemaPooler(Pooler): - """Runs the joint schema head on each finished prompt. - - It never raises, because it runs in the engine core: a prompt without a - valid schema, such as a profiling dummy, gets an empty output. - """ - - def __init__( - self, - head: JointSchemaHead, - embed_lexical: Callable[[torch.Tensor], torch.Tensor], - ) -> None: - super().__init__() - self.pooling = AllPool() - self.head = head - self.embed_lexical = embed_lexical - - def get_supported_tasks(self) -> Set[PoolingTask]: - return {"token_classify"} - - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: - return PoolingParamsUpdate(requires_token_ids=True) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled = self.pooling(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() - outputs: list[torch.Tensor | None] = [] - for hidden, token_ids, params in zip( - pooled, prompt_token_ids, pooling_metadata.pooling_params - ): - if hidden is None: - outputs.append(None) - continue - parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) - if parsed is None: - outputs.append(hidden.new_empty(0, dtype=torch.float32)) - continue - schema, lexical_ids, lexical_span_ids = self._indices( - *parsed, hidden.device - ) - lexical = segment_mean( - self.embed_lexical( - async_tensor_h2d(token_ids[lexical_ids], hidden.device) - ), - lexical_span_ids, - schema.span_lengths[len(parsed[0]) :], - ) - logits = self.head(hidden, lexical, schema).float() - if params.use_activation is not False: - mask = schema.option_mask - grid = logits.new_full((mask.numel(),), float("-inf")) - grid[schema.option_slots] = logits - probs = grid.view(mask.shape).softmax(dim=-1) - logits = probs.view(-1)[schema.option_slots] - outputs.append(logits) - return outputs - - @staticmethod - def _indices( - types: list[int], - spans: list[tuple[int, int]], - counts: list[int], - device: torch.device, - ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: - """The schema's indices, plus the option tokens' positions (on the CPU) - and their option spans, for the lexical embeddings.""" - num_questions = len(types) - lengths = torch.tensor([end - start for start, end in spans]) - span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) - positions = torch.cat([torch.arange(start, end) for start, end in spans]) - num_question_tokens = int(lengths[:num_questions].sum()) - - counts_cpu = torch.tensor(counts) - max_options = max(counts) - option_question = torch.repeat_interleave( - torch.arange(num_questions), counts_cpu - ) - option_slots = torch.cat( - [torch.arange(count) + q * max_options for q, count in enumerate(counts)] - ) - schema = SchemaIndices( - positions=async_tensor_h2d(positions, device), - span_ids=async_tensor_h2d(span_ids, device), - span_lengths=async_tensor_h2d(lengths, device), - question_types=async_tensor_h2d(types, device, dtype=torch.long), - option_question=async_tensor_h2d(option_question, device), - option_slots=async_tensor_h2d(option_slots, device), - option_mask=async_tensor_h2d( - torch.arange(max_options) < counts_cpu[:, None], device - ), - ) - lexical_span_ids = async_tensor_h2d( - span_ids[num_question_tokens:] - num_questions, device - ) - return schema, positions[num_question_tokens:], lexical_span_ids - - MAX_FRAME_PIXELS = 256 * 32 * 32 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 36f9bd6e6fcc..cd0557ac4585 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefFlashForJointSchema": ("clef_flash", "ClefFlashForJointSchema"), "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 147b65759dea..06a937697106 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -647,6 +647,27 @@ def __init__( _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3_5ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + if getattr(hf_config, "architectures", None) == [ + "Qwen3_5ForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefFlashForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -850,6 +871,7 @@ def derive_max_model_len_and_key(self) -> tuple[float, str | None]: "nemotron-nas": NemotronNasModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, + "qwen3_5": Qwen3_5ModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor, From 86b1c2ef159dcc268d219b074cb047340ea718aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Tue, 13 Oct 2026 07:51:22 +0200 Subject: [PATCH 12/23] [Attention][DCP] Support a2a DCP comm backend for FlashMLA sparse (#61285) Signed-off-by: NickLucche Co-authored-by: Claude Opus 5.5 --- tests/v1/attention/test_sparse_mla_backends.py | 16 ++++++++++++++++ .../v1/attention/backends/mla/flashmla_sparse.py | 8 +------- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index c21fc15a5278..f7c47fe6a3ec 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -4157,6 +4157,22 @@ def test_fp8_dcp_head_envelope_guard(local_heads, dcp_world_size, should_raise): assert local_pad == gathered_pad +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +@pytest.mark.parametrize("comm_backend", ["ag_rs", "a2a"]) +def test_fp8_dcp_accepts_comm_backend(comm_backend): + """The mixed-batch path returns (out, lse) for every token in the layout + both DCP combines consume, so neither comm backend is rejected.""" + device = torch.device(DEVICE_TYPE) + vllm_config = _build_sparse_dcp_vllm_config(16, 4, comm_backend) + kv_cache_spec = create_standard_kv_cache_spec(vllm_config) + builder_cls = FlashMLASparseBackend.get_builder_cls() + builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device) + assert builder.fp8_use_mixed_batch + + def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch): """A decode row whose top-k shard holds no local candidates (all -1) has undefined kernel out/lse; it must come back as (0, -inf), the identity of diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 3cf0c33c534d..80b8da179769 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -446,12 +446,6 @@ def __init__( ) if parallel_config.decode_context_parallel_size > 1: - if parallel_config.dcp_comm_backend != "ag_rs": - raise NotImplementedError( - "DCP for FlashMLA sparse is only validated with the " - "default 'ag_rs' DCP comm backend; got " - f"'{parallel_config.dcp_comm_backend}'" - ) if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": raise NotImplementedError( "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " @@ -1268,7 +1262,7 @@ def _forward_fp8_kv_mixed_batch( out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0) lse.masked_fill_(empty_rows.view(-1, 1), float("-inf")) # The head-padding slice above can leave `out` non-contiguous, and the - # merge feeds it to reduce_scatter. + # DCP combines (reduce_scatter, direct A2A) need packed heads. return out.contiguous(), lse def _fp8_flash_mla_kernel( From a6009ea85be08161f4b5954bbf317027f6ddfd82 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 12 Oct 2026 23:24:43 -0700 Subject: [PATCH 13/23] [Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer (#61149) Signed-off-by: Yongye Zhu Co-authored-by: Claude Opus 5.5 (1M context) --- cmake/external_projects/deepgemm.cmake | 12 +- csrc/cache_kernels.cu | 20 +- .../attention/test_deepgemm_attention.py | 189 +++++++++++------- .../test_rocm_paged_mxfp4_indexer.py | 1 + tests/kernels/test_compressor_kv_cache.py | 42 +++- .../test_fused_deepseek_v32_norm_rope.py | 15 +- .../test_fused_indexer_q_rope_quant.py | 5 + .../test_kpool_decode_update_batched.py | 8 +- .../layers/test_mla_short_prefill_indexer.py | 39 ++-- .../test_indexer_deepseek_v4_slot_mapping.py | 11 +- .../test_rocm_paged_mxfp4_indexer_plan.py | 1 + tools/build_deepgemm_C.py | 3 +- tools/check_wheel_deepgemm.py | 11 +- tools/install_deepgemm.sh | 8 +- .../layers/sparse_attn_indexer.py | 110 +++++++--- vllm/models/deepseek_v32/common/kernels.py | 14 +- .../common/ops/fused_compress_quant_cache.py | 12 +- .../common/ops/indexer_k_store.py | 15 +- .../glm5next/nvidia/ops/kpool_compress.py | 30 ++- vllm/models/glm5next/nvidia/sparse_indexer.py | 72 ++++--- vllm/utils/deep_gemm.py | 101 +++++++--- vllm/v1/attention/backends/mla/indexer.py | 4 + vllm/v1/attention/ops/common.py | 7 + 23 files changed, 519 insertions(+), 211 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 1dcee345f845..b796623b7ac1 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -29,12 +29,11 @@ if(DEEPGEMM_SRC_DIR) else() # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") - # Pinned to the fork's dev branch: upstream 2.8.0 plus the SM120 - # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU - # alpha/beta, and SiTU for FP8/FP4 Mega MoE. Also includes the CUDA 12.x - # layout header fix from vllm-project/DeepGEMM#12. - # TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. - set(_DEEPGEMM_UPSTREAM_TAG "1e1842a833699298f7afc02eefb2ed168fff6938") + # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) + # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, + # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 + # Mega MoE. TORCH_LIBRARY build with a CPython 3.10 stable-ABI floor. + set(_DEEPGEMM_UPSTREAM_TAG "6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -146,6 +145,7 @@ if(DEEPGEMM_ARCHS) # install(FILES "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py" + "${deepgemm_SOURCE_DIR}/deep_gemm/locality_domain.py" DESTINATION vllm/third_party/deep_gemm COMPONENT _deep_gemm_C) diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index d1c1e9df9896..24453ff1c154 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -630,7 +630,8 @@ __global__ void indexer_k_quant_and_cache_kernel( const int cache_block_size, // cache block size const int64_t cache_block_stride, // stride for each block in kv_cache - const bool use_ue8m0 // use ue8m0 scale format + const bool use_ue8m0, // use ue8m0 scale format + const bool use_mx_scale // store the ue8m0 scale as packed MX bytes ) { constexpr int VEC_SIZE = 4; const int64_t token_idx = blockIdx.x; @@ -679,7 +680,13 @@ __global__ void indexer_k_quant_and_cache_kernel( const int64_t dst_scale_idx = block_idx * cache_block_stride + cache_block_size * head_dim + (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size; - reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + if (use_mx_scale) { + // One UE8M0 byte per 32 elements, all equal to the shared scale's. + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = + ((__float_as_uint(scale) >> 23) & 0xFF) * 0x01010101u; + } else { + reinterpret_cast(kv_cache)[dst_scale_idx / 4] = scale; + } } } @@ -1912,7 +1919,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( reinterpret_cast(k.data_ptr()), \ reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), head_dim, quant_block_size, \ - cache_block_size, cache_block_stride, use_ue8m0); + cache_block_size, cache_block_stride, use_ue8m0, use_mx_scale); void indexer_k_quant_and_cache( torch::stable::Tensor& k, // [num_tokens, head_dim] @@ -1924,8 +1931,13 @@ void indexer_k_quant_and_cache( int head_dim = k.size(1); int cache_block_size = kv_cache.size(1); int64_t cache_block_stride = kv_cache.stride(0); - bool use_ue8m0 = scale_fmt == "ue8m0"; + // "mxfp8" stores the ue8m0 scale as four packed UE8M0 bytes, the MX scales + // DeepGEMM's SM100 MQA logits kernels read. + bool use_mx_scale = scale_fmt == "mxfp8"; + bool use_ue8m0 = use_mx_scale || scale_fmt == "ue8m0"; + STD_TORCH_CHECK(!use_mx_scale || (head_dim == 128 && quant_block_size == 128), + "mxfp8 scales require head_dim == quant_block_size == 128"); STD_TORCH_CHECK(k.device() == kv_cache.device(), "k and kv_cache must be on the same device"); STD_TORCH_CHECK(k.device() == slot_mapping.device(), diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index f332fb0fea21..109cc83830b1 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -5,6 +5,7 @@ import pytest import torch +from vllm.model_executor.layers.sparse_attn_indexer import as_fp32_logits from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( _ceil_to_ue8m0, @@ -13,18 +14,29 @@ fp8_fp4_paged_mqa_logits, get_num_sms, get_paged_mqa_logits_metadata, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, native_next_n_supported, ) from vllm.utils.import_utils import has_deep_gemm from vllm.utils.math_utils import cdiv +def encode_indexer_scales(sf: torch.Tensor) -> torch.Tensor: + """Per-token power-of-two fp32 scales as the indexer K cache stores them: + four equal packed UE8M0 bytes under MX scales, else the fp32 itself.""" + if not mqa_logits_use_mx_scales(): + return sf + ue8m0 = (sf.view(torch.int32) >> 23).to(torch.uint8) + return ue8m0[..., None].repeat_interleave(4, dim=-1).view(torch.int32)[..., 0] + + def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: # x: (num_blocks, block_size, 1, head_dim) num_blocks, block_size, num_heads, head_dim = x.shape assert num_heads == 1 x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) - sf = x_amax / 448.0 + sf = _ceil_to_ue8m0(x_amax / 448.0) x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) x_fp8 = torch.empty( (num_blocks, block_size * (head_dim + 4)), @@ -34,9 +46,9 @@ def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: x_fp8[:, : block_size * head_dim] = x_scaled.view( num_blocks, block_size * head_dim ).view(dtype=torch.uint8) - x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view( - dtype=torch.uint8 - ) + x_fp8[:, block_size * head_dim :] = encode_indexer_scales( + sf.view(num_blocks, block_size) + ).view(dtype=torch.uint8) return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) @@ -91,64 +103,72 @@ def _ref_fp8_mqa_logits( return logits +def _packed_requests_test_data(seq_len: int, seq_len_kv: int): + # Two causal prefill requests packed into one KV buffer, so the second + # request's rows start at a nonzero K position. + half, half_kv = seq_len // 2, seq_len_kv // 2 + rows = torch.arange(seq_len, dtype=torch.int, device="cuda") + ks = torch.where(rows < half, 0, half_kv).to(torch.int) + ke = ks + half_kv - half + 1 + rows % half + return ks, ke + + @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" ) -@pytest.mark.parametrize("clean_logits", [True, False]) -def test_deepgemm_fp8_mqa_logits(clean_logits: bool): +@pytest.mark.parametrize("layout", ["cp", "no_cp", "packed"]) +def test_deepgemm_fp8_mqa_logits(layout: str): + """FP8 indexer prefill logits, compressed to each row's valid K span.""" torch.manual_seed(0) random.seed(0) num_heads, head_dim = 32, 128 - for seq_len in (512,): - for seq_len_kv in (1024,): - for disable_cp in (False, True): - q = torch.randn( - seq_len, - num_heads, - head_dim, - device="cuda", - dtype=torch.bfloat16, - ) - kv = torch.randn( - seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16 - ) - weights = torch.randn( - seq_len, num_heads, device="cuda", dtype=torch.float32 - ) - - if disable_cp: - ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") - ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( - seq_len_kv - seq_len - ) - else: - ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) - - q_fp8 = q.to(torch.float8_e4m3fn) - kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) - logits = fp8_fp4_mqa_logits( - (q_fp8, None), kv_fp8, weights, ks, ke, clean_logits=clean_logits - ) - - ref_logits = _ref_fp8_mqa_logits( - q=q, - kv=kv, - weights=weights, - cu_seqlen_ks=ks, - cu_seqlen_ke=ke, - ) - ref_neginf_mask = ref_logits == float("-inf") - - if clean_logits: - neginf_mask = logits == float("-inf") - assert torch.equal(neginf_mask, ref_neginf_mask) - - ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) - logits = logits.masked_fill(ref_neginf_mask, 0) - diff = calc_diff(logits, ref_logits) - assert diff < 1e-3, f"{diff=}" + seq_len, seq_len_kv = 512, 1024 + q = torch.randn(seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(seq_len_kv, head_dim, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device="cuda", dtype=torch.float32) + + if layout == "cp": + ks, ke = _generate_cp_test_data(seq_len, seq_len_kv) + elif layout == "no_cp": + ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") + ke = torch.arange(seq_len, dtype=torch.int, device="cuda") + ( + seq_len_kv - seq_len + ) + else: + ks, ke = _packed_requests_test_data(seq_len, seq_len_kv) + + q_fp8 = q.to(torch.float8_e4m3fn) + kv_fp8, kv_sf = per_custom_dims_cast_to_fp8(kv, (0,), True) + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + max_seqlen_k = int((ke - ks).max()) + logits = fp8_fp4_mqa_logits( + (q_fp8, q_scale), + (kv_fp8, encode_indexer_scales(kv_sf)), + kernel_weights, + ks, + ke, + max_seqlen_k, + ) + assert logits.shape == (seq_len, max_seqlen_k) + logits = as_fp32_logits(logits, ke - ks) + + ref_logits = _ref_fp8_mqa_logits( + q=q, + kv=kv, + weights=kernel_weights.float(), + cu_seqlen_ks=ks, + cu_seqlen_ke=ke, + ) + # Compressed row i holds K positions [ks[i], ke[i]) from column zero. + cols = ks[:, None] + torch.arange(max_seqlen_k, device="cuda")[None, :] + valid = cols < ke[:, None] + ref_logits = ref_logits.gather(1, cols.clamp(max=seq_len_kv - 1).long()) + ref_logits = ref_logits.masked_fill(~valid, 0) + logits = logits.masked_fill(~valid, 0) + diff = calc_diff(logits, ref_logits) + assert diff < 1e-3, f"{diff=}" def _ref_fp8_fp4_paged_mqa_logits( @@ -212,9 +232,6 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): if not native_next_n_supported(next_n): pytest.skip(f"next_n={next_n} has no native kernel on this architecture") - # NOTE: clean_logits=True is incompatible with the 2D context_lens - # required by csrc/apis/attention.hpp; only the False path is exercised. - clean_logits = False torch.manual_seed(0) random.seed(0) @@ -273,26 +290,42 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens_2d = ( context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange ).contiguous() + q_scale, kernel_weights = mqa_logits_mx_inputs(q_fp8, None, weights) + kernel_q = (q_fp8, q_scale) + kernel_context_lens, kernel_block_tables = context_lens_2d, block_tables + indices = None + if mqa_logits_use_mx_scales(): + # SM100 only takes varlen rows: one Q token each, tied to its + # request by `indices`. + assert q_scale is not None + kernel_q = ( + q_fp8.flatten(0, 1)[:, None], + q_scale.flatten(0, 1)[:, None], + ) + kernel_context_lens = context_lens_2d.view(-1, 1) + kernel_block_tables = block_tables.repeat_interleave(next_n, dim=0) + indices = torch.arange(batch_size, device="cuda", dtype=torch.int32) + indices = indices.repeat_interleave(next_n) schedule_metadata = get_paged_mqa_logits_metadata( - context_lens_2d, - blocksize, - get_num_sms(), + kernel_context_lens, blocksize, get_num_sms(), indices=indices ) logits = fp8_fp4_paged_mqa_logits( - (q_fp8, None), + kernel_q, kv_cache_fp8, - weights, - context_lens_2d, - block_tables, + kernel_weights, + kernel_context_lens, + kernel_block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, + indices=indices, ) + if indices is not None: + logits = as_fp32_logits(logits, kernel_context_lens.view(-1)) ref_logits = _ref_fp8_fp4_paged_mqa_logits( q, kv_cache, - weights, + kernel_weights.float(), context_lens, block_tables, max_model_len, @@ -309,7 +342,7 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): context_lens[row_indices] - next_n + next_n_offset ).unsqueeze(1) - logits = logits.masked_fill(~mask, 0) + logits = logits.float().masked_fill(~mask, 0) ref_logits = ref_logits.masked_fill(~mask, 0) diff = calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" @@ -349,12 +382,28 @@ def test_deepgemm_paged_mqa_packed_manager_block_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) + q_scale, weights = mqa_logits_mx_inputs( + q, None, 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()) + indices = ( + torch.zeros(1, device="cuda", dtype=torch.int32) + if mqa_logits_use_mx_scales() + else None + ) + schedule = get_paged_mqa_logits_metadata( + context_lens, page_size, get_num_sms(), indices=indices + ) logits = [ fp8_fp4_paged_mqa_logits( - (q, None), cache, weights, context_lens, table, schedule, block_size, False + (q, q_scale), + cache, + weights, + context_lens, + table, + schedule, + block_size, + indices=indices, ) for cache, table in zip((compact_pages, packed_pages), tables) ] diff --git a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py index 311ee89a8d12..e270ae4a5813 100644 --- a/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py +++ b/tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py @@ -253,6 +253,7 @@ def _prefill_metadata( cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE), token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=t0, token_end=t1, num_reqs=req_hi - req_lo, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a931..e9924229055a 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -1422,7 +1422,10 @@ def test_dequantize_and_gather_k_cache( @pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) @pytest.mark.parametrize("block_size", [16, 64]) -def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): +@pytest.mark.parametrize("scale_fmt", ["ue8m0", "mxfp8"]) +def test_indexer_quant_cache_roundtrip( + num_tokens: int, block_size: int, scale_fmt: str +): """K → indexer_k_quant_and_cache → cp_gather_indexer_k_quant_cache → manual dequant → compare against original.""" HEAD_DIM = 128 @@ -1442,7 +1445,9 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): ) slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) - ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, "ue8m0") + ops.indexer_k_quant_and_cache( + k, kv_cache, slot_mapping, QUANT_BLOCK_SIZE, scale_fmt + ) # ── Gather ────────────────────────────────────────────────────────── max_blocks_per_seq = num_blocks @@ -1465,7 +1470,12 @@ def test_indexer_quant_cache_roundtrip(num_tokens: int, block_size: int): # ── Manual dequant ────────────────────────────────────────────────── k_fp8 = dst_k.view(torch.float8_e4m3fn).float() # [num_tokens, 128] - scale = dst_scale.view(torch.float32) # [num_tokens, 1] + if scale_fmt == "mxfp8": + # One UE8M0 byte per 32 elements, all equal to the token's scale. + assert (dst_scale == dst_scale[:, :1]).all() + scale = (dst_scale[:, :1].to(torch.int32) << 23).view(torch.float32) + else: + scale = dst_scale.view(torch.float32) # [num_tokens, 1] k_recovered = k_fp8 * scale # [num_tokens, 128] # ── Compare ───────────────────────────────────────────────────────── @@ -1749,8 +1759,13 @@ def _reference_kv_compress_norm_rope( ), ], ) -def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool): +@pytest.mark.parametrize("mx_scale", [False, True]) +def test_fused_kv_insert_indexer( + num_tokens: int, kv_block_size: int, use_fp4: bool, mx_scale: bool +): """Fused K compress+norm+rope+quant+insert for the indexer KV cache.""" + if use_fp4 and mx_scale: + pytest.skip("MX scales only change the FP8 cache") HEAD_DIM = 128 ROPE_DIM = 64 BLOCK_SIZE = 16 @@ -1833,6 +1848,7 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b TOKEN_STRIDE=TOKEN_STRIDE, SCALE_DIM=SCALE_DIM, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=mx_scale, num_warps=1, ) @@ -1876,12 +1892,18 @@ def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: b ), f"token {i}: FP8 bytes differ" scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM - actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view( - torch.float32 - ) - assert torch.equal(actual_scale, scale[i : i + 1]), ( - f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" - ) + actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM] + if mx_scale: + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale[i : i + 1].view(torch.int32) >> 23).to(torch.uint8) + assert torch.equal(actual_scale, ue8m0.repeat(SCALE_DIM)), ( + f"token {i}: mx scale {actual_scale.tolist()} != {ue8m0.item()}" + ) + else: + actual_scale = actual_scale.view(torch.float32) + assert torch.equal(actual_scale, scale[i : i + 1]), ( + f"token {i}: scale {actual_scale.item()} != {scale[i].item()}" + ) @pytest.mark.parametrize("compress_ratio", [4, 128]) diff --git a/tests/kernels/test_fused_deepseek_v32_norm_rope.py b/tests/kernels/test_fused_deepseek_v32_norm_rope.py index 63bbe0c95ff4..ec61b94abd5e 100644 --- a/tests/kernels/test_fused_deepseek_v32_norm_rope.py +++ b/tests/kernels/test_fused_deepseek_v32_norm_rope.py @@ -32,6 +32,7 @@ from vllm.models.deepseek_v32.common import kernels as K from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales FP8 = torch.float8_e4m3fn FP8_MAX = 448.0 @@ -116,6 +117,16 @@ def rope( return out +def indexer_scales(scale_bytes: torch.Tensor) -> torch.Tensor: + """Per-token fp32 scales from the indexer cache's 4 scale bytes per token.""" + if not mqa_logits_use_mx_scales(): + return scale_bytes.view(torch.float32) + # MX scales: the power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale_bytes.view(-1, 4) + assert (ue8m0 == ue8m0[:, :1]).all(), "MX scale bytes differ within a token" + return (ue8m0[:, 0].to(torch.int32) << 23).view(torch.float32) + + def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" amax = vals.float().abs().amax(dim=-1, keepdim=True) @@ -243,7 +254,7 @@ def test_fused_norm_rope( q_ref, s_ref = ue8m0_quant(ik_ref) flat = idx_cache[0].reshape(-1) vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim) - scales = flat[bs * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(flat[bs * cfg.index_head_dim :]) assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) @@ -318,7 +329,7 @@ def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig): .view(FP8) .view(block_size, cfg.index_head_dim) ) - scales = packed[block_size * cfg.index_head_dim :].view(torch.float32) + scales = indexer_scales(packed[block_size * cfg.index_head_dim :]) assert_fp8(values, q_ref, "packed indexer-K fp8") torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) assert (backing[block_size * idx_row : packed_block_stride] == 0).all() diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 0ce3a918513f..c6dfa11b0f13 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -26,6 +26,7 @@ ) from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.import_utils import has_cutedsl HEAD_DIM = 128 @@ -347,6 +348,10 @@ def _indexer_k_reference( exponent = k.float().abs().amax(dim=-1, keepdim=True).clamp(min=1e-4) / 448 scale = exponent.log2().ceil().exp2() values = (k.float() / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = (scale.view(torch.int32) >> 23).to(torch.uint8) + return values.view(torch.uint8), ue8m0.repeat(1, 4) return values.view(torch.uint8), scale.view(torch.uint8) diff --git a/tests/kernels/test_kpool_decode_update_batched.py b/tests/kernels/test_kpool_decode_update_batched.py index 4accf527f3d1..5dd9eff5f10a 100644 --- a/tests/kernels/test_kpool_decode_update_batched.py +++ b/tests/kernels/test_kpool_decode_update_batched.py @@ -29,6 +29,7 @@ import torch from vllm.platforms import current_platform +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.kpool_compress import ( @@ -206,7 +207,12 @@ def _torch_reference( k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM) s_off = page_base + k_region + loc_tok * 4 kv_flat[k_off] = quantized.view(torch.uint8) - kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8) + scale_bytes = scale.detach().reshape(1).view(torch.uint8) + if mqa_logits_use_mx_scales(): + # The power-of-two scale's exponent as 4 equal UE8M0 bytes. + ue8m0 = scale.reshape(1).view(torch.int32) >> 23 + scale_bytes = ue8m0.to(torch.uint8).repeat(4) + kv_flat[s_off : s_off + 4] = scale_bytes # stash -- gated on the TOKEN-granular tail slot, not on pos_valid. # pos_valid keys off the pool-granular cache_loc, which is -1 for diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index ba4ff895cbf6..7ca4cff4040d 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -771,6 +771,8 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): candidates = _make_candidates(lens, num_candidates, cbk).cuda() # Reference: dense bf16 logits masked to the candidates, then torch.topk. + # The logits are compressed: row r starts at K position ks[r]. + max_len = int(lens.max()) dense = deep_gemm.fp8_fp4_mqa_logits( ( q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), @@ -780,19 +782,13 @@ def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): weights, ks, ke, - clean_logits=False, - logits_dtype=torch.bfloat16, + max_len, ) - cols = torch.arange(total, device="cuda") - in_range = (cols - ks.unsqueeze(1) >= 0) & ( - cols - ks.unsqueeze(1) < lens.unsqueeze(1) - ) - masked = dense.masked_fill( - ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), - float("-inf"), - ).float() - ref_abs = _reference_topk(masked, topk) # packed-workspace positions - ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + in_range = torch.arange(max_len, device="cuda") < lens.unsqueeze(1) + # A row's last candidate block may extend past max_len. + in_cand = _candidate_token_mask(candidates, cbk, max_len + cbk)[:, :max_len] + masked = dense.masked_fill(~(in_range & in_cand), float("-inf")).float() + ref = _reference_topk(masked, topk) out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") scratch = _sparse_scratch(rows, num_candidates, topk) @@ -827,11 +823,11 @@ def run_sparse(kernel_metadata=None): graph.replay() else: kernel_metadata = run_sparse() - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) out.fill_(-1) run_sparse(kernel_metadata) - _assert_same_selection(dense, out, ref, row_starts=ks) + _assert_same_selection(dense, out, ref) def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): @@ -902,21 +898,22 @@ def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() def reference_topk(): - context_lens_2d = row_lens.view(batch, next_n) + # SM100 paged MQA logits take varlen rows: one Q token per row. + context_lens_2d = row_lens.view(rows, 1) + indices = torch.arange(batch, dtype=torch.int32, device="cuda") + indices = indices.repeat_interleave(next_n) sched = deep_gemm.get_paged_mqa_logits_metadata( - context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + context_lens_2d, page_kv, deep_gemm.get_num_sms(), indices ) logits = deep_gemm.fp8_fp4_paged_mqa_logits( - (q, q_scale), + (q.view(rows, 1, num_heads, -1), q_scale.view(rows, 1, num_heads)), kv_cache, weights, context_lens_2d, - block_table, + block_table.repeat_interleave(next_n, dim=0), sched, int(req_lens.max()), - False, - torch.bfloat16, - None, + indices, ) lens_flat = row_lens.flatten() cols = torch.arange(logits.shape[1], device="cuda") 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 464fce66f212..62c8bc32bff5 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -14,6 +14,7 @@ from vllm.models.deepseek_v41.sparse_mla import ( DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend, ) +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.mla import indexer as indexer_module from vllm.v1.attention.backends.mla.compressor_utils import ( @@ -113,10 +114,17 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): num_reqs, 4 ) num_sms = torch.cuda.get_device_properties(device).multi_processor_count + # SM100's paged MQA logits only take varlen rows. + indices = ( + torch.arange(num_reqs, dtype=torch.int32, device=device) + if mqa_logits_use_mx_scales() + else None + ) schedule_metadata = indexer_module.get_paged_mqa_logits_metadata( torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1), 64, num_sms, + indices=indices, ).clone() builder = object.__new__(DeepseekV32IndexerMetadataBuilder) @@ -139,6 +147,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device), requires_padding=False, schedule_metadata=schedule_metadata, + indices=indices, ), block_table=block_table, ) @@ -157,7 +166,7 @@ def test_indexer_draft_decode_metadata_update_is_capture_safe(): torch.testing.assert_close( schedule_metadata, indexer_module.get_paged_mqa_logits_metadata( - expected_decode_seq_lens.unsqueeze(1), 64, num_sms + expected_decode_seq_lens.unsqueeze(1), 64, num_sms, indices=indices ), ) assert metadata.decode is not None diff --git a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py index 0283ed99ee3a..a3dbec86dca5 100644 --- a/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py +++ b/tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py @@ -36,6 +36,7 @@ def _chunk(block_table, token_start, token_end, row_ends): cu_seq_lens=torch.zeros(1, dtype=torch.int32), token_to_seq=torch.zeros(1, dtype=torch.int32), total_seq_lens=1, + max_seqlen_k=int(ends.max()), token_start=token_start, token_end=token_end, num_reqs=block_table.shape[0], diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 36824318c051..dec8b0105ac8 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -35,8 +35,9 @@ env.pop("DG_SKIP_CUDA_BUILD", None) print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True) +# CMake decides when to rebuild; setuptools would only check python_api.cpp. subprocess.check_call( - [sys.executable, "setup.py", "build_ext", "--inplace"], + [sys.executable, "setup.py", "build_ext", "--inplace", "--force"], cwd=src, env=env, ) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index fcc35ead58c4..c8da1351bb54 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -4,10 +4,12 @@ """Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout. Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under -``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100 -deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. +``vllm.third_party.deep_gemm``, and that the package imports: vLLM treats an +import failure as DeepGEMM being absent. Run after vLLM is installed, e.g. the +H100 deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. """ +import importlib import importlib.util import sys from pathlib import Path @@ -38,6 +40,11 @@ def main() -> int: if missing: print(f"missing: {missing}", file=sys.stderr) return 1 + try: + importlib.import_module("vllm.third_party.deep_gemm") + except Exception as e: + print(f"vllm.third_party.deep_gemm failed to import: {e!r}", file=sys.stderr) + return 1 return 0 diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 2fb2351a0ca5..b617e6d5679f 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -7,10 +7,10 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" -# NOTE: This targets the vLLM fork's dev branch, which carries the sm120 -# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE -# activations, plus the CUDA 12.x layout header fix from vllm-project/DeepGEMM#12. -DEEPGEMM_GIT_REF="1e1842a833699298f7afc02eefb2ed168fff6938" +# NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release +# plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU +# Mega MoE activations, and NVFP4 Mega MoE. +DEEPGEMM_GIT_REF="6ce1c5ec6627a6bf11dbcc698b0b1e19c9248432" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 935e4c0d2d44..3af3a161ac8b 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -42,6 +42,8 @@ fp8_fp4_mqa_logits, fp8_fp4_paged_mqa_logits, has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, ) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( @@ -286,6 +288,42 @@ def fused_indexer_q_rope_quant( return q_fp8, weights_out +@triton.jit +def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr): + row = tl.program_id(0) + end = tl.load(row_ends_ptr + row) + src_ptr += row.to(tl.int64) * stride + dst_ptr += row.to(tl.int64) * stride + for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK): + offs = start + tl.arange(0, BLOCK) + mask = offs < end + vals = tl.load(src_ptr + offs, mask=mask) + tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask) + + +def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor: + """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels. + + Only columns below ``row_ends`` (int32, one per row) are converted, the + span the consumers read; DeepGEMM's aligned row stride is kept. + """ + if logits.dtype == torch.float32: + return logits + out = torch.empty_strided( + logits.shape, logits.stride(), dtype=torch.float32, device=logits.device + ) + num_rows, width = logits.shape + if num_rows == 0 or width == 0: + return out + assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32 + block = 1024 + splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows)) + _upcast_logits_kernel[(num_rows, splits)]( + logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block + ) + return out + + def _gather_workspace_shapes( total_seq_lens: int, head_dim: int, @@ -537,7 +575,7 @@ def sparse_attn_indexer( kv_cache, slot_mapping_for_cache, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) # The indexer and main MLA may classify the same short extend differently # because they use independent decode thresholds. Only the main MLA route @@ -649,6 +687,9 @@ def sparse_attn_indexer( q_scale_slice = q_scale[row_start:row_end] if q_scale is not None else None topk_indices = topk_indices_buffer[row_start:row_end, :topk_tokens] + # DeepGEMM returns compressed logits: row i starts at K position + # cu_seqlen_ks[i]. None means every row starts at column zero. + row_starts: torch.Tensor | None = None if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) @@ -662,7 +703,9 @@ def sparse_attn_indexer( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) if current_platform.is_xpu(): if q_scale_slice is not None: raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") @@ -674,14 +717,23 @@ def sparse_attn_indexer( cu_seqlen_ks, cu_seqlen_ke, ) + row_starts = cu_seqlen_ks + row_ends = cu_seqlen_ke else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) num_rows = logits.shape[0] if candidate_blocks is not None: @@ -694,8 +746,8 @@ def sparse_attn_indexer( if candidate_write: _select_candidate_blocks( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates.shape[1], candidate_block_size, chunk_candidates, @@ -703,15 +755,15 @@ def sparse_attn_indexer( else: _apply_candidate_mask( logits, - cu_seqlen_ks, - cu_seqlen_ke, + row_starts, + row_ends, chunk_candidates, candidate_block_size, ) ops.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends) if row_starts is None else row_starts, + row_ends, topk_indices, num_rows, logits.stride(0), @@ -729,7 +781,7 @@ def sparse_attn_indexer( dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, - row_starts=cu_seqlen_ks, + row_starts=row_starts, ) if row_shard is not None: @@ -822,16 +874,22 @@ def sparse_attn_indexer( max_model_len, ) else: - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_model_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_model_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] if candidate_blocks is not None: diff --git a/vllm/models/deepseek_v32/common/kernels.py b/vllm/models/deepseek_v32/common/kernels.py index 255b1f5aaf2f..f876120a6ef2 100644 --- a/vllm/models/deepseek_v32/common/kernels.py +++ b/vllm/models/deepseek_v32/common/kernels.py @@ -6,7 +6,9 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # Cache of tiny 1-element dummy tensors (per device, dtype) reused by the # has_indexer=False path so the indexer args don't allocate every call. @@ -87,6 +89,7 @@ def _fp8_quant_and_cache_write( SHUFFLE: tl.constexpr, BLOCK_TILE: tl.constexpr, HEAD_TILE: tl.constexpr, + MX_SCALE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FNUZ: tl.constexpr, ): @@ -111,6 +114,8 @@ def _fp8_quant_and_cache_write( tl.store(kv_cache_ptr + block_start + value_off, k_fp8, mask=mask) scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4 + if MX_SCALE: + scale = pack_ue8m0_x4(scale) tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale) @@ -167,6 +172,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE: tl.constexpr, INDEXER_CACHE_BLOCK_TILE: tl.constexpr, INDEXER_CACHE_HEAD_TILE: tl.constexpr, + INDEXER_CACHE_MX_SCALE: tl.constexpr, # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr) mla_cache_ptr, mla_cache_block_stride, @@ -501,6 +507,7 @@ def _fused_norm_rope_kernel( INDEXER_CACHE_SHUFFLE, INDEXER_CACHE_BLOCK_TILE, INDEXER_CACHE_HEAD_TILE, + INDEXER_CACHE_MX_SCALE, FP8_MAX, USE_FNUZ, ) @@ -571,7 +578,10 @@ def fused_norm_rope( assert slot_mapping is not None if indexer_slot_mapping is None: indexer_slot_mapping = slot_mapping - idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32) + idx_cache_mx_scale = mqa_logits_use_mx_scales() + idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view( + torch.int32 if idx_cache_mx_scale else torch.float32 + ) idx_cache_block_size = indexer_k_cache.shape[1] idx_cache_stride = indexer_k_cache.shape[2] # The caller's cache reports whether its reader expects the shuffled @@ -591,6 +601,7 @@ def fused_norm_rope( ) else: idx_cache_scale_view = None + idx_cache_mx_scale = False idx_cache_block_size = 1 idx_cache_stride = 0 idx_cache_shuffle = False @@ -701,6 +712,7 @@ def fused_norm_rope( idx_cache_shuffle, _INDEXER_CACHE_BLOCK_TILE, idx_cache_head_tile, + idx_cache_mx_scale, # MLA KV cache (uses same slot_mapping) mla_kv_cache, mla_block_stride, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 50ca86558e05..703169f4829f 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -33,7 +33,9 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales from vllm.utils.math_utils import round_up +from vllm.v1.attention.ops.common import pack_ue8m0_x4 if current_platform.is_rocm(): from vllm.platforms.rocm import _ON_GFX950 @@ -742,6 +744,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( TOKEN_STRIDE: tl.constexpr, # 128 for indexer SCALE_DIM: tl.constexpr, # 4 for indexer (1 float32) KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # store the scale as 4 packed UE8M0 bytes ): """Fused compress → RMSNorm → RoPE → FP8 quant → store. @@ -878,7 +881,10 @@ def _fused_kv_compress_norm_rope_insert_indexer_attn( # Single float32 scale scale_val = tl.exp2(exponent) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) + if MX_SCALE: + tl.store(scale_ptr.to(tl.pointer_type(tl.int32)), pack_ue8m0_x4(scale_val)) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), scale_val) # ============================================================================= @@ -927,6 +933,7 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn( TOKEN_STRIDE: tl.constexpr, # HEAD_SIZE // 2 = 64 packed bytes/token SCALE_DIM: tl.constexpr, # HEAD_SIZE // QUANT_BLOCK = 4 ue8m0 bytes/token KV_BLOCK_STRIDE: tl.constexpr, + MX_SCALE: tl.constexpr, # unused for MXFP4 (kept for signature parity) ): """Fused compress → RMSNorm → RoPE → MXFP4 quant → store. @@ -1102,6 +1109,7 @@ class CompileKey: block_size: int kv_cache_block_size: int kv_block_stride: int + mx_scale: bool # Canonical kernel for the launcher's arg_names mapping; the launched # variant (fp8/mxfp4, same arg_names) is picked per call in __call__. @@ -1181,6 +1189,7 @@ def dispatch( # type: ignore[override] block_size=state_block_size, kv_cache_block_size=kv_cache_block_size, kv_block_stride=kv_block_stride, + mx_scale=not use_fp4_cache and mqa_logits_use_mx_scales(), ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -1301,6 +1310,7 @@ def __call__( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), num_warps=1, **pdl_kwargs, ) diff --git a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py index ba8ddbd95a9a..fe203e917861 100644 --- a/vllm/models/deepseek_v41/common/ops/indexer_k_store.py +++ b/vllm/models/deepseek_v41/common/ops/indexer_k_store.py @@ -17,6 +17,8 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2 @@ -51,7 +53,8 @@ def indexer_k_norm_rope_store( kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip). compress_ratio: group size; keys are emitted at group boundaries. use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else - per-token FP8 with a single fp32 scale. + per-token FP8 with a single fp32 scale (four packed ue8m0 bytes + under `mqa_logits_use_mx_scales`). """ num_tokens = kv_slot_mapping.numel() @@ -111,6 +114,7 @@ def indexer_k_norm_rope_store( KV_BLOCK_STRIDE=k_cache.stride(0), FP8_MAX=448.0, USE_FP4=use_fp4_cache, + MX_SCALE=not use_fp4_cache and mqa_logits_use_mx_scales(), SHUFFLE=shuffle, BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE, HEAD_TILE_SIZE=_HEAD_TILE_SIZE, @@ -139,6 +143,7 @@ def _indexer_k_norm_rope_quant_store_kernel( KV_BLOCK_STRIDE: tl.constexpr, FP8_MAX: tl.constexpr, USE_FP4: tl.constexpr, + MX_SCALE: tl.constexpr, SHUFFLE: tl.constexpr, BLOCK_TILE_SIZE: tl.constexpr, HEAD_TILE_SIZE: tl.constexpr, @@ -252,4 +257,10 @@ def _indexer_k_norm_rope_quant_store_kernel( tl.store(val_ptr + tiled, x_uint8) else: tl.store(val_ptr + block, x_uint8) - tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) + if MX_SCALE: + tl.store( + scale_ptr.to(tl.pointer_type(tl.int32)), + pack_ue8m0_x4(tl.exp2(exponent)), + ) + else: + tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent)) diff --git a/vllm/models/glm5next/nvidia/ops/kpool_compress.py b/vllm/models/glm5next/nvidia/ops/kpool_compress.py index eb289c65cf68..310d60dfd41e 100644 --- a/vllm/models/glm5next/nvidia/ops/kpool_compress.py +++ b/vllm/models/glm5next/nvidia/ops/kpool_compress.py @@ -15,6 +15,8 @@ import torch from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import mqa_logits_use_mx_scales +from vllm.v1.attention.ops.common import pack_ue8m0_x4 # The GLM-5.3-Flash indexer head dimension is fixed at 128. INDEX_HEAD_DIM = 128 @@ -157,6 +159,7 @@ def _kpool_softmax_rotate_write_cache_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, HAS_WRITE_MASK: tl.constexpr, RETURN_COMPRESSED: tl.constexpr, WRITE_CACHE: tl.constexpr, @@ -245,7 +248,14 @@ def _kpool_softmax_rotate_write_cache_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask) - tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + mask=do_write, + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write) if RETURN_COMPRESSED: tl.store( @@ -324,6 +334,7 @@ def kpool_compress_and_write_cache( buf_fp8 = buf.view(torch.float8_e4m3fn) buf_fp32 = buf.view(torch.float32) + mx_scale = write_cache and mqa_logits_use_mx_scales() # bytes per page (last dim of kv_cache) viewed as uint8 buf_numel_per_page = buf.stride(0) s_offset_nbytes_in_page = page_size * head_dim @@ -361,7 +372,9 @@ def kpool_compress_and_write_cache( POOL_SIZE=slot_k.shape[1], HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mx_scale, + MX_SCALE=mx_scale, HAS_WRITE_MASK=has_write_mask, RETURN_COMPRESSED=return_compressed, WRITE_CACHE=write_cache, @@ -479,6 +492,7 @@ def _kpool_decode_update_batched_kernel( HEAD_DIM: tl.constexpr, S_OFFSET_NBYTES_IN_PAGE: tl.constexpr, ROUND_SCALE: tl.constexpr, + MX_SCALE: tl.constexpr, BLOCK_D: tl.constexpr, ): """One program per request; iterates its NEXT_N verify tokens in order. @@ -604,7 +618,13 @@ def _kpool_decode_update_batched_kernel( + loc_token_offset_in_page ) tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask) - tl.store(buf_fp32_ptr + out_s_offset, scale) + if MX_SCALE: + tl.store( + buf_fp32_ptr.to(tl.pointer_type(tl.int32)) + out_s_offset, + pack_ue8m0_x4(scale), + ) + else: + tl.store(buf_fp32_ptr + out_s_offset, scale) # Stash the current token AFTER any completion read so the completion # uses prior stashes (and the current token's own key/score via @@ -718,7 +738,9 @@ def kpool_decode_update_and_maybe_write_cache_batched( KPOOL_HEAD=tail_kv_cache.stride(1), HEAD_DIM=head_dim, S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, - ROUND_SCALE=round_scale, + # MX scales need power-of-two scales. + ROUND_SCALE=round_scale or mqa_logits_use_mx_scales(), + MX_SCALE=mqa_logits_use_mx_scales(), BLOCK_D=triton.next_power_of_2(head_dim), ) diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 154fdcda900e..1a3fbedb8387 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -16,6 +16,7 @@ PrefillRowShard, _merge_dcp_topk_global, _reserve_row_shard_exchange_transient, + as_fp32_logits, ) from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -29,7 +30,13 @@ ) from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops from vllm.platforms import current_platform -from vllm.utils.deep_gemm import has_deep_gemm +from vllm.utils.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + has_deep_gemm, + mqa_logits_mx_inputs, + mqa_logits_use_mx_scales, +) from vllm.utils.import_utils import has_cutedsl from vllm.utils.torch_utils import ( LayerNameType, @@ -253,7 +260,7 @@ def sparse_attn_indexer_kpool( kv_cache, slot_mapping, quant_block_size, - scale_fmt, + "mxfp8" if mqa_logits_use_mx_scales() else scale_fmt, ) topk_indices_buffer[: hidden_states.shape[0]] = -1 @@ -343,20 +350,29 @@ def sparse_attn_indexer_kpool( else: q_slice_cast = q_slice k_quant_cast = k_quant - k_scale_cast = k_scale.view(torch.float32).squeeze(-1) - from vllm.utils.deep_gemm import fp8_fp4_mqa_logits + k_scale_cast = k_scale.view( + torch.int32 if mqa_logits_use_mx_scales() else torch.float32 + ).squeeze(-1) num_rows = q_slice.shape[0] if chunk.local_total_seq_lens == 0: logits = q_slice.new_empty((num_rows, 0), dtype=torch.float32) else: - logits = fp8_fp4_mqa_logits( - (q_slice_cast, q_scale_slice), - (k_quant_cast, k_scale_cast), - weights[row_start:row_end], - cu_seqlen_ks, - cu_seqlen_ke, - clean_logits=False, + q_scale_slice, chunk_weights = mqa_logits_mx_inputs( + q_slice_cast, q_scale_slice, weights[row_start:row_end] + ) + # Compressed: row i starts at K position cu_seqlen_ks[i]. + row_ends = cu_seqlen_ke - cu_seqlen_ks + logits = as_fp32_logits( + fp8_fp4_mqa_logits( + (q_slice_cast, q_scale_slice), + (k_quant_cast, k_scale_cast), + chunk_weights, + cu_seqlen_ks, + cu_seqlen_ke, + chunk.max_seqlen_k, + ), + row_ends, ) # kpool: logits are pool-granular (compress_ratio == index_kpool), @@ -376,8 +392,8 @@ def sparse_attn_indexer_kpool( else: torch.ops._C.top_k_per_row_prefill( logits, - cu_seqlen_ks, - cu_seqlen_ke, + torch.zeros_like(row_ends), + row_ends, topk_dst, num_rows, logits.stride(0), @@ -392,7 +408,7 @@ def sparse_attn_indexer_kpool( dcp_rank, dcp_world_size, pool_cp_interleave, - row_starts=cu_seqlen_ks, + row_starts=None, ) if index_kpool > 1: @@ -594,18 +610,22 @@ def sparse_attn_indexer_kpool( if use_fp4_cache else padded_q_quant_decode_tokens ) - from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits - - logits = fp8_fp4_paged_mqa_logits( - (padded_q_quant_cast, padded_q_scale), - kv_cache, - padded_weights[:num_padded_tokens], - seq_lens, - decode_metadata.block_table, - decode_metadata.schedule_metadata, - max_model_len=max_pool_len, - clean_logits=False, - indices=decode_metadata.indices, + padded_q_scale, decode_weights = mqa_logits_mx_inputs( + padded_q_quant_cast, padded_q_scale, padded_weights[:num_padded_tokens] + ) + # SM100 decode is varlen: one row and context length per token. + logits = as_fp32_logits( + fp8_fp4_paged_mqa_logits( + (padded_q_quant_cast, padded_q_scale), + kv_cache, + decode_weights, + seq_lens, + decode_metadata.block_table, + decode_metadata.schedule_metadata, + max_model_len=max_pool_len, + indices=decode_metadata.indices, + ), + seq_lens.view(-1), ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools, diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 04c4508f0ba9..3dc9be5d9385 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -591,37 +591,79 @@ def transform_sf_into_required_layout(*args, **kwargs): ) +@functools.cache +def mqa_logits_use_mx_scales() -> bool: + """Whether the MQA logits kernels take only MX-scaled inputs. + + On SM100 they need MXFP8/MXFP4 Q and K with packed UE8M0 scales and BF16 + weights, and return BF16 logits. FP8 indexer K caches then store each + token's power-of-two scale as four equal UE8M0 bytes instead of an fp32. + """ + return current_platform.is_cuda() and current_platform.is_device_capability_family( + 100 + ) + + +# Four packed UE8M0 bytes encoding 1.0. +_UNIT_UE8M0_X4 = 0x7F7F7F7F +# Grown buffers are kept alive: CUDA graphs captured earlier still read them. +_unit_mx_scales: dict[torch.device, list[torch.Tensor]] = {} + + +def mqa_logits_mx_inputs( + q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor +) -> tuple[torch.Tensor | None, torch.Tensor]: + """Adapt indexer Q scales and weights to the SM100 MQA logits kernels. + + FP8 Q folds its per-(token, head) power-of-two scale into ``weights``, so + it is passed with unit MX scales, which is exact since relu commutes with + a positive scale. Returns ``(q_scale, weights)`` unchanged off SM100. + """ + if not mqa_logits_use_mx_scales(): + return q_scale, weights + if q_scale is None: + numel = q.numel() // q.shape[-1] + bufs = _unit_mx_scales.setdefault(q.device, []) + if not bufs or bufs[-1].numel() < numel: + size = max(numel, 2 * bufs[-1].numel() if bufs else 0) + bufs.append( + torch.full((size,), _UNIT_UE8M0_X4, dtype=torch.int32, device=q.device) + ) + q_scale = bufs[-1][:numel].view(q.shape[:-1]) + return q_scale, weights.to(torch.bfloat16) + + def fp8_fp4_mqa_logits( q: tuple[torch.Tensor, torch.Tensor | None], kv: tuple[torch.Tensor, torch.Tensor], weights: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, - clean_logits: bool, + max_seqlen_k: int, ) -> torch.Tensor: """Compute MQA logits for a single sequence without KV paging. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)`` where ``scales`` is None for FP8 Q - (per-token scale is folded into ``weights``) and a packed block-scale - tensor for MXFP4 Q. + ``q = (values, scales_or_None)``. SM100 takes MXFP8 or MXFP4 with packed + UE8M0 block scales; SM90/SM120 also take FP8 Q with ``q_scale=None`` + (per-token scale folded into ``weights``). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [M, H, D] - float8_e4m3fn and q_scale is None (per-token scale is folded - into ``weights``). FP4 path: q_values is packed uint8 and - q_scale is the companion block-scale tensor. - kv: Tuple `(k_packed, k_scales)` — FP8 layout is [N, D] - float8_e4m3fn plus fp32 scales [N]; FP4 layout is packed uint8. - weights: weights of shape [M, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape [M, H, D] (packed + to D/2 bytes for MXFP4) and q_scale int32 [M, H] or None. + kv: Tuple ``(k_values, k_scales)`` of shape [N, D] and [N]; k_scales + is int32 packed UE8M0 with MX scales, fp32 otherwise. + weights: Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere. cu_seqlen_ks: Start indices (inclusive) for valid K per query position, shape [M], dtype int32. cu_seqlen_ke: End indices (exclusive) for valid K per query position, shape [M], dtype int32. - clean_logits: Whether to clean the unfilled logits into `-inf`. + max_seqlen_k: Upper bound on ``cu_seqlen_ke - cu_seqlen_ks``. Returns: - Logits tensor of shape [M, N], dtype `torch.float32`. + Compressed logits of shape [M, max_seqlen_k]: row ``i`` holds K + positions ``[cu_seqlen_ks[i], cu_seqlen_ke[i])`` starting at column + zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -633,7 +675,7 @@ def fp8_fp4_mqa_logits( weights, cu_seqlen_ks, cu_seqlen_ke, - clean_logits=clean_logits, + max_seqlen_k, ) @@ -699,24 +741,25 @@ def fp8_fp4_paged_mqa_logits( block_tables: torch.Tensor, schedule_metadata: torch.Tensor, max_model_len: int, - clean_logits: bool, indices: torch.Tensor | None = None, ) -> torch.Tensor: """Compute MQA logits using a paged KV-cache. Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes - ``q = (values, scales_or_None)``; pass ``(q_tensor, None)`` for the FP8 - path and ``(q_values, q_scale)`` for MXFP4. + ``q = (values, scales_or_None)``; see `fp8_fp4_mqa_logits` for the + per-architecture formats. SM100 requires the varlen form (``indices`` + with next_n == 1). Args: - q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is - [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path: - q_values is packed uint8 and q_scale is the companion - block-scale tensor. - kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1, - D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos) - storing the float dequant scale. - weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. + q: Tuple ``(q_values, q_scale)``, q_values of shape + [B, next_n, H, D] (packed to D/2 bytes for MXFP4) and q_scale + int32 [B, next_n, H] or None. + kv_cache: Paged KV-cache of shape [num_blocks, block_size, 1, D+4] + (D/2+4 for MXFP4), dtype `torch.uint8`; the last 4 bytes per + (block, pos) hold the int32 packed UE8M0 scales with MX scales, + or the fp32 dequant scale otherwise. + weights: Tensor of shape [B * next_n, H], bf16 on SM100 and fp32 + elsewhere. context_lens: Tensor of shape [B], dtype int32; effective context length for each batch element. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical @@ -724,12 +767,11 @@ def fp8_fp4_paged_mqa_logits( schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. - clean_logits: Whether to clean the unfilled logits into `-inf`. indices: Optional request index for each varlen row. Returns: - Logits tensor of shape [B * next_n, max_model_len], dtype - `torch.float32`. + Logits tensor of shape [B * next_n, max_model_len]; columns past each + row's context length are not cleaned. bf16 on SM100, fp32 elsewhere. """ _lazy_init() @@ -751,7 +793,6 @@ def fp8_fp4_paged_mqa_logits( block_tables, schedule_metadata, max_model_len, - clean_logits=clean_logits, **kwargs, ) @@ -1052,6 +1093,8 @@ def should_use_deepgemm_for_fp8_linear( "fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits", "get_paged_mqa_logits_metadata", + "mqa_logits_mx_inputs", + "mqa_logits_use_mx_scales", "native_next_n_supported", "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2d797ae66283..0382a81fe044 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -467,6 +467,9 @@ class DeepseekV32IndexerPrefillChunkMetadata: cu_seq_lens: torch.Tensor token_to_seq: torch.Tensor total_seq_lens: int + # Host upper bound on cu_seqlen_ke - cu_seqlen_ks, sizing the compressed + # DeepGEMM prefill logits. + max_seqlen_k: int token_start: int token_end: int num_reqs: int @@ -1957,6 +1960,7 @@ def build_prefill_chunk_metadata( cu_seq_lens=cu_seq_lens, token_to_seq=token_to_seq, total_seq_lens=total_seq_lens, + max_seqlen_k=int(compressed_seq_lens_cpu[start_idx:end_idx].max()), block_table=block_table[start_idx:end_idx], token_start=token_start, token_end=token_end, diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 5c1677cac479..6936be9fa169 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -154,6 +154,13 @@ def __call__( ) +@triton.jit +def pack_ue8m0_x4(scale): + """Encode a power-of-two fp32 scale as four equal packed UE8M0 bytes, the + MX scales of a 128-wide row (see `mqa_logits_use_mx_scales`).""" + return ((scale.to(tl.int32, bitcast=True) >> 23) & 0xFF) * 0x01010101 + + def pack_seq_triton( x: torch.Tensor, lengths: torch.Tensor, From e2ad0ce5f5a944e10ed2ea664e4a24bf85791346 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 13 Oct 2026 14:25:16 +0800 Subject: [PATCH 14/23] [Bugfix][LoRA] Slice expert LoRA by the unpadded per-rank size under TP (#61235) Signed-off-by: aoshen524 Co-authored-by: aoshen524 Co-authored-by: Claude Opus 5.5 Co-authored-by: linitra24 --- vllm/lora/layers/fused_moe.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index c58c18e71ff7..19a1128ce145 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from transformers import PreTrainedConfig from vllm import envs @@ -110,6 +111,15 @@ def use_ep(self) -> bool: def intermediate_size_per_partition(self) -> int: return self.moe_config.intermediate_size_per_partition + @property + def _tp_shard_size(self) -> int: + """Per-rank stride of an adapter along the expert intermediate dim: the + base weights are sharded by the unpadded size even when a backend pads.""" + unpadded = self.moe_config.intermediate_size_per_partition_unpadded + if unpadded is None: + return self.intermediate_size_per_partition + return unpadded + def _init_lora_stream_context(self) -> None: self._lora_stream: torch.cuda.Stream | None = None self._events: tuple[torch.cuda.Event, ...] | None = None @@ -311,7 +321,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -322,7 +332,7 @@ def _slice_w2_a(self, w2_lora_a: torch.Tensor) -> torch.Tensor: if self.tp_size == 1: return w2_lora_a # w2_lora_a shape (num_experts,rank,input_size) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size @@ -536,7 +546,7 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): return w13_lora_b # w13_lora_b shape (num_experts,output_size,rank) - shard_size = self.intermediate_size_per_partition + shard_size = self._tp_shard_size start_idx = self.tp_rank * shard_size end_idx = (self.tp_rank + 1) * shard_size # HACK: Currently, only GPT-OSS is in interleaved order @@ -557,6 +567,11 @@ def _slice_w13_b(self, w13_lora_b: torch.Tensor): w3_lora_b = w13_lora_b[:, slice_size:, :] sliced_w1_lora_b = w1_lora_b[:, start_idx:end_idx, :] sliced_w3_lora_b = w3_lora_b[:, start_idx:end_idx, :] + # w3 starts at the allocated (possibly padded) per-rank size, as + # in the base w13 weight. + pad = self.intermediate_size_per_partition - sliced_w1_lora_b.shape[1] + if pad > 0: + sliced_w1_lora_b = F.pad(sliced_w1_lora_b, (0, 0, 0, pad)) return torch.cat([sliced_w1_lora_b, sliced_w3_lora_b], dim=1) From 4e4eab73ed2b4e00931a4636e926b8df1cb1fe50 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 12 Oct 2026 23:57:32 -0700 Subject: [PATCH 15/23] [CI] Shard Model Executor and L4 Pipeline + Context Parallelism automatically (#61202) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 --- .buildkite/test_areas/distributed.yaml | 1 + .buildkite/test_areas/model_executor.yaml | 1 + 2 files changed, 2 insertions(+) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index b30cbac0ced4..2f929b53b840 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -641,6 +641,7 @@ steps: - label: ":nvidia: (L4) Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus + automatic_shard: true timeout_in_minutes: 55 working_dir: "/vllm-workspace/tests" device: l4 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index a23359b768a4..c9f467941af9 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Model Executor" device: h200_35gb key: model-executor + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - vllm/engine/arg_utils.py From bc3ddb198a0d5eb864bf484c4719864b671cc93b Mon Sep 17 00:00:00 2001 From: trienebutdiene <59642603+kimseunghyun-kr@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:30:19 +0800 Subject: [PATCH 16/23] [Model][Qwen4Exp] FP8 QSA main KV cache on Ampere (sm_80, sm_86) (#60602) Signed-off-by: kimseunghyun Signed-off-by: kimseunghyun Signed-off-by: Thien Tran Co-authored-by: kimseunghyun Co-authored-by: Claude Opus 5.5 Co-authored-by: Thien Tran --- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 69 +++++++++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 21 +++++--- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 401048e02c55..c7fe6014497e 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -22,6 +22,28 @@ def _is_sm120() -> bool: return current_platform.get_device_capability() == (12, 0) +@triton.jit +def _maybe_dequant(tile): + if tile.dtype == tl.uint8: + # e4m3 bytes below sm_89 (no Triton fp8e4nv pointer type): sign to bit + # 15, exp+mantissa to bits 13..7 gives the fp16 of x * 2^-8. Exact for + # every finite code; the two NaN codes decode to +-480 after the host + # scale instead of NaN. + code = tile.to(tl.uint16) + bits = ((code & 0x80) << 8) | ((code & 0x7F) << 7) + return bits.to(tl.float16, bitcast=True) + elif tile.dtype == tl.float8e4nv: + return tile.to(tl.float16) + else: + return tile + + +@lru_cache(maxsize=1) +def _is_sm80() -> bool: + """True on sm_80 (A100/A30): selects the wide FP8 byte-decode tiles.""" + return current_platform.get_device_capability() == (8, 0) + + @lru_cache(maxsize=1) def _is_sm90() -> bool: """True on sm_90 (H100/H200/H20): selects the sm_90 tuning table.""" @@ -69,7 +91,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, - IS_FP8: tl.constexpr, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -141,6 +162,9 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[None, :], other=0.0, ) + # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK + # measured slower here and less accurate). + keys = _maybe_dequant(keys).to(query.dtype) values = tl.load( v_cache_ptr + safe_page[:, None] * stride_v_block @@ -150,10 +174,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( mask=valid[:, None], other=0.0, ) - if IS_FP8: - # e4m3 -> Q dtype is exact; keep the QK dot in Q's dtype (fp8 QK - # measured slower here and less accurate). - keys = keys.to(query.dtype) + # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online + # softmax) so fp16 has the range, its wider mantissa is more + # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. + values = _maybe_dequant(values) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16; for fp8 # caches the K dequant scale is already folded into softmax_scale on the @@ -165,11 +189,6 @@ def _qsa_sparse_paged_gqa_splitk_kernel( probabilities = tl.where( valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 ) - if IS_FP8: - # Dequant V to fp16 (not bf16) for the PV dot: P <= 1 (online - # softmax) so fp16 has the range, its wider mantissa is more - # accurate, and the fp8->fp16 upcast with an fp16 PV dot is faster. - values = values.to(tl.float16) accumulator = tl.dot( probabilities.to(values.dtype), values, @@ -585,6 +604,10 @@ def _select_config( BLOCK_N, target_splits, num_warps = 64, 4, 2 else: BLOCK_N, target_splits, num_warps = 64, 1, 2 + if is_fp8 and _is_sm80(): + # Wider tiles measured faster for the byte-decode path on A100; + # unmeasured on sm_86, which keeps the table's tiles. + BLOCK_N, num_warps = 128, 4 num_tiles = triton.cdiv(num_columns, BLOCK_N) # Never more splits than tiles, never empty. num_splits = min(target_splits, num_tiles) @@ -607,7 +630,9 @@ def qsa_sparse_paged_attention( ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 or FP8-e4m3 K/V caches. - With fp8 caches, k_scale/v_scale are the layer's per-tensor dequant scales + An fp8 cache is a native e4m3 view on sm_89+ and the raw uint8 storage + below sm_89 (no Triton fp8e4nv pointer type there); the kernel dequantizes + either on load. k_scale/v_scale are the layer's per-tensor dequant scales as host floats (e.g. layer._k_scale_float/_v_scale_float): k_scale is pre-multiplied into the softmax scale and v_scale becomes the kernel's output scale, so the kernel needs no device scale buffers. @@ -636,15 +661,21 @@ def qsa_sparse_paged_attention( assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 assert q.dtype == torch.bfloat16 assert k_cache.dtype == v_cache.dtype - is_fp8 = k_cache.dtype == torch.float8_e4m3fn + is_fp8 = k_cache.dtype != torch.bfloat16 if is_fp8: assert k_scale is not None and v_scale is not None # Host pre-multiply: fold the K dequant scale into the attention scale # and pass V's dequant scale as the kernel's output scale. softmax_scale = (head_dim**-0.5) * float(k_scale) output_scale = float(v_scale) + if k_cache.dtype == torch.uint8: + # _maybe_dequant's byte decode yields x * 2^-8; fold 2^8 into the + # scales so the decode needs no multiply. + softmax_scale *= 256.0 + output_scale *= 256.0 + else: + assert k_cache.dtype == torch.float8_e4m3fn else: - assert k_cache.dtype == torch.bfloat16 softmax_scale = head_dim**-0.5 output_scale = 1.0 assert logical_indices.dtype == block_table.dtype == torch.int32 @@ -731,7 +762,6 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=partial_warps, num_stages=2, ) @@ -768,9 +798,15 @@ def warmup_qsa_sparse_paged_attention( """Compile every production-reachable split-K/merge specialization.""" head_dim = kv_cache.shape[-1] // 2 key_cache, value_cache = kv_cache.transpose(1, 2).split(head_dim, dim=-1) - # An fp8 cache is allocated as uint8 and viewed as e4m3 at attention time. + # An fp8 cache is allocated as uint8 storage; warm the dtype the impl hands + # over: an e4m3 view on sm_89+, the raw bytes below. is_fp8 = kv_cache.dtype == torch.uint8 - cache_dtype = torch.float8_e4m3fn if is_fp8 else key_cache.dtype + fp8_dtype = ( + torch.float8_e4m3fn + if current_platform.has_device_capability(89) + else torch.uint8 + ) + cache_dtype = fp8_dtype if is_fp8 else key_cache.dtype num_kv_heads = key_cache.shape[2] group_size = num_query_heads // num_kv_heads block_m = triton.next_power_of_2(group_size) @@ -874,7 +910,6 @@ def warmup_qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, - IS_FP8=is_fp8, num_warps=warps, num_stages=2, grid=(num_rows, num_kv_heads, num_splits), diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 9f0934d5a0c6..9fb36673a5e7 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -297,10 +297,12 @@ def forward_qsa( key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) k_scale = v_scale = None if self.kv_cache_dtype in ("fp8", "fp8_e4m3"): - # The cache is allocated as uint8; reinterpret the e4m3 bytes - # (same itemsize, so shape and strides are preserved). - key_cache = key_cache.view(torch.float8_e4m3fn) - value_cache = value_cache.view(torch.float8_e4m3fn) + # The cache is allocated as uint8; on sm_89+ reinterpret the e4m3 + # bytes (same itemsize, so shape and strides are preserved). Below + # sm_89 the kernel byte-decodes the raw uint8 storage instead. + if current_platform.has_device_capability(89): + key_cache = key_cache.view(torch.float8_e4m3fn) + value_cache = value_cache.view(torch.float8_e4m3fn) # Host-side per-tensor dequant scales (Python floats), as used by # other host-scale backends; folded into the kernel's scales. k_scale = layer._k_scale_float @@ -308,6 +310,7 @@ def forward_qsa( if query.dtype != torch.bfloat16 or key_cache.dtype not in ( torch.bfloat16, torch.float8_e4m3fn, + torch.uint8, ): raise NotImplementedError( "Qwen4Exp QSA requires BF16 Q and BF16 or FP8-e4m3 K/V" @@ -492,9 +495,15 @@ def __init__( ) # One launch does the indexer prepare, the main QK-norm/RoPE/gate and # the main K/V cache write (see QSAIndexer.forward); otherwise all of - # them take the separate kernels. + # them take the separate kernels. The fused writer receives an fp8 + # pointer, which Triton cannot lower below sm_89. self.use_fused_qsa_prepare = ( - self.use_fused_qk_norm_rope_gate and self.indexer.use_fused_pre_indexer + self.use_fused_qk_norm_rope_gate + and self.indexer.use_fused_pre_indexer + and ( + self.kv_cache_torch_dtype != torch.uint8 + or current_platform.has_device_capability(89) + ) ) self.fuse_indexer_projection = vllm_config.lora_config is None if self.fuse_indexer_projection: From 2cc2e68d1c5652a5bc7ae10422697308e2193359 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 15:58:59 +0800 Subject: [PATCH 17/23] [Core][Mamba] Remove leftovers of mamba_cache_mode="all" (#60043) Signed-off-by: zjy0516 Signed-off-by: Jiangyun Zhu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/disaggregated.yaml | 2 +- docs/design/hybrid_kv_cache_manager.md | 2 +- docs/features/automatic_prefix_caching.md | 7 +- .../test_kimi_linear_context_parallel.py | 1 - .../configs/Qwen3.8-Flash-Next-FP8-rocm.yaml | 2 +- tests/kernels/mamba/test_mamba_mixer2.py | 10 +- .../test_mamba_eagle_resume_checkpoint.py | 4 +- tests/v1/core/test_contiguous_kv_packing.py | 9 +- tests/v1/core/test_prefix_caching.py | 90 +++------- .../v1/e2e/general/test_mamba_prefix_cache.py | 2 - tests/v1/e2e/test_hybrid_chunked_prefill.py | 2 - tests/v1/e2e/test_replayssm_decode.py | 4 - .../run_mamba_prefix_cache_test.sh | 2 - .../unit/test_offloading_connector.py | 1 - .../unit/test_remote_prefill_lifecycle.py | 3 + tests/v1/simple_kv_offload/test_scheduler.py | 74 ++++++-- vllm/config/cache.py | 28 ++- vllm/config/model.py | 4 - vllm/config/vllm.py | 10 +- vllm/engine/arg_utils.py | 18 +- .../layers/mamba/mamba_mixer2.py | 6 +- .../layers/mamba/ops/causal_conv1d.py | 164 +----------------- vllm/model_executor/models/config.py | 2 +- vllm/model_executor/models/falcon_h1.py | 2 - .../model_executor/models/granitemoehybrid.py | 2 - vllm/model_executor/models/interfaces.py | 25 +-- vllm/model_executor/models/jamba.py | 2 - vllm/model_executor/models/nemotron_h.py | 2 - vllm/model_executor/models/registry.py | 3 - vllm/model_executor/models/zamba2.py | 8 +- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/core/single_type_kv_cache_manager.py | 3 + 33 files changed, 161 insertions(+), 344 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d815ee9268b2..7f0a351524e1 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3523,7 +3523,7 @@ steps: - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' - export PREFILL_BLOCK_SIZE='2048' - export DECODE_BLOCK_SIZE='2048' - - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b9bd59583b82..4821fac0a20b 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -426,7 +426,7 @@ steps: # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. PREFILL_BLOCK_SIZE: "2048" DECODE_BLOCK_SIZE: "2048" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 9891c82f97c9..3984000c7d8f 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -204,7 +204,7 @@ The second question is the cache eviction policy. For now, we use one LRU queue ### Case 3: mamba models -The prefix caching support of the mamba model is work in progress. Once implemented, models with mamba layer + full attention layer can be supported via the full attention + X algorithm in case 2. +Models with mamba layers + full attention layers support prefix caching via the full attention + X algorithm in case 2, with the Mamba state cached at block boundaries. See [Automatic Prefix Caching](../features/automatic_prefix_caching.md). ## Implementation diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md index cfc8c19acb2b..bfd87e7a2c57 100644 --- a/docs/features/automatic_prefix_caching.md +++ b/docs/features/automatic_prefix_caching.md @@ -15,23 +15,22 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl ## Hybrid Mamba models -Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. +With prefix caching enabled, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold: -- `--mamba-cache-mode align` +- prefix caching enabled - EAGLE/MTP speculative decoding on the Mamba group - `--prefix-match-unit` smaller than the Mamba block size - the model does not use multi-module MTP ```bash vllm serve \ - --mamba-cache-mode align \ --prefix-match-unit 64 \ --enable-mamba-shared-prefix-checkpoint ``` -`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect. +`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. For hybrid Mamba models that is the block size itself, so no sub-block boundary exists and the flag has no effect. Choose a value that divides the block size of every prefix-cacheable KV cache group, and that is a multiple of the per-state compression ratio for models that use one, such as sparse MLA. vLLM validates both at startup and names the offending sizes in the error. Read the served block size from the startup log. 64 is a reasonable starting point. diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index 06348074bb14..8bddf6e609c0 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -180,7 +180,6 @@ def _run_k3_partial_prefix_reuse( block_size=block_size, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", ) as runner: assert ( runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml index db43869eb4ae..c35614cd5f8c 100644 --- a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml +++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml @@ -23,4 +23,4 @@ generation_config: extra_body: chat_template_kwargs: enable_thinking: true -server_args: "--tensor-parallel-size 4 --mamba-cache-mode align" +server_args: "--tensor-parallel-size 4" diff --git a/tests/kernels/mamba/test_mamba_mixer2.py b/tests/kernels/mamba/test_mamba_mixer2.py index 973e7885c680..db2274c3c933 100644 --- a/tests/kernels/mamba/test_mamba_mixer2.py +++ b/tests/kernels/mamba/test_mamba_mixer2.py @@ -96,15 +96,17 @@ def mixer2_gated_norm_tensor_parallel( gate_states = torch.randn(batch_size, seq_len, hidden_size) # create gated-norm with TP - mixer = Mixer2RMSNormGated( - full_hidden_size=hidden_size, - full_n_groups=n_groups, - ) + with ensure_current_vllm_config(): + mixer = Mixer2RMSNormGated( + full_hidden_size=hidden_size, + full_n_groups=n_groups, + ) mixer.weight.weight_loader(mixer.weight, weight) # load # create gated-norm without TP to compute reference # - utilize mock patching to disable TP when with ( + ensure_current_vllm_config(), unittest.mock.patch( "vllm.model_executor.layers.mamba.mamba_mixer2." "get_tensor_model_parallel_world_size", diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 324aef425b89..54ae43e8e232 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -38,9 +38,7 @@ def _manager( num_prefill_lookahead=0, ): init_none_hash(sha256) - config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] - ) + config = _make_hybrid_kv_cache_config(block_size, num_blocks, ["full", "mamba"]) if eagle_group is not None: groups = list(config.kv_cache_groups) groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 64cd55054015..f71fe160d83d 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -278,7 +278,14 @@ def test_scratch_group_survives_computed_block_truncation(self): neither divides the hit length nor bounds the (empty) block list.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True - groups = get_kv_cache_groups(config, _make_csa_linear_specs()) + config.cache_config.mamba_cache_mode = "align" + specs = { + name: replace(spec, mamba_cache_mode="align") + if isinstance(spec, MambaSpec) + else spec + for name, spec in _make_csa_linear_specs().items() + } + groups = get_kv_cache_groups(config, specs) kv_cache_config = get_kv_cache_config_from_groups( config, groups, available_memory=BYTES_PER_BLOCK * 64 ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 836286d127fa..52455c608305 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1617,6 +1617,7 @@ def make_kv_cache_config_hybrid_model( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) return KVCacheConfig( @@ -1652,6 +1653,7 @@ def make_kv_cache_config_three_types( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) elif third_spec_type == "sliding_window": third_spec = SlidingWindowSpec( @@ -1732,6 +1734,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): block_size=block_size, shapes=(1, 1), dtypes=(torch.float32,), + mamba_cache_mode="align", ), ), ], @@ -1757,10 +1760,13 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, len(req0.prompt_token_ids), 0, computed_blocks - ) - assert blocks is not None + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the shared prefix. + for num_new_tokens in (len(common_token_ids), 7): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + req0.num_computed_tokens += num_new_tokens + computed_blocks = None req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) @@ -2285,7 +2291,7 @@ def _make_hybrid_kv_cache_config( - "full": FullAttentionSpec - "sliding_window": SlidingWindowSpec with window=2*block_size - "sliding_window_large": SlidingWindowSpec with window=4*block_size - - "mamba": MambaSpec + - "mamba": MambaSpec in prefix-caching ("align") mode """ spec_map = { @@ -2313,11 +2319,6 @@ def _make_hybrid_kv_cache_config( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), - ), - "mamba_align": lambda: MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), mamba_cache_mode="align", ), } @@ -2427,14 +2428,16 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): assert not computed_blocks.blocks[0] # No cache hit initially assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, 55, len(computed_blocks.blocks[0]) * block_size, computed_blocks - ) - assert blocks is not None - # Should have blocks for all groups - assert len(blocks.get_block_ids()) == num_groups - - manager.new_step_starts() + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the common prefix. + for num_new_tokens in (len(common_token_ids), len(unique_token_ids)): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + # Should have blocks for all groups + assert len(blocks.get_block_ids()) == num_groups + req0.num_computed_tokens += num_new_tokens + computed_blocks = None + manager.new_step_starts() # Second request: should hit cached blocks for common prefix req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn) @@ -2549,7 +2552,7 @@ def test_prefill_hybrid_model_mamba_align(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = make_kv_cache_manager( kv_cache_config, @@ -2585,7 +2588,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): """ block_size = 16 manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, @@ -2660,7 +2663,7 @@ def test_hybrid_model_mamba_align_with_dynamic_draft_tokens(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = KVCacheManager( kv_cache_config, @@ -5896,49 +5899,6 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(): - """Manager-level check of the full wiring: a pinned shared-prefix boundary - (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style - detection) keeps its Mamba state block cached under - ``prefix_cache_retention_interval=0``, which otherwise retains only the - end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request.""" - block_size = 16 - - # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. - token_ids = [i for i in range(16) for _ in range(block_size)] - - def cached_mamba_blocks(shared_prefix_boundary): - # Fresh manager per scenario so cached blocks don't leak between runs. - manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"]), - max_model_len=8192, - enable_caching=True, - hash_block_size=block_size, - retention_interval=0, - ) - req = make_request("r", token_ids, block_size, sha256) - req.shared_prefix_boundary = shared_prefix_boundary - computed_blocks, num_computed, _ = manager.get_computed_blocks(req) - blocks = manager.allocate_slots( - req, len(token_ids), num_computed, computed_blocks - ) - assert blocks is not None - pool = manager.block_pool - return { - i - for i in range(16) - if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1]) - is not None - } - - # Without a pinned boundary, retention=0 keeps only the replay boundary (14). - assert cached_mamba_blocks(0) == {14} - # Pinning the shared prefix at token 96 (state block 5) retains it too, so a - # later request sharing that prefix can hit the Mamba state. - assert cached_mamba_blocks(96) == {5, 14} - - def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under @@ -5949,7 +5909,7 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): def last_req_hit(retention, pin): manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 50f34737085d..6c208e19fb5d 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -898,7 +898,6 @@ def _run_mamba_prefix_cache_mrv1( load_format="dummy", enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, @@ -1158,7 +1157,6 @@ def fake_sample( skip_tokenizer_init=True, enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index 66c8d5b9b76c..c63a4ebc7893 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -88,7 +88,6 @@ def test_mtp_speculative_mixed_batch_short_prefill( trust_remote_code=True, enable_chunked_prefill=True, enable_prefix_caching=enable_prefix_caching, - mamba_cache_mode="align" if enable_prefix_caching else "none", ) as llm: sampling_params = SamplingParams( temperature=0.0, @@ -161,7 +160,6 @@ def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch): enforce_eager=True, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", speculative_config={"method": "mtp", "num_speculative_tokens": 2}, ) as runner: llm = runner.get_llm() diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 89797db4cdd7..f75151414afd 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -60,7 +60,6 @@ def _check_replayssm_parity( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", tensor_parallel_size=tensor_parallel_size, mamba_backend=mamba_backend, ) @@ -130,7 +129,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", speculative_config={ "method": "ngram", "num_speculative_tokens": 3, @@ -182,7 +180,6 @@ def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner) max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", mamba_backend="flashinfer", tensor_parallel_size=2, disable_log_stats=False, @@ -250,7 +247,6 @@ def _check_replayssm_prefix_caching_parity( trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, - mamba_cache_mode="align", disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c6f79f37211c..6615a5c103b1 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -60,7 +60,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & @@ -78,7 +77,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 76d219746913..53e7378a7b5c 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -583,7 +583,6 @@ def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int): kv_transfer_config=kv_transfer_config, language_model_only=True, enable_prefix_caching=True, - mamba_cache_mode="align", # Use lossless state storage so this exact-equality regression isolates # offload boundary selection from low-precision Mamba checkpointing. mamba_cache_dtype="float32", diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index b72de811c0f9..74de3c3a8cf1 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -632,6 +632,8 @@ def test_p_side_chunked_prefill_mamba(mock_platform): block_size=BLOCK_SIZE, ) vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + # mamba_cache_mode "none" only runs without prefix caching. + vllm_config.cache_config.enable_prefix_caching = False kv_cache_config = make_kv_cache_config( block_size=BLOCK_SIZE, @@ -855,6 +857,7 @@ def test_async_load_does_not_double_reserve_mamba_spec_blocks( disable_hybrid_kv_cache_manager=False, ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.enable_prefix_caching = mamba_cache_mode == "align" kv_cache_config = make_kv_cache_config( block_size, mamba_enabled=True, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 9f818c866fff..527f7bee05c2 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -49,6 +49,7 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, KVCacheTensor, MambaSpec, SlidingWindowSpec, @@ -2355,8 +2356,6 @@ def _make_hybrid_attention_mamba_scheduler( enable_kv_cache_events: bool = False, ) -> SchedulerFixture: """Build a scheduler for one attention group plus one Mamba group.""" - scheduler_block_size = scheduler_block_size or block_size - hash_block_size = hash_block_size or block_size attention_spec = FullAttentionSpec( block_size=attention_block_size, num_kv_heads=NUM_KV_HEADS, @@ -2369,9 +2368,38 @@ def _make_hybrid_attention_mamba_scheduler( dtypes=(torch.float32,), mamba_cache_mode=mamba_cache_mode, ) + return _make_two_group_scheduler( + attention_spec, + ("mamba", mamba_spec), + num_cpu_blocks=num_cpu_blocks, + num_gpu_blocks=num_gpu_blocks, + scheduler_block_size=scheduler_block_size or block_size, + hash_block_size=hash_block_size or block_size, + dcp_world_size=dcp_world_size, + lazy=lazy, + mamba_cache_mode=mamba_cache_mode, + enable_kv_cache_events=enable_kv_cache_events, + ) + + +def _make_two_group_scheduler( + attention_spec: FullAttentionSpec, + second_group: tuple[str, KVCacheSpec], + *, + num_cpu_blocks: int, + num_gpu_blocks: int, + scheduler_block_size: int, + hash_block_size: int, + dcp_world_size: int, + lazy: bool = False, + mamba_cache_mode: MambaCacheMode = "align", + enable_kv_cache_events: bool = False, +) -> SchedulerFixture: + """Build a scheduler for one attention group plus one named second group.""" + second_name, second_spec = second_group groups = [ KVCacheGroupSpec(["attention"], attention_spec), - KVCacheGroupSpec(["mamba"], mamba_spec), + KVCacheGroupSpec([second_name], second_spec), ] tensors = [ KVCacheTensor( @@ -2380,7 +2408,7 @@ def _make_hybrid_attention_mamba_scheduler( layer_stride=spec.page_size_bytes * num_gpu_blocks, block_stride=spec.page_size_bytes, ) - for group, spec in zip(groups, (attention_spec, mamba_spec)) + for group, spec in zip(groups, (attention_spec, second_spec)) ] kv_cache_config = KVCacheConfig( num_blocks=num_gpu_blocks, @@ -2428,34 +2456,48 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: silently offloads only a fraction of them. """ attention_block_size = BLOCK_SIZE - mamba_block_size = 4 * BLOCK_SIZE + replicated_block_size = 4 * BLOCK_SIZE dcp_world_size = 2 - fix = _make_hybrid_attention_mamba_scheduler( + fix = _make_two_group_scheduler( + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ), + ( + "replicated_attention", + FullAttentionSpec( + block_size=replicated_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + dcp_sharded=False, + ), + ), num_cpu_blocks=32, num_gpu_blocks=32, - attention_block_size=attention_block_size, - block_size=mamba_block_size, + scheduler_block_size=replicated_block_size, # Every prefix-cacheable group's resolved block size must be a multiple # of the hash block, so it cannot exceed the smallest of them. hash_block_size=BLOCK_SIZE, dcp_world_size=dcp_world_size, - mamba_cache_mode="none", ) sched = fix.scheduler gpu_pool = fix.gpu_block_pool - attention_size, mamba_size = sched.group_block_sizes + attention_size, replicated_size = sched.group_block_sizes assert attention_size == attention_block_size * dcp_world_size - assert mamba_size == mamba_block_size + assert replicated_size == replicated_block_size confirmed = 2 * sched.block_size req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) attention_blocks = _allocate_cp_gpu_blocks( gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 ) - mamba_blocks = _allocate_cp_gpu_blocks( - gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 + replicated_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // replicated_size, replicated_size, group_id=1 ) - kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, replicated_blocks)) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) sched.build_connector_meta( make_scheduler_output( @@ -2476,11 +2518,11 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: # count. assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ confirmed // attention_size, - confirmed // mamba_size, + confirmed // replicated_size, ] assert set(meta.store_gpu_blocks) == { *(b.block_id for b in attention_blocks), - *(b.block_id for b in mamba_blocks), + *(b.block_id for b in replicated_blocks), } diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 5eb26726c718..9bb716621594 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -176,8 +176,7 @@ def device_memory_utilization(self, value: float) -> None: page.""" mamba_block_size: int | None = Field(default=None, gt=0) """Size of a contiguous cache block in number of tokens for mamba cache. - Can be set only when prefix caching is enabled. - Value must be a multiple of 8 to align with causal_conv1d kernel.""" + Can be set only when prefix caching is enabled.""" mamba_cache_dtype: MambaDType = "auto" """The data type to use for the Mamba cache (both the conv as well as the ssm state). If set to 'auto', the data type will be inferred from the model @@ -187,19 +186,19 @@ def device_memory_utilization(self, value: float) -> None: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers: + """The cache strategy for Mamba layers, derived from `enable_prefix_caching`: - - "none": set when prefix caching is disabled. - - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. This is the default when prefix - caching is enabled. + - "none": prefix caching is disabled. + - "align": prefix caching is enabled. Only cache the mamba state of the last + token of each scheduler step and when the token is at position + i * block_size. """ enable_mamba_shared_prefix_checkpoint: bool = False """Also register a Mamba "align" checkpoint at the shared-prefix junction -- where an EAGLE/MTP sibling was observed to resume -- instead of only at the - prompt tail. Off by default; only takes effect with `mamba_cache_mode` - "align", EAGLE on the Mamba group, and a prefix match unit smaller than the - Mamba block size.""" + prompt tail. Off by default; only takes effect with prefix caching enabled, + EAGLE on the Mamba group, and a prefix match unit smaller than the Mamba + block size.""" replayssm_buffer_len: int = Field(default=16, gt=0) """ReplaySSM logical history length B for Mamba2. Triton uses B physical rows and FlashInfer uses B+T, where T is the target verification length. @@ -207,11 +206,10 @@ def device_memory_utilization(self, value: float) -> None: use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. - Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton - or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer - and mamba_cache_mode 'none'. In align mode flushes are most efficient when - mamba_block_size is a multiple of replayssm_buffer_len, but this is not - required.""" + Requires the Triton or FlashInfer mamba backend; FlashInfer does not support + prefix caching yet. Mamba2 speculative decode requires FlashInfer. With + prefix caching, flushes are most efficient when mamba_block_size is a + multiple of replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 87759f46061b..c4016e82bc3d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2038,10 +2038,6 @@ def has_noops(self) -> bool: def has_inner_state(self): return self._model_info.has_inner_state - @property - def supports_mamba_prefix_caching(self) -> bool: - return self._model_info.supports_mamba_prefix_caching - @property def supports_replayssm(self) -> bool: return self._model_info.supports_replayssm diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f27c8418ab4d..d6e834861a19 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3360,7 +3360,7 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": and not self.use_v2_model_runner ): raise ValueError( - "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + "RecoverSSM with prefix caching requires VLLM_USE_V2_MODEL_RUNNER=1" ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError( @@ -3371,8 +3371,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif use_mamba_replayssm_spec: if self.cache_config.mamba_cache_mode != "none": raise ValueError( - "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none" + "FlashInfer ReplaySSM speculative decoding does not support " + "prefix caching; use --no-enable-prefix-caching" ) query_len = 1 + self.num_speculative_tokens if self.cache_config.replayssm_buffer_len < query_len: @@ -3390,8 +3390,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: if self.cache_config.mamba_cache_mode == "align": raise ValueError( - "FlashInfer ReplaySSM does not support " - "--mamba-cache-mode align yet; use none" + "FlashInfer ReplaySSM does not support prefix caching yet; " + "use --no-enable-prefix-caching" ) elif self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 86199e5f007b..c9aa433162d2 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ def get_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: ) +_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE = ( + "--mamba-cache-mode is deprecated and will be removed in v0.33.0. " + "It has no effect: Mamba prefix caching follows --enable-prefix-caching." +) + + _DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( "--disable-cascade-attn is deprecated and has no effect since cascade " "attention has been removed. It will be removed in v0.33.0." @@ -798,7 +804,7 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") - mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + mamba_cache_mode: MambaCacheMode | None = None enable_mamba_shared_prefix_checkpoint: bool = ( CacheConfig.enable_mamba_shared_prefix_checkpoint ) @@ -1403,7 +1409,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--prefix-match-unit", **cache_kwargs["prefix_match_unit"] ) cache_group.add_argument( - "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] + "--mamba-cache-mode", + **{ + **cache_kwargs["mamba_cache_mode"], + "default": None, + "help": _MAMBA_CACHE_MODE_DEPRECATION_MESSAGE, + }, ) cache_group.add_argument( "--enable-mamba-shared-prefix-checkpoint", @@ -2263,6 +2274,8 @@ def create_engine_config( assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + if self.mamba_cache_mode is not None: + logger.warning_once(_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] @@ -2282,7 +2295,6 @@ def create_engine_config( mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, - mamba_cache_mode=self.mamba_cache_mode, enable_mamba_shared_prefix_checkpoint=( self.enable_mamba_shared_prefix_checkpoint ), diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index bfa1103bb368..14c483122ab1 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -812,10 +812,8 @@ def conv_ssm_forward( # - It will read the initial states for every sequence, # that has "has_initial_states_p" == True, # from "cache_indices", using "state_indices_tensor_p". - # - It updates the "conv_state" cache in positions pointed - # to by "state_indices_tensor_p". - # In particular, it will always write the state at the - # sequence end. + # - It writes the state at the sequence end to the "conv_state" + # cache in positions pointed to by "state_indices_tensor_p". x = hidden_states_B_C_p.transpose( 0, 1 ) # this is the form that causal-conv see diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index e8254e4f34eb..2094df902a27 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -20,17 +20,11 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptr, # (dim, width) bias_ptr, initial_states_ptr, # conv_states_ptr - cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains - # the block indices relevant for each sequence - # plus potential 0-padding at the beginning and at the end + cache_indices_ptr, # (batch,) has_initial_states_ptr, query_start_loc_ptr, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, # (batch,) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) - num_computed_tokens, # (batch,) o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, @@ -46,7 +40,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching stride_cache_indices: tl.constexpr, stride_o_dim: tl.constexpr, stride_o_token: tl.int64, - stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M # others pad_slot_id: tl.constexpr, null_block_id: tl.constexpr, @@ -54,7 +47,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching HAS_BIAS: tl.constexpr, KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, @@ -93,40 +85,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching # find the actual sequence length seqlen = sequence_end_index - sequence_start_index - B_size: tl.constexpr = stride_block_m * BLOCK_M - - if IS_APC_ENABLED: - # Handle the case if prefix caching is enabled. - # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" - - # Get the length of the completed sequence so far and compute the offset. - current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - sequence_completed_index = tl.load(num_computed_tokens + idx_seq) - - # Compute the offset where the first stride_block_m-aligned first full block is - # Value in "token-space" - sequence_completed_offset_token = sequence_completed_index % B_size - seq_completed_offset = B_size - sequence_completed_offset_token - seq_end_offset = (seqlen - seq_completed_offset) % B_size - last_full_block_token_index = sequence_end_index - seq_end_offset - # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one - if seq_end_offset == 0: - last_full_block_token_index = last_full_block_token_index - B_size - - # Get the number of blocks to be filled for the current sequence - # If n_block_to_fill = 0, then only the state at the sequence end is stored - n_block_to_fill = current_last_index - current_first_index - - # Get the index of the init block - conv_state_init_index = tl.load(initial_state_idx + idx_seq) - else: - n_block_to_fill = 0 - current_last_index = 0 - conv_state_init_index = 0 - current_first_index = 0 - last_full_block_token_index = 0 - token_offset = BLOCK_M * chunk_offset segment_len = min(BLOCK_M, seqlen - token_offset) @@ -137,7 +95,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + conv_state_indices_ptr + idx_seq * stride_cache_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -224,20 +182,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching loaded_x = tl.load(x_ptrs, mask_x, 0.0) idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - # Compute the offset where the last block should be written in the conv_states - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_last_index - ).to(tl.int64) - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] tl.debug_barrier() # NOTE: use this due to bug in Triton compiler @@ -346,52 +294,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") - # Store intermediate states aligned with stride_block_m - # The additional states are cached starting from the last stride_block_m. - # For example: - # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. - # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last - # stride_block_m are cached. - # For example chunk_offset = n_block_to_fill stores the state at last_full_block - if (chunk_offset - 1) < n_block_to_fill: - # Store the states at the chunk boundaries from the start of the sequence - idx_tokens_last = ( - last_full_block_token_index - - (n_block_to_fill - chunk_offset) * B_size - - state_len - ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] - x_ptrs = ( - x_ptr - + (idx_tokens_last * stride_x_token)[:, None] - + (idx_feats * stride_x_dim)[None, :] - ) # [BLOCK_M,BLOCK_N,] - - mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ - None, : - ] # token-index # token-index # feature-index - loaded_x = tl.load(x_ptrs, mask_x, 0.0) - idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - - # cache_idx - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_first_index - + (chunk_offset - 1) - ).to(tl.int64) - - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] - - mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] - tl.debug_barrier() # NOTE: use this due to bug in Triton compiler - tl.store(conv_states_ptrs_target, loaded_x, mask) - if HAS_BIAS: bias = bias_ptr + idx_feats mask_bias = idx_feats < dim @@ -489,11 +391,6 @@ def causal_conv1d_fn( activation: str | None = "silu", pad_slot_id: int = PAD_SLOT_ID, null_block_id: int = NULL_BLOCK_ID, - block_idx_first_scheduled_token: torch.Tensor | None = None, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, - num_computed_tokens: torch.Tensor | None = None, - block_size_to_align=0, metadata=None, validate_data=False, ): @@ -536,16 +433,6 @@ def causal_conv1d_fn( for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] in this case, the kernel will not process entries at indices 0 and 3 - block_idx_first_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the first cache block to be filled is located. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into cache_indices, where the cache block containing the initial state is located. - num_computed_tokens: (batch,), dtype int32 - The number of tokens already completed for each sequence - block_size_to_align: int - The block size to align the cached states to out: same shape as `x` """ if isinstance(activation, bool) and activation: @@ -631,12 +518,6 @@ def causal_conv1d_fn( assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" - if block_size_to_align is not None and block_size_to_align > 0: - assert (block_size_to_align % BLOCK_M) == 0, ( - "The mamba block size needs to be divisible by the BLOCK_M" - ) - else: - block_size_to_align = BLOCK_M if metadata is None: @@ -720,10 +601,6 @@ def grid(META): query_start_loc, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, - block_idx_last_scheduled_token, - initial_state_idx, - num_computed_tokens, out, # Matrix dimensions dim, @@ -739,7 +616,6 @@ def grid(META): stride_cache_indices, stride_o_dim, stride_o_token, - block_size_to_align // BLOCK_M, # others pad_slot_id, null_block_id, @@ -747,7 +623,6 @@ def grid(META): HAS_BIAS=bias is not None, KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, HAS_NULL_BLOCK=null_block_id is not None, NP2_STATELEN=np2_statelen, # launch_cooperative_grid=True @@ -769,8 +644,6 @@ def _causal_conv1d_update_kernel( conv_state_indices_ptr, num_accepted_tokens_ptr, query_start_loc_ptr, # (batch + 1) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) o_ptr, # (batch, dim, seqlen) # Matrix dimensions batch: int, @@ -798,7 +671,6 @@ def _causal_conv1d_update_kernel( KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, IS_VARLEN: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, @@ -818,17 +690,9 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - if IS_APC_ENABLED: - # Get the state from the initial_state_idx - conv_state_init = tl.load(initial_state_idx + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - else: - conv_state_init = 0 - current_last_index = 0 - # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + conv_state_indices_ptr + idx_seq * stride_state_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -941,14 +805,9 @@ def _causal_conv1d_update_kernel( new_conv_state = tl.where(mask, conv_state, loaded_x) - # Get the state from the initial_state_idx - # cache_idx - conv_states_offset = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index - ).to(tl.int64) conv_state_ptrs_target = ( conv_state_ptr - + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (conv_states_input_coord * stride_conv_state_seq) # Offset from seq + (idx_feats * stride_conv_state_dim) )[None, :] + ( # [BLOCK_N,] idx_tokens * stride_conv_state_tok @@ -1104,8 +963,6 @@ def causal_conv1d_update( query_start_loc: torch.Tensor | None = None, max_query_len: int = -1, null_block_id: int = NULL_BLOCK_ID, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, validate_data=False, out: torch.Tensor | None = None, ): @@ -1123,10 +980,6 @@ def causal_conv1d_update( If not None, the conv_state is a larger tensor along the batch dim, and we are selecting the batch coords specified by conv_state_indices. Useful for a continuous batching scenario. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into conv_state_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into conv_state_indices, where the cache block containing the initial state is located. num_accepted_tokens: (batch,), dtype int32 If not None, it indicates the number of accepted tokens for each sequence in the batch. @@ -1237,8 +1090,6 @@ def grid(META): conv_state_indices, num_accepted_tokens, query_start_loc, - block_idx_last_scheduled_token, - initial_state_idx, out, # Matrix dimensions batch, @@ -1266,7 +1117,6 @@ def grid(META): KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], IS_VARLEN=query_start_loc is not None, - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, IS_SPEC_DECODING=num_accepted_tokens is not None, NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d64e51bfa8bf..69ca643a3a97 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -692,7 +692,7 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None: model_config.architecture, ) assert vllm_config.scheduler_config.enable_chunked_prefill, ( - "Chunked prefill is required for mamba cache mode 'align'." + "Chunked prefill is required for Mamba prefix caching." ) # By default, mamba block size will be set to max_model_len (see # below). When enabling prefix caching, we align mamba block size diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 203a3f3ec393..7aff819dfcd7 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -499,7 +498,6 @@ class FalconH1ForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 979d272c13d1..040f7453fb5b 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -42,7 +42,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, ) @@ -587,7 +586,6 @@ class GraniteMoeHybridForCausalLM( SupportsPP, IsHybrid, SupportsQuant, - SupportsMambaPrefixCaching, ): packed_modules_mapping: dict[str, list[str]] = { "qkv_proj": [ diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 4dc67ff34d27..a780cbe23cf6 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1134,13 +1134,10 @@ def has_noops( @runtime_checkable class SupportsMambaPrefixCaching(Protocol): - """The interface for models whose mamba layers support prefix caching. - - This is currently experimental. + """The interface for attention-free Mamba models that support prefix + caching. Hybrid models get the same methods from `IsHybrid`. """ - supports_mamba_prefix_caching: ClassVar[Literal[True]] = True - @classmethod def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: """Return copy functions for the model's Mamba states.""" @@ -1156,24 +1153,6 @@ def get_mamba_state_copy_funcs( return {mamba_type: copy_funcs for mamba_type in mamba_types} -@overload -def supports_mamba_prefix_caching( - model: type[object], -) -> TypeIs[type[SupportsMambaPrefixCaching]]: ... - - -@overload -def supports_mamba_prefix_caching( - model: object, -) -> TypeIs[SupportsMambaPrefixCaching]: ... - - -def supports_mamba_prefix_caching( - model: type[object] | object, -) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]: - return getattr(model, "supports_mamba_prefix_caching", False) - - @runtime_checkable class SupportsReplaySSM(Protocol): """The interface for models whose recurrent layers support ReplaySSM diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index fe9349c30b10..315c821d6a0e 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -44,7 +44,6 @@ HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -380,7 +379,6 @@ class JambaForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"}, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ae653950ac48..050a13431ddc 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -74,7 +74,6 @@ SupportsEagle, SupportsEagle3, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, SupportsReplaySSM, @@ -725,7 +724,6 @@ class NemotronHForCausalLM( IsHybrid, SupportsQuant, MixtureOfExperts, - SupportsMambaPrefixCaching, SupportsReplaySSM, ): # Relevant only if self.has_moe is True diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cd0557ac4585..09e5025c4e0e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -50,7 +50,6 @@ is_hybrid, requires_raw_input_tokens, supports_hc_sp, - supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -881,7 +880,6 @@ class _ModelInfo: is_attention_free: bool is_hybrid: bool has_noops: bool - supports_mamba_prefix_caching: bool supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -911,7 +909,6 @@ def from_model_cls(model: type[nn.Module]) -> "_ModelInfo": has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), - supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index 2e90da86ea01..35e972818726 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -45,7 +45,7 @@ ) from vllm.sequence import IntermediateTensors -from .interfaces import HasInnerState, IsHybrid, SupportsMambaPrefixCaching +from .interfaces import HasInnerState, IsHybrid from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -833,7 +833,7 @@ def forward( return hidden_states -class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching): +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid): """Zamba2 model with causal language modeling head. This class wraps the core Zamba2 model and adds: @@ -911,10 +911,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: LoRA and scheduler settings prefix: Optional prefix for parameter names - Raises: - AssertionError: If prefix caching is enabled - (not supported by Mamba) - """ config = vllm_config.model_config.hf_config diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 3e7ae905a30d..9820e8660e55 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -245,13 +245,10 @@ def _compute_chunk_metadata( The code below carefully constructs the chunks such that: 1. Chunks contain tokens from a *single* sequence only. - 2. For every sequence, we are guaranteed that we can - retrieve the mamba state *every* chunk_size tokens. + 2. Chunk boundaries fall on multiples of chunk_size in each + sequence's absolute token positions, so the chunking does not + depend on how chunked prefill split the sequence. Constraint (1) dramatically simplifies the mamba kernels. - Constraint (2) dramatically simplifies the implementation - of prefix caching for mamba (wip). We need to take care - of the interaction with chunked prefill in order to - satisfy constraint (2). `checkpoint_offsets_p[i]` forces an extra chunk end that many tokens into row i's query (0 = none), because the SSD scan only materializes diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2e4fe55b96c2..28a7a1106f77 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1426,6 +1426,9 @@ def __init__( ) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode + assert self.mamba_cache_mode == "align" or not self.enable_caching, ( + "Mamba prefix caching requires mamba_cache_mode='align'" + ) self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks self.has_prefill_checkpoint_blocks = ( self.mamba_cache_mode == "align" From 3ad096faf81b4734ae995aa66f4ca0c904cea978 Mon Sep 17 00:00:00 2001 From: zhouyou Date: Tue, 13 Oct 2026 15:59:48 +0800 Subject: [PATCH 18/23] [Bugfix][EC Connector] Don't create hash directories on cache-miss lookups in ECExampleConnector (#60826) Signed-off-by: zhouyou9505 --- .../ec_connector/unit/test_ec_example_connector.py | 12 ++++++++++++ .../ec_transfer/ec_connector/example_connector.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index 725bb99eb055..a940fd4fcf5e 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -246,6 +246,18 @@ def test_has_cache_item_partial_exist( assert result[1] # Second exists assert not result[2] # Third doesn't exist + def test_has_cache_item_miss_creates_no_folder( + self, mock_vllm_config_consumer, temp_storage + ): + """Test a cache miss leaves no empty folder in the shared storage.""" + connector = ECExampleConnector( + vllm_config=mock_vllm_config_consumer, + role=ECConnectorRole.SCHEDULER, + ) + + assert not connector.has_cache_item("missing_hash") + assert os.listdir(temp_storage) == [] + class TestStateManagement: """Test connector state management.""" diff --git a/vllm/distributed/ec_transfer/ec_connector/example_connector.py b/vllm/distributed/ec_transfer/ec_connector/example_connector.py index e080f09b6bef..97075cdf38c9 100644 --- a/vllm/distributed/ec_transfer/ec_connector/example_connector.py +++ b/vllm/distributed/ec_transfer/ec_connector/example_connector.py @@ -202,7 +202,7 @@ def request_finished( def _found_match_for_mm_data(self, mm_hash) -> bool: """Check if the cache is hit for the request.""" - filename = self._generate_filename_debug(mm_hash) + filename = self._generate_filename_debug(mm_hash, create_folder=False) return os.path.exists(filename) def _generate_foldername_debug( @@ -219,11 +219,13 @@ def _generate_foldername_debug( os.makedirs(foldername, exist_ok=True) return foldername - def _generate_filename_debug(self, mm_hash: str) -> str: + def _generate_filename_debug( + self, + mm_hash: str, + create_folder: bool = True, + ) -> str: """Return the full path of the safetensors file for this mm_hash. - Ensures the parent directory exists because - `_generate_foldername_debug` is called with its default - (`create_folder=True`). + If `create_folder` is True (default) the parent directory is created. """ - foldername = self._generate_foldername_debug(mm_hash) # <- folder auto-created + foldername = self._generate_foldername_debug(mm_hash, create_folder) return os.path.join(foldername, "encoder_cache.safetensors") From ff305934f1c6f98edf7c499464a3c2be5f6e3f04 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 16:15:47 +0800 Subject: [PATCH 19/23] [Tests][Mamba] Wait for accepted counts in fake MTP proposer (#61342) Signed-off-by: zjy0516 --- tests/v1/e2e/general/test_mamba_prefix_cache.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 6c208e19fb5d..cf8aee141bc5 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -129,6 +129,10 @@ def fake_propose_draft_token_ids_fn( first_token_id_index = ( num_computed_tokens + 1 ) # bonus token isn't considered as computed + # num_accepted_tokens_cpu is filled by a non-blocking D2H copy in the + # same step; wait for it or a lagging GPU yields stale counts. + if self.num_accepted_tokens_event is not None: + self.num_accepted_tokens_event.synchronize() first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item() proposed_draft_token_ids = [ prompt_token_ids[ From 24aa1191f433638722178fed646f0fd52e96a587 Mon Sep 17 00:00:00 2001 From: stefankoncarevic Date: Tue, 13 Oct 2026 10:35:57 +0200 Subject: [PATCH 20/23] =?UTF-8?q?[ROCm][Bugfix]=20Run=20the=20AITER=20AR+R?= =?UTF-8?q?MSNorm=20fallback=20test=20on=20the=20platform=E2=80=A6=20(#613?= =?UTF-8?q?50)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Stefan Koncarevic Signed-off-by: Stefan Koncarevic --- tests/distributed/test_aiter_rms_fallback.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/distributed/test_aiter_rms_fallback.py b/tests/distributed/test_aiter_rms_fallback.py index 4a0fc67873a1..211ff15e6e0d 100644 --- a/tests/distributed/test_aiter_rms_fallback.py +++ b/tests/distributed/test_aiter_rms_fallback.py @@ -11,6 +11,7 @@ import vllm._aiter_ops as aiter_ops import vllm.distributed as vllm_distributed import vllm.kernels # noqa: F401 +from vllm.platforms import current_platform def test_rocm_aiter_gemma_rmsnorm_falls_back_after_allreduce( @@ -72,9 +73,14 @@ def fake_rmsnorm2d_fwd_with_add( fake_aiter.rmsnorm2d_fwd_with_add = fake_rmsnorm2d_fwd_with_add monkeypatch.setitem(sys.modules, "aiter", fake_aiter) - inp = torch.randn(8192, 16, dtype=torch.bfloat16) + # On the platform's device, not the CPU. The fallback leaves Python through + # `torch.ops.vllm_aiter.fused_add_rms_norm`, which is registered for the + # platform's dispatch key alone, so CPU tensors have nothing to dispatch to + # on a GPU build. The kernel behind it is still the fake above. + device = current_platform.device_type + inp = torch.randn(8192, 16, dtype=torch.bfloat16, device=device) residual = torch.randn_like(inp) - weight = torch.randn(16, dtype=torch.bfloat16) + weight = torch.randn(16, dtype=torch.bfloat16, device=device) out, residual_out = aiter_ops._rocm_aiter_fused_allreduce_rmsnorm_impl( inp, residual, weight, 1e-6, gemma_norm=True From 6f11c220b5b0f3b45b2a14c49ebbca2df12d1c80 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:41:23 +0100 Subject: [PATCH 21/23] [Transformers backend] Fuse MoE routers that score after the top-k (#57410) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) --- tests/models/transformers/fusers/test_moe.py | 70 ++++++++++++++++++ .../models/transformers/fusers/moe.py | 73 +++++++++++++++---- .../models/transformers/fx_utils.py | 30 +++++++- .../model_executor/models/transformers/moe.py | 31 ++++---- vllm/transformers_utils/config.py | 36 +++++++++ .../model_arch_config_convertor.py | 33 ++------- 6 files changed, 213 insertions(+), 60 deletions(-) diff --git a/tests/models/transformers/fusers/test_moe.py b/tests/models/transformers/fusers/test_moe.py index 7de4589eef91..f0e70c896baa 100644 --- a/tests/models/transformers/fusers/test_moe.py +++ b/tests/models/transformers/fusers/test_moe.py @@ -29,6 +29,42 @@ def forward(self, hidden_states): return logits, value, index +class SelectionScoredRouter(TopKRouter): + """Scores the selected logits only (GraniteMoE/Aria style), in fp32.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight).float() + value, index = torch.topk(logits, self.top_k, dim=-1) + value = ( + torch.sigmoid(value) if self.sigmoid else F.softmax(value, dim=-1) + ).type_as(hidden_states) + return logits, value, index + + +class EpsilonRenormRouter(TopKRouter): + """Renormalizes via an epsilon-offset denominator (DeepSeek-V3 style).""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = torch.sigmoid(logits) + value, index = torch.topk(scores, self.top_k, dim=-1) + denominator = value.sum(dim=-1, keepdim=True) + 1e-20 + value /= denominator + return logits, value, index + + +class UntraceableRouter(TopKRouter): + """Scores and selects, then loses the trace before the renormalization.""" + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight) + scores = F.softmax(logits, dim=-1) + value, index = torch.topk(scores, self.top_k, dim=-1) + if value.sum() > 0: # data-dependent branch: tracing stops here + value = value / value.sum(dim=-1, keepdim=True) + return logits, value, index + + class ScaledRouter(TopKRouter): """Greedy router scaling its top-k weights (DeepSeek `routed_scaling_factor`).""" @@ -268,6 +304,19 @@ def test_moe_fuser_detects_router(sigmoid): assert fuser.shared_name is None and fuser.shared_gate_name is None +@pytest.mark.parametrize("sigmoid", [False, True]) +def test_moe_fuser_matches_router_scoring_the_selection(sigmoid): + """Scoring the selection rather than all logits still allows the fuser to match.""" + with torch.device("meta"): + block = MoEBlock(lambda: SelectionScoredRouter(sigmoid=sigmoid)) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.gate_name == "gate" + assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax") + assert fuser.renormalize is (not sigmoid) + assert fuser.router_dtype == torch.float32 + + def test_moe_fuser_matches_scaled_router(): """Weight scaling after the top-k (DeepSeek style) does not break matching.""" with torch.device("meta"): @@ -275,6 +324,27 @@ def test_moe_fuser_matches_scaled_router(): assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser) +@pytest.mark.parametrize( + "router_cls,expected", + [ + (TopKRouter, True), # divides by the sum of the selection + (EpsilonRenormRouter, True), # ... with an epsilon in the denominator + (ScaledRouter, False), # scales the weights instead of renormalizing + (UntraceableRouter, None), # trace stops before it could say + ], +) +def test_moe_fuser_reads_renormalize_from_the_gate(router_cls, expected): + """The router's graph, not the config, decides whether vLLM renormalizes: it + is the code being replaced, and it already bakes in `norm_topk_prob` because + the flag is a Python bool that fx resolves at trace time. Absence of the + division only means "no" on a complete trace.""" + with torch.device("meta"): + block = MoEBlock(router_cls) + fuser = MoEBlockFuser.match(block, "experts") + assert isinstance(fuser, MoEBlockFuser) + assert fuser.renormalize is expected + + def test_moe_fuser_matches_grouped_router(): """Group masking between the score and the routing top-k still matches: the scorer is anchored on the routing (last) top-k, not the group one.""" diff --git a/vllm/model_executor/models/transformers/fusers/moe.py b/vllm/model_executor/models/transformers/fusers/moe.py index af3a101a0ce1..06fdac109ca4 100644 --- a/vllm/model_executor/models/transformers/fusers/moe.py +++ b/vllm/model_executor/models/transformers/fusers/moe.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models.transformers.fx_utils import ( find_node, + is_complete, is_op, peel, trace, @@ -107,6 +108,24 @@ def _reaches(node: fx.Node, key: str) -> set[fx.Node]: return seen +def _renormalizes(topk: fx.Node) -> bool: + """Are the top-k weights divided by their own sum? + + Every HF router spells renormalization as a division by a sum over the + selection, optionally offset by an epsilon.""" + downstream = _reaches(topk, "users") + for node in downstream: + if not (is_op(node, "div") or is_op(node, "truediv")) or len(node.args) < 2: + continue + denominator = node.args[1] + if not isinstance(denominator, fx.Node): + continue + cone = _reaches(denominator, "all_input_nodes") + if any(is_op(n, "sum") and n in downstream for n in cone): + return True + return False + + class SharedExpertMLP(nn.Module): """Wraps an HF shared expert, applying the output gating it is paired with.""" @@ -147,12 +166,18 @@ class MoEBlockFuser: shared_name: str | None shared_gate_name: str | None router_dtype: torch.dtype | None = None + renormalize: bool | None = None @staticmethod - def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: - """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`. - - Returns the scoring function and the dtype the router computes in.""" + def _match_router( + gate: nn.Module, + ) -> tuple[str, torch.dtype | None, bool | None] | None: + """Matches `topk(score(linear(x)))` or `score(topk(linear(x)))`, `score` being + `softmax`/`sigmoid`. + + Returns the scoring function, the dtype the router computes in, and + whether the top-k weights are renormalized (`None` if the trace was cut + short before it could say).""" state = {name for name, _ in named_state(gate)} if "weight" not in state or state - {"weight", "e_score_correction_bias"}: return None @@ -164,20 +189,31 @@ def _match_router(gate: nn.Module) -> tuple[str, torch.dtype | None] | None: if not topks: return None topk = topks[-1] - # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear. - scorers = [ - n - for n in _reaches(topk, "all_input_nodes") - if is_op(n, "softmax") or is_op(n, "sigmoid") - ] + is_scorer = lambda n: is_op(n, "softmax") or is_op(n, "sigmoid") + post_topk = False + # Find all scoring functions that reach the top-k node. + scorers = [n for n in _reaches(topk, "all_input_nodes") if is_scorer(n)] + if not scorers: + # Find all scoring functions that are users of the top-k node. + scorers = [n for n in _reaches(topk, "users") if is_scorer(n)] + post_topk = True if len(scorers) != 1: return None scorer = scorers[0] logits_cone = _reaches(scorer, "all_input_nodes") if not any(is_op(n, "linear") for n in logits_cone): return None - scoring_func = "softmax" if is_op(scorer, "softmax") else "sigmoid" - return scoring_func, _forced_dtype(logits_cone) + is_softmax = is_op(scorer, "softmax") + # vLLM always scores every expert, so matching a router that softmaxes + # the selection alone takes renormalizing. Otherwise look for the + # division, reading its absence as "no" only from a complete trace. + renormalize: bool | None + if (post_topk and is_softmax) or _renormalizes(topk): + renormalize = True + else: + renormalize = False if is_complete(graph) else None + scoring_func = "softmax" if is_softmax else "sigmoid" + return scoring_func, _forced_dtype(logits_cone), renormalize @staticmethod def _match_shared_experts( @@ -221,14 +257,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None if _returns_tuple(type(moe_block)): return None # Router: the child that scores + top-k selects. - gate_name = scoring_func = router_dtype = None + gate_name = scoring_func = router_dtype = renormalize = None for name, child in moe_block.named_children(): if ( name != experts_name and (router := cls._match_router(child)) is not None ): gate_name = name - scoring_func, router_dtype = router + scoring_func, router_dtype, renormalize = router break if gate_name is None or scoring_func is None: return None @@ -256,7 +292,14 @@ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None for name, child in moe_block.named_children(): if name not in accounted and next(named_state(child), None) is not None: return None - return cls(gate_name, scoring_func, shared_name, shared_gate_name, router_dtype) + return cls( + gate_name, + scoring_func, + shared_name, + shared_gate_name, + router_dtype, + renormalize, + ) def gate( self, moe_block: nn.Module, prefix: str, out_dtype: torch.dtype | None = None diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index cb630866a390..818e1ca7bb4f 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -35,6 +35,21 @@ """The unpatched `nn.Module.__call__`. During tracing fx patches it to record `call_module` nodes; meta execution must call modules for real.""" +_MODULE_GETATTR = nn.Module.__getattr__ +"""The unpatched `nn.Module.__getattr__`. During tracing fx patches it to record +`get_attr` nodes; reading one while inferring a `get_attr`'s own value would +recurse until the stack runs out, truncating the graph.""" + + +@contextlib.contextmanager +def _unpatched_modules(): + """Restores the module hooks fx patches, so meta execution records nothing.""" + with ( + mock.patch.object(nn.Module, "__call__", _MODULE_CALL), + mock.patch.object(nn.Module, "__getattr__", _MODULE_GETATTR), + ): + yield + def is_leaf_call(node: object) -> TypeIs[fx.Node]: """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" @@ -132,7 +147,8 @@ def _infer_meta(self, kind: str, target: object, args: tuple, kwargs: dict): ) return _UNKNOWN if kind == "get_attr": - value = operator.attrgetter(str(target))(self.root) + with _unpatched_modules(): + value = operator.attrgetter(str(target))(self.root) if isinstance(value, torch.Tensor): value = torch.empty_like(value, device="meta") return value @@ -159,8 +175,6 @@ def meta_of(arg: object) -> object: if kind == "call_module": # Run the child's forward with all its state on "meta", without # mutating it (at match time params may be meta but buffers real). - # fx patches `nn.Module.__call__` while tracing; restore the real - # one so this execution is not itself recorded. child = self.root.get_submodule(str(target)) state = { name: torch.empty_like(tensor, device="meta") @@ -168,7 +182,7 @@ def meta_of(arg: object) -> object: child.named_parameters(), child.named_buffers() ) } - with mock.patch.object(nn.Module, "__call__", _MODULE_CALL): + with _unpatched_modules(): return torch.func.functional_call(child, state, meta_args, meta_kwargs) return _UNKNOWN @@ -686,6 +700,14 @@ def visit(self, node: ast.AST) -> ast.AST: _Replacer().visit(module) +def is_complete(graph: fx.Graph) -> bool: + """Did the trace run to completion? A partial graph has no `output` node. + + Only a complete graph can be read for the *absence* of an op; in a partial + one the tail is exactly what is missing.""" + return any(node.op == "output" for node in graph.nodes) + + def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None: """The first node in `graph` matching `predicate`, or `None`.""" return next((n for n in graph.nodes if predicate(n)), None) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 71f5bd066c8d..953459172d36 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -39,6 +39,7 @@ from vllm.model_executor.models.transformers.fusers.glu import GLUFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix +from vllm.transformers_utils.config import CONFIG_FIELD_ALIASES from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op from .base import Base @@ -180,23 +181,18 @@ def recursive_replace(self): text_config = self.text_config # Positional arguments - num_experts = self.model_config.get_num_experts() - top_k = getattr_iter(text_config, ["num_experts_per_tok", "top_k"], None) + num_experts_names = CONFIG_FIELD_ALIASES["num_experts"] + num_experts = getattr_iter(text_config, num_experts_names, None) + top_k_names = CONFIG_FIELD_ALIASES["num_experts_per_token"] + top_k = getattr_iter(text_config, top_k_names, None) assert top_k is not None hidden_size = text_config.hidden_size - intermediate_size = getattr_iter( - text_config, ["moe_intermediate_size", "intermediate_size"], None - ) + intermediate_size_names = CONFIG_FIELD_ALIASES["moe_intermediate_size"] + intermediate_size = getattr_iter(text_config, intermediate_size_names, None) assert intermediate_size is not None - num_shared_experts = getattr_iter( - text_config, - [ - "n_shared_experts", # DeepSeek, Docs, GLM - "moe_num_shared_experts", # Aria, Ernie - ], - 0, - ) + num_shared_experts_names = CONFIG_FIELD_ALIASES["num_shared_experts"] + num_shared_experts = getattr_iter(text_config, num_shared_experts_names, 0) # Common kwargs norm_topk_prob = getattr(text_config, "norm_topk_prob", None) @@ -275,6 +271,9 @@ def _recursive_replace(module: nn.Module, prefix: str): if "bias" in experts_param_name: has_bias = True break + # Whether the expert weights are stored in transposed format. + # Comes from Transformers `use_experts_implementation` decorator. + is_transposed = getattr(experts, "is_transposed", False) # If the config does not specify num_shared_experts, but # the model has shared experts, we assume there is one. if self.num_shared_experts == 0: @@ -308,6 +307,7 @@ def _recursive_replace(module: nn.Module, prefix: str): enable_eplb=enable_eplb, num_redundant_experts=num_redundant_experts, has_bias=has_bias, + is_fused_checkpoint_transposed=is_transposed, routed_experts_cls=TransformersRoutedExperts, ) fuser = MoEBlockFuser.match(moe_block, experts_name) @@ -346,6 +346,11 @@ def _recursive_replace(module: nn.Module, prefix: str): gate = fuser.gate(moe_block, prefix, router_dtype) kwargs |= dict( scoring_func=fuser.scoring_func, + renormalize=( + layer_renormalize + if fuser.renormalize is None + else fuser.renormalize + ), is_sequence_parallel=( self.parallel_config.use_sequence_parallel_moe ), diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index e046c301d23a..1d93a04e662a 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -52,6 +52,42 @@ logger = init_logger(__name__) +CONFIG_FIELD_ALIASES: dict[str, tuple[str, ...]] = { + "total_num_kv_heads": ( + "n_head_kv", # Falcon + "num_kv_heads", + "num_key_value_heads", # LLaMA-2 + "multi_query_group_num", # ChatGLM + "num_attention_groups", # Step3p5 + ), + "num_experts": ( + "num_experts", # Jamba + "moe_num_experts", # Dbrx + "n_routed_experts", # DeepSeek + "num_local_experts", # Mixtral + ), + "num_experts_per_token": ( + "num_experts_per_tok", + "num_experts_per_token", + "top_k_experts", + "moe_topk", # Hunyuan + "moe_top_k", + "top_k", + ), + "moe_intermediate_size": ( + "moe_intermediate_size", + "intermediate_size", + ), + "num_shared_experts": ( + "n_shared_experts", # DeepSeek, Docs, GLM + "moe_num_shared_experts", # Aria, Ernie + ), +} +"""HF config fields that models spell differently, keyed by the name vLLM uses. + +`ModelArchConfigConvertorBase` reads a scalar from these; a caller that needs the +raw value (a per-layer `list`, say) passes the aliases to `getattr_iter` itself.""" + _ST_POOLING_MODULE_TYPES = { "sentence_transformers.models.Pooling", "sentence_transformers.sentence_transformer.modules.pooling.Pooling", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 06a937697106..6290f6c6645b 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -14,6 +14,7 @@ from vllm.config.utils import getattr_iter from vllm.logger import init_logger from vllm.transformers_utils.config import ( + CONFIG_FIELD_ALIASES, ConfigFormat, get_safetensors_params_metadata, ) @@ -148,23 +149,11 @@ def _get_qk_rope_head_dim(self) -> int: return qk_rope_head_dim def get_total_num_kv_heads(self) -> int: - attributes = [ - # For Falcon: - "n_head_kv", - "num_kv_heads", - # For LLaMA-2: - "num_key_value_heads", - # For ChatGLM: - "multi_query_group_num", - # For Step3p5: - "num_attention_groups", - ] + names = CONFIG_FIELD_ALIASES["total_num_kv_heads"] # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. default_factory = self.get_total_num_attention_heads - return getattr_iter( - self.hf_text_config, attributes, default_factory=default_factory - ) + return getattr_iter(self.hf_text_config, names, default_factory=default_factory) def get_num_experts_from_block_configs(self) -> int: """Check block_configs for heterogeneous models (e.g., NemotronH). @@ -188,13 +177,7 @@ def get_num_experts_from_block_configs(self) -> int: def get_num_experts(self) -> int: """Returns the number of experts in the model.""" - num_expert_names = [ - "num_experts", # Jamba - "moe_num_experts", # Dbrx - "n_routed_experts", # DeepSeek - "num_local_experts", # Mixtral - ] - + num_expert_names = CONFIG_FIELD_ALIASES["num_experts"] num_experts = getattr_iter(self.hf_text_config, num_expert_names, 0) if isinstance(num_experts, list): # Ernie VL's remote code uses list[int]... @@ -206,13 +189,7 @@ def get_num_experts(self) -> int: return num_experts def get_num_experts_per_token(self) -> int: - names = [ - "num_experts_per_tok", - "num_experts_per_token", - "top_k_experts", - "moe_topk", - "moe_top_k", - ] + names = CONFIG_FIELD_ALIASES["num_experts_per_token"] num_experts_per_token = getattr_iter(self.hf_text_config, names, 0) if isinstance(num_experts_per_token, list): return max(num_experts_per_token, default=0) From 183625d4f9f49968218ffbe43242877ead72d569 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Tue, 13 Oct 2026 12:48:13 +0400 Subject: [PATCH 22/23] [Bugfix] Fall back to auto for unquantized MoE layers when --moe-backend is marlin (#60219) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- .../moe/test_unquantized_backend_selection.py | 14 ++++++++++++++ .../layers/fused_moe/oracle/unquantized.py | 7 ++++--- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/kernels/moe/test_unquantized_backend_selection.py b/tests/kernels/moe/test_unquantized_backend_selection.py index b64e909fbd45..e3a83ec3c5c7 100644 --- a/tests/kernels/moe/test_unquantized_backend_selection.py +++ b/tests/kernels/moe/test_unquantized_backend_selection.py @@ -626,3 +626,17 @@ def test_select_explicit_triton_backend(is_lora_enabled): assert selected_backend == UnquantizedMoeBackend.TRITON assert experts_cls is not None + + +@skipif_not_cuda_rocm +@pytest.mark.parametrize("moe_backend", ["humming", "marlin"]) +def test_quantization_only_backend_falls_back_to_auto(moe_backend): + """A quantization-only --moe-backend also reaches the unquantized MoE + layers of a quantized model (e.g. a BF16 MTP draft), which must use the + auto-selected backend instead of failing.""" + moe_config = make_dummy_moe_config() + moe_config.moe_backend = moe_backend + + assert select_unquantized_moe_backend( + moe_config=moe_config + ) == select_unquantized_moe_backend(moe_config=make_dummy_moe_config()) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 1af0450dfd13..2f05074c643c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -303,9 +303,10 @@ def _return_or_raise( ) runner_backend = moe_config.moe_backend - # 'humming' is quantization-only; an unquantized layer (e.g. excluded via - # modules_to_not_convert) falls through to auto instead of erroring. - if runner_backend not in ["auto", "humming"]: + # 'humming' and 'marlin' are quantization-only; an unquantized layer (e.g. + # excluded via modules_to_not_convert, or an unquantized MTP draft of a + # quantized model) falls through to auto instead of erroring. + if runner_backend not in ["auto", "humming", "marlin"]: requested_backend = map_unquantized_backend(runner_backend) if ( activation_format == mk.FusedMoEActivationFormat.BatchedExperts From b1a08980092b6c29eb66d9f2652030958d533704 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Tue, 13 Oct 2026 08:58:38 +0000 Subject: [PATCH 23/23] [Perf][GLM-5.3-Flash] Fill short-context decode causal top-k rows with one Triton launch The kpool indexer's short-context decode path (max_seq_len <= index_topk) filled the causal top-k rows with five kernels per call. Replace them with a single Triton kernel on GPU; CPU tensors keep the torch formulation. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Jared Wen --- .../test_sparse_indexer_decode_seq_lens.py | 15 +++++++++++ vllm/models/glm5next/common/sparse_indexer.py | 25 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py index d87a82c37b2b..585e0157ef02 100644 --- a/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py +++ b/tests/v1/attention/test_sparse_indexer_decode_seq_lens.py @@ -18,6 +18,7 @@ kernel is documented to replicate. """ +import pytest import torch # Bootstrap the glm5next package before importing kpool_compress directly: its @@ -34,6 +35,7 @@ from vllm.models.glm5next.common.sparse_indexer import ( _decode_topk_seq_lens, + _fill_causal_indices, _fill_short_decode_causal_indices, ) from vllm.platforms import current_platform @@ -186,3 +188,16 @@ def test_tail_expansion_legacy_vs_fixed(): assert legacy_tail[10, :2].tolist() == [556, 557] assert not torch.equal(legacy_tail, fixed_tail) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("width", [8, 1000, 2176]) +def test_fill_causal_indices_kernel_matches_torch(width: int): + """The single-launch GPU fill must match the torch formulation (CPU path), + including rows whose position exceeds the row width.""" + positions = torch.tensor([0, 3, 7, width - 1, width + 5], dtype=torch.int64) + ref = torch.full((positions.numel(), width), 99, dtype=torch.int32) + _fill_causal_indices(ref, positions) + got = torch.full_like(ref, 99, device="cuda") + _fill_causal_indices(got, positions.cuda()) + assert torch.equal(got.cpu(), ref) diff --git a/vllm/models/glm5next/common/sparse_indexer.py b/vllm/models/glm5next/common/sparse_indexer.py index ad3845457063..45585474e09d 100644 --- a/vllm/models/glm5next/common/sparse_indexer.py +++ b/vllm/models/glm5next/common/sparse_indexer.py @@ -4,6 +4,8 @@ import torch +from vllm.triton_utils import tl, triton + RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. @@ -106,7 +108,30 @@ def _decode_topk_seq_lens( return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail +@triton.jit +def _fill_causal_indices_kernel( + rows_ptr, + rows_stride, + positions_ptr, + width, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + pos = tl.load(positions_ptr + row).to(tl.int32) + vals = tl.where(cols <= pos, cols.to(tl.int32), -1) + tl.store(rows_ptr + row * rows_stride + cols, vals, mask=cols < width) + + def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None: + """rows[t, j] = j for j <= positions[t], else -1 (one launch on GPU).""" + if rows.is_cuda and rows.shape[0] > 0: + block = 1024 + grid = (rows.shape[0], triton.cdiv(rows.shape[1], block)) + _fill_causal_indices_kernel[grid]( + rows, rows.stride(0), positions, rows.shape[1], BLOCK=block + ) + return causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32) positions = positions.to(torch.int32) rows[:] = causal_range[None, :]