From 62acb3dab87ba2a51459b435bdb2efb9f695431b Mon Sep 17 00:00:00 2001 From: MOVZX Date: Tue, 22 Sep 2026 04:12:35 -0400 Subject: [PATCH] [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. Co-Authored-By: Hinata Signed-off-by: MOVZX --- 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 | 697 +++++++++++++++++- 4 files changed, 718 insertions(+), 10 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5fc880c2e3..139d104d4d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -694,6 +694,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 ee3e519a6e..a011898e07 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 b8c44b16a9..d0b25b22a9 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -90,6 +90,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_recoverssm else None ) @@ -108,6 +111,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. @@ -168,11 +174,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 33a2b8b480..52ec5fdf55 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -7,6 +7,7 @@ from typing import Any, NamedTuple import torch +import vllm.envs as envs from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( @@ -26,14 +27,299 @@ from vllm.v1.kv_cache_interface import ( UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker import guard_telemetry from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch logger = init_logger(__name__) +# Bias so a negative row/index survives atomic_max on a zeroed counter +# -- it MUST be a tl.constexpr: a plain module-level int is not readable +# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died +# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. +_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) + +# Production value for the align post-process's block-table row rule. True +# means "the row is the request's own state slot", which is the only correct +# rule for the PP-deferred launch (see run_fused_postprocess_align). It is a +# module-level name rather than a literal at the launch site for exactly one +# reason: the failing control for the slot-ownership check below has to be able +# to flip it. NOTHING reads it from the environment, so there is no way to turn +# it off on a serving host; tests monkeypatch it. +ALIGN_BT_ROW_IS_REQ = True + +# Host-side mamba state-slot ownership check (PLE diagnostics, llama.cpp#29092 class / +# vllm#55506 failure mode). off | warn | raise, default off, so it costs +# nothing in production. It is a DIAGNOSTIC, not a device guard: it copies +# decision arrays and block-table rows to the host, which synchronizes. +# +# WHY IT EXISTS. Mamba state blocks are the one KV class with no zeroing and no +# ownership instrument: +# * single_type_kv_cache_manager.py:96-99 records newly allocated blocks for +# zeroing only when the spec `isinstance(kv_cache_spec, AttentionSpec)`. +# `MambaSpec` is not an `AttentionSpec`, so mamba state blocks are NEVER +# zeroed by `KVBlockZeroer` (gpu/model_runner.py:1213). +# * the host-side KV-zeroer ownership check (sched/scheduler.py:1508-1535) only +# inspects `new_block_ids_to_zero`, a list mamba blocks are never on. +# * a genuinely fresh request is still safe: gdn_attn.py:401 sets +# `has_initial_state = context_lens_tensor > 0`, so at context_len 0 the +# GDN kernel ignores whatever bytes the reused block holds. +# * a RESUMING request is not. With `has_initial_state` true the kernel +# consumes the block's contents, and the only thing that puts the right +# contents there is this align copy. Everything else we have bounds +# indices (the bounds guards): a bounds check catches an +# out-of-range row, never a wrong-but-in-range one, which is the whole of +# this defect class. +# Read through the envs registry, not os.getenv: the registry owns the +# normalization, so the value reported by +# `vllm.envs` and the value this check applies cannot drift apart. It is +# registered as a non-compile factor there, and an unrecognised value raises +# on first read rather than silently disarming the instrument. +_SLOT_OWNERSHIP_MODE: str = envs.VLLM_E2_MAMBA_SLOT_OWNERSHIP +# Liveness cadence: first check and every Nth are logged with a running count +# (see MambaSpecDecodeGPUContext.slot_ownership_checks_run). +_E2_CHECK_HEARTBEAT = 4096 + +_NULL_BLOCK_ID = 0 + + +def mamba_slot_ownership_violations( + *, + num_reqs: int, + idx_mapping: list[int], + bt_row_is_req: bool, + block_tables: list[list[list[int]]], + num_rows: int, + num_accepted_tokens: list[int], + state_idx: list[int], + new_num_computed_tokens: list[int], + block_size: int, + registered_ptrs: list[int] | None = None, + view_ptrs: list[int] | None = None, +) -> list[str]: + """ALIGN ROW-RULE AND SIMULTANEOUS BLOCK-COLLISION CHECK. + + Read the name literally. It was narrowed from "slot ownership": + with the production row rule ``row_used == req_idx``, so the + destination comparison below reduces to ``table[req_idx][dst_col]`` + against itself and is TAUTOLOGICAL on a healthy run. This function derives + what it expects and what it observes from THE SAME TABLE, and a table + cannot testify about its own staleness. What it does and does not + establish, stated so that zero findings are not over-read: + + ESTABLISHES + + 1. **Row attribution.** The row the kernel will use is the request's own + state slot. This is the property vllm#55506 violates: under PP the + deferred launch carries an ``idx_mapping`` captured ``pp_size`` steps + earlier, so ``batch_idx`` names a different request. The + ``BT_ROW_IS_REQ=False`` control exercises exactly this and nothing + wider. + 2. **Simultaneous block collision.** Within ONE step, no two live request + slots reach the same state block -- counting both the source and the + destination column, and counting slots that are live but not copying, + because a writer colliding with a non-copying live owner is the same + corruption. ``NULL_BLOCK_ID`` is excluded: block id 0 is the pool's + placeholder (``core/block_pool.py:188-191``), owned by no request. + 3. **Pointer-registration agreement**, when the caller supplies + ``registered_ptrs``: the tables read here are the tables whose addresses + the kernel was handed. That is the one genuinely INDEPENDENT source of + truth available at this layer, and without it a wrong pointer array + would be invisible -- the checker would faithfully validate a table the + kernel never reads. + + DOES NOT ESTABLISH -- do not quote zero findings as covering these: + + * **Stale ownership across steps.** The block-table row could itself be + stale; nothing here compares it against the scheduler's allocation. + * **Request-generation reuse.** A slot re-issued to a different request + with the same index is indistinguishable here: there is no generation + id in any tensor this function sees. + * **Cross-step source staleness.** A source block correct for this step but + written by a previous request's state is invisible. + * **Anything about a captured replay**, which does not re-enter the caller. + + Closing those needs an independent allocation/generation map from the + scheduler side, which is a plumbing change, not a diagnostic. + + Pure and side-effect free; every argument is a plain Python list or int, so + this runs with no GPU and no CUDA context. The caller owns the device->host + copies and the mode. + + ``block_tables[g][row][col]`` is the PERSISTENT per-req-slot table of mamba + group ``g``. ``registered_ptrs`` and ``view_ptrs``, when given, are the + kernel's ``persistent_block_table_ptrs`` values and the ``data_ptr()`` of + the tables actually read, compared elementwise. + + Returns one string per violation, empty when clean. + """ + violations: list[str] = [] + # (group, block id) -> (req slot, which column class claimed it). Both the + # source and the destination column are entered, and live-but-not-copying + # slots are entered too, so a writer colliding with a non-copying owner is + # caught (the destination-only map missed that case). + block_owner: dict[tuple[int, int], tuple[int, str]] = {} + + # (3) pointer-registration agreement -- the one independent source of truth + # available here. Without it the checker can faithfully validate a table + # the kernel never reads. + if registered_ptrs is not None and view_ptrs is not None: + if len(registered_ptrs) != len(view_ptrs): + violations.append( + f"pointer registration: {len(registered_ptrs)} registered " + f"block-table pointer(s) but {len(view_ptrs)} retained view(s)" + ) + for group_idx, (registered, view) in enumerate( + zip(registered_ptrs, view_ptrs) + ): + if int(registered) != int(view): + violations.append( + f"pointer registration: group {group_idx} kernel reads " + f"0x{int(registered):x} but the table checked here is at " + f"0x{int(view):x}. Every other finding in this run is " + f"about the wrong tensor." + ) + + def _claim(group_idx: int, block: int, req_idx: int, kind: str) -> None: + """Record one live claim on a state block and report a collision.""" + if block == _NULL_BLOCK_ID: + return # the pool's placeholder; owned by no request + key = (group_idx, block) + previous = block_owner.get(key) + if previous is not None and previous[0] != req_idx: + violations.append( + f"block collision: group {group_idx}, state block {block} is " + f"claimed in one step by request slot {previous[0]} " + f"({previous[1]}) and request slot {req_idx} ({kind}) " + f"(llama.cpp#29092 class: a prior request's recurrent state in " + f"a slot another request resumes into)." + ) + elif previous is None: + block_owner[key] = (req_idx, kind) + + def _cell(group_idx: int, table, row: int, col: int, what: str) -> int | None: + """Bounded read. Returns None and reports rather than indexing wrongly. + + A NEGATIVE column silently indexes a Python list from the end, so an + unbounded read here invents a block id and can report an unrelated + collision (example: accepted 1 with computed 0 yields dst_col -1, + which the device guard masks and this reader would turn into the LAST + column). Both bounds, both indices, every time. + """ + if not 0 <= row < len(table): + violations.append( + f"malformed metadata: group {group_idx}, {what} row {row} is " + f"outside the table's [0, {len(table)}) rows" + ) + return None + if not 0 <= col < len(table[row]): + violations.append( + f"malformed metadata: group {group_idx}, {what} column {col} " + f"is outside row {row}'s [0, {len(table[row])}) columns. The " + f"device guard masks this access; it is reported, not read." + ) + return None + return int(table[row][col]) + + # Pass 1: every LIVE slot's current state block, whether or not it copies. + # A slot that does not copy this step still owns its block, and a copying + # writer that lands on it is the corruption -- so the claims have to be + # entered before any writer is judged. + for batch_idx in range(min(num_reqs, len(idx_mapping))): + slot = int(idx_mapping[batch_idx]) + if slot < 0 or slot >= len(state_idx): + continue + for group_idx, table in enumerate(block_tables): + block = _cell(group_idx, table, slot, int(state_idx[slot]), "live src") + if block is not None: + _claim(group_idx, block, slot, "live source") + + for batch_idx in range(num_reqs): + if batch_idx >= len(idx_mapping): + violations.append( + f"idx_mapping is shorter ({len(idx_mapping)}) than num_reqs " + f"({num_reqs}); batch row {batch_idx} has no request slot" + ) + continue + req_idx = int(idx_mapping[batch_idx]) + if req_idx < 0: + continue # the kernel's own inactive-row sentinel + if req_idx >= min( + len(num_accepted_tokens), len(state_idx), len(new_num_computed_tokens) + ): + violations.append( + f"slot bound: batch row {batch_idx} maps to request slot " + f"{req_idx}, past the per-slot decision arrays " + f"({len(num_accepted_tokens)}, {len(state_idx)}, " + f"{len(new_num_computed_tokens)} entries)" + ) + continue + + # Mirror of the kernel's decision block (postprocess_mamba_fused_kernel, + # the PRECOMPUTED_NEW_COMPUTED branch). Every array is read at req_idx, + # never at batch_idx -- that asymmetry is the bug's hiding place. + num_accepted = int(num_accepted_tokens[req_idx]) + src_col = int(state_idx[req_idx]) + new_num_computed = int(new_num_computed_tokens[req_idx]) + num_tokens_running_state = new_num_computed - num_accepted + 1 + aligned_new_computed = (new_num_computed // block_size) * block_size + if aligned_new_computed < num_tokens_running_state: + continue # needs_copy is False: the block table is never read + accept_token_bias = aligned_new_computed - num_tokens_running_state + dst_col = aligned_new_computed // block_size - 1 + if src_col == dst_col and accept_token_bias == 0: + continue # no-op self-copy: the kernel returns before the copy + + row_used = req_idx if bt_row_is_req else batch_idx + + # (1) row attribution + if row_used != req_idx: + violations.append( + f"row attribution: batch row {batch_idx} is request slot " + f"{req_idx}, but the align post-process will read block-table " + f"row {row_used}. Under PP the idx_mapping was captured " + f"pp_size steps before this launch, so row {row_used} belongs " + f"to a different micro-batch's request (vllm#55506)." + ) + if not 0 <= row_used < num_rows: + violations.append( + f"row bound: batch row {batch_idx} resolves to block-table row " + f"{row_used}, outside [0, {num_rows})" + ) + continue + + for group_idx, table in enumerate(block_tables): + owner_block = _cell(group_idx, table, req_idx, dst_col, "owner dst") + used_block = _cell(group_idx, table, row_used, dst_col, "used dst") + if owner_block is None or used_block is None: + continue + # (2) destination agreement. NOTE: with the production row rule + # row_used == req_idx, so this comparison is tautological on a + # healthy run and proves nothing on its own. It earns + # its place only in the wrong-row world, where it reports the + # defect in BLOCK IDS rather than in indices -- which is what says + # the copy really would land on another request's state. + if used_block != owner_block: + violations.append( + f"destination disagreement: group {group_idx}, request " + f"slot {req_idx} would resume into state block " + f"{used_block} (row {row_used}, col {dst_col}) but its own " + f"row holds block {owner_block} there. The copy would " + f"write this request's recurrent state over another " + f"request's block." + ) + _claim(group_idx, used_block, req_idx, f"align destination col {dst_col}") + + return violations + # 16 saturates HBM on H100/GB200 across the reqs=8..128 range in # microbenchmarks _TEMPORAL_TILES = 16 +# Row bound used when no trustworthy row count is available: large enough never +# to reject a real request slot, so the guard degrades to "reject negatives" +# rather than to "reject correct metadata". See initialize_from_forward_context. +_UNBOUNDED_ROWS = 1 << 30 + @triton.jit(do_not_specialize=["num_requests"]) def get_aligned_state_indices_multi_group_kernel( @@ -187,6 +473,11 @@ def _reinterpret_u64_as_i64(value: int) -> int: return value if value < (1 << 63) else value - (1 << 64) +def _reinterpret_i64_as_u64(value: int) -> int: + """Inverse of ``_reinterpret_u64_as_i64``, for comparing pointers host side.""" + return value + (1 << 64) if value < 0 else value + + @triton.jit def _copy_mamba_state_block( state_idx, @@ -196,6 +487,8 @@ def _copy_mamba_state_block( token_bias, block_table_ptrs_ptr, block_table_stride_req, + block_table_num_rows, + guard_ctr, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, @@ -243,12 +536,39 @@ def _copy_mamba_state_block( group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) + # Bounds-guard instrumentation: bound the ROW as well as the column. #50021 bounds which + # column is read within a row; this bounds which row. CAN fire on + # postprocess_mamba_fused_kernel, where bt_row_idx comes from a device + # idx_mapping value. CANNOT fire on precopy_mamba_align_fused_kernel, whose + # bt_row_idx is program_id(0) bounded by num_reqs <= rows -- that path has a + # TEMPORAL staleness problem, fixed separately, not a spatial one. + if (bt_row_idx < 0) or (bt_row_idx >= block_table_num_rows): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, _GUARD_BIAS + (-1)) + return block_table_base = block_table_typed + bt_row_idx * block_table_stride_req # Widen block ids to int64 before they reach `block_id * state_block_stride` # below: state_block_stride can exceed 2**31 bytes for large mamba caches, # and Triton would otherwise do the multiply in int32 and wrap. - dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64) + # PR #50021 (upstream :178): these columns are derived from per-request + # accepted-token counts. A wrong count walks the read past this request's + # row; the loaded int32 is then multiplied by state_block_stride and + # dereferenced, so an out-of-range column becomes a wild (and possibly + # misaligned) address rather than merely a wrong copy. Only the + # out-of-range case is counted: `<= 0` also matches NULL_BLOCK_ID (0), + # which is the ordinary unallocated marker and not a fault. + dst_col_ok = (dst_col >= 0) & (dst_col < block_table_stride_req) + if (dst_col < 0) or (dst_col >= block_table_stride_req): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (dst_col).to(tl.int64) + _GUARD_BIAS) + dest_block_id = tl.load( + block_table_base + dst_col, mask=dst_col_ok, other=-1 + ).to(tl.int64) + if dest_block_id <= 0: + return dst_addr = state_base_addr + dest_block_id * state_block_stride is_conv_state = conv_width > 0 @@ -259,7 +579,16 @@ def _copy_mamba_state_block( if tile_idx > 0: return # DS conv layout: state_len is the slide axis; copy per dim row. - src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + src_col_ok = (src_col >= 0) & (src_col < block_table_stride_req) + if (src_col < 0) or (src_col >= block_table_stride_req): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (src_col).to(tl.int64) + _GUARD_BIAS) + src_block_id = tl.load( + block_table_base + src_col, mask=src_col_ok, other=-1 + ).to(tl.int64) + if src_block_id <= 0: + return dim_rows = tl.load(state_dim_row_count_ptr + state_idx) row_stride = tl.load(state_dim_row_stride_ptr + state_idx) src_block_addr = state_base_addr + src_block_id * state_block_stride @@ -310,7 +639,16 @@ def _copy_mamba_state_block( # SD conv: copy # state[bt[src_col], token_bias:] -> # state[bt[dst_col], :conv_width - token_bias] - src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + src_col_ok = (src_col >= 0) & (src_col < block_table_stride_req) + if (src_col < 0) or (src_col >= block_table_stride_req): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (src_col).to(tl.int64) + _GUARD_BIAS) + src_block_id = tl.load( + block_table_base + src_col, mask=src_col_ok, other=-1 + ).to(tl.int64) + if src_block_id <= 0: + return src_block_addr = state_base_addr + src_block_id * state_block_stride token_bytes = state_inner_size * state_elem_size num_dst_tokens = conv_width - token_bias @@ -349,7 +687,17 @@ def _copy_mamba_state_block( # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] # Body u64 range is partitioned across TEMPORAL_TILES CTAs to keep the # SMs filled at small batch. - actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64) + tmp_col = src_col + token_bias + tmp_col_ok = (tmp_col >= 0) & (tmp_col < block_table_stride_req) + if (tmp_col < 0) or (tmp_col >= block_table_stride_req): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (bt_row_idx).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (tmp_col).to(tl.int64) + _GUARD_BIAS) + actual_src_block_id = tl.load( + block_table_base + tmp_col, mask=tmp_col_ok, other=-1 + ).to(tl.int64) + if actual_src_block_id <= 0: + return src_addr = state_base_addr + actual_src_block_id * state_block_stride # Use natural block data size (inner_size * elem_size), NOT # state_block_stride which is the page stride and can exceed the @@ -378,6 +726,8 @@ def postprocess_mamba_fused_kernel( # block table. block_table_ptrs_ptr, block_table_stride_req: tl.int64, # stride between requests (in elements) + block_table_num_rows: tl.int64, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] # Mamba state metadata, flattened across all layer/state tensors. Different # Mamba groups may contribute different numbers of state tensors. state_base_addrs_ptr, # base address of each state tensor @@ -409,6 +759,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 @@ -476,7 +834,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: # noqa: SIM108 -- a ternary would drop this comment + # 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, @@ -485,6 +851,8 @@ def postprocess_mamba_fused_kernel( accept_token_bias, block_table_ptrs_ptr, block_table_stride_req, + block_table_num_rows, + guard_ctr, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, @@ -556,6 +924,8 @@ def precopy_mamba_align_fused_kernel( # Same flattened state-layout metadata as postprocess_mamba_fused_kernel block_table_ptrs_ptr, block_table_stride_req: tl.int64, + block_table_num_rows: tl.int64, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, @@ -616,6 +986,8 @@ def precopy_mamba_align_fused_kernel( token_bias, block_table_ptrs_ptr, block_table_stride_req, + block_table_num_rows, + guard_ctr, state_base_addrs_ptr, state_block_strides_ptr, state_elem_sizes_ptr, @@ -810,6 +1182,31 @@ 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 + + # The same tensors `persistent_block_table_ptrs` addresses, retained so the + # host-side slot-ownership check can READ the rows the kernel dereferences + # rather than model them. Holding these costs no allocation: they are the + # runner's own `BlockTables.block_tables[i].gpu`, alive for the engine's + # lifetime, and nothing here writes to them. + persistent_block_table_views: list[torch.Tensor] | None = None + # 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 @@ -828,6 +1225,18 @@ class MambaSpecDecodeGPUContext: # Flag to track if metadata has been populated is_initialized: bool = False + # Last result of the host-side slot-ownership check (always [] unless + # VLLM_E2_MAMBA_SLOT_OWNERSHIP is set). Readable by a test so the control + # does not have to parse log records. + last_slot_ownership_violations: list[str] = dataclasses.field( + default_factory=list + ) + # Liveness for the same check (diagnostics on 0.30): how many launches it has + # judged and how many had findings. Logged on the first check and every + # _E2_CHECK_HEARTBEAT after, so a zero-findings run can prove it ran. + slot_ownership_checks_run: int = 0 + slot_ownership_checks_with_findings: int = 0 + @classmethod def create( cls, @@ -900,6 +1309,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), @@ -924,6 +1336,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. @@ -954,10 +1367,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 @@ -968,6 +1385,7 @@ class MambaSpecDecodeGPUContext: forward_context, mamba_state_copy_funcs, block_tables, + persistent_block_tables, ) def _populate_metadata( @@ -976,6 +1394,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): @@ -1084,9 +1503,90 @@ class MambaSpecDecodeGPUContext: f"all mamba block tables must share stride(0), got {strides}" ) self.block_table_stride_req = int(next(iter(strides))) + # Bounds-guard instrumentation: row bound for the guard in _copy_mamba_state_block. + # + # NOT from `block_tables`. Every caller passes a per-batch VIEW -- + # get_device_tensor(num_reqs) at :1690 and :1838, a num_reqs_padded + # slice at mamba_hybrid.py:196 -- and this function runs once, so a + # bound taken from shape[0] here freezes at the first batch's request + # count and rejects every higher row in every later batch. The guard + # would then skip real state copies and report masked faults that never + # happened. (stride(0) is unaffected: it is max_num_blocks and survives + # a row slice, which is why the existing stride assert never caught + # this.) + # + # The persistent tables ARE [max_num_reqs, ...], so use them when they + # are registered. When they are not there is no correct row count at + # this point, and a wrong bound is worse than a loose one: fall back to + # a sentinel that still rejects negative rows -- the case the poison + # guard row exercises and the only one that reads outside the tensor -- + # while never rejecting a real request slot. + rows = {bt.shape[0] for bt in block_tables} + assert len(rows) == 1, ( + f"all mamba block tables must share shape[0], got {rows}" + ) + self.block_table_num_rows = _UNBOUNDED_ROWS for i, bt in enumerate(block_tables): self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) + # Same capture for the persistent per-req-slot tables. `input_block_tables` + # is `torch.zeros_like` of these (block_table.py:74-76), so stride(0) is + # identical and `block_table_stride_req` covers both; assert it rather + # than assume. + if persistent_block_tables is not None: + assert len(persistent_block_tables) == self.num_groups, ( + f"expected {self.num_groups} persistent block tables, " + f"got {len(persistent_block_tables)}" + ) + persistent_strides = {bt.stride(0) for bt in persistent_block_tables} + assert persistent_strides == strides, ( + "persistent block tables must share stride(0) with the gathered " + f"tables, got {persistent_strides} vs {strides}" + ) + assert self.persistent_block_table_ptrs is not None + for i, bt in enumerate(persistent_block_tables): + self.persistent_block_table_ptrs[i] = _reinterpret_u64_as_i64( + bt.data_ptr() + ) + # The align paths index the PERSISTENT tables by req_idx while + # every other geometry scalar is measured from the GATHERED ones + # (by design). Two different scalars, two different treatments: + # + # * the STRIDE is load-bearing -- it is the multiplier in + # `block_table_typed + bt_row_idx * block_table_stride_req`, so + # a mismatch reads the wrong row entirely. It is ASSERTED equal + # just above; there is no correct behaviour if it differs. + # * the ROW COUNT is not required to be equal (the fixture in + # test_pp_deferred_mamba_block_table.py builds a 2-row gathered + # table against an 8-row persistent one on purpose), so it is + # CARRIED separately. Anything bounding a row against the + # persistent tables must use this scalar; using + # block_table_num_rows would falsely reject a valid request slot + # whenever the gathered buffer is narrower, and a guard that + # fires on correct metadata is worse than no guard. + persistent_rows = {bt.shape[0] for bt in persistent_block_tables} + assert len(persistent_rows) == 1, ( + f"all persistent block tables must share shape[0], got " + f"{persistent_rows}" + ) + self.persistent_block_table_num_rows = int(next(iter(persistent_rows))) + # ...and this is the only trustworthy row count in this object, so + # the same-step guard takes it too. + self.block_table_num_rows = self.persistent_block_table_num_rows + self.has_persistent_block_tables = True + self.persistent_block_table_views = list(persistent_block_tables) + else: + # Legal for V1 callers, which never reach the V2 align + # post-process. If a V2 align launch does happen on such a context, + # `run_fused_postprocess_align` raises rather than falling back to + # the batch-order gather: that fallback is exactly the earlier + # wrong-row bug. + logger.warning_once( + "Mamba align context initialized without the persistent " + "per-req-slot block tables. The V1 paths are unaffected; a V2 " + "align post-process launch on this context will raise." + ) + self.is_initialized = True def compute_aligned_state_indices( @@ -1168,6 +1668,10 @@ class MambaSpecDecodeGPUContext: num_draft_tokens_gpu, self.block_table_ptrs, self.block_table_stride_req, + self.block_table_num_rows, + guard_telemetry.counter( + guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device + ), self.state_base_addrs, self.state_block_strides, self.state_elem_sizes, @@ -1215,6 +1719,10 @@ class MambaSpecDecodeGPUContext: token_bias_gpu, self.block_table_ptrs, self.block_table_stride_req, + self.block_table_num_rows, + guard_telemetry.counter( + guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device + ), self.state_base_addrs, self.state_block_strides, self.state_elem_sizes, @@ -1252,11 +1760,49 @@ 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 num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) + if _SLOT_OWNERSHIP_MODE != "off": + self._check_slot_ownership( + num_reqs=num_reqs, + num_accepted_tokens=num_accepted_tokens_snapshot, + state_idx_gpu=state_idx_gpu, + new_num_computed_tokens_gpu=new_num_computed_tokens_gpu, + idx_mapping=idx_mapping, + ) + total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) postprocess_mamba_fused_kernel[grid]( @@ -1265,8 +1811,12 @@ class MambaSpecDecodeGPUContext: None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED new_num_computed_tokens_gpu, None, # num_draft: unused under PRECOMPUTED_NEW_COMPUTED - self.block_table_ptrs, + block_table_ptrs, self.block_table_stride_req, + self.persistent_block_table_num_rows, + guard_telemetry.counter( + guard_telemetry.SITE_MAMBA_COPY, self.state_base_addrs.device + ), self.state_base_addrs, self.state_block_strides, self.state_elem_sizes, @@ -1283,9 +1833,136 @@ class MambaSpecDecodeGPUContext: CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), HAS_IDX_MAPPING=True, PRECOMPUTED_NEW_COMPUTED=True, + BT_ROW_IS_REQ=ALIGN_BT_ROW_IS_REQ, TEMPORAL_TILES=_TEMPORAL_TILES, ) + def _check_slot_ownership( + self, + *, + num_reqs: int, + num_accepted_tokens: torch.Tensor, + state_idx_gpu: torch.Tensor, + new_num_computed_tokens_gpu: torch.Tensor, + idx_mapping: torch.Tensor, + ) -> list[str]: + """Host-side mamba state-slot ownership check, off by default. + + VLLM_E2_MAMBA_SLOT_OWNERSHIP=warn|raise. Same discipline as the KV + zeroer's ownership check (VLLM_E2_ZERO_OWNERSHIP_CHECK): a diagnostic in the + scheduler-visible sense -- host side, opt-in, never a device guard. + Unlike that check it lives here, because the thing at risk is not visible to + the scheduler at all: `MambaSpec` is not an `AttentionSpec`, so mamba + state blocks never appear on `new_block_ids_to_zero` and that check never + sees them. + + COST, and it is not only wall-clock. Four device->host copies plus one + block-table slice per call, each a synchronize. That SERIALIZES work + that is normally overlapped, so an instrumented boot has materially + different timing from a production one: zero findings under this check + cannot by itself exclude a timing-sensitive defect. Run the + instrumented arm for diagnosis and an uninstrumented arm for the + concurrency qualification; do not merge their conclusions. + + CAPTURE CONTRACT, stated explicitly rather than assumed. These copies + are illegal inside a stream capture, so the function asserts it is not + capturing instead of relying on "this is engine-thread Python" -- Python + also runs while a graph is being captured, and the caller's own contract + is the thing being checked here. A captured REPLAY does not re-enter + this function at all, which is a real scope limit of the instrument, + not a safety argument. + + Returns the violation list (also stored on the context so a test can + read it without parsing logs). + """ + if torch.cuda.is_current_stream_capturing(): + # The D2H copies below would abort the capture. Report rather than + # silently skip: a diagnostic reached from a path it cannot serve + # is a contract error in the caller. + logger.warning_once( + "VLLM_E2_MAMBA_SLOT_OWNERSHIP is set but the align " + "post-process was reached during a cudagraph capture, where " + "the check's device->host copies are illegal. Skipped for " + "this launch; the caller's outside-capture contract does not " + "hold." + ) + self.last_slot_ownership_violations = [] + return [] + tables = self.persistent_block_table_views + if not tables: + # Unreachable from run_fused_postprocess_align, which raises above + # when the persistent tables were never registered. Kept so the + # check degrades to silence rather than to an exception if some + # future caller reaches it first. + logger.warning_once( + "VLLM_E2_MAMBA_SLOT_OWNERSHIP is set but the persistent " + "per-req-slot block tables were never registered, so slot " + "ownership cannot be checked." + ) + self.last_slot_ownership_violations = [] + return [] + + num_cols = min(int(t.shape[1]) for t in tables) + violations = mamba_slot_ownership_violations( + num_reqs=int(num_reqs), + idx_mapping=idx_mapping.detach().to("cpu").reshape(-1).tolist(), + bt_row_is_req=bool(ALIGN_BT_ROW_IS_REQ), + block_tables=[ + t[:, :num_cols].detach().to("cpu").tolist() for t in tables + ], + num_rows=int(self.persistent_block_table_num_rows), + # The independent evidence: what the kernel was handed, against + # what is being read here. `persistent_block_table_ptrs` holds the + # addresses reinterpreted as signed, so undo that before comparing. + registered_ptrs=( + [ + _reinterpret_i64_as_u64(int(v)) + for v in self.persistent_block_table_ptrs.detach() + .to("cpu") + .reshape(-1) + .tolist()[: len(tables)] + ] + if self.persistent_block_table_ptrs is not None + else None + ), + view_ptrs=[int(t.data_ptr()) for t in tables], + num_accepted_tokens=num_accepted_tokens.detach() + .to("cpu") + .reshape(-1) + .tolist(), + state_idx=state_idx_gpu.detach().to("cpu").reshape(-1).tolist(), + new_num_computed_tokens=new_num_computed_tokens_gpu.detach() + .to("cpu") + .reshape(-1) + .tolist(), + block_size=int(self.block_size), + ) + self.last_slot_ownership_violations = violations + self.slot_ownership_checks_run += 1 + if violations: + self.slot_ownership_checks_with_findings += 1 + runs = self.slot_ownership_checks_run + if runs == 1 or runs % _E2_CHECK_HEARTBEAT == 0: + logger.info( + "Mamba slot-ownership check (mode=%s): %d checks run, %d with " + "findings", + _SLOT_OWNERSHIP_MODE, + runs, + self.slot_ownership_checks_with_findings, + ) + if not violations: + return violations + message = ( + f"Mamba state-slot ownership violation on the align " + f"post-process ({len(violations)} finding(s), num_reqs=" + f"{num_reqs}, BT_ROW_IS_REQ={ALIGN_BT_ROW_IS_REQ}): " + + " | ".join(violations) + ) + if _SLOT_OWNERSHIP_MODE == "raise": + raise RuntimeError(message) + logger.warning("%s", message) + return violations + @dataclasses.dataclass class MambaBuffers: