# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass from typing import Any from unittest.mock import MagicMock, patch import numpy as np import pytest import torch from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncsByType, get_conv_copy_spec, get_temporal_copy_spec, ) from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, MambaSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, _reinterpret_u64_as_i64, batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, get_mamba_groups, preprocess_mamba, stage_postprocess_inputs_to_gpu, ) # Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect. _DEFAULT_COPY_FUNCS: tuple[MambaStateCopyFunc, ...] = ( get_conv_copy_spec, get_temporal_copy_spec, ) _COPY_FUNCS: MambaStateCopyFuncsByType = { MambaAttentionBackendEnum.MAMBA2: _DEFAULT_COPY_FUNCS, MambaAttentionBackendEnum.GDN_ATTN: _DEFAULT_COPY_FUNCS, MambaAttentionBackendEnum.SHORT_CONV: (get_conv_copy_spec,), } def postprocess_mamba( scheduler_output: "SchedulerOutput", kv_cache_config: "KVCacheConfig", input_batch: Any, requests: dict[str, Any], forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, copy_bufs: "MambaCopyBuffers", ): """CPU reference for the align-mode postprocess. Used as a golden against the GPU fused kernel (``postprocess_mamba_align_gpu``). Mirrors what the production code did before the fused kernel replaced it; kept here because production no longer has a CPU implementation. """ assert input_batch.mamba_state_idx_cpu is not None num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu mamba_state_idx_cpu = input_batch.mamba_state_idx_cpu mamba_group_ids = copy_bufs.mamba_group_ids mamba_spec = copy_bufs.mamba_spec copy_bufs.offset = 0 for i, req_id in enumerate(input_batch.req_ids): req_state = requests[req_id] num_computed_tokens = req_state.num_computed_tokens num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, [])) num_scheduled_tokens = num_scheduled_tokens_dict[req_id] num_accepted_tokens = num_accepted_tokens_cpu[i] num_tokens_running_state = ( num_computed_tokens + num_scheduled_tokens - num_draft_tokens ) new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1 aligned_new_computed_tokens = ( new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size ) if aligned_new_computed_tokens >= num_tokens_running_state: accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state src_block_idx = mamba_state_idx_cpu[i] dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1 collect_mamba_copy_meta( copy_bufs, kv_cache_config, mamba_state_copy_funcs, mamba_group_ids, src_block_idx, dest_block_idx, accept_token_bias, req_state, forward_context, ) if src_block_idx == dest_block_idx: num_accepted_tokens_cpu[i] = 1 do_mamba_copy_block(copy_bufs) def _make_scheduler_output( finished_req_ids: set[str], preempted_req_ids: set[str] | None, resumed_req_ids: set[str], ) -> SchedulerOutput: cached = CachedRequestData.make_empty() cached.resumed_req_ids = resumed_req_ids return SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, num_scheduled_tokens={}, total_num_scheduled_tokens=0, scheduled_spec_decode_tokens={}, scheduled_encoder_inputs={}, num_common_prefix_blocks=[], finished_req_ids=finished_req_ids, free_encoder_mm_hashes=[], preempted_req_ids=preempted_req_ids, ) def test_resumed_req_ids_cleared_from_mamba_state_idx(): """When a request is force-preempted (e.g. reset_prefix_cache), it appears in resumed_req_ids but NOT in preempted_req_ids. preprocess_mamba must still clear its mamba_state_idx entry, otherwise stale indices can point beyond the new block allocation. """ spec = MagicMock(block_size=64, num_speculative_blocks=0) cache_config = MagicMock(enable_prefix_caching=True) input_batch = MagicMock(req_ids=[]) copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec) mamba_state_idx: dict[str, int] = { "finished": 1, "preempted": 2, "resumed": 3, # only in resumed_req_ids, NOT in preempted "keep": 99, } sched = _make_scheduler_output( finished_req_ids={"finished"}, preempted_req_ids={"preempted"}, resumed_req_ids={"resumed"}, ) with patch( "vllm.v1.worker.mamba_utils.get_mamba_groups", return_value=([0], spec), ): preprocess_mamba( sched, MagicMock(), # kv_cache_config cache_config, mamba_state_idx, input_batch, {}, # requests {}, # forward_context {}, # mamba_state_copy_funcs copy_bufs, ) assert mamba_state_idx == {"keep": 99} # ----------------------------------------------------------------------------- # Golden tests for postprocess_mamba_fused_kernel # ----------------------------------------------------------------------------- @dataclass class _TestConfig: """Common test configuration for fused kernel tests.""" block_size: int = 16 num_blocks: int = 32 num_layers: int = 2 num_reqs: int = 4 max_num_reqs: int = 8 # Conv state shape: [num_blocks, conv_width, inner_dim] conv_width: int = 4 conv_inner_dim: int = 64 # Temporal state shape: [num_blocks, state_dim] temporal_state_dim: int = 128 dtype: torch.dtype = torch.float16 class _MockCpuGpuBuffer: """Mock CpuGpuBuffer for testing without pinned memory.""" def __init__(self, size: int, dtype: torch.dtype, device: torch.device): self.cpu = torch.zeros(size, dtype=dtype, device="cpu") self.gpu = torch.zeros(size, dtype=dtype, device=device) self.np = self.cpu.numpy() def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: if n is None: return self.gpu.copy_(self.cpu, non_blocking=True) return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True) class _FakeDataPtrTensor: """Tensor wrapper that exposes a controlled data_ptr for metadata tests.""" def __init__(self, tensor: torch.Tensor, data_ptr: int): self._tensor = tensor self._data_ptr = data_ptr self.shape = tensor.shape def data_ptr(self) -> int: return self._data_ptr def dim(self) -> int: return self._tensor.dim() def stride(self, *args): return self._tensor.stride(*args) def numel(self) -> int: return self._tensor.numel() def element_size(self) -> int: return self._tensor.element_size() def size(self, *args): return self._tensor.size(*args) def __getitem__(self, item): return self._tensor[item] def test_reinterpret_u64_as_i64_preserves_pointer_bits(): ptrs = [ 0, 1, (1 << 63) - 1, 1 << 63, (1 << 63) + 1234, (1 << 64) - 1, ] ptr_tensor = torch.zeros(len(ptrs), dtype=torch.int64) for idx, ptr in enumerate(ptrs): ptr_tensor[idx] = _reinterpret_u64_as_i64(ptr) assert ptr_tensor.numpy().view(np.uint64).tolist() == ptrs def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): cfg = _TestConfig(num_layers=1) device = torch.device("cpu") kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) conv_ptr = 1 << 63 temporal_ptr = (1 << 64) - 8 block_table_ptr = (1 << 63) + 42 conv_state = _FakeDataPtrTensor( torch.empty( cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim, dtype=cfg.dtype, ), conv_ptr, ) temporal_state = _FakeDataPtrTensor( torch.empty(cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype), temporal_ptr, ) block_table = _FakeDataPtrTensor( torch.empty(1, 4, dtype=torch.int32), block_table_ptr, ) forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} gpu_ctx.initialize_from_forward_context( kv_cache_config, forward_context, _COPY_FUNCS, [block_table] ) assert gpu_ctx.state_base_addrs.tolist() == [ _reinterpret_u64_as_i64(conv_ptr), _reinterpret_u64_as_i64(temporal_ptr), ] assert gpu_ctx.block_table_ptrs.tolist() == [ _reinterpret_u64_as_i64(block_table_ptr) ] def _make_postprocess_scheduler_output( req_ids: list[str], num_scheduled_tokens: dict[str, int], scheduled_spec_decode_tokens: dict[str, list] | None = None, ) -> SchedulerOutput: """Create a minimal SchedulerOutput for postprocess testing.""" cached = CachedRequestData.make_empty() return SchedulerOutput( scheduled_new_reqs=[], scheduled_cached_reqs=cached, num_scheduled_tokens=num_scheduled_tokens, total_num_scheduled_tokens=sum(num_scheduled_tokens.values()), scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {}, scheduled_encoder_inputs={}, num_common_prefix_blocks=[], finished_req_ids=set(), free_encoder_mm_hashes=[], preempted_req_ids=set(), ) def _make_mock_attention( conv_state: torch.Tensor, temporal_state: torch.Tensor ) -> MagicMock: """Create a mock attention object with kv_cache.""" attention = MagicMock() attention.kv_cache = [conv_state, temporal_state] return attention def _make_dual_states( cfg: "_TestConfig", layer_names: list[str], device: torch.device, *, num_blocks: int | None = None, ) -> tuple[ list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], dict[str, MagicMock], dict[str, MagicMock], ]: """Allocate conv+temporal state tensors for the Python path, clone them for the GPU path, and build matching ``forward_context`` dicts for both. Returns ``(conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu)`` where the four state lists are parallel to ``layer_names``. """ n_blocks = num_blocks if num_blocks is not None else cfg.num_blocks conv_py = [ torch.randn( n_blocks, cfg.conv_width, cfg.conv_inner_dim, dtype=cfg.dtype, device=device, ) for _ in layer_names ] temporal_py = [ torch.randn(n_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device) for _ in layer_names ] conv_gpu = [s.clone() for s in conv_py] temporal_gpu = [s.clone() for s in temporal_py] fwd_py = { name: _make_mock_attention(c, t) for name, c, t in zip(layer_names, conv_py, temporal_py) } fwd_gpu = { name: _make_mock_attention(c, t) for name, c, t in zip(layer_names, conv_gpu, temporal_gpu) } return conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu def _make_dual_layer_state( cfg: "_TestConfig", device: torch.device, *, num_blocks: int | None = None, layer_name: str = "layer_0", ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, MagicMock], dict[str, MagicMock], ]: """Single-layer convenience form of ``_make_dual_states``.""" conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu = _make_dual_states( cfg, [layer_name], device, num_blocks=num_blocks ) return conv_py[0], temporal_py[0], conv_gpu[0], temporal_gpu[0], fwd_py, fwd_gpu def _make_kv_cache_config(cfg: _TestConfig, layer_names: list[str]) -> KVCacheConfig: """Create a KVCacheConfig with mamba groups.""" mamba_spec = MambaSpec( block_size=cfg.block_size, shapes=( (cfg.conv_width, cfg.conv_inner_dim), (cfg.temporal_state_dim,), ), dtypes=(cfg.dtype, cfg.dtype), mamba_cache_mode="all", ) group = KVCacheGroupSpec( layer_names=layer_names, kv_cache_spec=mamba_spec, ) return KVCacheConfig( num_blocks=cfg.num_blocks, kv_cache_tensors=[], kv_cache_groups=[group], ) def _make_input_batch( req_ids: list[str], num_accepted_tokens: list[int], mamba_state_idx: list[int], ) -> MagicMock: """Create a mock GPUInputBatch.""" batch = MagicMock() batch.req_ids = req_ids batch.req_id_to_index = {rid: i for i, rid in enumerate(req_ids)} # Use numpy arrays so modifications persist batch.num_accepted_tokens_cpu = np.array(num_accepted_tokens, dtype=np.int32) batch.mamba_state_idx_cpu = np.array(mamba_state_idx, dtype=np.int32) return batch def _make_requests( req_ids: list[str], num_computed_tokens: list[int], block_ids_per_req: list[list[int]], ) -> dict[str, MagicMock]: """Create mock CachedRequestState objects.""" requests = {} for i, req_id in enumerate(req_ids): req = MagicMock() req.num_computed_tokens = num_computed_tokens[i] req.block_ids = {0: block_ids_per_req[i]} # group_id=0 requests[req_id] = req return requests def _make_copy_bufs( cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device, copy_funcs: MambaStateCopyFuncsByType = _COPY_FUNCS, ) -> MambaCopyBuffers: """Create MambaCopyBuffers for the Python path.""" def make_buffer(n, dtype): return _MockCpuGpuBuffer(n, dtype, device) return MambaCopyBuffers.create( max_num_reqs=cfg.max_num_reqs, kv_cache_config=kv_cache_config, copy_funcs=copy_funcs, make_buffer=make_buffer, ) def _make_gpu_ctx( cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device, copy_funcs: MambaStateCopyFuncsByType = _COPY_FUNCS, ) -> MambaSpecDecodeGPUContext: """Create MambaSpecDecodeGPUContext for the GPU path.""" def make_buffer(n, dtype): return _MockCpuGpuBuffer(n, dtype, device) return MambaSpecDecodeGPUContext.create( max_num_reqs=cfg.max_num_reqs, kv_cache_config=kv_cache_config, copy_funcs=copy_funcs, device=device, make_buffer=make_buffer, ) def test_mamba_groups_support_different_state_specs(): gdn_spec = MambaSpec( block_size=16, shapes=((4, 4), (2, 4, 4)), dtypes=(torch.float16, torch.float32), mamba_type=MambaAttentionBackendEnum.GDN_ATTN, mamba_cache_mode="align", num_speculative_blocks=3, ) short_conv_spec = MambaSpec( block_size=16, shapes=((12, 8),), dtypes=(torch.float16,), mamba_type=MambaAttentionBackendEnum.SHORT_CONV, mamba_cache_mode="align", num_speculative_blocks=3, ) kv_cache_config = KVCacheConfig( num_blocks=8, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec(["gdn.0", "gdn.1"], gdn_spec), KVCacheGroupSpec(["ple.0"], short_conv_spec), ], ) mamba_groups = get_mamba_groups(kv_cache_config) assert mamba_groups == {gdn_spec: [0], short_conv_spec: [1]} model_state = object.__new__(MambaHybridModelState) model_state._mamba_group_ids = [] model_state._mamba_spec = None group_ids, representative_spec = model_state._get_mamba_group_info(kv_cache_config) assert group_ids == [0, 1] assert representative_spec is gdn_spec cfg = _TestConfig() copy_bufs = _make_copy_bufs(cfg, kv_cache_config, torch.device("cpu")) assert copy_bufs.src_ptrs.cpu.numel() == cfg.max_num_reqs * 5 assert copy_bufs.dst_ptrs.cpu.numel() == cfg.max_num_reqs * 5 assert copy_bufs.sizes.cpu.numel() == cfg.max_num_reqs * 5 ctx = _make_gpu_ctx(cfg, kv_cache_config, torch.device("cpu")) assert ctx.num_states == 5 forward_context = {} gdn_spec_layer_names = ["gdn.0", "gdn.1"] for layer_name in gdn_spec_layer_names: attention = MagicMock() attention.kv_cache = ( torch.empty(8, 4, 4, dtype=torch.float16), torch.empty(8, 2, 4, 4, dtype=torch.float32), ) forward_context[layer_name] = attention ple_attention = MagicMock() ple_attention.kv_cache = (torch.empty(8, 12, 8, dtype=torch.float16),) forward_context["ple.0"] = ple_attention assert gdn_spec_layer_names == kv_cache_config.kv_cache_groups[0].layer_names block_tables = [ torch.zeros(4, 8, dtype=torch.int32), torch.zeros(4, 8, dtype=torch.int32), ] ctx.initialize_from_forward_context( kv_cache_config, forward_context, _COPY_FUNCS, block_tables, ) assert ctx.is_initialized assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1] assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] def test_mamba_groups_support_mixed_specs_in_uniform_group(): gdn_spec = MambaSpec( block_size=16, shapes=((4, 4), (2, 4, 4)), dtypes=(torch.float16, torch.float32), mamba_type=MambaAttentionBackendEnum.GDN_ATTN, mamba_cache_mode="align", num_speculative_blocks=3, ) short_conv_spec = MambaSpec( block_size=16, shapes=((12, 8),), dtypes=(torch.float16,), mamba_type=MambaAttentionBackendEnum.SHORT_CONV, mamba_cache_mode="align", num_speculative_blocks=3, ) layer_specs = { "gdn.0": gdn_spec, "gdn.1": gdn_spec, "ple.0": short_conv_spec, } kv_cache_config = KVCacheConfig( num_blocks=8, kv_cache_tensors=[], kv_cache_groups=[ KVCacheGroupSpec( list(layer_specs), UniformTypeKVCacheSpecs( block_size=16, kv_cache_specs=layer_specs, ), ) ], ) assert kv_cache_config.has_mamba_layers assert get_mamba_groups(kv_cache_config) == { gdn_spec: [0], short_conv_spec: [0], } # The group is listed once per spec, so its id must not be repeated. model_state = object.__new__(MambaHybridModelState) model_state._mamba_group_ids = [] model_state._mamba_spec = None group_ids, representative_spec = model_state._get_mamba_group_info(kv_cache_config) assert group_ids == [0] assert representative_spec is gdn_spec cfg = _TestConfig() copy_bufs = _make_copy_bufs(cfg, kv_cache_config, torch.device("cpu")) assert copy_bufs.mamba_group_ids == [0] assert copy_bufs.src_ptrs.cpu.numel() == cfg.max_num_reqs * 5 ctx = _make_gpu_ctx(cfg, kv_cache_config, torch.device("cpu")) assert ctx.mamba_group_ids == [0] assert ctx.num_states == 5 forward_context = { "gdn.0": MagicMock( kv_cache=( torch.empty(8, 4, 4, dtype=torch.float16), torch.empty(8, 2, 4, 4, dtype=torch.float32), ) ), "gdn.1": MagicMock( kv_cache=( torch.empty(8, 4, 4, dtype=torch.float16), torch.empty(8, 2, 4, 4, dtype=torch.float32), ) ), "ple.0": MagicMock(kv_cache=(torch.empty(8, 12, 8, dtype=torch.float16),)), } ctx.initialize_from_forward_context( kv_cache_config, forward_context, _COPY_FUNCS, [torch.zeros(4, 8, dtype=torch.int32)], ) assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 0] assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] # ----------------------------------------------------------------------------- # stage_postprocess_inputs_to_gpu: single-pass staging into pinned views # ----------------------------------------------------------------------------- def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: """Build a MambaSpecDecodeGPUContext stand-in exposing only the four per-request staging buffers touched by stage_postprocess_inputs_to_gpu.""" ctx = MagicMock() ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_computed_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_draft_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) return ctx def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): """stage_postprocess_inputs_to_gpu writes each request's state_idx and scheduled/computed/draft counts into slot `i` for `req_ids[i]`, leaves slots past `num_reqs` untouched, and mirrors the values to the GPU buffers.""" device = torch.device("cpu") max_num_reqs = 8 ctx = _make_staging_ctx(max_num_reqs, device) # Any negative int32 works as a sentinel: all staged values (state_idx, # scheduled/computed/draft token counts) are non-negative, so a negative # entry surviving in slots [num_reqs:] proves the loop didn't overrun. sentinel = np.int32(-1) bufs = ( ctx.mamba_state_idx_buf, ctx.num_scheduled_tokens_buf, ctx.num_computed_tokens_buf, ctx.num_draft_tokens_buf, ) for buf in bufs: buf.np[:] = sentinel req_ids = ["req_a", "req_b", "req_c"] num_reqs = len(req_ids) scheduler_output = _make_postprocess_scheduler_output( req_ids=req_ids, num_scheduled_tokens={"req_a": 5, "req_b": 12, "req_c": 3}, # req_b intentionally absent -> defaults to 0 drafts. scheduled_spec_decode_tokens={ "req_a": [1, 2], "req_c": [1, 2, 3, 4], }, ) requests = _make_requests( req_ids=req_ids, num_computed_tokens=[10, 20, 30], block_ids_per_req=[[0], [0], [0]], ) mamba_state_idx = {"req_a": 100, "req_b": 200, "req_c": 300} # A trailing entry past num_reqs must not be read. req_ids_padded = req_ids + ["not_scheduled"] stage_postprocess_inputs_to_gpu( ctx, scheduler_output, req_ids_padded, num_reqs, requests, mamba_state_idx, ) np.testing.assert_array_equal( ctx.mamba_state_idx_buf.np[:num_reqs], [100, 200, 300] ) np.testing.assert_array_equal( ctx.num_scheduled_tokens_buf.np[:num_reqs], [5, 12, 3] ) np.testing.assert_array_equal( ctx.num_computed_tokens_buf.np[:num_reqs], [10, 20, 30] ) np.testing.assert_array_equal(ctx.num_draft_tokens_buf.np[:num_reqs], [2, 0, 4]) for buf in bufs: assert (buf.np[num_reqs:] == sentinel).all() assert torch.equal( ctx.mamba_state_idx_buf.gpu[:num_reqs], torch.tensor([100, 200, 300], dtype=torch.int32), ) assert torch.equal( ctx.num_draft_tokens_buf.gpu[:num_reqs], torch.tensor([2, 0, 4], dtype=torch.int32), ) def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): """If preprocess_mamba didn't populate mamba_state_idx for a req in the batch, staging must fail loudly rather than silently writing a stale index.""" device = torch.device("cpu") ctx = _make_staging_ctx(max_num_reqs=4, device=device) scheduler_output = _make_postprocess_scheduler_output( req_ids=["req_a"], num_scheduled_tokens={"req_a": 1}, ) requests = _make_requests(["req_a"], [0], [[0]]) # mamba_state_idx has an entry for a *different* request but not for # req_a (the one being staged). This is the realistic failure mode: # preprocess_mamba ran, populated some reqs, but missed this one. # The safety assert in stage_postprocess_inputs_to_gpu must fire on # req_a's missing entry rather than silently writing None. mamba_state_idx: dict[str, int] = {"other_req": 42} with pytest.raises(AssertionError, match="mamba_state_idx missing entry"): stage_postprocess_inputs_to_gpu( ctx, scheduler_output, ["req_a"], 1, requests, mamba_state_idx, ) def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: device = torch.device("cpu") config = _TestConfig(num_layers=1) layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(config, layer_names) conv_state = torch.empty( config.num_blocks, config.conv_width, config.conv_inner_dim, dtype=config.dtype, ) temporal_state = torch.empty( config.num_blocks, config.temporal_state_dim, dtype=config.dtype ) attention = MagicMock() attention.kv_cache = [ conv_state, temporal_state, *(torch.empty(config.num_blocks, 1) for _ in range(4)), ] context = _make_gpu_ctx(config, kv_cache_config, device) context.initialize_from_forward_context( kv_cache_config, {"layer_0": attention}, _COPY_FUNCS, [torch.zeros(1, 1, dtype=torch.int32)], ) assert context.is_initialized assert context.state_base_addrs.tolist() == [ conv_state.data_ptr(), temporal_state.data_ptr(), ] def _run_gpu_postprocess( gpu_ctx: MambaSpecDecodeGPUContext, *, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], copy_funcs: MambaStateCopyFuncsByType, block_table: torch.Tensor, req_ids: list[str], num_accepted_tokens: list[int], mamba_state_idx: list[int], num_scheduled_tokens: dict[str, int], num_computed_tokens: list[int], num_draft_tokens: dict[str, int], device: torch.device, ) -> None: """Initialize the GPU context against `block_table`, run the fused postprocess kernel for `req_ids`, and synchronize.""" def t(values): return torch.tensor(values, dtype=torch.int32, device=device) gpu_ctx.initialize_from_forward_context( kv_cache_config, forward_context, copy_funcs, [block_table] ) gpu_ctx.run_fused_postprocess( num_reqs=len(req_ids), num_accepted_tokens_gpu=t(num_accepted_tokens), mamba_state_idx_gpu=t(mamba_state_idx), num_scheduled_tokens_gpu=t([num_scheduled_tokens[r] for r in req_ids]), num_computed_tokens_gpu=t(num_computed_tokens), num_draft_tokens_gpu=t([num_draft_tokens.get(r, 0) for r in req_ids]), ) torch.accelerator.synchronize() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") class TestPostprocessMambaFusedKernel: """Tests for postprocess_mamba_fused_kernel comparing GPU vs CPU paths.""" @pytest.fixture def device(self): return torch.device("cuda:0") @pytest.fixture 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. This test: 1. Sets up identical initial state for both paths 2. Runs Python postprocess_mamba (modifies states via batch_memcpy) 3. Runs GPU fused kernel (modifies states directly) 4. Compares resulting state tensors and num_accepted_tokens """ cfg = test_config torch.manual_seed(42) # Test scenario: 4 requests with different copy conditions # Copy needed when: aligned_new_computed >= num_tokens_running_state # where: num_tokens_running_state = num_computed + num_scheduled - num_draft # new_num_computed = num_tokens_running_state + num_accepted - 1 # aligned_new_computed = (new_num_computed // block_size) * block_size req_ids = ["req_0", "req_1", "req_2", "req_3"] # Configure requests so some need copies, some don't # block_size = 16 # req_0: running=60+5-2=63, new=63+3-1=65, aligned=64 >= 63 -> COPY # req_1: running=30+3-0=33, new=33+2-1=34, aligned=32 < 33 -> NO COPY # req_2: running=45+8-3=50, new=50+4-1=53, aligned=48 < 50 -> NO COPY # req_3: running=10+6-0=16, new=16+2-1=17, aligned=16 >= 16 -> COPY num_computed_tokens = [60, 30, 45, 10] num_scheduled_tokens = {"req_0": 5, "req_1": 3, "req_2": 8, "req_3": 6} num_draft_tokens = {"req_0": 2, "req_1": 0, "req_2": 3, "req_3": 0} num_accepted_tokens = [3, 2, 4, 2] mamba_state_idx = [3, 1, 2, 0] # source block indices # Block IDs for each request (simulate block table) block_ids_per_req = [ list(range(8)), # req_0: blocks 0-7 list(range(8, 16)), # req_1: blocks 8-15 list(range(16, 24)), # req_2: blocks 16-23 list(range(24, 32)), # req_3: blocks 24-31 ] layer_names = [f"layer_{i}" for i in range(cfg.num_layers)] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_states_py, temporal_states_py, conv_states_gpu, temporal_states_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_states(cfg, layer_names, device) # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) max_blocks = max(len(b) for b in block_ids_per_req) block_table_gpu = torch.zeros( num_reqs, max_blocks, dtype=torch.int32, device=device ) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32 ) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Compare results --- # 1. Compare state tensors for i in range(cfg.num_layers): torch.testing.assert_close( conv_states_gpu[i], conv_states_py[i], msg=f"Conv state mismatch at layer {i}", ) torch.testing.assert_close( temporal_states_gpu[i], temporal_states_py[i], msg=f"Temporal state mismatch at layer {i}", ) # 2. Compare num_accepted_tokens updates expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch", ) def test_no_copy_when_not_needed(self, device, test_config): """Kernel should not modify state when no copy is needed.""" cfg = test_config torch.manual_seed(123) # Single request where no copy is needed: # running = 30 + 3 = 33, new = 33 + 1 - 1 = 33, aligned = 32 < 33 req_ids = ["req_0"] num_computed_tokens = [30] num_scheduled_tokens = {"req_0": 3} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [1] mamba_state_idx = [1] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) # Create state tensor conv_state = torch.randn( cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim, dtype=cfg.dtype, device=device, ) temporal_state = torch.randn( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) # Clone to verify no modification conv_state_orig = conv_state.clone() temporal_state_orig = temporal_state.clone() forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # State should be unchanged torch.testing.assert_close(conv_state, conv_state_orig) torch.testing.assert_close(temporal_state, temporal_state_orig) @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16]) def test_various_batch_sizes(self, device, test_config, num_reqs): """Verify kernel works correctly with different batch sizes.""" cfg = _TestConfig(max_num_reqs=max(16, num_reqs)) torch.manual_seed(456) req_ids = [f"req_{i}" for i in range(num_reqs)] # All requests will trigger a copy num_computed_tokens = [60] * num_reqs num_scheduled_tokens = {r: 5 for r in req_ids} num_draft_tokens = {r: 0 for r in req_ids} num_accepted_tokens = [3] * num_reqs mamba_state_idx = [3] * num_reqs # Each request gets unique blocks block_ids_per_req = [list(range(i * 8, (i + 1) * 8)) for i in range(num_reqs)] # Ensure we have enough blocks total_blocks = num_reqs * 8 cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs)) layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # Run Python path scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # Run GPU path gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) max_blocks_per_req = 8 block_table_gpu = torch.zeros( num_reqs, max_blocks_per_req, dtype=torch.int32, device=device ) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32 ) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # Compare results torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="Conv state mismatch" ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="Temporal state mismatch" ) def test_block_table_with_realistic_stride(self, device, test_config): """ Test kernel with realistic block table strides. In real usage, the block table is pre-allocated with shape [max_num_reqs, max_num_blocks_per_req] and then sliced to [:num_reqs]. This means stride(0) = max_num_blocks_per_req, which is typically much larger than the actual blocks used. This test verifies the kernel handles non-tight strides correctly, catching bugs where stride is incorrectly treated as bytes vs elements. """ cfg = test_config torch.manual_seed(789) # Use multiple requests to exercise stride-based indexing num_reqs = 4 req_ids = [f"req_{i}" for i in range(num_reqs)] # All requests trigger copies (same setup as test_various_batch_sizes) num_computed_tokens = [60] * num_reqs num_scheduled_tokens = {r: 5 for r in req_ids} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [3] * num_reqs mamba_state_idx = [3] * num_reqs # Each request uses only 8 blocks, but we allocate much more blocks_used_per_req = 8 block_ids_per_req = [ list(range(i * blocks_used_per_req, (i + 1) * blocks_used_per_req)) for i in range(num_reqs) ] total_blocks = num_reqs * blocks_used_per_req cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs)) layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # Run Python path scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # Run GPU path with REALISTIC block table stride gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) # KEY DIFFERENCE: Create a large block table like real code does # Real system has max_num_blocks_per_req >> blocks actually used max_num_reqs_full = 16 max_blocks_per_req_full = 512 # Much larger than blocks_used_per_req=8 # Allocate full-size table (simulates pre-allocated CpuGpuBuffer) block_table_full = torch.zeros( max_num_reqs_full, max_blocks_per_req_full, dtype=torch.int32, device=device ) # Fill in actual block IDs (only first few columns used) for i, block_ids in enumerate(block_ids_per_req): block_table_full[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32 ) # Slice like real code: block_table.gpu[:num_reqs] # This preserves stride(0) = 512, not 8! block_table_gpu = block_table_full[:num_reqs] # Verify stride is large (the key property we're testing) assert block_table_gpu.stride(0) == max_blocks_per_req_full, ( f"Expected stride {max_blocks_per_req_full}, " f"got {block_table_gpu.stride(0)}" ) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # Compare results - this will fail if stride handling is incorrect torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="Conv state mismatch - possible stride bug in kernel", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="Temporal state mismatch - possible stride bug in kernel", ) def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( self, device, test_config ): """ Test the ``src_addr == dst_addr`` early-return path in postprocess_mamba_fused_kernel matches Python behavior. When src_addr == dst_addr (source and destination memory addresses are identical), both implementations should: 1. Skip the copy (state unchanged) 2. Set num_accepted_tokens to 1 This condition occurs when: - src_block_idx == dest_block_idx (same logical block) - accept_token_bias == 0 (no offset within the block) Python reference (collect_mamba_copy_meta): if src_block_idx == dest_block_idx and accept_token_bias == 0: return # No copy added Python reference (postprocess_mamba): if src_block_idx == dest_block_idx: num_accepted_tokens_cpu[i] = 1 Test setup (block_size=16): - num_tokens_running_state = 30 + 2 - 0 = 32 - new_num_computed = 32 + 1 - 1 = 32 - aligned_new_computed = 32 - accept_token_bias = 32 - 32 = 0 - dest_block_idx = 32 // 16 - 1 = 1 - src_block_idx = 1 (set explicitly) """ cfg = test_config torch.manual_seed(1001) req_ids = ["req_0"] num_computed_tokens = [30] num_scheduled_tokens = {"req_0": 2} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [5] # Initial value, should become 1 mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # Also clone to verify no modification conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Verify Python behavior (ground truth) --- # State should be unchanged (no copy when src_addr == dst_addr) torch.testing.assert_close( conv_state_py, conv_state_orig, msg="Python: Conv state should be unchanged when src==dst", ) torch.testing.assert_close( temporal_state_py, temporal_state_orig, msg="Python: Temporal state should be unchanged when src==dst", ) # num_accepted_tokens should be 1 assert input_batch_py.num_accepted_tokens_cpu[0] == 1, ( f"Python: num_accepted_tokens should be 1, " f"got {input_batch_py.num_accepted_tokens_cpu[0]}" ) # --- Verify GPU matches Python --- torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="GPU conv state should match Python", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="GPU temporal state should match Python", ) expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( self, device, test_config ): """ Test the ``src_block_idx == dest_block_idx`` post-copy update in postprocess_mamba_fused_kernel matches Python behavior. When src_block_idx == dest_block_idx but accept_token_bias > 0, both implementations should: 1. Perform the copy (src_addr != dst_addr due to offset) 2. Set num_accepted_tokens to 1 AFTER the copy Python reference (postprocess_mamba): if src_block_idx == dest_block_idx: num_accepted_tokens_cpu[i] = 1 For conv states: copies state[block, offset:] to state[block, :] (shifted window) For temporal states: copies state[block_ids[src_idx + offset]] to state[block_ids[dest_idx]] Test setup (block_size=16): - num_tokens_running_state = 30 + 1 - 0 = 31 - new_num_computed = 31 + 2 - 1 = 32 - aligned_new_computed = 32 - accept_token_bias = 32 - 31 = 1 (> 0, so copy happens) - dest_block_idx = 32 // 16 - 1 = 1 - src_block_idx = 1 (set explicitly, == dest_block_idx) """ cfg = test_config torch.manual_seed(1002) req_ids = ["req_0"] num_computed_tokens = [30] num_scheduled_tokens = {"req_0": 1} 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 block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # Clone to verify modification conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1 # 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:] ) 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) src_block_id_temporal = block_ids_per_req[0][2] # actual_src_block_idx = 2 dest_block_id_temporal = block_ids_per_req[0][1] # dest_block_idx = 1 torch.testing.assert_close( temporal_state_py[dest_block_id_temporal], temporal_state_orig[src_block_id_temporal], msg="Python: Temporal state copy should have happened", ) # num_accepted_tokens should be 1 assert input_batch_py.num_accepted_tokens_cpu[0] == 1, ( f"Python: num_accepted_tokens should be 1, " f"got {input_batch_py.num_accepted_tokens_cpu[0]}" ) # --- Verify GPU matches Python --- torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="GPU conv state should match Python", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="GPU temporal state should match Python", ) expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) def test_different_block_idx_copies_without_setting_accepted_to_1( self, device, test_config ): """ Test that neither special-case path triggers when src_block_idx != dest_block_idx, and GPU matches Python behavior. When copying between different blocks: 1. src_addr != dst_addr (different blocks = different addresses) 2. src_block_idx != dest_block_idx Therefore: - The ``src_addr == dst_addr`` early-return does NOT trigger - The ``src_block_idx == dest_block_idx`` post-copy update does NOT trigger - Copy happens normally - num_accepted_tokens remains UNCHANGED Test setup (block_size=16): - num_tokens_running_state = 60 + 3 - 0 = 63 - new_num_computed = 63 + 3 - 1 = 65 - aligned_new_computed = 64 - accept_token_bias = 64 - 63 = 1 - dest_block_idx = 64 // 16 - 1 = 3 - src_block_idx = 2 (set explicitly, != dest_block_idx) """ cfg = test_config torch.manual_seed(1003) req_ids = ["req_0"] num_computed_tokens = [60] num_scheduled_tokens = {"req_0": 3} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [3] # Should remain 3, NOT set to 1 mamba_state_idx = [2] # src_block_idx = 2, dest_block_idx will be 3 block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # Clone to verify modification conv_state_orig = conv_state_py.clone() # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][3] # dest_block_idx = 3 # Copy DID happen (dest block should be modified) conv_changed = not torch.allclose( conv_state_py[dest_block_id], conv_state_orig[dest_block_id] ) assert conv_changed, "Python: Conv state copy should have happened" # num_accepted_tokens should NOT be changed to 1 assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], ( f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, " f"got {input_batch_py.num_accepted_tokens_cpu[0]}" ) # --- Verify GPU matches Python --- torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="GPU conv state should match Python", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="GPU temporal state should match Python", ) expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) def test_prefix_caching_shared_block_does_not_set_accepted_to_1( self, device, test_config ): """ Regression test: with prefix caching, different logical block indices can map to the same physical block. The kernel must NOT set num_accepted_tokens to 1 in that case. When src_block_idx != dest_block_idx but block_table maps both to the same physical block ID, src_addr == dst_addr. The copy is correctly skipped (self-copy is a no-op), but num_accepted_tokens must be preserved — only logical-index equality justifies setting it to 1. Test setup (block_size=16): - num_tokens_running_state = 30 + 2 - 0 = 32 - new_num_computed = 32 + 3 - 1 = 34 - aligned_new_computed = 32 - accept_token_bias = 32 - 32 = 0 - dest_block_idx = 32 // 16 - 1 = 1 - src_block_idx = 0 (set explicitly, != dest_block_idx) - block_ids = [5, 5, ...] — prefix caching: both logical indices map to the same physical block 5 """ cfg = test_config torch.manual_seed(2001) req_ids = ["req_0"] num_computed_tokens = [30] num_scheduled_tokens = {"req_0": 2} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [3] # Must stay 3, NOT become 1 mamba_state_idx = [0] # src_block_idx = 0, dest_block_idx will be 1 # Prefix caching: logical blocks 0 and 1 share physical block 5 block_ids_per_req = [[5, 5, 2, 3, 4, 6, 7, 8]] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Verify Python behavior (ground truth) --- # Copy is self-to-self (same physical block), state unchanged torch.testing.assert_close( conv_state_py, conv_state_orig, msg="Python: Conv state should be unchanged (self-copy)", ) torch.testing.assert_close( temporal_state_py, temporal_state_orig, msg="Python: Temporal state should be unchanged (self-copy)", ) # num_accepted_tokens must NOT be set to 1 (src_block_idx != dest_block_idx) assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], ( f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, " f"got {input_batch_py.num_accepted_tokens_cpu[0]}" ) # --- Verify GPU matches Python --- torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="GPU conv state should match Python", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="GPU temporal state should match Python", ) expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python (must NOT be 1)", ) def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_config): """ Regression test: non-sequential physical block IDs under prefix caching with the needs_copy boundary at exact equality. Under PC, the block allocator assigns physical block IDs in arbitrary order (e.g., [17, 3, 42, 9] instead of [0, 1, 2, 3]). The needs_copy condition is purely token-count based and must evaluate identically regardless of the physical block IDs assigned. This test verifies that the kernel's address arithmetic (block_table lookup, stride computation) produces correct copies when physical IDs are non-sequential. Two requests exercise different boundary behaviors: - req_0: aligned_new_computed == num_tokens_running_state (exact boundary) This is the tightest edge: one fewer accepted token and no copy needed. - req_1: aligned_new_computed == num_tokens_running_state (exact boundary) Different block layout, src!=dest, real copy happens. Both use non-sequential block IDs typical of PC reuse patterns. Test setup (block_size=16): req_0: - num_tokens_running_state = 48 + 0 - 0 = 48 - new_num_computed = 48 + 1 - 1 = 48 - aligned_new_computed = 48 - needs_copy = (48 >= 48) = True (exact boundary!) - accept_token_bias = 48 - 48 = 0 - dest_block_idx = 48 // 16 - 1 = 2 - src_block_idx = 2 (same as dest -> num_accepted = 1) req_1: - num_tokens_running_state = 31 + 1 - 0 = 32 - new_num_computed = 32 + 3 - 1 = 34 - aligned_new_computed = 32 - needs_copy = (32 >= 32) = True (exact boundary!) - accept_token_bias = 32 - 32 = 0 - dest_block_idx = 32 // 16 - 1 = 1 - src_block_idx = 0 (diff from dest -> num_accepted unchanged) """ cfg = test_config torch.manual_seed(4001) req_ids = ["req_0", "req_1"] num_computed_tokens = [48, 31] num_scheduled_tokens = {"req_0": 0, "req_1": 1} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [1, 3] mamba_state_idx = [2, 0] # Non-sequential block IDs typical of prefix caching allocation block_ids_per_req = [ [17, 3, 42, 9, 25, 11, 30, 2], # req_0: scattered physical blocks [41, 7, 22, 15, 38, 19, 4, 28], # req_1: different scattered blocks ] layer_names = [f"layer_{i}" for i in range(cfg.num_layers)] # Need enough physical blocks for the scattered IDs num_blocks = 50 local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs) kv_cache_config = _make_kv_cache_config(local_cfg, layer_names) ( conv_states_py, temporal_states_py, conv_states_gpu, temporal_states_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_states(local_cfg, layer_names, device) # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device) num_reqs = len(req_ids) max_blocks = max(len(b) for b in block_ids_per_req) block_table_gpu = torch.zeros( num_reqs, max_blocks, dtype=torch.int32, device=device ) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32 ) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Compare results --- for i in range(cfg.num_layers): torch.testing.assert_close( conv_states_gpu[i], conv_states_py[i], msg=f"Conv state mismatch at layer {i} with non-sequential block IDs", ) torch.testing.assert_close( temporal_states_gpu[i], temporal_states_py[i], msg=( f"Temporal state mismatch at layer {i} " f"with non-sequential block IDs" ), ) expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch with non-sequential block IDs", ) # Verify req_0 had num_accepted set to 1 (src==dest) and req_1 unchanged assert input_batch_py.num_accepted_tokens_cpu[0] == 1 assert input_batch_py.num_accepted_tokens_cpu[1] == num_accepted_tokens[1] def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_config): """ Regression test: mixed batch under prefix caching where some requests have shared physical blocks (aliased) and others have distinct blocks, with the needs_copy boundary at various positions. This tests the interaction between: 1. PC block aliasing (src and dest map to same physical block) 2. The needs_copy boundary (exact equality vs well-past vs no-copy) 3. Non-sequential physical block IDs Batch of 4 requests: - req_0: needs_copy=True, src!=dest, shared physical block (PC aliased) -> copy skipped (src_addr==dst_addr), num_accepted PRESERVED - req_1: needs_copy=True, src!=dest, distinct blocks, non-sequential IDs -> real copy happens, num_accepted PRESERVED - req_2: needs_copy=False (below boundary) -> no action at all - req_3: needs_copy=True, src==dest (exact boundary, zero bias) -> copy skipped (self-copy), num_accepted SET TO 1 Test setup (block_size=16): req_0: running=30+2-0=32, new=32+3-1=34, aligned=32, 32>=32 -> COPY bias=0, dest=32//16-1=1, src=0 (!=dest) block_ids=[5,5,...] -> same physical -> skip, keep accepted=3 req_1: running=60+5-2=63, new=63+3-1=65, aligned=64, 64>=63 -> COPY bias=1, dest=64//16-1=3, src=2 (!=dest) block_ids=[41,7,22,15,...] -> distinct -> real copy, keep accepted=3 req_2: running=30+3-0=33, new=33+1-1=33, aligned=32, 32<33 -> NO COPY req_3: running=48+0-0=48, new=48+1-1=48, aligned=48, 48>=48 -> COPY bias=0, dest=48//16-1=2, src=2 (==dest) block_ids=[10,20,30,...] -> distinct IDs, same logical idx -> self-copy (src_addr==dst_addr), set accepted=1 """ cfg = test_config torch.manual_seed(5001) req_ids = ["req_0", "req_1", "req_2", "req_3"] num_computed_tokens = [30, 60, 30, 48] num_scheduled_tokens = {"req_0": 2, "req_1": 5, "req_2": 3, "req_3": 0} num_draft_tokens = {"req_1": 2} num_accepted_tokens = [3, 3, 1, 1] mamba_state_idx = [0, 2, 1, 2] # Block IDs with various PC patterns: # req_0: shared blocks (PC alias: logical 0 and 1 -> physical 5) # req_1: distinct non-sequential blocks # req_2: doesn't matter (no copy) # req_3: distinct sequential blocks (no aliasing) block_ids_per_req = [ [5, 5, 12, 18, 23, 31, 44, 2], # req_0: blocks 0,1 share phys 5 [41, 7, 22, 15, 38, 19, 4, 28], # req_1: all distinct [10, 20, 30, 40, 1, 6, 8, 14], # req_2: irrelevant [10, 20, 30, 40, 1, 6, 8, 14], # req_3: distinct, dest=src=idx 2 ] layer_names = [f"layer_{i}" for i in range(cfg.num_layers)] num_blocks = 50 local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs) kv_cache_config = _make_kv_cache_config(local_cfg, layer_names) ( conv_states_py, temporal_states_py, conv_states_gpu, temporal_states_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_states(local_cfg, layer_names, device) # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device) num_reqs = len(req_ids) max_blocks = max(len(b) for b in block_ids_per_req) block_table_gpu = torch.zeros( num_reqs, max_blocks, dtype=torch.int32, device=device ) for i, block_ids in enumerate(block_ids_per_req): block_table_gpu[i, : len(block_ids)] = torch.tensor( block_ids, dtype=torch.int32 ) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Compare all state tensors --- for i in range(cfg.num_layers): torch.testing.assert_close( conv_states_gpu[i], conv_states_py[i], msg=( f"Conv state mismatch at layer {i} — " f"mixed PC batch with shared/distinct blocks" ), ) torch.testing.assert_close( temporal_states_gpu[i], temporal_states_py[i], msg=( f"Temporal state mismatch at layer {i} — " f"mixed PC batch with shared/distinct blocks" ), ) # --- Compare num_accepted_tokens --- expected_accepted = torch.tensor( input_batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch in mixed PC batch", ) # Verify per-request expectations: # req_0: src!=dest, shared block -> preserved (3) assert input_batch_py.num_accepted_tokens_cpu[0] == 3 # req_1: src!=dest, distinct blocks -> preserved (3) assert input_batch_py.num_accepted_tokens_cpu[1] == 3 # req_2: no copy -> preserved (1) assert input_batch_py.num_accepted_tokens_cpu[2] == 1 # req_3: src==dest -> set to 1 assert input_batch_py.num_accepted_tokens_cpu[3] == 1 def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr( self, device, test_config ): """ Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return guard must compare logical block indices, not physical addresses. Under prefix caching, different logical blocks (src_block_idx=0, dest_block_idx=1) can map to the same physical block. When accept_token_bias=0, this makes src_addr == dst_addr for BOTH conv and temporal states. A buggy guard `if src_addr == dst_addr` would incorrectly set num_accepted_tokens=1; the correct guard is `if src_block_idx == dest_block_idx and accept_token_bias == 0`. The Python reference only sets num_accepted_tokens=1 when src_block_idx == dest_block_idx (line 79 of postprocess_mamba). With src_block_idx=0, dest_block_idx=1, num_accepted_tokens must be preserved even though the physical addresses match. Test setup (block_size=16): - num_tokens_running_state = 30 + 2 - 0 = 32 - new_num_computed = 32 + 3 - 1 = 34 - aligned_new_computed = 32 - needs_copy = (32 >= 32) = True - accept_token_bias = 32 - 32 = 0 - dest_block_idx = 32 // 16 - 1 = 1 - src_block_idx = 0 (explicitly, != dest_block_idx) - block_ids = [7, 7, ...] -> physical aliasing via prefix caching Expected: num_accepted_tokens stays 3 (not set to 1). Bug (959ca0fd): kernel saw src_addr == dst_addr, set it to 1. """ cfg = test_config torch.manual_seed(6001) req_ids = ["req_0"] num_computed_tokens = [30] num_scheduled_tokens = {"req_0": 2} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [3] mamba_state_idx = [0] # src_block_idx = 0 # Prefix caching: logical blocks 0 and 1 both map to physical block 7. block_ids_per_req = [[7, 7, 10, 11, 12, 13, 14, 15]] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, forward_context_py, forward_context_gpu, ) = _make_dual_layer_state(cfg, device) # --- Run Python path --- scheduler_output = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) input_batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( scheduler_output, kv_cache_config, input_batch_py, requests, forward_context_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # Python reference: src_block_idx(0) != dest_block_idx(1) -> no change assert input_batch_py.num_accepted_tokens_cpu[0] == 3, ( f"Python: num_accepted_tokens should remain 3, " f"got {input_batch_py.num_accepted_tokens_cpu[0]}" ) # --- Run GPU path --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=forward_context_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # The critical assertion: kernel must NOT set num_accepted_tokens to 1 # when src_block_idx != dest_block_idx, even though src_addr == dst_addr # due to prefix caching aliasing. # # Old kernel (959ca0fd): `if src_addr == dst_addr` -> FAILS here (sets 1) # Fixed kernel (6466ce0d): `if src_block_idx == dest_block_idx and # accept_token_bias == 0` -> PASSES (preserves 3) kernel_accepted = gpu_ctx.num_accepted_tokens_out[0].item() assert kernel_accepted == 3, ( f"Kernel set num_accepted_tokens to {kernel_accepted} but expected 3. " f"The early-return guard likely compared physical addresses " f"(src_addr == dst_addr) instead of logical block indices " f"(src_block_idx == dest_block_idx). Under prefix caching, " f"different logical blocks can share the same physical block." ) # Also verify state tensors match Python torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="GPU conv state should match Python", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg="GPU temporal state should match Python", ) def test_as_strided_temporal_copy_size(self, device, test_config): """ Regression test for 240723d46: temporal copy_size must be inner_size * elem_size, not state_block_stride. In production (gpu_model_runner.py), conv and temporal states share a raw buffer via torch.as_strided where stride(0) equals page_size_bytes / elem_size — larger than either state's natural element count. Using stride(0) as copy_size for temporal states overwrites into the next block's conv region. Layout per page (384 float16 elements = 768 bytes): [conv: 256 elems | temporal: 128 elems] The test triggers a temporal copy from block 4 to block 3. With the bug the kernel copies 768 bytes (page stride) instead of 256 bytes (128 * 2), overwriting conv_state[4] with conv_state[5]'s data. Test setup (block_size=16): - running = 60 + 5 - 2 = 63 - new = 63 + 3 - 1 = 65 - aligned = 64 >= 63 -> COPY needed - accept_token_bias = 64 - 63 = 1 - dest_block_idx = 64 // 16 - 1 = 3 - temporal: actual_src_block_idx = 3 + 1 = 4 (block_ids[4] = 4) """ cfg = test_config torch.manual_seed(7001) req_ids = ["req_0"] num_computed_tokens = [60] num_scheduled_tokens = {"req_0": 5} num_draft_tokens = {"req_0": 2} num_accepted_tokens = [3] mamba_state_idx = [3] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) # --- Production-like packed layout (mirrors gpu_model_runner.py) --- conv_shape = (cfg.conv_width, cfg.conv_inner_dim) temporal_shape = (cfg.temporal_state_dim,) dtype = cfg.dtype elem_size = torch.tensor([], dtype=dtype).element_size() conv_natural_elems = cfg.conv_width * cfg.conv_inner_dim temporal_natural_elems = cfg.temporal_state_dim page_size_bytes = (conv_natural_elems + temporal_natural_elems) * elem_size num_element_per_page = page_size_bytes // elem_size assert num_element_per_page > temporal_natural_elems, ( "Test requires padded stride; page must be larger than one state" ) raw_py = torch.randn( cfg.num_blocks * num_element_per_page, dtype=dtype, device=device ) raw_gpu = raw_py.clone() def make_views(raw): conv_tgt = (cfg.num_blocks, *conv_shape) conv_nat_stride = torch.empty(conv_tgt).stride() conv = torch.as_strided( raw, size=conv_tgt, stride=(num_element_per_page, *conv_nat_stride[1:]), storage_offset=0, ) temp_tgt = (cfg.num_blocks, *temporal_shape) temp_nat_stride = torch.empty(temp_tgt).stride() temp = torch.as_strided( raw, size=temp_tgt, stride=(num_element_per_page, *temp_nat_stride[1:]), storage_offset=conv_natural_elems, ) return conv, temp conv_py, temp_py = make_views(raw_py) conv_gpu, temp_gpu = make_views(raw_gpu) fwd_py = {"layer_0": _make_mock_attention(conv_py, temp_py)} fwd_gpu = {"layer_0": _make_mock_attention(conv_gpu, temp_gpu)} # --- Python reference --- sched = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( sched, kv_cache_config, batch_py, requests, fwd_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- GPU fused kernel --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = 1 block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=fwd_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- Assertions --- # With the bug (pre-240723d46), the kernel copies page_size_bytes # (768) for temporal state instead of 256 bytes, overwriting # conv_state[4] with conv_state[5]'s data. torch.testing.assert_close( conv_gpu, conv_py, msg=( "Conv state corrupted: temporal copy_size was likely " "state_block_stride instead of inner_size * elem_size" ), ) torch.testing.assert_close( temp_gpu, temp_py, msg="Temporal state mismatch", ) expected_accepted = torch.tensor( batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch", ) def test_temporal_copy_with_bias_ge_2(self, device, test_config): """ Coverage test for the temporal-state block-table stride arithmetic when ``accept_token_bias >= 2``. The kernel computes, for temporal (non-conv) states:: actual_src_block_idx = src_block_idx + accept_token_bias actual_src_block_id = block_table[req, actual_src_block_idx] 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 - new = 30 + 3 - 1 = 32 - aligned = 32 >= 30 -> COPY needed - bias = 32 - 30 = 2 (key: >= 2) - dest_idx = 32 // 16 - 1 = 1 - src_idx = 1 (same as dest -> exercises post-copy accepted=1 write) - temporal actual_src_block_idx = 1 + 2 = 3 (reads block_table[0, 3]) 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 mismatch against the untouched snapshot. """ cfg = test_config torch.manual_seed(7002) req_ids = ["req_0"] num_computed_tokens = [28] num_scheduled_tokens = {"req_0": 2} num_draft_tokens: dict[str, int] = {} num_accepted_tokens = [3] # -> accept_token_bias = 2 mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) ( conv_state_py, temporal_state_py, conv_state_gpu, temporal_state_gpu, 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 --- sched = _make_postprocess_scheduler_output( req_ids, num_scheduled_tokens, {k: [None] * v for k, v in num_draft_tokens.items() if v > 0}, ) batch_py = _make_input_batch( req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy() ) requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req) copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device) postprocess_mamba( sched, kv_cache_config, batch_py, requests, fwd_py, _COPY_FUNCS, copy_bufs, ) torch.accelerator.synchronize() # --- GPU fused kernel --- gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) num_reqs = 1 block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) _run_gpu_postprocess( gpu_ctx, kv_cache_config=kv_cache_config, forward_context=fwd_gpu, copy_funcs=_COPY_FUNCS, block_table=block_table, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) # --- 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" ), ) # --- GPU kernel must match Python byte-for-byte --- torch.testing.assert_close( conv_state_gpu, conv_state_py, msg="Conv state mismatch at accept_token_bias=2", ) torch.testing.assert_close( temporal_state_gpu, temporal_state_py, msg=( "Temporal state mismatch at accept_token_bias=2: the kernel " "likely read the wrong slot of the block table " "(actual_src_block_idx stride arithmetic)" ), ) expected_accepted = torch.tensor( batch_py.num_accepted_tokens_cpu[:num_reqs], dtype=torch.int32, device=device, ) torch.testing.assert_close( gpu_ctx.num_accepted_tokens_out[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch at accept_token_bias=2", ) @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, ): """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"] # 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 = [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"] kv_cache_config = _make_kv_cache_config(cfg, layer_names) num_reqs = len(req_ids) block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device) block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32) # Same logical conv state in SD and DS layouts. sd_source_conv = torch.randn( cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim, dtype=cfg.dtype, device=device, ) ds_source_conv = sd_source_conv.permute(0, 2, 1).contiguous() sd_source_temporal = torch.randn( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) # 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() forward_context_sd = { "layer_0": _make_mock_attention(sd_conv, sd_temporal), } gpu_ctx_sd = _make_gpu_ctx(cfg, kv_cache_config, device) _run_gpu_postprocess( gpu_ctx_sd, kv_cache_config=kv_cache_config, forward_context=forward_context_sd, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) torch.accelerator.synchronize() 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. monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS") model_mamba_utils.get_conv_state_layout.cache_clear() try: ds_conv = ds_source_conv.clone() ds_temporal = sd_source_temporal.clone() forward_context_ds = { "layer_0": _make_mock_attention(ds_conv, ds_temporal), } gpu_ctx_ds = _make_gpu_ctx(cfg, kv_cache_config, device) _run_gpu_postprocess( gpu_ctx_ds, kv_cache_config=kv_cache_config, forward_context=forward_context_ds, copy_funcs=_COPY_FUNCS, block_table=block_table_gpu, req_ids=req_ids, num_accepted_tokens=num_accepted_tokens, mamba_state_idx=mamba_state_idx, num_scheduled_tokens=num_scheduled_tokens, num_computed_tokens=num_computed_tokens, num_draft_tokens=num_draft_tokens, device=device, ) torch.accelerator.synchronize() finally: # Reset the lru cache so other tests see the default layout again. model_mamba_utils.get_conv_state_layout.cache_clear() # 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_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, 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], expected_accepted_tensor, rtol=0, atol=0, msg="DS num_accepted_tokens result is wrong", )