diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index d37e19d0bd..06348074bb 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -2,12 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from pathlib import Path import pytest from tests.conftest import VllmRunner from tests.utils import create_new_process_for_each_test from vllm import SamplingParams, TokensPrompt +from vllm.config import CUDAGraphMode +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" @@ -126,3 +129,111 @@ def test_kimi_linear_dcp_tiny( logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) assert max(logprob_drifts) <= 1e-2 + + +def _make_tiny_k3_config(model_dir: Path) -> str: + config = KimiK3Config( + text_config=_make_tiny_overrides(), + vision_config={ + "vt_num_attention_heads": 2, + "vt_num_hidden_layers": 1, + "vt_hidden_size": 32, + "vt_intermediate_size": 64, + "qkv_hidden_size": 48, + }, + architectures=["KimiK3ForConditionalGeneration"], + ) + config.save_pretrained(model_dir) + return str(model_dir) + + +def _run_k3_partial_prefix_reuse( + vllm_runner: type[VllmRunner], model_name: str, dcp_size: int +) -> tuple[list[int], list[float], int]: + block_size = 256 + common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)] + prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17]) + replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19]) + sampling_params = SamplingParams( + temperature=0, + max_tokens=8, + seed=0, + ignore_eos=True, + logprobs=20, + ) + + with vllm_runner( + model_name=model_name, + skip_tokenizer_init=True, + load_format="dummy", + language_model_only=True, + tensor_parallel_size=2, + decode_context_parallel_size=dcp_size, + cp_kv_cache_interleave_size=1, + distributed_executor_backend="mp", + dtype="bfloat16", + seed=0, + max_model_len=512, + max_num_seqs=4, + max_num_batched_tokens=256, + gpu_memory_utilization=0.85, + 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 + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + runner.llm.generate( + [prime_prompt], + sampling_params, + use_tqdm=False, + ) + replay_output = runner.llm.generate( + [replay_prompt], + sampling_params, + use_tqdm=False, + )[0] + + completion = replay_output.outputs[0] + token_ids = list(completion.token_ids) + assert len(token_ids) == sampling_params.max_tokens + assert completion.logprobs is not None + assert len(completion.logprobs) == sampling_params.max_tokens + selected_logprobs = [ + step_logprobs[token_id].logprob + for token_id, step_logprobs in zip(token_ids, completion.logprobs) + ] + return token_ids, selected_logprobs, replay_output.num_cached_tokens + + +@create_new_process_for_each_test() +@pytest.mark.distributed(num_gpus=2) +def test_kimi_k3_dcp_partial_prefix_reuse( + vllm_runner: type[VllmRunner], + num_gpus_available: int, + tmp_path: Path, +) -> None: + if num_gpus_available < 2: + pytest.skip("Need at least 2 GPUs") + + model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3") + baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=1 + ) + dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=2 + ) + + assert baseline_cached == 256 + assert dcp_cached == 256 + assert dcp_tokens == baseline_tokens + assert len(dcp_logprobs) == len(baseline_logprobs) == 8 + logprob_drifts = [] + for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs): + assert math.isfinite(baseline_logprob) + assert math.isfinite(dcp_logprob) + logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) + assert max(logprob_drifts) <= 1e-2 diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 72773adcb6..cd164250a2 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ from vllm.model_executor.layers.attention.mla_attention import ( MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform @@ -93,6 +94,48 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( assert spec.page_size_bytes == 64 * 656 +@pytest.mark.cpu_test +def test_dcp_chunked_context_accepts_non_virtual_block_aligned_prefix(): + context_lens_cpu = torch.tensor([65, 96], dtype=torch.int32) + prefill_query_start_loc_cpu = torch.tensor([0, 1, 2], dtype=torch.int32) + + metadata = build_mla_chunked_context_metadata( + context_lens_cpu=context_lens_cpu, + prefill_query_start_loc_cpu=prefill_query_start_loc_cpu, + chunked_prefill_workspace=torch.empty(0), + chunked_prefill_workspace_size=128, + block_size=64, + align_chunk_to_block=True, + device=torch.device("cpu"), + dcp_world_size=4, + dcp_local_block_size=1, + dcp_virtual_block_size=4, + ) + + assert metadata is not None + assert metadata.context_lens.tolist() == [65, 96] + assert [chunk.seq_lens.tolist() for chunk in metadata.chunks] == [[65], [96]] + assert [chunk.starts.tolist() for chunk in metadata.chunks] == [[0], [0]] + assert [chunk.local_context_lens_allranks for chunk in metadata.chunks] == [ + [[17, 16, 16, 16]], + [[24, 24, 24, 24]], + ] + assert [chunk.padded_local_seq_lens for chunk in metadata.chunks] == [ + [17], + [24], + ] + assert [chunk.padded_local_cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 17], + [0, 24], + ] + assert [chunk.cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 65], + [0, 96], + ] + assert [chunk.num_context_tokens for chunk in metadata.chunks] == [65, 96] + assert [chunk.num_local_context_tokens for chunk in metadata.chunks] == [17, 24] + + # Remove sm100 backends from the list if not using sm100 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10: BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index d237ce9f34..da5bd5b5a1 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -4,6 +4,7 @@ "align") models: scheduler chunk splitting, partial tail registration, CoW on partial hits, and same-step deferral.""" +from math import lcm from types import SimpleNamespace from unittest.mock import MagicMock @@ -68,12 +69,62 @@ def test_capable_connector_uses_divergent_partial_hit_lookup(): manager.get_computed_blocks.assert_not_called() -def test_mamba_align_split_partial_tail_schedule(): +def make_full_mamba_manager( + *, + dcp_world_size: int, + hash_block_size: int = 2, + full_block_size: int = 4, + mamba_block_size: int = 4, + num_blocks: int = 32, + use_eagle: bool = False, +): + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=full_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + scheduler_block_size = lcm( + full_block_size * dcp_world_size, + mamba_block_size, + ) + return make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + use_eagle=use_eagle, + ) + + +@pytest.mark.parametrize("dcp_world_size", [1, 4]) +def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: 0 -> 8192 -> 9728 -> 9984 -> 10000.""" block_size = 512 + scheduler_block_size = block_size * dcp_world_size hash_block_size = 32 mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), @@ -81,6 +132,8 @@ def test_mamba_align_split_partial_tail_schedule(): scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=hash_block_size, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, mamba_partial_cache_hit=True, ) split = Scheduler._mamba_block_aligned_split @@ -109,6 +162,7 @@ def test_mamba_align_split_partial_tail_schedule(): # stops at the next block boundary (10240), later chunk ends re-align. req2 = make_request("1", [0] * 12000, hash_block_size, sha256) req2.num_computed_tokens = 9984 + assert req2.num_computed_tokens % scheduler_block_size != 0 assert split(self=mock, request=req2, num_new_tokens=2016) == 256 req2.num_computed_tokens = 10240 assert split(self=mock, request=req2, num_new_tokens=1000) == 512 @@ -1499,3 +1553,151 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert num_computed == mamba_block_size assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed + + +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): + hash_block_size = 2 + physical_block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=hash_block_size, + full_block_size=physical_block_size, + mamba_block_size=physical_block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + partial_hash = req0.block_hashes[2] + partial_full_block = manager.block_pool.get_cached_block(partial_hash, [0]) + partial_mamba_block = manager.block_pool.get_cached_block(partial_hash, [1]) + assert partial_full_block is not None + assert partial_mamba_block is not None + + req1 = make_request("dcp-replay", [0, 0, 1, 1, 2, 2, 3, 3], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 6 + full_block_size = physical_block_size * dcp_world_size + assert [len(group) for group in computed_blocks.blocks] == [ + (6 + full_block_size - 1) // full_block_size, + 2, + ] + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + mamba_new_block_id = new_blocks.get_block_ids()[1][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block[0].block_id, full_new_block_id) in copies + assert ( + KVCacheBlockCopy(partial_mamba_block[0].block_id, mamba_new_block_id) in copies + ) + manager.block_pool.free_blocks(retained) + + +@pytest.mark.parametrize("dcp_world_size", [2, 4]) +def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( + dcp_world_size: int, +): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + assert manager.coordinator._cache_hit_alignment_tokens == block_size + assert manager.coordinator.single_type_managers[0].block_size == ( + block_size * dcp_world_size + ) + assert manager.coordinator.single_type_managers[1].block_size == block_size + + prefix = list(range(12)) + req0 = make_request("snapshot-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + req1 = make_request( + "snapshot-replay", prefix + list(range(12, 16)), block_size, sha256 + ) + + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [ + (12 + block_size * dcp_world_size - 1) // (block_size * dcp_world_size), + 3, + ] + partial_full_block = computed_blocks.blocks[0][-1] + new_blocks = manager.allocate_slots(req1, 4, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block.block_id, full_new_block_id) in copies + assert all( + copy.src_block_id != computed_blocks.blocks[1][-1].block_id for copy in copies + ) + manager.block_pool.free_blocks(retained) + + +def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + prefix = list(range(12)) + req0 = make_request("joint-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None + manager.new_step_starts() + + replay_tokens = prefix + list(range(12, 16)) + req1 = make_request("joint-before", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 8 + assert [len(group) for group in computed_blocks.blocks] == [1, 2] + + req0.num_computed_tokens = 8 + assert manager.allocate_slots(req0, 4) is not None + manager.new_step_starts() + + req2 = make_request("joint-after", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req2) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [2, 3] + + +def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit(): + hash_block_size = 2 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=hash_block_size, + full_block_size=4, + mamba_block_size=4, + use_eagle=True, + ) + + req0 = make_request("eagle-owner", [7] * 6, hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 4 + manager.new_step_starts() + assert manager.allocate_slots(req0, 2) is not None + req0.num_computed_tokens = 6 + manager.new_step_starts() + + req1 = make_request("eagle-replay", [7] * 6 + [9] * 2, 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 4 + assert [len(group) for group in computed_blocks.blocks] == [1, 1] + assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py index 225d87e867..bbb7d26260 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py @@ -55,9 +55,9 @@ def cache_full_block_and_partial_tail( token_ids: list[int], *, enable_kv_cache_events: bool = False, + block_size: int = 6, ) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]: hash_block_size = 2 - block_size = 6 kv_cache_group_id = 0 req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( @@ -389,27 +389,34 @@ def test_reset_prefix_cache_clears_partial_entry_metadata(): assert pool.cached_block_hashes_by_block == {} -def test_evict_cached_block_removes_full_hash_and_partial_entry(): - pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail( - [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_evict_cached_block_removes_full_hash_and_partial_entry( + dcp_world_size: int, +): + block_size = 6 * dcp_world_size + partial_num_tokens = 2 * block_size - 2 + pool, req, blocks, partial_hash = cache_full_block_and_partial_tail( + list(range(partial_num_tokens)), block_size=block_size ) - full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0] + full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, block_size)[0] assert pool.get_cached_block(full_hash, [0]) == [blocks[0]] - assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [0]) == [blocks[1]] pool.evict_blocks({blocks[0].block_id, blocks[1].block_id}) assert pool.get_cached_block(full_hash, [0]) is None - assert pool.get_cached_block(partial_hash_10, [0]) is None + assert pool.get_cached_block(partial_hash, [0]) is None assert pool.cached_block_hashes_by_block == {} -def test_partial_block_promotes_to_direct_full_block_hash(): +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_partial_block_promotes_to_direct_full_block_hash(dcp_world_size: int): hash_block_size = 2 - block_size = 6 + block_size = 6 * dcp_world_size kv_cache_group_id = 0 - token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] + partial_num_tokens = 2 * block_size - hash_block_size + token_ids = list(range(partial_num_tokens)) req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( num_gpu_blocks=3, @@ -426,27 +433,22 @@ def test_partial_block_promotes_to_direct_full_block_hash(): block_size=block_size, kv_cache_group_id=kv_cache_group_id, ) - partial_hash_10 = boundary_hash(req, hash_block_size, 10) + partial_hash = boundary_hash(req, hash_block_size, partial_num_tokens) assert pool.cache_partial_block( request=req, block=blocks[1], - num_tokens=10, + num_tokens=partial_num_tokens, kv_cache_group_id=kv_cache_group_id, block_size=block_size, ) - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) == [blocks[1]] - req.append_output_token_ids([5, 5]) + req.append_output_token_ids(list(range(partial_num_tokens, 2 * block_size))) full_hashes = BlockHashListWithBlockSize( req.block_hashes, hash_block_size, block_size ) promoted_full_hash = full_hashes[1] - # The promoted full-block hash is the fine hash at the 12-token boundary, - # not a concatenation of the fine hashes inside the block. - assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1] - assert promoted_full_hash != BlockHash( - req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5] - ) + assert promoted_full_hash == req.block_hashes[2 * block_size // hash_block_size - 1] pool.cache_full_blocks( request=req, @@ -457,4 +459,4 @@ def test_partial_block_promotes_to_direct_full_block_hash(): kv_cache_group_id=kv_cache_group_id, ) assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]] - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) is None diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py new file mode 100644 index 0000000000..59c11e79c4 --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + +@pytest.mark.parametrize( + ("mamba_cache_mode", "num_speculative_blocks", "expected"), + [ + pytest.param("align", 0, 65_536, id="align-prefix-cache"), + pytest.param("none", 7, 8, id="no-prefix-cache-with-speculation"), + ], +) +def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table( + monkeypatch, + mamba_cache_mode: str, + num_speculative_blocks: int, + expected: int, +): + """Mamba/GDN block-table rows index global positions, unlike DCP KV.""" + + max_model_len = 1_048_576 + attention_block_size = 1_536 + mamba_block_size = 16 + dcp_size = 8 + full_attention_spec = FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + shapes=((1,),), + dtypes=(torch.bfloat16,), + block_size=mamba_block_size, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], full_attention_spec), + KVCacheGroupSpec(["kda"], mamba_spec), + ], + ) + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size), + cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode), + ) + runner = SimpleNamespace( + max_model_len=max_model_len, + is_encoder_decoder=False, + vllm_config=vllm_config, + ) + + class _CapturedWidths(Exception): + pass + + captured: list[int] = [] + + def capture_width(max_num_blocks: int, *_args, **_kwargs) -> int: + captured.append(max_num_blocks) + if len(captured) == 2: + raise _CapturedWidths + return max_num_blocks + + monkeypatch.setattr(model_runner_module, "get_block_table_width", capture_width) + + with pytest.raises(_CapturedWidths): + GPUModelRunner.initialize_kv_cache(runner, kv_cache_config) + + # Attention KV is local to one of eight DCP ranks; KDA state is replicated + # and therefore needs one table entry for every global 16-token page. + assert captured == [86, expected] + + +def test_append_block_ids_rejects_write_past_row_capacity(): + """Reject an oversized staged write before it can corrupt the next row.""" + + class _BlockTable: + gpu = torch.empty((2, 4), dtype=torch.int32) + + def stage_write(self, *_args): + pytest.fail("an oversized write must not be staged") + + block_tables = BlockTables.__new__(BlockTables) + block_tables.num_kv_cache_groups = 1 + block_tables.blocks_per_kv_block = [1] + block_tables.block_tables = [_BlockTable()] + block_tables.num_blocks = SimpleNamespace( + np=torch.tensor([[0, 3]], dtype=torch.int32) + ) + + with pytest.raises( + RuntimeError, + match=r"request 1, group 0 exceeds row capacity \(5 > 4\)", + ): + block_tables.append_block_ids( + req_index=1, + new_block_ids=([4, 5],), + overwrite=False, + ) + + assert block_tables.num_blocks.np[0, 1] == 3 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 0829ec0d2b..767ea00334 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -619,15 +619,22 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): "full-attention and Mamba groups, got: " f"{type(g.kv_cache_spec).__name__}." ) - # Fine-grained hash hits require Mamba "align", no context - # parallelism, and compatible cache managers in every group. + # Fine-grained hash hits require Mamba "align" and compatible cache + # managers in every group. TP needs hashing finer than the Mamba block; + # DCP accepts equality because it scales the effective full-attention + # block instead. has_partial_mamba_group = any( isinstance(g.kv_cache_spec, MambaSpec) and g.kv_cache_spec.mamba_cache_mode == "align" - and g.kv_cache_spec.block_size > hash_block_size + and ( + (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) + or ( + dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size + ) + ) for g in kv_cache_config.kv_cache_groups ) - self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group + self.enable_partial_hash_hits = has_partial_mamba_group if self.enable_partial_hash_hits: unsupported_partial_hit_managers = { type(manager).__name__ diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 90fc104dc2..fbc4f383cb 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -116,8 +116,15 @@ class BlockTables: bpk = self.blocks_per_kv_block[i] if bpk > 1: block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] + end = start + len(block_ids) + row_capacity = self.block_tables[i].gpu.shape[1] + if end > row_capacity: + raise RuntimeError( + f"Block table write for request {req_index}, group {i} exceeds " + f"row capacity ({end} > {row_capacity})" + ) self.block_tables[i].stage_write(req_index, start, block_ids) - self.num_blocks.np[i, req_index] = start + len(block_ids) + self.num_blocks.np[i, req_index] = end def apply_staged_writes(self) -> None: if self.num_kv_cache_groups == 0: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 77ae59f0ee..4f65769c95 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -60,7 +60,6 @@ from vllm.multimodal.encoder_budget import ( ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask -from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -517,17 +516,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - # When using DCP, each request's KV cache is sharded among different ranks. - # As a result, one block on the current rank covers `block_size * cp_size` - # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + # Let each cache type account for CP. Attention KV is DCP-sharded, + # while Mamba/GDN recurrent state is replicated across DCP ranks. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len ) - # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens + # Preserve each cache type's alignment requirements after applying + # its topology-aware block-table width. if isinstance(spec, MambaSpec): - max_num_blocks = ( - max_num_blocks if self.cache_config.enable_prefix_caching else 1 - ) + spec.num_speculative_blocks max_num_blocks = get_block_table_width( max_num_blocks, spec.block_size, token_alignment=None )