From 638e92f28cc3076b6306228709ca96690429b718 Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:12:35 -0400 Subject: [PATCH 007/103] worker: index the persistent per-request-slot block table under PP Under pipeline parallelism the Mamba/GDN state-copy path indexed a block-table row gathered for a different step. The align post-process now indexes the persistent per-request-slot table, and fails closed when it is absent. --- tests/v1/worker/gpu/__init__.py | 2 + .../gpu/test_pp_deferred_mamba_block_table.py | 706 ++++++++++++++++++ .../worker/gpu/test_pp_deferred_mamba_gpu.py | 542 ++++++++++++++ tests/v1/worker/test_mamba_utils.py | 25 +- vllm/v1/worker/gpu/model_runner.py | 4 + vllm/v1/worker/gpu/model_states/interface.py | 8 + .../worker/gpu/model_states/mamba_hybrid.py | 19 + vllm/v1/worker/mamba_utils.py | 83 +- 8 files changed, 1381 insertions(+), 8 deletions(-) create mode 100644 tests/v1/worker/gpu/__init__.py create mode 100644 tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py create mode 100644 tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py diff --git a/tests/v1/worker/gpu/__init__.py b/tests/v1/worker/gpu/__init__.py new file mode 100644 index 0000000..208f01a --- /dev/null +++ b/tests/v1/worker/gpu/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py new file mode 100644 index 0000000..30bed0a --- /dev/null +++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_block_table.py @@ -0,0 +1,706 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which block-table row the align post-process kernel reads. + +Under pipeline parallelism the Mamba align post-process is *deferred*: it is +enqueued by ``update_pp_decode_requests`` at the top of ``execute_model`` +(``model_runner.py:1644``), before this step's ``gather_block_tables`` +(``model_runner.py:1442`` -> ``block_table.py:149``), and the ``idx_mapping`` it +carries was captured ``pp_size`` steps earlier on a ``PendingRecv`` +(``pp_utils.py:57-58, 258-271``). There is exactly one gathered destination +buffer, reused in place for the engine's lifetime (``block_table.py:74-76``), so +at that moment it holds the batch of step T-1 while the mapping describes the +batch of step T-2. With two-way PP micro-batch alternation those are different +micro-batches by construction, so ``bt_row_idx = batch_idx`` +(``mamba_utils.py:481``) resolved a *different request's* row. + +These tests model the row selection on CPU: they drive the production plumbing +(``MambaHybridModelState._ensure_align_ctx`` -> +``MambaSpecDecodeGPUContext.initialize_from_forward_context`` -> +``run_fused_postprocess_align``), capture the launch instead of running Triton, +and then replay the kernel's index arithmetic against the two candidate tables. + +No GPU is required. +""" + +import inspect +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFuncsByType, + get_conv_copy_spec, + get_temporal_copy_spec, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec +from vllm.v1.worker import mamba_utils as mamba_utils_mod +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext + +MAX_NUM_REQS = 8 +MAX_NUM_BLOCKS = 6 +MAMBA_BLOCK_SIZE = 16 +NUM_PHYS_BLOCKS = 64 +CONV_WIDTH = 4 +CONV_INNER = 8 +TEMPORAL_DIM = 16 +DTYPE = torch.float16 + +_COPY_FUNCS: MambaStateCopyFuncsByType = { + MambaAttentionBackendEnum.GDN_ATTN: (get_conv_copy_spec, get_temporal_copy_spec), +} + +# Captured before any monkeypatching so the replay binds against the real +# kernel signature; a launch that stops matching it fails loudly here. +_POSTPROCESS_SIG = inspect.signature( + mamba_utils_mod.postprocess_mamba_fused_kernel.fn +) +_PRECOPY_SIG = inspect.signature( + mamba_utils_mod.precopy_mamba_align_fused_kernel.fn +) + + +class _MockCpuGpuBuffer: + """CpuGpuBuffer stand-in that needs no 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: + self.gpu.copy_(self.cpu) + return self.gpu + self.gpu[:n].copy_(self.cpu[:n]) + return self.gpu[:n] + + +class _RecordingKernel: + """Stands in for a triton.jit kernel and records the launch.""" + + def __init__(self) -> None: + self.launches: list[tuple[tuple[int, ...], tuple, dict]] = [] + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.launches.append((grid, args, kwargs)) + + return launch + + +def _ptr_to_int(value: int) -> int: + """Undo ``_reinterpret_u64_as_i64`` (mamba_utils.py:186-189).""" + return value + (1 << 64) if value < 0 else value + + +def _gather_rows( + persistent: torch.Tensor, + num_blocks: list[int], + idx_mapping: list[int], + num_reqs_padded: int, + dst: torch.Tensor, +) -> None: + """CPU port of ``_gather_block_tables_kernel`` (block_table.py:236-273). + + Row ``batch_idx`` of ``dst`` receives ``num_blocks[req_idx]`` entries of + persistent row ``req_idx``; rows in ``[num_reqs, num_reqs_padded)`` are + zeroed whole; rows at or beyond ``num_reqs_padded`` are left holding the + previous gather's content. + """ + num_reqs = len(idx_mapping) + for batch_idx in range(num_reqs_padded): + if batch_idx >= num_reqs: + dst[batch_idx].zero_() + continue + req_idx = idx_mapping[batch_idx] + nb = num_blocks[req_idx] + dst[batch_idx, :nb] = persistent[req_idx, :nb] + + +def _make_kv_cache_config() -> KVCacheConfig: + spec = MambaSpec( + block_size=MAMBA_BLOCK_SIZE, + shapes=((CONV_WIDTH, CONV_INNER), (TEMPORAL_DIM,)), + dtypes=(DTYPE, DTYPE), + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="align", + num_speculative_blocks=3, + ) + return KVCacheConfig( + num_blocks=NUM_PHYS_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer_0"], spec)], + ) + + +def _make_forward_context() -> dict: + attention = MagicMock() + attention.kv_cache = [ + torch.zeros(NUM_PHYS_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE), + torch.zeros(NUM_PHYS_BLOCKS, TEMPORAL_DIM, dtype=DTYPE), + ] + return {"layer_0": attention} + + +def _register_block_tables(state: MambaHybridModelState, block_tables) -> bool: + """Hand the runner's persistent ``BlockTables`` to the model state. + + Deliberately tolerant: on a tree without the fix the hook does not exist, and + the point of these tests is to report *which row the kernel reads*, not to + die on a missing attribute. + """ + setter = getattr(state, "set_block_tables", None) + if setter is None: + return False + setter(block_tables) + return True + + +def _build_scenario( + device: torch.device = torch.device("cpu"), + slot_mapping: list[int] = (0, 1, 2), + t1_mapping: list[int] = (2,), + register: bool = True, +): + """T-2 gathers ``slot_mapping``; T-1 re-gathers ``t1_mapping`` into the same + buffer; the deferred post-process then runs with T-2's mapping. + + ``slot_mapping`` is the batch-row -> req-slot mapping carried on the + PendingRecv. Making it non-identity is what separates "the kernel used the + persistent table" from "the kernel used req_idx as the row": with an + identity mapping ``batch_idx == req_idx`` and the two are indistinguishable. + """ + kv_cache_config = _make_kv_cache_config() + forward_context = _make_forward_context() + + # Persistent per-req-slot table (block_table.py:53-59). Row r owns block + # ids 10*(r+1) + col, so a row is identifiable from any single entry. + persistent = torch.zeros( + MAX_NUM_REQS, MAX_NUM_BLOCKS, dtype=torch.int32, device=device + ) + for r in range(MAX_NUM_REQS): + for c in range(MAX_NUM_BLOCKS): + persistent[r, c] = 10 * (r + 1) + c + num_blocks = [3] * MAX_NUM_REQS + + # The single gathered destination buffer (block_table.py:74-76). + slot_mapping = list(slot_mapping) + t1_mapping = list(t1_mapping) + gathered = torch.zeros_like(persistent) + _gather_rows(persistent, num_blocks, slot_mapping, len(slot_mapping) + 1, gathered) + _gather_rows(persistent, num_blocks, t1_mapping, len(t1_mapping) + 1, gathered) + + ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=MAX_NUM_REQS, + kv_cache_config=kv_cache_config, + copy_funcs=_COPY_FUNCS, + device=device, + make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), + ) + + state = object.__new__(MambaHybridModelState) + state.device = device + state.max_num_reqs = MAX_NUM_REQS + state._align_mode = True + state._mamba_ctx = ctx + state._mamba_group_ids = [] + state._mamba_spec = None + state._mamba_state_copy_funcs = _COPY_FUNCS + # object.__new__ skips __init__, so mirror the attribute it would set. + state._runner_block_tables = None + state.vllm_config = SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context=forward_context) + ) + registered = ( + _register_block_tables( + state, + SimpleNamespace(block_tables=[SimpleNamespace(gpu=persistent)]), + ) + if register + else False + ) + + # What the runner hands prepare_attn/preprocess_state: the T-1 gather view. + gathered_view = gathered[: len(t1_mapping) + 1] + state._ensure_align_ctx(kv_cache_config, [0], (gathered_view,)) + + tables_by_ptr = { + persistent.data_ptr(): ("persistent", persistent), + gathered.data_ptr(): ("gathered", gathered), + } + return SimpleNamespace( + ctx=ctx, + state=state, + slot_mapping=slot_mapping, + persistent=persistent, + gathered=gathered, + tables_by_ptr=tables_by_ptr, + registered=registered, + ) + + +def _replay_row(scenario, launch, batch_idx: int) -> tuple[str, int, dict]: + """Replay the kernel's row selection (mamba_utils.py:437-490) for one row. + + Returns ``(table_name, bt_row_idx, predicate)``. + """ + grid, args, kwargs = launch + bound = _POSTPROCESS_SIG.bind(*args, **kwargs) + bound.apply_defaults() + a = bound.arguments + + num_reqs = a["num_reqs"] + assert grid[0] == num_reqs, (grid, num_reqs) + assert batch_idx < num_reqs + + has_idx_mapping = bool(a["HAS_IDX_MAPPING"]) + # Absent on a tree without the fix; the kernel then always uses batch_idx. + bt_row_is_req = bool(a.get("BT_ROW_IS_REQ", False)) + + if has_idx_mapping: + req_idx = int(a["idx_mapping_ptr"][batch_idx]) + if req_idx < 0: + return ("skipped", -1, {}) + else: + req_idx = batch_idx + + num_accepted = int(a["num_accepted_tokens_ptr"][req_idx]) + src_block_idx = int(a["mamba_state_idx_ptr"][req_idx]) + assert bool(a["PRECOMPUTED_NEW_COMPUTED"]), "V2 align path only" + new_num_computed = int(a["num_computed_tokens_ptr"][req_idx]) + block_size = int(a["block_size"]) + + num_tokens_running_state = new_num_computed - num_accepted + 1 + aligned_new_computed = (new_num_computed // block_size) * block_size + needs_copy = aligned_new_computed >= num_tokens_running_state + accept_token_bias = aligned_new_computed - num_tokens_running_state + dest_block_idx = aligned_new_computed // block_size - 1 + predicate = { + "needs_copy": needs_copy, + "num_accepted": num_accepted, + "new_num_computed_mod_block": new_num_computed % block_size, + "src_col": src_block_idx, + "dst_col": dest_block_idx, + "token_bias": accept_token_bias, + "src_eq_dst": src_block_idx == dest_block_idx, + } + + if bt_row_is_req: + bt_row_idx = req_idx + else: + bt_row_idx = batch_idx if has_idx_mapping else req_idx + + group_idx = int(a["state_group_indices_ptr"][0]) + ptr = _ptr_to_int(int(a["block_table_ptrs_ptr"][group_idx])) + name, _table = scenario.tables_by_ptr[ptr] + return (name, bt_row_idx, predicate) + + +def _block_ids(scenario, launch, batch_idx: int) -> dict: + name, bt_row_idx, predicate = _replay_row(scenario, launch, batch_idx) + table = next( + tbl for label, tbl in scenario.tables_by_ptr.values() if label == name + ) + return { + "table": name, + "row": bt_row_idx, + "src_block_id": int(table[bt_row_idx, predicate["src_col"]]), + "dst_block_id": int(table[bt_row_idx, predicate["dst_col"]]), + "temporal_src_block_id": int( + table[bt_row_idx, predicate["src_col"] + predicate["token_bias"]] + ), + **predicate, + } + + +def _run_deferred(scenario, monkeypatch: pytest.MonkeyPatch) -> _RecordingKernel: + recorder = _RecordingKernel() + monkeypatch.setattr( + mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder + ) + device = scenario.persistent.device + + def t(values): + return torch.tensor(values, dtype=torch.int32, device=device) + + # Per-req-SLOT decision arrays (mamba_utils.py:439-458 reads them at + # req_idx, never at batch_idx), so they are keyed by the slot mapping. + copier, *bystanders = scenario.slot_mapping + num_accepted = torch.ones(MAX_NUM_REQS, dtype=torch.int32, device=device) + state_idx = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + new_num_computed = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + # Only the first batch row crosses a block boundary inside its accepted + # window: 32 % 16 = 0 < num_accepted 2. + num_accepted[copier] = 2 + new_num_computed[copier] = 32 + for slot in bystanders: + num_accepted[slot] = 2 + new_num_computed[slot] = 40 # 40 % 16 = 8 >= 2 -> no copy + + scenario.ctx.run_fused_postprocess_align( + num_reqs=len(scenario.slot_mapping), + num_accepted_tokens_gpu=num_accepted, + state_idx_gpu=state_idx, + new_num_computed_tokens_gpu=new_num_computed, + idx_mapping=t(scenario.slot_mapping), # captured at T-2 + ) + return recorder + + +def test_gather_port_matches_production_semantics() -> None: + """Instrument check: the CPU gather port must reproduce block_table.py.""" + persistent = torch.arange( + MAX_NUM_REQS * MAX_NUM_BLOCKS, dtype=torch.int32 + ).reshape(MAX_NUM_REQS, MAX_NUM_BLOCKS) + dst = torch.full_like(persistent, -7) + num_blocks = [2] * MAX_NUM_REQS + + _gather_rows(persistent, num_blocks, [3, 1], 3, dst) + + # Gathered rows carry only the first num_blocks entries... + assert dst[0, :2].tolist() == persistent[3, :2].tolist() + assert dst[1, :2].tolist() == persistent[1, :2].tolist() + # ...the tail of a gathered row is untouched (block_table.py:269-273), + assert dst[0, 2:].tolist() == [-7] * (MAX_NUM_BLOCKS - 2) + # ...padded rows in [num_reqs, num_reqs_padded) are zeroed whole, + assert dst[2].tolist() == [0] * MAX_NUM_BLOCKS + # ...and rows at or beyond num_reqs_padded keep the previous content. + assert dst[3].tolist() == [-7] * MAX_NUM_BLOCKS + + +def test_deferred_postprocess_resolves_the_requests_own_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """r0's deferred copy must land on r0's blocks, not on r2's. + + T-2 gathered [r0, r1, r2]; T-1 re-gathered [r2] into the same buffer, so + gathered row 0 now holds r2. The deferred launch still carries T-2's + mapping, so ``batch_idx = 0`` means r0. + """ + scenario = _build_scenario() + recorder = _run_deferred(scenario, monkeypatch) + assert len(recorder.launches) == 1 + launch = recorder.launches[0] + + row0 = _block_ids(scenario, launch, 0) + why = ( + f"resolved={row0} " + f"runner_hook_registered={scenario.registered} " + f"(a missing runner hook leaves the kernel on the batch-order gather)" + ) + + # The copy predicate really fires for this row, so the kernel does + # dereference the block table here (mamba_utils.py:462-479). + assert row0["needs_copy"] is True, why + assert row0["src_eq_dst"] is False, why + assert (row0["src_col"], row0["dst_col"], row0["token_bias"]) == (0, 1, 1), why + + # The gathered buffer holds r2 at batch row 0 after the T-1 gather, so a + # kernel that still indexes it by batch_idx corrupts r2 on r0's behalf. + assert scenario.gathered[0, 0].item() == 30 + + # r0's persistent row is 10, 11, 12, ... + assert row0["table"] == "persistent", why + assert row0["row"] == 0, why + assert row0["src_block_id"] == 10, why + assert row0["dst_block_id"] == 11, why + assert row0["temporal_src_block_id"] == 11, why + assert scenario.registered, why + + +def test_deferred_postprocess_skips_rows_without_a_boundary_crossing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The other two rows do not reach the block table at all.""" + scenario = _build_scenario() + recorder = _run_deferred(scenario, monkeypatch) + launch = recorder.launches[0] + + for batch_idx in (1, 2): + _name, _row, predicate = _replay_row(scenario, launch, batch_idx) + assert predicate["needs_copy"] is False, (batch_idx, predicate) + + +def test_precopy_still_uses_the_same_step_gather( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The pre-copy site must keep batch-order rows. + + ``precopy_mamba_align_fused_kernel`` (mamba_utils.py:613-615) runs from + ``preprocess_state``, five lines after the gather that feeds it + (``model_runner.py:1718`` -> ``:1723``), so batch row i *is* this request. + Switching it to req rows would break it. + """ + scenario = _build_scenario() + recorder = _RecordingKernel() + monkeypatch.setattr( + mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder + ) + device = scenario.persistent.device + z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + + scenario.ctx.run_fused_precopy( + num_reqs=3, + state_idx_gpu=z.clone(), + src_col_gpu=z.clone(), + token_bias_gpu=z.clone(), + idx_mapping=torch.tensor([0, 1, 2], dtype=torch.int32, device=device), + ) + + assert len(recorder.launches) == 1 + _grid, args, kwargs = recorder.launches[0] + bound = _PRECOPY_SIG.bind(*args, **kwargs) + bound.apply_defaults() + ptr = _ptr_to_int(int(bound.arguments["block_table_ptrs_ptr"][0])) + assert scenario.tables_by_ptr[ptr][0] == "gathered" + assert bound.arguments["HAS_IDX_MAPPING"] is True + assert "BT_ROW_IS_REQ" not in bound.arguments + + +def test_deferred_postprocess_uses_req_rows_under_a_non_identity_mapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Batch row 0 -> request slot 5, so batch_idx and req_idx disagree. + + The identity-mapping case above pins *which table*; this one pins *which + row*. Every wrong combination lands on a different set of block ids: + + ========================== ========== ================== + table / row resolves block ids + ========================== ========== ================== + persistent[req_idx=5] correct 60, 61 + persistent[batch_idx=0] wrong row 10, 11 + gathered[batch_idx=0] old 30, 31 (r2, from T-1) + gathered[req_idx=5] wrong both 0, 0 (never written) + ========================== ========== ================== + """ + scenario = _build_scenario(slot_mapping=[5, 1, 2], t1_mapping=[2]) + recorder = _run_deferred(scenario, monkeypatch) + assert len(recorder.launches) == 1 + launch = recorder.launches[0] + + _grid, args, kwargs = launch + bound = _POSTPROCESS_SIG.bind(*args, **kwargs) + bound.apply_defaults() + # The constexpr itself, not just its effect. + assert bound.arguments["BT_ROW_IS_REQ"] is True + assert bound.arguments["HAS_IDX_MAPPING"] is True + + row0 = _block_ids(scenario, launch, 0) + why = f"resolved={row0} runner_hook_registered={scenario.registered}" + + assert row0["needs_copy"] is True, why + assert row0["src_eq_dst"] is False, why + assert (row0["src_col"], row0["dst_col"], row0["token_bias"]) == (0, 1, 1), why + + # The row index is the request slot, not the batch row. + assert row0["table"] == "persistent", why + assert row0["row"] == 5, why + assert row0["src_block_id"] == 60, why + assert row0["dst_block_id"] == 61, why + assert row0["temporal_src_block_id"] == 61, why + + # None of the three wrong coordinates. + assert scenario.persistent[0, 0].item() == 10 # persistent, batch row + assert scenario.gathered[0, 0].item() == 30 # gathered, batch row + assert scenario.gathered[5, 0].item() == 0 # gathered, req row: never written + assert row0["src_block_id"] not in (10, 30, 0), why + + +def test_align_postprocess_fails_closed_without_persistent_tables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No silent fallback to the batch-order gather. + + A context that never received the persistent tables is a caller bug: the + only other thing the kernel could index is the stale gather, which is the + corruption this change exists to remove. It must raise, and it must raise + before mutating the accepted-token snapshot. + """ + scenario = _build_scenario(slot_mapping=[5, 1, 2], t1_mapping=[2], register=False) + assert scenario.registered is False + assert scenario.ctx.has_persistent_block_tables is False + + recorder = _RecordingKernel() + monkeypatch.setattr( + mamba_utils_mod, "postprocess_mamba_fused_kernel", recorder + ) + device = scenario.persistent.device + before = scenario.ctx.num_accepted_tokens_out.clone() + num_accepted = torch.full( + (MAX_NUM_REQS,), 7, dtype=torch.int32, device=device + ) + + with pytest.raises(RuntimeError, match="persistent per-req-slot block tables"): + scenario.ctx.run_fused_postprocess_align( + num_reqs=3, + num_accepted_tokens_gpu=num_accepted, + state_idx_gpu=torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device), + new_num_computed_tokens_gpu=torch.full( + (MAX_NUM_REQS,), 32, dtype=torch.int32, device=device + ), + idx_mapping=torch.tensor([5, 1, 2], dtype=torch.int32, device=device), + ) + + assert recorder.launches == [] + torch.testing.assert_close(scenario.ctx.num_accepted_tokens_out, before) + + +def test_precopy_still_works_without_persistent_tables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The V1 pre-copy contract is untouched by the fail-closed rule. + + Only the V2 align post-process needs per-req-slot rows; the pre-copy reads + the same-step gather and must keep launching on a context that never + received the persistent tables. + """ + scenario = _build_scenario(register=False) + assert scenario.ctx.has_persistent_block_tables is False + recorder = _RecordingKernel() + monkeypatch.setattr( + mamba_utils_mod, "precopy_mamba_align_fused_kernel", recorder + ) + device = scenario.persistent.device + z = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=device) + + scenario.ctx.run_fused_precopy( + num_reqs=1, + state_idx_gpu=z.clone(), + src_col_gpu=torch.full_like(z, -1), + token_bias_gpu=z.clone(), + idx_mapping=None, + ) + + assert len(recorder.launches) == 1 + _grid, args, kwargs = recorder.launches[0] + bound = _PRECOPY_SIG.bind(*args, **kwargs) + bound.apply_defaults() + ptr = _ptr_to_int(int(bound.arguments["block_table_ptrs_ptr"][0])) + assert scenario.tables_by_ptr[ptr][0] == "gathered" + + +# ---------------------------------------------------- generation masking (CPU) + + +def mask_from_generation_change( + idx_mapping: list[int], + reused_slot_rows: list[int], +) -> list[int]: + """Run the real generation masking and return the mapping it produces. + + Drives ``PPHandler.get_prev_sampled_outputs`` (pp_utils.py:132-140) rather + than hand-writing a -1: a slot whose generation counter moved between + ``receive`` and consumption has been freed and handed to a new request, and + the handler must replace its row with -1 so the kernel skips it + (mamba_utils.py:441-442). + + Exported so the GPU suite can consume a mapping that a real reuse produced. + """ + from collections import deque + + from vllm.v1.worker.gpu import pp_utils as pp_utils_mod + + num_reqs = len(idx_mapping) + idx_np = np.asarray(idx_mapping, dtype=np.intp) + + handler = object.__new__(pp_utils_mod.PPHandler) + handler.device = torch.device("cpu") + handler.main_stream = SimpleNamespace(wait_event=lambda event: None) + # Slot generation counters as they stand now. + handler.req_idx_gen_np = np.zeros(max(int(idx_np.max()) + 1, 1), dtype=np.int32) + # Snapshot taken at receive time; a reused slot has since been bumped. + gen_at_receive = handler.req_idx_gen_np[idx_np].copy() + for row in reused_slot_rows: + handler.req_idx_gen_np[idx_np[row]] += 1 + + slot = pp_utils_mod.PendingRecv( + event=object(), + sampled_tokens=torch.zeros(num_reqs, 1, dtype=torch.int64), + num_sampled=torch.ones(num_reqs, dtype=torch.int32), + num_rejected=torch.zeros(num_reqs, dtype=torch.int32), + idx_mapping=torch.as_tensor(idx_np, dtype=torch.int32), + idx_mapping_np=idx_np, + need_sampled_mask=np.ones(num_reqs, dtype=bool), + gen_at_receive_np=gen_at_receive, + ) + handler.queue = deque([slot]) + + captured: dict[str, object] = {} + + def fake_async_copy_to_gpu(array, device=None): + captured["np"] = np.asarray(array).copy() + return torch.as_tensor(np.asarray(array), dtype=torch.int32) + + original = pp_utils_mod.async_copy_to_gpu + pp_utils_mod.async_copy_to_gpu = fake_async_copy_to_gpu + try: + outputs = handler.get_prev_sampled_outputs() + finally: + pp_utils_mod.async_copy_to_gpu = original + + assert outputs is not None, "the whole batch was excluded" + return [int(v) for v in outputs["idx_mapping"].tolist()] + + +def test_generation_change_masks_only_the_reused_row(): + """A slot handed to a new request between capture and use becomes -1.""" + assert mask_from_generation_change([5, 1, 2], [1]) == [5, -1, 2] + assert mask_from_generation_change([5, 1, 2], [0, 2]) == [-1, 1, -1] + # No reuse: the handler returns the original tensor untouched. + assert mask_from_generation_change([5, 1, 2], []) == [5, 1, 2] + + +def test_kernel_rejects_negative_req_idx_before_any_load_that_uses_it(): + """Structural check: the guard precedes every load indexed by req_idx. + + ``mamba_utils.py:441-442`` returns on ``req_idx < 0``; the decision loads at + ``:445-446`` index by ``req_idx``. If a future edit reordered them, a masked + row would dereference index -1 before being rejected. The GPU suite plants + poison at index -1 to catch that at runtime; this catches it on CPU, in the + source, with no device. + """ + import ast + import textwrap + + source = textwrap.dedent( + inspect.getsource(mamba_utils_mod.postprocess_mamba_fused_kernel.fn) + ) + tree = ast.parse(source) + func = tree.body[0] + + guard_lines = [ + node.lineno + for node in ast.walk(func) + if isinstance(node, ast.If) + and isinstance(node.test, ast.Compare) + and isinstance(node.test.left, ast.Name) + and node.test.left.id == "req_idx" + and any(isinstance(op, ast.Lt) for op in node.test.ops) + and any(isinstance(stmt, ast.Return) for stmt in node.body) + ] + assert guard_lines, "no `if req_idx < 0: return` guard found in the kernel" + guard_line = min(guard_lines) + + # Every expression that adds req_idx to a pointer, i.e. every load or store + # that would dereference a negative index. + uses = [ + node.lineno + for node in ast.walk(func) + if isinstance(node, ast.BinOp) + and isinstance(node.op, ast.Add) + and isinstance(node.right, ast.Name) + and node.right.id == "req_idx" + ] + assert uses, "no pointer arithmetic on req_idx found; has the kernel moved?" + assert min(uses) > guard_line, ( + f"a req_idx-indexed access at line {min(uses)} precedes the " + f"`req_idx < 0` guard at line {guard_line} (offsets are relative to the " + f"kernel body)" + ) diff --git a/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py new file mode 100644 index 0000000..f036227 --- /dev/null +++ b/tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py @@ -0,0 +1,542 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU correctness for the deferred-PP align post-process row selection. + +The CPU model in ``test_pp_deferred_mamba_block_table.py`` pins which pointer +and which row the launcher passes; it replaces the Triton body with a recorder, +so it proves nothing about the bytes the kernel moves. These tests launch the +real kernel and check the destination state, including that the request the +first-cut coordinate would have hit is left alone. + +Scenario, shared by every case and identical to the CPU model's: + + step T-2 gather [r5, r1, r2] into the shared destination buffer + step T-1 gather [r2] into the same buffer + step T the deferred post-process runs with T-2's mapping + +so batch row 0 means request slot 5, while buffer row 0 holds r2. The four +possible coordinates resolve to four different physical blocks: + + persistent[req_idx=5] correct + persistent[batch_idx=0] wrong row -> r0's blocks + gathered[batch_idx=0] first cut -> r2's blocks (the victim) + gathered[req_idx=5] wrong both -> never written, block 0 + +Requires CUDA **and** an explicit opt-in: set ``VLLM_E2_GPU_TESTS=1``. The reference host +always reports CUDA available, and the reference host spends most of its time running +long serving runs and sanitizer runs that should not share a device with a +test suite, so `torch.cuda.is_available()` is not a sufficient gate here. Run +this when the GPU is otherwise idle: + + export VLLM_E2_GPU_TESTS=1 + export PYTHONPATH=: + python -m pytest -q tests/v1/worker/gpu/test_pp_deferred_mamba_gpu.py +""" + +import os +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFuncsByType, + get_conv_copy_spec, + get_temporal_copy_spec, + is_conv_state_dim_first, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec +from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + os.environ.get("VLLM_E2_GPU_TESTS") != "1", + reason="set VLLM_E2_GPU_TESTS=1 to run GPU tests (shared GPU host)", + ), +] + +MAX_NUM_REQS = 8 +BLOCKS_PER_REQ = 6 +NUM_PHYS_BLOCKS = MAX_NUM_REQS * BLOCKS_PER_REQ +# Two extra physical blocks, reachable only through the block-table guard row +# that a negative req_idx would address: CANARY_DST sits at the poison +# destination column and CANARY_SRC everywhere else, so a negative-index load +# would produce a visible copy between them. See _build()'s backing tensors. +# Block ids are 1-based (see phys()), so the canaries sit above the last real +# id, NUM_PHYS_BLOCKS, rather than on it. +CANARY_DST = NUM_PHYS_BLOCKS + 1 +CANARY_SRC = NUM_PHYS_BLOCKS + 2 +TOTAL_BLOCKS = NUM_PHYS_BLOCKS + 3 +# Decision values planted at index -1 of every per-req array. They describe a +# copy (new=48, a=4 -> aligned=48, running=45, bias=3, dst_col=2, src_col=0), +# so a missing `req_idx < 0` guard is a write, not a silent no-op. +POISON = (4, 0, 48) # (num_accepted, state_idx, new_num_computed) +POISON_DST_COL = 2 +MAMBA_BLOCK_SIZE = 16 +NUM_SPEC_BLOCKS = 3 +CONV_WIDTH = 4 +CONV_INNER = 64 +TEMPORAL_DIM = 128 +DTYPE = torch.float16 + +_COPY_FUNCS: MambaStateCopyFuncsByType = { + MambaAttentionBackendEnum.GDN_ATTN: (get_conv_copy_spec, get_temporal_copy_spec), +} + + +def phys(req_slot: int, col: int) -> int: + """Physical block id owned by ``req_slot`` at logical column ``col``. + + Disjoint per request, so a stray write is attributable. + + ONE-BASED, and that is load-bearing rather than cosmetic. Block id 0 is + NULL_BLOCK_ID (v1/attention/backends/utils.py:46); the block pool pops it + out of the free queue at construction "to represent a placeholder block + with block_id=0" (v1/core/block_pool.py:188-191), so no request ever owns + it. A fixture that hands out id 0 as a real block is asking the kernel to + copy from the null block, and the bounds guards -- correctly, and following + upstream #50021 -- refuse. Numbering from 0 made this fixture the one place + in the tree where a legitimate block id was 0. + """ + return req_slot * BLOCKS_PER_REQ + col + 1 + + +class _MockCpuGpuBuffer: + 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: + self.gpu.copy_(self.cpu) + return self.gpu + self.gpu[:n].copy_(self.cpu[:n]) + return self.gpu[:n] + + +def _make_kv_cache_config() -> KVCacheConfig: + spec = MambaSpec( + block_size=MAMBA_BLOCK_SIZE, + shapes=((CONV_WIDTH, CONV_INNER), (TEMPORAL_DIM,)), + dtypes=(DTYPE, DTYPE), + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="align", + num_speculative_blocks=NUM_SPEC_BLOCKS, + ) + return KVCacheConfig( + num_blocks=TOTAL_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer_0"], spec)], + ) + + +def _gather_rows(persistent, num_blocks, idx_mapping, num_reqs_padded, dst) -> None: + """Host port of ``_gather_block_tables_kernel`` (block_table.py:236-273).""" + num_reqs = len(idx_mapping) + for batch_idx in range(num_reqs_padded): + if batch_idx >= num_reqs: + dst[batch_idx].zero_() + continue + req_idx = idx_mapping[batch_idx] + nb = num_blocks[req_idx] + dst[batch_idx, :nb] = persistent[req_idx, :nb] + + +@dataclass +class _Fixture: + device: torch.device + ctx: MambaSpecDecodeGPUContext + conv: torch.Tensor + temporal: torch.Tensor + conv0: torch.Tensor + temporal0: torch.Tensor + persistent: torch.Tensor + gathered: torch.Tensor + + +def _build(slot_mapping=(5, 1, 2), t1_mapping=(2,)) -> _Fixture: + device = torch.device("cuda:0") + torch.manual_seed(1234) + + conv = torch.randn( + TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device + ) + temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) + attention = MagicMock() + attention.kv_cache = [conv, temporal] + + # The block tables and the per-req decision arrays each get one extra + # leading row/element that the kernel is never supposed to reach. Because + # the real tensors are views starting one row/element in, index -1 in + # pointer arithmetic lands on that guard row -- so a kernel that failed to + # reject req_idx < 0 would load the poison decisions below and write into + # canaries. See test_freed_slot_is_rejected_before_any_negative_index_load. + persistent_backing = torch.zeros( + MAX_NUM_REQS + 1, BLOCKS_PER_REQ, dtype=torch.int32, device=device + ) + persistent_backing[0, :] = CANARY_SRC + persistent_backing[0, POISON_DST_COL] = CANARY_DST + persistent = persistent_backing[1:] + for r in range(MAX_NUM_REQS): + for c in range(BLOCKS_PER_REQ): + persistent[r, c] = phys(r, c) + num_blocks = [BLOCKS_PER_REQ] * MAX_NUM_REQS + + gathered_backing = torch.zeros_like(persistent_backing) + gathered_backing[0, :] = CANARY_SRC + gathered_backing[0, POISON_DST_COL] = CANARY_DST + gathered = gathered_backing[1:] + slot_mapping = [s for s in slot_mapping if s >= 0] + _gather_rows(persistent, num_blocks, slot_mapping, len(slot_mapping) + 1, gathered) + _gather_rows( + persistent, num_blocks, list(t1_mapping), len(t1_mapping) + 1, gathered + ) + + kv_cache_config = _make_kv_cache_config() + ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=MAX_NUM_REQS, + kv_cache_config=kv_cache_config, + copy_funcs=_COPY_FUNCS, + device=device, + make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), + ) + ctx.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [gathered], + persistent_block_tables=[persistent], + ) + assert ctx.has_persistent_block_tables + + return _Fixture( + device=device, + ctx=ctx, + conv=conv, + temporal=temporal, + conv0=conv.clone(), + temporal0=temporal.clone(), + persistent=persistent, + gathered=gathered, + ) + + +def _launch(fx: _Fixture, mapping, per_slot): + """Run the deferred align post-process. + + ``per_slot`` maps req slot -> (num_accepted, state_idx, new_num_computed). + """ + device = fx.device + + def backed(fill: int, poison: int) -> torch.Tensor: + """A [MAX_NUM_REQS] view whose index -1 holds ``poison``.""" + backing = torch.full( + (MAX_NUM_REQS + 1,), fill, dtype=torch.int32, device=device + ) + backing[0] = poison + return backing[1:] + + poison_acc, poison_sidx, poison_computed = POISON + num_accepted = backed(1, poison_acc) + state_idx = backed(0, poison_sidx) + new_computed = backed(0, poison_computed) + for slot, (acc, sidx, computed) in per_slot.items(): + num_accepted[slot] = acc + state_idx[slot] = sidx + new_computed[slot] = computed + fx.ctx.run_fused_postprocess_align( + num_reqs=len(mapping), + num_accepted_tokens_gpu=num_accepted, + state_idx_gpu=state_idx, + new_num_computed_tokens_gpu=new_computed, + idx_mapping=torch.tensor(mapping, dtype=torch.int32, device=device), + ) + torch.accelerator.synchronize() + return num_accepted + + +def _assert_bitwise(actual: torch.Tensor, expected: torch.Tensor, what: str) -> None: + """Compare the integer view, so "bit-identical" means what it says. + + A zero-tolerance float compare establishes numeric equality, not bit + identity: +0.0 and -0.0 compare equal and NaN payloads compare unequal. + These are byte copies, so the integer view is the right instrument. + + """ + assert actual.dtype == expected.dtype + lhs = actual.contiguous().view(torch.int16) + rhs = expected.contiguous().view(torch.int16) + assert torch.equal(lhs, rhs), ( + f"{what}: {int((lhs != rhs).sum().item())} of {lhs.numel()} 16-bit " + f"words differ" + ) + + +def _assert_only_blocks_changed(fx: _Fixture, changed: set[int]) -> None: + """Every physical block outside ``changed`` must be bit-identical. + + The canary blocks (``CANARY_DST``, ``CANARY_SRC``) are included: they are + only reachable through the guard row a negative ``req_idx`` would select. + """ + for block in range(TOTAL_BLOCKS): + if block in changed: + continue + _assert_bitwise(fx.conv[block], fx.conv0[block], f"conv[{block}]") + _assert_bitwise(fx.temporal[block], fx.temporal0[block], f"temporal[{block}]") + + +@pytest.mark.skipif( + is_conv_state_dim_first(), + reason="expected-bytes model assumes the SD conv layout", +) +def test_positive_bias_copy_lands_on_the_requests_own_blocks(): + """batch row 0 -> req slot 5, num_accepted 4, so token_bias is 3. + + new=48, a=4 -> aligned=48, running=45, needs_copy; bias=3; dst_col=2. + state_idx=0, so conv reads column 0 and temporal reads column + ``src_col + bias`` = 3 -- three distinct columns, so a wrong row cannot be + masked by a coincidental self-copy. + """ + fx = _build() + _launch(fx, [5, 1, 2], {5: (4, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) + + dst = phys(5, 2) + conv_src = phys(5, 0) + temporal_src = phys(5, 3) + + # conv: state[bt[src_col], bias:] -> state[bt[dst_col], :width - bias] + _assert_bitwise(fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[conv_src][3:], "conv dst") + # the tail of the destination conv window is not written + _assert_bitwise( + fx.conv[dst][CONV_WIDTH - 3 :], fx.conv0[dst][CONV_WIDTH - 3 :], "conv tail" + ) + # temporal: state[bt[src_col + bias]] -> state[bt[dst_col]] + _assert_bitwise(fx.temporal[dst], fx.temporal0[temporal_src], "temporal dst") + + # Nothing else moved -- in particular not the victim r2 (the first-cut row) and + # not r0 (the "persistent table but batch row" mistake). + _assert_only_blocks_changed(fx, {dst}) + victim = phys(2, 2) # gathered[batch row 0] column 2 -- the first-cut destination + _assert_bitwise(fx.conv[victim], fx.conv0[victim], "victim conv") + _assert_bitwise(fx.temporal[victim], fx.temporal0[victim], "victim temporal") + wrong_row = phys(0, 2) # persistent[batch row 0] column 2 + _assert_bitwise(fx.conv[wrong_row], fx.conv0[wrong_row], "wrong-row conv") + + +@pytest.mark.skipif( + is_conv_state_dim_first(), + reason="expected-bytes model assumes the SD conv layout", +) +def test_zero_bias_migration_copies_the_whole_window(): + """new=48, a=1 -> aligned=48, running=48, bias=0, dst_col=2, src_col=0. + + src != dst so the self-copy exit does not fire, and the whole conv window + plus the temporal block migrate. + """ + fx = _build() + _launch(fx, [5, 1, 2], {5: (1, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) + + dst = phys(5, 2) + src = phys(5, 0) + _assert_bitwise(fx.conv[dst], fx.conv0[src], "conv dst") + _assert_bitwise(fx.temporal[dst], fx.temporal0[src], "temporal dst") + _assert_only_blocks_changed(fx, {dst}) + + +def test_self_copy_exits_without_touching_state_and_resets_num_accepted(): + """src_block_idx == dest_block_idx and bias == 0 -> return (mamba_utils:478). + + new=33, a=2 -> aligned=32, running=32, needs_copy true, bias=0, dst_col=1; + state_idx=1 makes src == dst. Seeding num_accepted at 2 rather than 1 is + what makes the reset store observable: the kernel must write + num_accepted_tokens_out[req] = 1 (mamba_utils.py:473-475) *before* the early + return, so deleting that store now fails the test. (A prior version seeded 1 and + asserted 1, which a deleted store would have passed.) + + Limit, stated rather than papered over: unchanged bytes cannot distinguish + "returned at the guard" from "fell through and memcpy'd a block onto + itself", because those are byte-identical by construction. What the reset + store does establish is that the kernel got past `needs_copy` and evaluated + `src == dst`, so it is the *guard* that stopped it and not an earlier + return. The conjunction itself is covered by the next test. + """ + fx = _build() + num_accepted = _launch( + fx, [5, 1, 2], {5: (2, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} + ) + + _assert_only_blocks_changed(fx, set()) + assert int(num_accepted[5].item()) == 1 + + +@pytest.mark.skipif( + is_conv_state_dim_first(), + reason="expected-bytes model assumes the SD conv layout", +) +def test_same_column_with_positive_bias_copies_and_still_resets(): + """The early-exit guard is a conjunction, and this is the other conjunct. + + new=33, a=4 -> aligned=32, running=30, needs_copy; bias = 4-1-1 = 2; + dst_col = 1; state_idx = 1, so src_block_idx == dest_block_idx but + token_bias != 0. The kernel must NOT return: conv takes the same-block + memmove branch (mamba_utils.py:320-336) and temporal reads + ``src_col + token_bias`` = column 3, a different physical block. + + Together with the previous test this makes the exit observable in both + directions: same column and zero bias moves nothing, same column and + positive bias moves bytes. + """ + fx = _build() + num_accepted = _launch( + fx, [5, 1, 2], {5: (4, 1, 33), 1: (2, 0, 40), 2: (2, 0, 40)} + ) + + blk = phys(5, 1) # both the conv source and the destination + temporal_src = phys(5, 3) # src_col + bias = 3 + + # conv, same block, shifted left by token_bias = 2 + _assert_bitwise(fx.conv[blk][: CONV_WIDTH - 2], fx.conv0[blk][2:], "conv shift") + _assert_bitwise( + fx.conv[blk][CONV_WIDTH - 2 :], fx.conv0[blk][CONV_WIDTH - 2 :], "conv tail" + ) + # temporal, from a different block + _assert_bitwise(fx.temporal[blk], fx.temporal0[temporal_src], "temporal dst") + + _assert_only_blocks_changed(fx, {blk}) + # The reset store fires here too: it is guarded only on src == dst. + assert int(num_accepted[5].item()) == 1 + + +@pytest.mark.skipif( + is_conv_state_dim_first(), + reason="expected-bytes model assumes the SD conv layout", +) +def test_freed_slot_is_rejected_before_any_negative_index_load(): + """A -1 row must be rejected *before* the decision arrays are read. + + ``mamba_utils.py:441-442`` returns on ``req_idx < 0`` before the loads at + ``:445-446``. A prior version of this test seeded decisions at some other + slot, which the kernel would never have touched either way, so it proved + nothing about the guard. This one plants the poison + where a missing guard would actually look: index -1 of every per-req array + and row -1 of the block table, which are real memory owned by the backing + tensors in ``_build``. + + If the guard were deleted, the kernel would read + ``num_accepted=4, state_idx=0, new_num_computed=48`` -> a copy with bias 3 + into block-table row -1, writing CANARY_DST from CANARY_SRC. Both canaries + are covered by ``_assert_only_blocks_changed``, so that write fails the + test. + + The row after the -1 must also still copy correctly, which is the + positioning check. + """ + fx = _build(slot_mapping=(5, 1, 2)) + _launch(fx, [-1, 5, 2], {5: (4, 0, 48), 2: (2, 0, 40)}) + + dst = phys(5, 2) + _assert_bitwise( + fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(5, 0)][3:], "conv dst" + ) + _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(5, 3)], "temporal dst") + _assert_only_blocks_changed(fx, {dst}) + + +def test_generation_masked_reuse_is_skipped_end_to_end(): + """The -1 above, but produced by the real generation masking. + + ``mask_from_generation_change`` in the CPU sibling test drives + ``PPHandler.get_prev_sampled_outputs`` (pp_utils.py:132-140) and returns + the mapping it computes; here that mapping is fed to the kernel. Slot 1 is + the one whose generation moved -- it was freed and its slot handed to a new + request between capture and use -- so its row must copy nothing even though + slot 1 now carries decisions that would copy. + """ + from tests.v1.worker.gpu.test_pp_deferred_mamba_block_table import ( + mask_from_generation_change, + ) + + mapping = mask_from_generation_change( + idx_mapping=[5, 1, 2], reused_slot_rows=[1] + ) + assert mapping == [5, -1, 2], mapping + + fx = _build(slot_mapping=(5, 1, 2)) + _launch(fx, mapping, {5: (4, 0, 48), 1: (4, 0, 48), 2: (2, 0, 40)}) + + dst = phys(5, 2) + _assert_bitwise( + fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(5, 0)][3:], "conv dst" + ) + _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(5, 3)], "temporal dst") + # Slot 1's own blocks are untouched despite its copy-shaped decisions. + _assert_only_blocks_changed(fx, {dst}) + + +@pytest.mark.skipif( + is_conv_state_dim_first(), + reason="expected-bytes model assumes the SD conv layout", +) +def test_same_step_identity_mapping_is_unaffected(): + """The last-rank / no-PP path: batch order == req order, gather is current. + + req_idx == batch_idx here, so the change must be a no-op relative to the first cut. + """ + fx = _build(slot_mapping=(0, 1, 2), t1_mapping=(0, 1, 2)) + _launch(fx, [0, 1, 2], {0: (4, 0, 48), 1: (2, 0, 40), 2: (2, 0, 40)}) + + dst = phys(0, 2) + _assert_bitwise( + fx.conv[dst][: CONV_WIDTH - 3], fx.conv0[phys(0, 0)][3:], "conv dst" + ) + _assert_bitwise(fx.temporal[dst], fx.temporal0[phys(0, 3)], "temporal dst") + _assert_only_blocks_changed(fx, {dst}) + + +def test_align_launch_raises_without_persistent_tables(): + """Fail closed on real hardware too, before any kernel is enqueued.""" + device = torch.device("cuda:0") + conv = torch.randn( + TOTAL_BLOCKS, CONV_WIDTH, CONV_INNER, dtype=DTYPE, device=device + ) + temporal = torch.randn(TOTAL_BLOCKS, TEMPORAL_DIM, dtype=DTYPE, device=device) + attention = MagicMock() + attention.kv_cache = [conv, temporal] + gathered = torch.zeros( + MAX_NUM_REQS, BLOCKS_PER_REQ, dtype=torch.int32, device=device + ) + + kv_cache_config = _make_kv_cache_config() + ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=MAX_NUM_REQS, + kv_cache_config=kv_cache_config, + copy_funcs=_COPY_FUNCS, + device=device, + make_buffer=lambda n, dtype: _MockCpuGpuBuffer(n, dtype, device), + ) + ctx.initialize_from_forward_context( + kv_cache_config, {"layer_0": attention}, _COPY_FUNCS, [gathered] + ) + assert ctx.has_persistent_block_tables is False + + snapshot = conv.clone() + with pytest.raises(RuntimeError, match="persistent per-req-slot block tables"): + ctx.run_fused_postprocess_align( + num_reqs=1, + num_accepted_tokens_gpu=torch.full( + (MAX_NUM_REQS,), 4, dtype=torch.int32, device=device + ), + state_idx_gpu=torch.zeros( + MAX_NUM_REQS, dtype=torch.int32, device=device + ), + new_num_computed_tokens_gpu=torch.full( + (MAX_NUM_REQS,), 48, dtype=torch.int32, device=device + ), + idx_mapping=torch.tensor([5], dtype=torch.int32, device=device), + ) + torch.accelerator.synchronize() + _assert_bitwise(conv, snapshot, "conv unchanged") diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a162c39..2ab23ce 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -277,8 +277,14 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): ) forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)} + # V1: batch order == req-slot order, so the same tensor is both the + # gathered and the persistent table. gpu_ctx.initialize_from_forward_context( - kv_cache_config, forward_context, _COPY_FUNCS, [block_table] + kv_cache_config, + forward_context, + _COPY_FUNCS, + [block_table], + persistent_block_tables=[block_table], ) assert gpu_ctx.state_base_addrs.tolist() == [ @@ -548,6 +554,7 @@ def test_mamba_groups_support_different_state_specs(): forward_context, _COPY_FUNCS, block_tables, + persistent_block_tables=block_tables, ) assert ctx.is_initialized assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1] @@ -628,11 +635,13 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): ), "ple.0": MagicMock(kv_cache=(torch.empty(8, 12, 8, dtype=torch.float16),)), } + group_table = torch.zeros(4, 8, dtype=torch.int32) ctx.initialize_from_forward_context( kv_cache_config, forward_context, _COPY_FUNCS, - [torch.zeros(4, 8, dtype=torch.int32)], + [group_table], + persistent_block_tables=[group_table], ) assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 0] assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] @@ -777,11 +786,13 @@ def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: ] context = _make_gpu_ctx(config, kv_cache_config, device) + stub_table = torch.zeros(1, 1, dtype=torch.int32) context.initialize_from_forward_context( kv_cache_config, {"layer_0": attention}, _COPY_FUNCS, - [torch.zeros(1, 1, dtype=torch.int32)], + [stub_table], + persistent_block_tables=[stub_table], ) assert context.is_initialized @@ -812,8 +823,14 @@ def _run_gpu_postprocess( def t(values): return torch.tensor(values, dtype=torch.int32, device=device) + # V1: batch order == req-slot order, so the same tensor is both the + # gathered and the persistent table. gpu_ctx.initialize_from_forward_context( - kv_cache_config, forward_context, copy_funcs, [block_table] + kv_cache_config, + forward_context, + copy_funcs, + [block_table], + persistent_block_tables=[block_table], ) gpu_ctx.run_fused_postprocess( num_reqs=len(req_ids), diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f162cb5..e3479c2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -677,6 +677,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, ) + # Mamba "align" indexes a block table from the PP-deferred post-process, + # which runs pp_size steps after the gather that produced its + # idx_mapping; give the model state the persistent per-req-slot tables. + self.model_state.set_block_tables(self.block_tables) self.pcp_manager = pcp.maybe_build_pcp_manager( self.vllm_config, self.device, diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index a5c8d6f..8ce5121 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -126,6 +126,14 @@ class ModelState(ABC): def apply_staged_writes(self) -> None: return None + def set_block_tables(self, block_tables: Any) -> None: + """Register the runner's `BlockTables` once it exists. + + Model states that index a block table outside the step that gathered it + (mamba "align" + the PP-deferred post-process) need the persistent + per-req-slot tables, not the batch-order gather. No-op by default.""" + return None + def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: """Cudagraph support of attention groups this ModelState builds outside ``init_attn_backend`` (e.g. encoder-only layers). diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 0ac3f2f..0a8b5d7 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -94,6 +94,9 @@ class MambaHybridModelState(DefaultModelState): # kernel reusing the postprocess copy machinery, so the per-step src # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" + # Set by the runner via `set_block_tables`; the align post-process needs + # the persistent per-req-slot tables, not the per-step gather. + self._runner_block_tables: Any = None self.recoverssm = ( RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ) @@ -112,6 +115,9 @@ class MambaHybridModelState(DefaultModelState): self._mamba_spec: MambaSpec | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None + def set_block_tables(self, block_tables: Any) -> None: + self._runner_block_tables = block_tables + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) # Must reset the speculative acceptance count in this idx which could be stale. @@ -172,11 +178,24 @@ class MambaHybridModelState(DefaultModelState): # block_tables are batch-order slices of the persistent # input_block_tables (stable data_ptr), so the metadata is captured # once here and reused across steps. + # + # Also hand over the per-req-slot tables the gather reads from + # (block_table.py:53-59). The align post-process can run pp_size + # steps after the gather that produced its idx_mapping, at which + # point the batch-order buffer belongs to another micro-batch; it + # indexes these by req slot instead. + persistent_block_tables = None + if self._runner_block_tables is not None: + persistent_block_tables = [ + self._runner_block_tables.block_tables[gid].gpu + for gid in mamba_group_ids + ] ctx.initialize_from_forward_context( kv_cache_config, forward_context, copy_funcs, [block_tables[gid] for gid in mamba_group_ids], + persistent_block_tables=persistent_block_tables, ) return ctx diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 1deb89b..621efcb 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -409,6 +409,14 @@ def postprocess_mamba_fused_kernel( # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds # the post-step new_num_computed value (V2 supplies the advanced count). PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, + # BT_ROW_IS_REQ: when True, block_table_ptrs_ptr points at the *persistent* + # per-req-slot block tables (BlockTables.block_tables[i].gpu, + # block_table.py:53-59), whose rows are request state slots, so the row + # index is req_idx rather than batch_idx. When False it points at the + # batch-order gather (BlockTables.input_block_tables, block_table.py:74-76) + # and batch_idx is the row. Required by the PP-deferred post-process, which + # runs pp_size steps after the gather that produced its idx_mapping. + BT_ROW_IS_REQ: tl.constexpr = False, # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves @@ -477,7 +485,15 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return - bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx + if BT_ROW_IS_REQ: + # Persistent per-req-slot table: the row is the request's own slot, so + # it stays correct however many steps separate the gather from this + # launch (the PP-deferred post-process is enqueued at the top of + # execute_model, before this step's gather, with an idx_mapping + # captured pp_size steps earlier). + bt_row_idx = req_idx + else: + bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx _copy_mamba_state_block( state_idx, bt_row_idx, @@ -812,6 +828,24 @@ class MambaSpecDecodeGPUContext: block_table_ptrs: torch.Tensor block_table_stride_req: int = 0 + # Per-group *persistent* per-req-slot block-table addresses: int64[num_groups]. + # `block_table_ptrs` above addresses the batch-order gather + # (BlockTables.input_block_tables, block_table.py:74-76), which is rewritten + # in place every step and is only valid for the step that gathered it. These + # address BlockTables.block_tables[i].gpu (block_table.py:53-59), the + # persistent [max_num_reqs, max_num_blocks] tables the gather reads from + # (their data_ptrs are already held device-side at block_table.py:94-96). + # Rows there are request state slots, so they can be indexed by req_idx from + # any later step. Used only by the align post-process (BT_ROW_IS_REQ). + persistent_block_table_ptrs: torch.Tensor | None = None + has_persistent_block_tables: bool = False + # shape[0] of the persistent tables; see the note in set_block_tables. + # -1, not 0: this is the ROW BOUND the align post-process uses, and a 0 + # default bounds EVERY row out, so a context that never registered the + # persistent tables would report guard hits on correct metadata instead of + # copying state. -1 is unusable on purpose and is asserted before use. + persistent_block_table_num_rows: int = -1 + # persistent output for the once-per-step, all-group aligned-index launch. # shape: [num_groups, max_num_reqs, 1 + num_speculative_blocks]. aligned_state_indices: torch.Tensor | None = None @@ -902,6 +936,9 @@ class MambaSpecDecodeGPUContext: block_table_ptrs=torch.zeros( len(mamba_group_ids), dtype=torch.int64, device=device ), + persistent_block_table_ptrs=torch.zeros( + len(mamba_group_ids), dtype=torch.int64, device=device + ), aligned_state_indices=torch.empty( ( len(mamba_group_ids), @@ -926,6 +963,7 @@ class MambaSpecDecodeGPUContext: forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], + persistent_block_tables: list[torch.Tensor] | None = None, ) -> None: """ Extract and cache memory layout metadata from Mamba state tensors. @@ -957,9 +995,14 @@ class MambaSpecDecodeGPUContext: have a `kv_cache` attribute containing the list of state tensors. mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to copy functions. - block_tables: per-mamba-group persistent block-table tensors, in - the same order as `mamba_group_ids`. Their `data_ptr()` / - `stride(0)` are captured once for the kernel to index into. + block_tables: per-mamba-group *gathered*, batch-order block-table + tensors, in the same order as `mamba_group_ids`. Their + `data_ptr()` / `stride(0)` are captured once for the kernel to + index into by batch row. + persistent_block_tables: per-mamba-group *persistent* per-req-slot + block-table tensors (`BlockTables.block_tables[i].gpu`), same + order. Captured separately so the align post-process can index + by req slot; see `persistent_block_table_ptrs`. """ if self.is_initialized: return @@ -970,6 +1013,7 @@ class MambaSpecDecodeGPUContext: forward_context, mamba_state_copy_funcs, block_tables, + persistent_block_tables, ) def _populate_metadata( @@ -978,6 +1022,7 @@ class MambaSpecDecodeGPUContext: forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, block_tables: list[torch.Tensor], + persistent_block_tables: list[torch.Tensor] | None = None, ) -> None: idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): @@ -1253,6 +1298,35 @@ class MambaSpecDecodeGPUContext: if num_reqs == 0 or not self.is_initialized: return + # Fail closed. Under PP this launch is the *deferred* post-process: it + # is enqueued by `update_pp_decode_requests` at the top of + # `execute_model` (model_runner.py:1644), before that step's + # `prepare_attn` (model_runner.py:1718) and the + # `gather_block_tables` inside it (:1442 / block_table.py:149), and + # `idx_mapping` was captured pp_size steps earlier + # (pp_utils.py:57-58, 258-271). The single batch-order gather buffer + # therefore holds a different micro-batch's rows and `batch_idx` + # resolves to the wrong request. This path MUST index the persistent + # per-req-slot tables by `req_idx`; there is no safe fallback, so a + # context without them is a bug in the caller, not a degraded mode. + # (`req_idx` is also correct on the same-step last-rank / no-PP path.) + if not self.has_persistent_block_tables: + raise RuntimeError( + "MambaSpecDecodeGPUContext was initialized without the " + "persistent per-req-slot block tables, so the V2 align " + "post-process cannot index them by req slot. Falling back to " + "the batch-order gather is the known deferred-PP wrong-row " + "corruption, so this raises instead. The runner must call " + "ModelState.set_block_tables() before the first align step." + ) + block_table_ptrs = self.persistent_block_table_ptrs + assert self.persistent_block_table_num_rows > 0, ( + "persistent_block_table_num_rows was never set, so the row bound " + "in _copy_mamba_state_block would reject every row and report " + "guard hits on correct metadata. set_block_tables() must run " + "before the first align post-process." + ) + # V2 reads non-contiguous idx_mapping positions, so snapshot the whole # decision buffer rather than only [:num_reqs]. num_accepted_tokens_snapshot = self.num_accepted_tokens_out @@ -1284,6 +1358,7 @@ class MambaSpecDecodeGPUContext: CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), HAS_IDX_MAPPING=True, PRECOMPUTED_NEW_COMPUTED=True, + BT_ROW_IS_REQ=True, TEMPORAL_TILES=_TEMPORAL_TILES, ) -- 2.43.0