diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a338b934b5..2ba1ce1c59 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -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", ) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 9b89af0812..82247e3ee4 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -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)