[Bugfix][Mamba] Fix overlapping state copy race (#50729)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Aug 17, 2026 at 01:29 UTC
a02cfccbc6187344325e364f09f6d8c33c4b253b
2 files changed
+228
-63
tests/v1/worker/test_mamba_utils.py
+152
-37
@@ -18,6 +18,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpe
18
from vllm.v1.worker.mamba_utils import (
19
MambaCopyBuffers,
20
MambaSpecDecodeGPUContext,
21
+ batch_memcpy,
22
collect_mamba_copy_meta,
23
do_mamba_copy_block,
24
preprocess_mamba,
@@ -536,6 +537,34 @@ class TestPostprocessMambaFusedKernel:
537
def test_config(self):
538
return _TestConfig()
539
540
+ def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device):
541
+ batch = 128
542
+ row_bytes = 32 * 1024
543
+ shift = 16
544
+ copy_size = row_bytes - shift
545
+
546
+ pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to(
547
+ torch.uint8
548
+ )
549
+ state = pattern.expand(batch, -1).clone()
550
+ snapshot = state.clone()
551
+
552
+ row_stride_bytes = state.stride(0) * state.element_size()
553
+ row_offsets = (
554
+ torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes
555
+ )
556
+ dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64)
557
+ src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64)
558
+ sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device)
559
+
560
+ expected = snapshot.clone()
561
+ expected[:, :copy_size].copy_(snapshot[:, shift:])
562
+ for _ in range(10):
563
+ state.copy_(snapshot)
564
+ batch_memcpy(src_ptrs, dst_ptrs, sizes)
565
+ torch.accelerator.synchronize()
566
+ torch.testing.assert_close(state, expected, rtol=0, atol=0)
567
+
568
def test_matches_python_postprocess_mamba(self, device, test_config):
569
"""
570
Golden test: GPU kernel produces identical results to Python impl.
@@ -1190,12 +1219,27 @@ class TestPostprocessMambaFusedKernel:
1219
# --- Verify Python behavior (ground truth) ---
1220
dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1
1221
1193
- # Conv state should be modified (shifted copy within block)
1194
- conv_changed = not torch.allclose(
1195
- conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
1222
+ # This is an overlapping in-place left shift, so comparing only the
1223
+ # Python and fused paths can hide the same memcpy race in both. Build
1224
+ # the memmove result from the untouched snapshot and check each path
1225
+ # independently.
1226
+ expected_conv_state = conv_state_orig.clone()
1227
+ expected_conv_state[dest_block_id, :-1].copy_(
1228
+ conv_state_orig[dest_block_id, 1:]
1229
)
1197
- assert conv_changed, (
1198
- "Python: Conv state should be modified when accept_token_bias > 0"
1230
+ torch.testing.assert_close(
1231
+ conv_state_py,
1232
+ expected_conv_state,
1233
+ rtol=0,
1234
+ atol=0,
1235
+ msg="Python: overlapping conv copy should have memmove semantics",
1236
+ )
1237
+ torch.testing.assert_close(
1238
+ conv_state_gpu,
1239
+ expected_conv_state,
1240
+ rtol=0,
1241
+ atol=0,
1242
+ msg="GPU: overlapping conv copy should have memmove semantics",
1243
)
1244
1245
# Temporal state should be modified (copy from different block)
@@ -2122,13 +2166,10 @@ class TestPostprocessMambaFusedKernel:
2166
actual_src_block_idx = src_block_idx + accept_token_bias
2167
actual_src_block_id = block_table[req, actual_src_block_idx]
2168
2125
- All prior regression tests exercise only ``bias == 1``, i.e. they
2126
- only ever read one slot ahead of ``src_block_idx`` in the block
2127
- table. An off-by-one (or missing scale) in the address computation
2128
- on line 143 of ``mamba_utils.py`` would be invisible to every
2129
- existing test but would silently read the wrong physical block on
2130
- any speculative-decode cycle that accepts multiple tokens across a
2131
- block boundary, feeding a stale hidden state forward one step.
2169
+ A ``bias == 1`` case only reads one slot ahead of ``src_block_idx``
2170
+ in the block table. This test isolates the larger-stride case, where
2171
+ an off-by-one would read the wrong physical block after multiple
2172
+ tokens are accepted across a block boundary.
2173
2174
Setup (block_size=16):
2175
- running = 28 + 2 - 0 = 30
@@ -2141,8 +2182,7 @@ class TestPostprocessMambaFusedKernel:
2182
2183
With identity block_ids = [0,1,2,3,...], an off-by-one that used
2184
bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3,
2144
- producing a clear state-value mismatch against the Python
2145
- reference.
2185
+ producing a clear mismatch against the untouched snapshot.
2186
"""
2187
cfg = test_config
2188
torch.manual_seed(7002)
@@ -2166,6 +2206,7 @@ class TestPostprocessMambaFusedKernel:
2206
fwd_py,
2207
fwd_gpu,
2208
) = _make_dual_layer_state(cfg, device)
2209
+ conv_state_orig = conv_state_py.clone()
2210
temporal_state_orig = temporal_state_py.clone()
2211
2212
# --- Python reference ---
@@ -2212,12 +2253,22 @@ class TestPostprocessMambaFusedKernel:
2253
device=device,
2254
)
2255
2215
- # --- Ground truth: Python must have sourced temporal from block 3 ---
2256
+ # --- Ground truth from untouched snapshots ---
2257
actual_src_block_id = block_ids_per_req[0][3] # == 3
2258
dest_block_id = block_ids_per_req[0][1] # == 1
2259
+ expected_conv_state = conv_state_orig.clone()
2260
+ expected_conv_state[dest_block_id, :-2].copy_(
2261
+ conv_state_orig[dest_block_id, 2:]
2262
+ )
2263
+ torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0)
2264
+ torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0)
2265
+
2266
+ # Python must have sourced temporal from block 3.
2267
torch.testing.assert_close(
2268
temporal_state_py[dest_block_id],
2269
temporal_state_orig[actual_src_block_id],
2270
+ rtol=0,
2271
+ atol=0,
2272
msg=(
2273
"Python reference did not copy from block_ids[src+bias]=3; "
2274
"test preconditions are wrong"
@@ -2251,21 +2302,44 @@ class TestPostprocessMambaFusedKernel:
2302
msg="num_accepted_tokens mismatch at accept_token_bias=2",
2303
)
2304
2254
- def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd(
2255
- self, device, test_config, monkeypatch
2305
+ @pytest.mark.parametrize(
2306
+ "same_physical_block", [True, False], ids=["same", "distinct"]
2307
+ )
2308
+ @pytest.mark.parametrize("accept_token_bias", [1, 2, 3])
2309
+ @pytest.mark.parametrize(
2310
+ "dtype",
2311
+ [torch.float16, torch.float32, torch.float64],
2312
+ ids=["fp16", "fp32", "fp64"],
2313
+ )
2314
+ def test_sd_and_ds_conv_layouts_match_snapshot(
2315
+ self,
2316
+ device,
2317
+ test_config,
2318
+ monkeypatch,
2319
+ accept_token_bias,
2320
+ same_physical_block,
2321
+ dtype,
2322
):
2257
- """DS conv postprocess should match SD when accept_token_bias > 0."""
2323
+ """SD and DS copies should independently match memmove semantics."""
2324
from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils
2325
2326
cfg = test_config
2327
+ cfg.dtype = dtype
2328
torch.manual_seed(38898)
2329
2330
req_ids = ["req_0"]
2264
- num_computed_tokens = [30]
2265
- num_scheduled_tokens = {"req_0": 1}
2331
+ # Keep new_num_computed on an aligned boundary while varying how far
2332
+ # below it the running state starts. This makes the copy bias exactly
2333
+ # ``accept_token_bias`` for each case. The 32 boundary keeps source and
2334
+ # destination in logical block 1; the 64 boundary copies block 2 -> 3.
2335
+ aligned_boundary = 32 if same_physical_block else 64
2336
+ num_computed_tokens = [aligned_boundary - 2 * accept_token_bias]
2337
+ num_scheduled_tokens = {"req_0": accept_token_bias}
2338
num_draft_tokens: dict[str, int] = {}
2267
- num_accepted_tokens = [2] # Results in accept_token_bias = 1
2268
- mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx
2339
+ num_accepted_tokens = [accept_token_bias + 1]
2340
+ dest_block_idx = aligned_boundary // cfg.block_size - 1
2341
+ src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1
2342
+ mamba_state_idx = [src_block_idx]
2343
block_ids_per_req = [list(range(8))]
2344
2345
layer_names = ["layer_0"]
@@ -2288,7 +2362,8 @@ class TestPostprocessMambaFusedKernel:
2362
cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device
2363
)
2364
2291
- # SD GPU path. Default layout is SD.
2365
+ # SD GPU path.
2366
+ monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False)
2367
model_mamba_utils.get_conv_state_layout.cache_clear()
2368
sd_conv = sd_source_conv.clone()
2369
sd_temporal = sd_source_temporal.clone()
@@ -2312,9 +2387,32 @@ class TestPostprocessMambaFusedKernel:
2387
)
2388
torch.accelerator.synchronize()
2389
2315
- # Sanity: SD path actually modified the state (copy was performed).
2316
- assert not torch.equal(sd_conv, sd_source_conv), (
2317
- "SD baseline did not modify conv state; test setup is wrong"
2390
+ src_block_id = block_ids_per_req[0][src_block_idx]
2391
+ dest_block_id = block_ids_per_req[0][dest_block_idx]
2392
+ expected_conv = sd_source_conv.clone()
2393
+ expected_conv[dest_block_id, :-accept_token_bias].copy_(
2394
+ sd_source_conv[src_block_id, accept_token_bias:]
2395
+ )
2396
+ torch.testing.assert_close(
2397
+ sd_conv,
2398
+ expected_conv,
2399
+ rtol=0,
2400
+ atol=0,
2401
+ msg="SD conv copy did not match the untouched source snapshot",
2402
+ )
2403
+
2404
+ actual_temporal_src_idx = src_block_idx + accept_token_bias
2405
+ actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx]
2406
+ expected_temporal = sd_source_temporal.clone()
2407
+ expected_temporal[dest_block_id].copy_(
2408
+ sd_source_temporal[actual_temporal_src_id]
2409
+ )
2410
+ torch.testing.assert_close(
2411
+ sd_temporal,
2412
+ expected_temporal,
2413
+ rtol=0,
2414
+ atol=0,
2415
+ msg="SD temporal copy did not match the untouched source snapshot",
2416
)
2417
2418
# DS GPU path on the DS twin.
@@ -2346,22 +2444,39 @@ class TestPostprocessMambaFusedKernel:
2444
# Reset the lru cache so other tests see the default layout again.
2445
model_mamba_utils.get_conv_state_layout.cache_clear()
2446
2349
- # DS bytes, un-permuted, should match the SD result.
2447
+ # Validate DS independently against the snapshot; otherwise a shared
2448
+ # SD/DS bug would remain invisible.
2449
+ ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous()
2450
torch.testing.assert_close(
2351
- ds_conv.permute(0, 2, 1).contiguous(),
2352
- sd_conv,
2353
- msg=(
2354
- "DS conv post-kernel does not match SD baseline; the DS "
2355
- "row-loop in postprocess_mamba_fused_kernel is wrong."
2356
- ),
2451
+ ds_conv_sd_layout,
2452
+ expected_conv,
2453
+ rtol=0,
2454
+ atol=0,
2455
+ msg="DS conv copy did not match the untouched source snapshot",
2456
)
2457
torch.testing.assert_close(
2458
ds_temporal,
2360
- sd_temporal,
2361
- msg="DS temporal state diverged from SD",
2459
+ expected_temporal,
2460
+ rtol=0,
2461
+ atol=0,
2462
+ msg="DS temporal copy did not match the untouched source snapshot",
2463
+ )
2464
+
2465
+ expected_accepted = 1 if same_physical_block else accept_token_bias + 1
2466
+ expected_accepted_tensor = torch.tensor(
2467
+ [expected_accepted], dtype=torch.int32, device=device
2468
)
2469
torch.testing.assert_close(
2364
- gpu_ctx_ds.num_accepted_tokens_out[:num_reqs],
2470
gpu_ctx_sd.num_accepted_tokens_out[:num_reqs],
2366
- msg="DS num_accepted_tokens diverged from SD",
2471
+ expected_accepted_tensor,
2472
+ rtol=0,
2473
+ atol=0,
2474
+ msg="SD num_accepted_tokens result is wrong",
2475
+ )
2476
+ torch.testing.assert_close(
2477
+ gpu_ctx_ds.num_accepted_tokens_out[:num_reqs],
2478
+ expected_accepted_tensor,
2479
+ rtol=0,
2480
+ atol=0,
2481
+ msg="DS num_accepted_tokens result is wrong",
2482
)
vllm/v1/worker/mamba_utils.py
+76
-26
@@ -189,19 +189,46 @@ def _copy_mamba_state_block(
189
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
190
dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
191
row_stride = tl.load(state_dim_row_stride_ptr + state_idx)
192
- per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size
193
- bias_bytes = token_bias.to(tl.int64) * state_elem_size
192
src_block_addr = state_base_addr + src_block_id * state_block_stride
193
offsets = tl.arange(0, COPY_BLOCK_SIZE)
196
- for d in range(0, dim_rows):
197
- row_src = src_block_addr + d * row_stride + bias_bytes
198
- row_dst = dst_addr + d * row_stride
199
- for i in range(0, per_row_bytes, COPY_BLOCK_SIZE):
200
- mask = (i + offsets) < per_row_bytes
201
- curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8))
202
- curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8))
203
- data = tl.load(curr_src, mask=mask)
204
- tl.store(curr_dst, data, mask=mask)
194
+
195
+ # Stable row-to-lane ownership makes left shifts memmove-safe while
196
+ # exposing the dimension rows in parallel. All addresses retain
197
+ # state_elem_size alignment: tensor strides and token offsets are
198
+ # measured in whole elements before conversion to bytes.
199
+ num_dst_tokens = conv_width - token_bias
200
+ for token_idx in range(0, num_dst_tokens):
201
+ for row_base in range(0, dim_rows, COPY_BLOCK_SIZE):
202
+ rows = row_base + offsets
203
+ mask = rows < dim_rows
204
+ src_byte_addr = (
205
+ src_block_addr
206
+ + rows * row_stride
207
+ + (token_idx + token_bias) * state_elem_size
208
+ )
209
+ dst_byte_addr = (
210
+ dst_addr + rows * row_stride + token_idx * state_elem_size
211
+ )
212
+ if state_elem_size == 2:
213
+ src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16))
214
+ dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16))
215
+ data_u16 = tl.load(src_u16, mask=mask)
216
+ tl.store(dst_u16, data_u16, mask=mask)
217
+ elif state_elem_size == 4:
218
+ src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32))
219
+ dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32))
220
+ data_u32 = tl.load(src_u32, mask=mask)
221
+ tl.store(dst_u32, data_u32, mask=mask)
222
+ else:
223
+ for byte_idx in range(0, state_elem_size):
224
+ src_u8 = (src_byte_addr + byte_idx).to(
225
+ tl.pointer_type(tl.uint8)
226
+ )
227
+ dst_u8 = (dst_byte_addr + byte_idx).to(
228
+ tl.pointer_type(tl.uint8)
229
+ )
230
+ data_u8 = tl.load(src_u8, mask=mask)
231
+ tl.store(dst_u8, data_u8, mask=mask)
232
return
233
234
if is_conv_state:
@@ -210,22 +237,40 @@ def _copy_mamba_state_block(
237
# SD conv: copy
238
# state[bt[src_col], token_bias:] ->
239
# state[bt[dst_col], :conv_width - token_bias]
213
- # Small per-block bytes (~60-80 KiB) make tiling degenerate, so
214
- # conv runs as a single-CTA memcpy (NUM_TILES=1).
240
src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
216
- src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size
217
- src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
218
- copy_size = (
219
- (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size
220
- )
221
- _memcpy_u64_tiled(
222
- src_addr,
223
- dst_addr,
224
- copy_size,
225
- tile_idx,
226
- COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
227
- NUM_TILES=1,
228
- )
241
+ src_block_addr = state_base_addr + src_block_id * state_block_stride
242
+ token_bytes = state_inner_size * state_elem_size
243
+ num_dst_tokens = conv_width - token_bias
244
+
245
+ # Distinct blocks and exact self-copies cannot have a destructive
246
+ # overlap, so retain the u64-vectorized single-CTA copy.
247
+ if src_block_id != dest_block_id or token_bias == 0:
248
+ src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes
249
+ copy_size = num_dst_tokens.to(tl.int64) * token_bytes
250
+ _memcpy_u64_tiled(
251
+ src_addr,
252
+ dst_addr,
253
+ copy_size,
254
+ tile_idx,
255
+ COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
256
+ NUM_TILES=1,
257
+ )
258
+ return
259
+
260
+ # Copy tokens from low to high. Each token-sized source and destination
261
+ # region is disjoint, so same-block left shifts are memmove-safe
262
+ # without a barrier.
263
+ for token_idx in range(0, num_dst_tokens):
264
+ src_token = src_block_addr + (token_idx + token_bias) * token_bytes
265
+ dst_token = dst_addr + token_idx * token_bytes
266
+ _memcpy_u64_tiled(
267
+ src_token,
268
+ dst_token,
269
+ token_bytes,
270
+ tile_idx,
271
+ COPY_BLOCK_SIZE=COPY_BLOCK_SIZE,
272
+ NUM_TILES=1,
273
+ )
274
return
275
276
# 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):
567
src_ptr = tl.load(src_ptrs + pid)
568
dst_ptr = tl.load(dst_ptrs + pid)
569
size = tl.load(sizes + pid)
570
+ is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr
571
572
offsets = tl.arange(0, BLOCK_SIZE)
573
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):
577
curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8))
578
579
data = tl.load(curr_src_ptr, mask=mask)
580
+ if is_left_overlap:
581
+ # Preserve each lane's source before a lower-address lane stores
582
+ # over it. The condition is uniform within the program.
583
+ tl.debug_barrier()
584
tl.store(curr_dst_ptr, data, mask=mask)
585
586