[Bugfix][Model Runner V2] Preserve Mamba block table capacity under DCP

[jongukc]

Purpose

Model Runner V2 currently applies DCP scaling when sizing every KV cache group's block table. This is correct for attention KV, which is sharded across DCP ranks, but Mamba state is replicated and retains its original block size.

With max_model_len=512, block_size=16, and DCP=2, the worker allocates 16 Mamba entries while the scheduler can allocate 32. A valid 257-token request receives 17 Mamba blocks, writes its last block ID into the adjacent row, and then triggers a CUDA device-side assertion when Mamba selects column 16 from the 16-column row.

Apply DCP scaling only to AttentionSpec. Also reject staged block table writes that exceed the destination row before launching the fused writer.

Test Plan

# On the base commit:
CUDA_LAUNCH_BLOCKING=1 PYTHONPATH=. python poc.py

# On this branch, using the capacity produced by the fixed model runner:
CUDA_LAUNCH_BLOCKING=1 PYTHONPATH=. python poc.py --fixed-capacity

tests/v1/worker/test_gpu_block_table.py

The PoC uses the real KVCacheManager allocation, fused GPU block table writer, block table gather, and Mamba state-index selection. It requires one CUDA GPU; the full DCP deployment requires at least two GPUs.

PoC Script

import argparse

import torch

from vllm.sampling_params import SamplingParams
from vllm.utils.hashing import sha256
from vllm.utils.math_utils import cdiv
from vllm.v1.core.kv_cache_manager import KVCacheManager
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    MambaSpec,
)
from vllm.v1.request import Request
from vllm.v1.worker.gpu.block_table import BlockTables

block_size = 16
max_model_len = 512
dcp_size = 2
worker_capacity = cdiv(max_model_len, block_size * dcp_size)
worker_capacity = cdiv(worker_capacity, 128 // block_size) * (128 // block_size)
scheduler_capacity = cdiv(max_model_len, block_size)

parser = argparse.ArgumentParser()
parser.add_argument("--fixed-capacity", action="store_true")
args = parser.parse_args()
mamba_capacity = scheduler_capacity if args.fixed_capacity else worker_capacity
assert (worker_capacity, scheduler_capacity) == (16, 32)

init_none_hash(sha256)
block_hasher = get_request_block_hasher(block_size, sha256)


def make_request(request_id: str, num_tokens: int) -> Request:
    return Request(
        request_id=request_id,
        prompt_token_ids=[1] * num_tokens,
        sampling_params=SamplingParams(max_tokens=1),
        pooling_params=None,
        block_hasher=block_hasher,
    )


config = KVCacheConfig(
    num_blocks=128,
    kv_cache_tensors=[],
    kv_cache_groups=[
        KVCacheGroupSpec(
            ["attention"],
            FullAttentionSpec(
                block_size=block_size,
                num_kv_heads=1,
                head_size=16,
                dtype=torch.float16,
            ),
        ),
        KVCacheGroupSpec(
            ["mamba"],
            MambaSpec(
                block_size=block_size,
                shapes=((1,),),
                dtypes=(torch.float16,),
                mamba_cache_mode="align",
            ),
        ),
    ],
)
manager = KVCacheManager(
    kv_cache_config=config,
    max_model_len=max_model_len,
    scheduler_block_size=block_size,
    hash_block_size=block_size,
    enable_caching=True,
    dcp_world_size=dcp_size,
)

neighbor = make_request("neighbor", 1)
long_request = make_request("long", worker_capacity * block_size + 1)
neighbor_blocks = manager.allocate_slots(neighbor, neighbor.num_tokens)
long_blocks = manager.allocate_slots(long_request, long_request.num_tokens)
assert neighbor_blocks is not None and long_blocks is not None
neighbor_ids = neighbor_blocks.get_block_ids()
long_ids = long_blocks.get_block_ids()
assert [len(group) for group in long_ids] == [9, 17]

tables = BlockTables(
    block_sizes=[block_size, block_size],
    max_num_reqs=2,
    max_num_batched_tokens=32,
    max_num_blocks_per_group=[worker_capacity, mamba_capacity],
    device=torch.device("cuda:0"),
    kernel_block_sizes=[block_size, block_size],
)
tables.append_block_ids(1, neighbor_ids, overwrite=True)
tables.apply_staged_writes()
torch.cuda.synchronize()

mamba_table = tables.block_tables[1].gpu
before = int(mamba_table[1, 0])
tables.append_block_ids(0, long_ids, overwrite=True)
tables.apply_staged_writes()
torch.cuda.synchronize()
after = int(mamba_table[1, 0])

print("Mamba row capacity:", mamba_capacity)
print("scheduler blocks:", [len(group) for group in long_ids])
print("neighbor block before/after:", before, after)
assert (before == after) == args.fixed_capacity

rows = torch.tensor([1, 0], dtype=torch.int32, device="cuda")
forward_table = tables.gather_block_tables(rows, num_reqs_padded=2)[1]
print("selecting long-request column 16", flush=True)
selected = torch.gather(
    forward_table,
    1,
    torch.tensor([[0], [worker_capacity]], device="cuda"),
)
torch.cuda.synchronize()
print("selected states:", selected.tolist())

Test Result

Before the fix:

Mamba row capacity: 16
scheduler blocks: [9, 17]
neighbor block before/after: 2 12
selecting long-request column 16
ScatterGatherKernel.cu:163: Assertion `idx_dim >= 0 && idx_dim < index_size` failed.
torch.AcceleratorError: CUDA error: device-side assert triggered

With the corrected capacity:

Mamba row capacity: 32
scheduler blocks: [9, 17]
neighbor block before/after: 2 2
selecting long-request column 16
selected states: [[2], [12]]

On the patched branch, deliberately retaining the undersized 16-entry row instead raises RuntimeError: Block table write for request 0, group 1 exceeds row capacity (17 > 16) before launching the GPU write. Existing GPU block table tests pass (2 passed).

AI Assistance

OpenAI Codex was used to assist with investigation, implementation, testing, and PR preparation. The human submitter reviewed the changes and must be prepared to explain and defend them end-to-end.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jongukc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jongukc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork