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

[Bugfix][Mamba] Fix overlapping state copy race (#50729)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Andreas Karatzas
2026-08-17 08:29:41 +02:00
committed by GitHub
co-authored by OpenAI Codex
parent 93550cc4cd
commit a02cfccbc6
2 changed files with 228 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_mamba_utils.py
+152 -37
View File
@@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpe
from vllm.v1.worker.mamba_utils import (
MambaCopyBuffers,
MambaSpecDecodeGPUContext,
batch_memcpy,
collect_mamba_copy_meta,
do_mamba_copy_block,
preprocess_mamba,
@@ -536,6 +537,34 @@ class TestPostprocessMambaFusedKernel:
def test_config(self):
return _TestConfig()
def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device):
batch = 128
row_bytes = 32 * 1024
shift = 16
copy_size = row_bytes - shift
pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to(
torch.uint8
)
state = pattern.expand(batch, -1).clone()
snapshot = state.clone()
row_stride_bytes = state.stride(0) * state.element_size()
row_offsets = (
torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes
)
dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64)
src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64)
sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device)
expected = snapshot.clone()
expected[:, :copy_size].copy_(snapshot[:, shift:])
for _ in range(10):
state.copy_(snapshot)
batch_memcpy(src_ptrs, dst_ptrs, sizes)
torch.accelerator.synchronize()
torch.testing.assert_close(state, expected, rtol=0, atol=0)
def test_matches_python_postprocess_mamba(self, device, test_config):
"""
Golden test: GPU kernel produces identical results to Python impl.
@@ -1190,12 +1219,27 @@ class TestPostprocessMambaFusedKernel:
# --- Verify Python behavior (ground truth) ---
dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1
# Conv state should be modified (shifted copy within block)
conv_changed = not torch.allclose(
conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
# This is an overlapping in-place left shift, so comparing only the
# Python and fused paths can hide the same memcpy race in both. Build
# the memmove result from the untouched snapshot and check each path
# independently.
expected_conv_state = conv_state_orig.clone()
expected_conv_state[dest_block_id, :-1].copy_(
conv_state_orig[dest_block_id, 1:]
)
assert conv_changed, (
"Python: Conv state should be modified when accept_token_bias > 0"
torch.testing.assert_close(
conv_state_py,
expected_conv_state,
rtol=0,
atol=0,
msg="Python: overlapping conv copy should have memmove semantics",
)
torch.testing.assert_close(
conv_state_gpu,
expected_conv_state,
rtol=0,
atol=0,
msg="GPU: overlapping conv copy should have memmove semantics",
)
# Temporal state should be modified (copy from different block)
@@ -2122,13 +2166,10 @@ class TestPostprocessMambaFusedKernel:
actual_src_block_idx = src_block_idx + accept_token_bias
actual_src_block_id = block_table[req, actual_src_block_idx]
All prior regression tests exercise only ``bias == 1``, i.e. they
only ever read one slot ahead of ``src_block_idx`` in the block
table. An off-by-one (or missing scale) in the address computation
on line 143 of ``mamba_utils.py`` would be invisible to every
existing test but would silently read the wrong physical block on
any speculative-decode cycle that accepts multiple tokens across a
block boundary, feeding a stale hidden state forward one step.
A ``bias == 1`` case only reads one slot ahead of ``src_block_idx``
in the block table. This test isolates the larger-stride case, where
an off-by-one would read the wrong physical block after multiple
tokens are accepted across a block boundary.
Setup (block_size=16):
- running = 28 + 2 - 0 = 30
@@ -2141,8 +2182,7 @@ class TestPostprocessMambaFusedKernel:
With identity block_ids = [0,1,2,3,...], an off-by-one that used
bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3,
producing a clear state-value mismatch against the Python
reference.
producing a clear mismatch against the untouched snapshot.
"""
cfg = test_config
torch.manual_seed(7002)
@@ -2166,6 +2206,7 @@ class TestPostprocessMambaFusedKernel:
fwd_py,
fwd_gpu,
) = _make_dual_layer_state(cfg, device)
conv_state_orig = conv_state_py.clone()
temporal_state_orig = temporal_state_py.clone()
# --- Python reference ---
@@ -2212,12 +2253,22 @@ class TestPostprocessMambaFusedKernel:
device=device,
)
# --- Ground truth: Python must have sourced temporal from block 3 ---
# --- Ground truth from untouched snapshots ---
actual_src_block_id = block_ids_per_req[0][3] # == 3
dest_block_id = block_ids_per_req[0][1] # == 1
expected_conv_state = conv_state_orig.clone()
expected_conv_state[dest_block_id, :-2].copy_(
conv_state_orig[dest_block_id, 2:]
)
torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0)
torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0)
# Python must have sourced temporal from block 3.
torch.testing.assert_close(
temporal_state_py[dest_block_id],
temporal_state_orig[actual_src_block_id],
rtol=0,
atol=0,
msg=(
"Python reference did not copy from block_ids[src+bias]=3; "
"test preconditions are wrong"
@@ -2251,21 +2302,44 @@ class TestPostprocessMambaFusedKernel:
msg="num_accepted_tokens mismatch at accept_token_bias=2",
)
def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd(
self, device, test_config, monkeypatch
@pytest.mark.parametrize(
"same_physical_block", [True, False], ids=["same", "distinct"]
)
@pytest.mark.parametrize("accept_token_bias", [1, 2, 3])
@pytest.mark.parametrize(
"dtype",
[torch.float16, torch.float32, torch.float64],
ids=["fp16", "fp32", "fp64"],
)
def test_sd_and_ds_conv_layouts_match_snapshot(
self,
device,
test_config,
monkeypatch,
accept_token_bias,
same_physical_block,
dtype,
):
"""DS conv postprocess should match SD when accept_token_bias > 0."""
"""SD and DS copies should independently match memmove semantics."""
from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils
cfg = test_config
cfg.dtype = dtype
torch.manual_seed(38898)
req_ids = ["req_0"]
num_computed_tokens = [30]
num_scheduled_tokens = {"req_0": 1}
# Keep new_num_computed on an aligned boundary while varying how far
# below it the running state starts. This makes the copy bias exactly
# ``accept_token_bias`` for each case. The 32 boundary keeps source and
# destination in logical block 1; the 64 boundary copies block 2 -> 3.
aligned_boundary = 32 if same_physical_block else 64
num_computed_tokens = [aligned_boundary - 2 * accept_token_bias]
num_scheduled_tokens = {"req_0": accept_token_bias}
num_draft_tokens: dict[str, int] = {}
num_accepted_tokens = [2] # Results in accept_token_bias = 1
mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx
num_accepted_tokens = [accept_token_bias + 1]
dest_block_idx = aligned_boundary // cfg.block_size - 1
src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1
mamba_state_idx = [src_block_idx]
block_ids_per_req = [list(range(8))]
layer_names = ["layer_0"]
@@ -2288,7 +2362,8 @@ class TestPostprocessMambaFusedKernel:
cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device
)
# SD GPU path. Default layout is SD.
# SD GPU path.
monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False)
model_mamba_utils.get_conv_state_layout.cache_clear()
sd_conv = sd_source_conv.clone()
sd_temporal = sd_source_temporal.clone()
@@ -2312,9 +2387,32 @@ class TestPostprocessMambaFusedKernel:
)
torch.accelerator.synchronize()
# Sanity: SD path actually modified the state (copy was performed).
assert not torch.equal(sd_conv, sd_source_conv), (
"SD baseline did not modify conv state; test setup is wrong"
src_block_id = block_ids_per_req[0][src_block_idx]
dest_block_id = block_ids_per_req[0][dest_block_idx]
expected_conv = sd_source_conv.clone()
expected_conv[dest_block_id, :-accept_token_bias].copy_(
sd_source_conv[src_block_id, accept_token_bias:]
)
torch.testing.assert_close(
sd_conv,
expected_conv,
rtol=0,
atol=0,
msg="SD conv copy did not match the untouched source snapshot",
)
actual_temporal_src_idx = src_block_idx + accept_token_bias
actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx]
expected_temporal = sd_source_temporal.clone()
expected_temporal[dest_block_id].copy_(
sd_source_temporal[actual_temporal_src_id]
)
torch.testing.assert_close(
sd_temporal,
expected_temporal,
rtol=0,
atol=0,
msg="SD temporal copy did not match the untouched source snapshot",
)
# DS GPU path on the DS twin.
@@ -2346,22 +2444,39 @@ class TestPostprocessMambaFusedKernel:
# Reset the lru cache so other tests see the default layout again.
model_mamba_utils.get_conv_state_layout.cache_clear()
# DS bytes, un-permuted, should match the SD result.
# Validate DS independently against the snapshot; otherwise a shared
# SD/DS bug would remain invisible.
ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous()
torch.testing.assert_close(
ds_conv.permute(0, 2, 1).contiguous(),
sd_conv,
msg=(
"DS conv post-kernel does not match SD baseline; the DS "
"row-loop in postprocess_mamba_fused_kernel is wrong."
),
ds_conv_sd_layout,
expected_conv,
rtol=0,
atol=0,
msg="DS conv copy did not match the untouched source snapshot",
)
torch.testing.assert_close(
ds_temporal,
sd_temporal,
msg="DS temporal state diverged from SD",
expected_temporal,
rtol=0,
atol=0,
msg="DS temporal copy did not match the untouched source snapshot",
)
expected_accepted = 1 if same_physical_block else accept_token_bias + 1
expected_accepted_tensor = torch.tensor(
[expected_accepted], dtype=torch.int32, device=device
)
torch.testing.assert_close(
gpu_ctx_sd.num_accepted_tokens_out[:num_reqs],
expected_accepted_tensor,
rtol=0,
atol=0,
msg="SD num_accepted_tokens result is wrong",
)
torch.testing.assert_close(
gpu_ctx_ds.num_accepted_tokens_out[:num_reqs],
gpu_ctx_sd.num_accepted_tokens_out[:num_reqs],
msg="DS num_accepted_tokens diverged from SD",
expected_accepted_tensor,
rtol=0,
atol=0,
msg="DS num_accepted_tokens result is wrong",
)
vllm/v1/worker/mamba_utils.py
+76 -26
View File
@@ -189,19 +189,46 @@ def _copy_mamba_state_block(
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
row_stride = tl.load(state_dim_row_stride_ptr + state_idx)
per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size
bias_bytes = token_bias.to(tl.int64) * state_elem_size
src_block_addr = state_base_addr + src_block_id * state_block_stride
offsets = tl.arange(0, COPY_BLOCK_SIZE)
for d in range(0, dim_rows):
row_src = src_block_addr + d * row_stride + bias_bytes
row_dst = dst_addr + d * row_stride
for i in range(0, per_row_bytes, COPY_BLOCK_SIZE):
mask = (i + offsets) < per_row_bytes
curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8))
curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8))
data = tl.load(curr_src, mask=mask)
tl.store(curr_dst, data, mask=mask)
# Stable row-to-lane ownership makes left shifts memmove-safe while
# exposing the dimension rows in parallel. All addresses retain
# state_elem_size alignment: tensor strides and token offsets are
# measured in whole elements before conversion to bytes.
num_dst_tokens = conv_width - token_bias
for token_idx in range(0, num_dst_tokens):
for row_base in range(0, dim_rows, COPY_BLOCK_SIZE):
rows = row_base + offsets
mask = rows < dim_rows
src_byte_addr = (
src_block_addr
+ rows * row_stride
+ (token_idx + token_bias) * state_elem_size
)
dst_byte_addr = (
dst_addr + rows * row_stride + token_idx * state_elem_size
)
if state_elem_size == 2:
src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16))
dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16))
data_u16 = tl.load(src_u16, mask=mask)
tl.store(dst_u16, data_u16, mask=mask)
elif state_elem_size == 4:
src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32))
dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32))
data_u32 = tl.load(src_u32, mask=mask)
tl.store(dst_u32, data_u32, mask=mask)
else:
for byte_idx in range(0, state_elem_size):
src_u8 = (src_byte_addr + byte_idx).to(
tl.pointer_type(tl.uint8)
)
dst_u8 = (dst_byte_addr + byte_idx).to(
tl.pointer_type(tl.uint8)
)
data_u8 = tl.load(src_u8, mask=mask)
tl.store(dst_u8, data_u8, mask=mask)
return
if is_conv_state:
@@ -210,22 +237,40 @@ def _copy_mamba_state_block(
# SD conv: copy
# state[bt[src_col], token_bias:] ->
# state[bt[dst_col], :conv_width - token_bias]
# Small per-block bytes (~60-80 KiB) make tiling degenerate, so
# conv runs as a single-CTA memcpy (NUM_TILES=1).
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size
src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
copy_size = (
(conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size
)
_memcpy_u64_tiled(
src_addr,
dst_addr,
copy_size,
tile_idx,
COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
NUM_TILES=1,
)
src_block_addr = state_base_addr + src_block_id * state_block_stride
token_bytes = state_inner_size * state_elem_size
num_dst_tokens = conv_width - token_bias
# Distinct blocks and exact self-copies cannot have a destructive
# overlap, so retain the u64-vectorized single-CTA copy.
if src_block_id != dest_block_id or token_bias == 0:
src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes
copy_size = num_dst_tokens.to(tl.int64) * token_bytes
_memcpy_u64_tiled(
src_addr,
dst_addr,
copy_size,
tile_idx,
COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
NUM_TILES=1,
)
return
# Copy tokens from low to high. Each token-sized source and destination
# region is disjoint, so same-block left shifts are memmove-safe
# without a barrier.
for token_idx in range(0, num_dst_tokens):
src_token = src_block_addr + (token_idx + token_bias) * token_bytes
dst_token = dst_addr + token_idx * token_bytes
_memcpy_u64_tiled(
src_token,
dst_token,
token_bytes,
tile_idx,
COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
NUM_TILES=1,
)
return
# Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
@@ -522,6 +567,7 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr):
src_ptr = tl.load(src_ptrs + pid)
dst_ptr = tl.load(dst_ptrs + pid)
size = tl.load(sizes + pid)
is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr
offsets = tl.arange(0, BLOCK_SIZE)
for i in range(0, size, BLOCK_SIZE):
@@ -531,6 +577,10 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr):
curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8))
data = tl.load(curr_src_ptr, mask=mask)
if is_left_overlap:
# Preserve each lane's source before a lower-address lane stores
# over it. The condition is uniform within the program.
tl.debug_barrier()
tl.store(curr_dst_ptr, data, mask=mask)