From eb372b59cefa7c19e68aa47007fdee23e1e776a9 Mon Sep 17 00:00:00 2001 From: Jonguk Cheong Date: Mon, 27 Jul 2026 14:25:30 +0900 Subject: [PATCH] [Bugfix][Model Runner V2] Preserve Mamba block table capacity under DCP Co-authored-by: OpenAI Codex Signed-off-by: Jonguk Cheong --- vllm/v1/worker/gpu/block_table.py | 9 ++++++++- vllm/v1/worker/gpu/model_runner.py | 7 +++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index df5d2be629d2..c8de8407c171 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -114,10 +114,17 @@ def append_block_ids( start = self.num_blocks.np[i, req_index] if not overwrite else 0 block_ids = new_block_ids[i] bpk = self.blocks_per_kv_block[i] + end = start + len(block_ids) * bpk + capacity = self.block_tables[i].gpu.shape[1] + if end > capacity: + raise RuntimeError( + f"Block table write for request {req_index}, group {i} " + f"exceeds row capacity ({end} > {capacity})" + ) if bpk > 1: block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] 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 == 1: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f0da049ce46d..70ecf5fe3c07 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -54,7 +54,7 @@ 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 -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, MambaSpec from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu import pcp_manager as pcp @@ -441,9 +441,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: # 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 - ) + dcp_size = self.dcp_size if isinstance(spec, AttentionSpec) else 1 + max_num_blocks = cdiv(block_table_max_model_len, spec.block_size * dcp_size) # Align to a multiple of (128 / block_size) as required by some attention # backends such as TRTLLM (#39324) if spec.block_size <= 128: