iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Index Mamba align state copies by request slot in MRV2 (#60210)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Karl0007 <44745209+Karl0007@users.noreply.github.com>
Co-authored-by: Yongqin Wang <yongqinwang@roblox.com>
This commit is contained in:
Jiangyun Zhu
2026-10-07 02:32:46 +00:00
committed by GitHub
co-authored by Karl0007 Yongqin Wang
parent 956960f85f
commit 8a26869bca
6 changed files with 195 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/mamba/test_precopy_mamba_align.py
+78 -2
View File
@@ -31,7 +31,11 @@ import torch
from vllm.model_executor.layers.mamba import mamba_utils as layer_mamba_utils
from vllm.platforms import current_platform
from vllm.v1.worker import mamba_utils as worker_mamba_utils
from vllm.v1.worker.mamba_utils import _TEMPORAL_TILES, precopy_mamba_align_fused_kernel
from vllm.v1.worker.mamba_utils import (
_TEMPORAL_TILES,
postprocess_mamba_fused_kernel,
precopy_mamba_align_fused_kernel,
)
_parametrize: Callable[..., Callable[[Any], Any]]
@@ -195,7 +199,8 @@ def test_precopy_matches_v1_copy_specs(
convs, ssms, device, conv_state_dim_first
)
bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
# Non-identity mapping: block-table rows are request slots, not batch rows.
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device).flip(0)
grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles)
precopy_mamba_align_fused_kernel[grid](
dst_col,
@@ -225,6 +230,77 @@ def test_precopy_matches_v1_copy_specs(
torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)
@_parametrize("conv_state_dim_first", [False, True])
@_cuda_required
def test_v2_postprocess_indexes_block_tables_by_req_slot(conv_state_dim_first):
"""Under PP the V2 postprocess runs after later steps have re-gathered the
batch-ordered tables, so it must index the request-slot tables by slot."""
device = torch.device("cuda")
torch.manual_seed(0)
num_reqs = 4
block_size = 8
num_blocks = num_reqs * MAX_COLS + 1
bt = torch.arange(1, num_blocks, dtype=torch.int32, device=device).view(
num_reqs, MAX_COLS
)
# slot 0: copy col 1 -> 0; slot 1: src == dst; slot 2: not aligned;
# slot 3: copy col 2 -> 0 with one accepted draft token.
num_accepted = torch.tensor([1, 1, 1, 2], dtype=torch.int32, device=device)
state_idx = torch.tensor([1, 0, 1, 2], dtype=torch.int32, device=device)
new_num_computed = torch.tensor([8, 8, 5, 8], dtype=torch.int32, device=device)
ref_src = torch.tensor([1, 0, -1, 2], dtype=torch.int32)
ref_dst = torch.zeros(num_reqs, dtype=torch.int32)
ref_bias = torch.tensor([0, 0, 0, 1], dtype=torch.int32)
convs, ssms = _build_state(num_blocks, device, conv_state_dim_first)
conv_ref, ssm_ref = _reference(
convs,
ssms,
bt.cpu(),
ref_src,
ref_dst,
ref_bias,
num_reqs,
conv_state_dim_first,
)
base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
convs, ssms, device, conv_state_dim_first
)
bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
idx_mapping = torch.tensor([2, 0, 3, 1], dtype=torch.int32, device=device)
postprocess_mamba_fused_kernel[(num_reqs, NUM_LAYERS * 2, 1)](
num_accepted.clone(),
state_idx,
None,
new_num_computed,
None,
bt_ptrs,
bt.stride(0),
base,
blk_stride,
elem,
inner,
width,
group,
drc,
drs,
num_accepted,
idx_mapping,
num_reqs,
block_size=block_size,
COPY_BLOCK_SIZE=1024,
CONV_STATE_DIM_FIRST=conv_state_dim_first,
HAS_IDX_MAPPING=True,
PRECOMPUTED_NEW_COMPUTED=True,
TEMPORAL_TILES=1,
)
torch.accelerator.synchronize()
for layer in range(NUM_LAYERS):
torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0)
torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)
def test_ds_conv_copy_spec_reproduces_multi_accept_assert(monkeypatch):
monkeypatch.setattr(
layer_mamba_utils,
tests/v1/worker/test_mamba_utils.py
+46
View File
@@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import (
MambaSpec,
UniformTypeKVCacheSpecs,
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
from vllm.v1.worker.mamba_utils import (
MambaCopyBuffers,
@@ -638,6 +639,51 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group():
assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12]
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
def test_v2_align_ctx_binds_request_slot_tables_for_state_copies():
"""V2 state copies may run steps after their batch (PP), so they must read
the per-request-slot tables; aligned state indices are computed for the
current batch and keep reading the gathered tables."""
cfg = _TestConfig(num_layers=1)
device = torch.device("cuda")
kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"])
forward_context = {
"layer_0": _make_mock_attention(
torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim),
torch.empty(cfg.num_blocks, cfg.temporal_state_dim),
)
}
block_tables = BlockTables(
block_sizes=[cfg.block_size],
max_num_reqs=cfg.max_num_reqs,
max_num_batched_tokens=64,
max_num_blocks_per_group=[4],
device=device,
kernel_block_sizes=[cfg.block_size],
)
state = object.__new__(MambaHybridModelState)
state._align_mode = True
state._mamba_group_ids = []
state._mamba_spec = None
state._mamba_state_copy_funcs = _COPY_FUNCS
state.max_num_reqs = cfg.max_num_reqs
state.device = device
state.vllm_config = MagicMock()
state.vllm_config.compilation_config.static_forward_context = forward_context
state.initialize_kv_cache(kv_cache_config, block_tables)
ctx = state._mamba_ctx
assert ctx is not None and ctx.is_initialized
assert ctx.block_table_ptrs.tolist() == [
block_tables.block_tables[0].gpu.data_ptr()
]
assert ctx.aligned_index_block_table_ptrs.tolist() == [
block_tables.input_block_tables[0].data_ptr()
]
# -----------------------------------------------------------------------------
# stage_postprocess_inputs_to_gpu: single-pass staging into pinned views
# -----------------------------------------------------------------------------
vllm/v1/worker/gpu/model_runner.py
+1
View File
@@ -778,6 +778,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.kv_caches = [
cache for cache in kv_caches_dict.values() if cache.device == self.device
]
self.model_state.initialize_kv_cache(self.kv_cache_config, self.block_tables)
if is_profiling:
self.kv_connector = NO_OP_KV_CONNECTOR
else:
vllm/v1/worker/gpu/model_states/interface.py
+6
View File
@@ -127,6 +127,12 @@ class ModelState(ABC):
def apply_staged_writes(self) -> None:
return None
def initialize_kv_cache(
self, kv_cache_config: KVCacheConfig, block_tables: BlockTables
) -> None:
"""Hook run after the KV cache tensors are allocated and bound."""
return None
def capture_inner_cudagraphs(
self,
input_buffers: InputBuffers,
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+36 -41
View File
@@ -21,6 +21,7 @@ from vllm.v1.core.sched.output import NewRequestData
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
from vllm.v1.utils import CpuGpuBuffer
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.input_batch import InputBatch
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
@@ -134,47 +135,42 @@ class MambaHybridModelState(DefaultModelState):
self._mamba_spec = mamba_spec
return self._mamba_group_ids, self._mamba_spec
def _ensure_align_ctx(
self,
kv_cache_config: KVCacheConfig,
mamba_group_ids: list[int],
block_tables: tuple[torch.Tensor, ...],
) -> MambaSpecDecodeGPUContext:
def initialize_kv_cache(
self, kv_cache_config: KVCacheConfig, block_tables: BlockTables
) -> None:
if not self._align_mode:
return
mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config)
if self._mamba_state_copy_funcs is None:
mamba_groups = get_mamba_groups(kv_cache_config)
mamba_types = {spec.mamba_type for spec in mamba_groups}
copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types)
validate_mamba_state_copy_funcs(mamba_groups, copy_funcs)
self._mamba_state_copy_funcs = copy_funcs
copy_funcs = self._mamba_state_copy_funcs
if self._mamba_ctx is None:
# Both SD and DS conv layouts support a >0 spec-decode shift: the
# fused pre-copy kernel (``_copy_mamba_state_block``) applies the
# ``token_bias = num_accepted - 1`` window shift per conv layout
# (SD: contiguous slice; DS: per-dim-row strided slice), matching
# the V1 ``get_conv_copy_spec`` semantics.
self._mamba_ctx = MambaSpecDecodeGPUContext.create(
max_num_reqs=self.max_num_reqs,
kv_cache_config=kv_cache_config,
copy_funcs=copy_funcs,
device=self.device,
make_buffer=lambda n, dtype: CpuGpuBuffer(
n, dtype=dtype, device=self.device
),
)
ctx = self._mamba_ctx
if not ctx.is_initialized:
forward_context = self.vllm_config.compilation_config.static_forward_context
# block_tables are batch-order slices of the persistent
# input_block_tables (stable data_ptr), so the metadata is captured
# once here and reused across steps.
ctx.initialize_from_forward_context(
kv_cache_config,
forward_context,
copy_funcs,
[block_tables[gid] for gid in mamba_group_ids],
)
return ctx
# Both SD and DS conv layouts support a >0 spec-decode shift: the
# fused pre-copy kernel (``_copy_mamba_state_block``) applies the
# ``token_bias = num_accepted - 1`` window shift per conv layout
# (SD: contiguous slice; DS: per-dim-row strided slice), matching
# the V1 ``get_conv_copy_spec`` semantics.
self._mamba_ctx = MambaSpecDecodeGPUContext.create(
max_num_reqs=self.max_num_reqs,
kv_cache_config=kv_cache_config,
copy_funcs=self._mamba_state_copy_funcs,
device=self.device,
make_buffer=lambda n, dtype: CpuGpuBuffer(
n, dtype=dtype, device=self.device
),
)
# Under PP a step's postprocess runs after later steps re-gathered the
# batch-ordered input tables, so state copies read the per-request-slot
# tables. Aligned state indices are for the current batch only.
self._mamba_ctx.initialize_from_forward_context(
kv_cache_config,
self.vllm_config.compilation_config.static_forward_context,
self._mamba_state_copy_funcs,
[block_tables.block_tables[gid].gpu for gid in mamba_group_ids],
[block_tables.input_block_tables[gid] for gid in mamba_group_ids],
)
def preprocess_state(
self,
@@ -194,8 +190,9 @@ class MambaHybridModelState(DefaultModelState):
num_reqs = input_batch.num_reqs
if num_reqs == 0:
return
mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config)
ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables)
_, mamba_spec = self._get_mamba_group_info(kv_cache_config)
ctx = self._mamba_ctx
assert ctx is not None
# The state-advance + pre-copy kernels run every step; they fast-exit per
# request when src_col < 0 or src_col == dst_col, so no copy happens on
@@ -301,10 +298,8 @@ class MambaHybridModelState(DefaultModelState):
if hasattr(builder, "mamba_aligned_state_indices"):
aligned_index_builders.append((group_idx, builder))
if aligned_index_builders:
ctx = self._ensure_align_ctx(
kv_cache_config, mamba_group_ids, block_tables
)
all_group_indices = ctx.compute_aligned_state_indices(
assert self._mamba_ctx is not None
all_group_indices = self._mamba_ctx.compute_aligned_state_indices(
input_batch.seq_lens, num_reqs
)
for group_idx, builder in aligned_index_builders:
vllm/v1/worker/mamba_utils.py
+28 -14
View File
@@ -392,8 +392,8 @@ def postprocess_mamba_fused_kernel(
# Output: num_accepted_tokens update (for src==dst case)
num_accepted_tokens_out_ptr,
# Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The
# per-request decision arrays are in req-state-slot order; the block table
# is in batch order, so HAS_IDX_MAPPING splits the two indexings.
# per-request decision arrays and the block tables are both in
# req-state-slot order.
idx_mapping_ptr,
# Runtime parameter (varies per batch - NOT constexpr to avoid recompilation)
num_reqs,
@@ -476,10 +476,9 @@ def postprocess_mamba_fused_kernel(
if src_block_idx == dest_block_idx and accept_token_bias == 0:
return
bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx
_copy_mamba_state_block(
state_idx,
bt_row_idx,
req_idx,
src_block_idx,
dest_block_idx,
accept_token_bias,
@@ -610,7 +609,7 @@ def precopy_mamba_align_fused_kernel(
token_bias = tl.load(token_bias_ptr + req_idx)
_copy_mamba_state_block(
state_idx,
batch_idx,
req_idx,
src_col,
dst_col,
token_bias,
@@ -806,8 +805,10 @@ class MambaSpecDecodeGPUContext:
# Per-group block-table base addresses: int64[num_groups]. Populated in
# initialize_from_forward_context from the persistent per-group block
# table tensors (whose data_ptr is stable across steps).
# table tensors (whose data_ptr is stable across steps). Rows are request
# state slots for the state copies and batch rows for the aligned indices.
block_table_ptrs: torch.Tensor
aligned_index_block_table_ptrs: torch.Tensor
block_table_stride_req: int = 0
# persistent output for the once-per-step, all-group aligned-index launch.
@@ -900,6 +901,9 @@ class MambaSpecDecodeGPUContext:
block_table_ptrs=torch.zeros(
len(mamba_group_ids), dtype=torch.int64, device=device
),
aligned_index_block_table_ptrs=torch.zeros(
len(mamba_group_ids), dtype=torch.int64, device=device
),
aligned_state_indices=torch.empty(
(
len(mamba_group_ids),
@@ -924,6 +928,7 @@ class MambaSpecDecodeGPUContext:
forward_context: dict[str, Any],
mamba_state_copy_funcs: MambaStateCopyFuncsByType,
block_tables: list[torch.Tensor],
aligned_index_block_tables: list[torch.Tensor] | None = None,
) -> None:
"""Extract and cache memory layout metadata from Mamba state tensors.
@@ -955,8 +960,12 @@ class MambaSpecDecodeGPUContext:
mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to
copy functions.
block_tables: per-mamba-group persistent block-table tensors, in
the same order as `mamba_group_ids`. Their `data_ptr()` /
`stride(0)` are captured once for the kernel to index into.
the same order as `mamba_group_ids`, with one row per request
state slot. Their `data_ptr()` / `stride(0)` are captured once
for the state-copy kernels to index into.
aligned_index_block_tables: per-mamba-group persistent tables
read by `compute_aligned_state_indices`, with one row per
batch row. Defaults to `block_tables`.
"""
if self.is_initialized:
@@ -968,6 +977,7 @@ class MambaSpecDecodeGPUContext:
forward_context,
mamba_state_copy_funcs,
block_tables,
aligned_index_block_tables or block_tables,
)
def _populate_metadata(
@@ -976,6 +986,7 @@ class MambaSpecDecodeGPUContext:
forward_context: dict[str, Any],
mamba_state_copy_funcs: MambaStateCopyFuncsByType,
block_tables: list[torch.Tensor],
aligned_index_block_tables: list[torch.Tensor],
) -> None:
idx = 0
for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids):
@@ -1076,16 +1087,19 @@ class MambaSpecDecodeGPUContext:
# `block_tables[i]` is the persistent 2D int32 block-table tensor for
# `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the
# engine's lifetime, so we capture them once here.
assert len(block_tables) == self.num_groups, (
f"expected {self.num_groups} block tables, got {len(block_tables)}"
)
strides = {bt.stride(0) for bt in block_tables}
assert len(block_tables) == len(aligned_index_block_tables) == self.num_groups
strides = {bt.stride(0) for bt in (*block_tables, *aligned_index_block_tables)}
assert len(strides) == 1, (
f"all mamba block tables must share stride(0), got {strides}"
)
self.block_table_stride_req = int(next(iter(strides)))
for i, bt in enumerate(block_tables):
for i, (bt, aligned_bt) in enumerate(
zip(block_tables, aligned_index_block_tables)
):
self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr())
self.aligned_index_block_table_ptrs[i] = _reinterpret_u64_as_i64(
aligned_bt.data_ptr()
)
self.is_initialized = True
@@ -1107,7 +1121,7 @@ class MambaSpecDecodeGPUContext:
block_rows = 32
grid = (triton.cdiv(num_reqs, block_rows),)
get_aligned_state_indices_multi_group_kernel[grid](
self.block_table_ptrs,
self.aligned_index_block_table_ptrs,
seq_lens,
self.aligned_state_indices,
self.block_table_stride_req,