iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Worker] Index the persistent per-request-slot block table under PP

Browse Source 
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 <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 19:43:54 +07:00
co-authored by Hinata
parent 41f6691c26
commit 62acb3dab8
4 changed files with 718 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/worker/gpu/model_runner.py
+4
View File
@@ -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,
vllm/v1/worker/gpu/model_states/interface.py
+8
View File
@@ -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).
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+19
View File
@@ -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
vllm/v1/worker/mamba_utils.py
+687 -10
View File
@@ -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: