iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
50dabc9989bc9d7b5625f8256196d9935de28ade
vLLM-4x-RTX-3090/patches/0027-worker-debug-gated-diagnostics-for-Mamba-slot-owners.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

1595 lines
72 KiB
Diff
Raw Blame History

From 2ae218c5e7b045e42c55fb26e929d8698728e11c Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:18:16 -0400
Subject: [PATCH 027/103] worker: debug-gated diagnostics for Mamba slot
ownership and guard encoding
- Host-side Mamba slot-ownership check (off|warn|raise, default off) with a
failing control that proves it detects a wrong block-table row.
- Guard telemetry: int64 encoding limits and the extreme-value pair closed by
test; a decoded-None counter field is now named in the warning; a
captured-replay sentinel probe.
- Both knobs registered in envs.py.
---
.../test_host_embed_capture_sentinel.py | 202 +++++++++
.../worker/gpu/test_guard_telemetry_limits.py | 328 ++++++++++++++
.../worker/gpu/test_mamba_slot_ownership.py | 366 +++++++++++++++
vllm/envs.py | 48 ++
vllm/v1/worker/guard_telemetry.py | 76 +++-
vllm/v1/worker/mamba_utils.py | 416 +++++++++++++++++-
6 files changed, 1433 insertions(+), 3 deletions(-)
create mode 100644 tests/models/qwen4_exp/test_host_embed_capture_sentinel.py
create mode 100644 tests/v1/worker/gpu/test_guard_telemetry_limits.py
create mode 100644 tests/v1/worker/gpu/test_mamba_slot_ownership.py
diff --git a/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py b/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py
new file mode 100644
index 0000000..aac4684
--- /dev/null
+++ b/tests/models/qwen4_exp/test_host_embed_capture_sentinel.py
@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The captured-replay host sentinel for the host-embedding layer (reference-host tier).
+
+The host-embedding work's bytewise check passed: on the real Merlin-e1 BF16
+table, TP2 over real NCCL, all 36 cases equal and the negative control (one
+flipped host byte) detected on both ranks. Its own title records what it did
+NOT settle:
+
+ "compiled-forward/captured-replay sentinel not yet run"
+ "that the SERVING path uses the same storage after compile and CUDA-graph
+ capture: the CPU-sentinel check through compiled forward and captured
+ replay, with residency recorded at each lifecycle stage"
+
+That is not a restatement of the layer-level gate. The layer-level gate proves
+the UVA read is byte-exact when Python calls it. This proves the READ IS STILL
+THE UVA READ after capture -- that graph capture did not bake a device-resident
+copy of the table into the replay. Those are different claims, and only the
+second one is about the serving path.
+
+WHY THIS IS THE RIGHT SHAPE OF CHECK. Equality alone cannot distinguish them:
+a captured replay that silently reads a stale device copy produces IDENTICAL
+bytes for identical inputs, forever. So equality is the first arm and a HOST
+SENTINEL is the second: mutate one row in the pinned host table between capture
+and replay and require the replay to produce the EXACT expected post-mutation
+bytes. If it does not, the replay is not reading live host memory, and every
+byte-equality result taken under capture was evidence about a copy
+(an instrument must discriminate both worlds).
+
+WHAT THE SENTINEL DOES *NOT* PROVE. Observing the
+mutation proves the replay DEPENDS ON MUTABLE HOST DATA. It does not prove
+exclusive UVA residency: a replay that copies the host row into a device buffer
+before reading it, on every launch, would also observe the mutation and pass.
+Ruling that out needs storage-identity and residency evidence
+(`host_memory_stats()` per worker, device allocation deltas across the
+lifecycle), which is not collected and is not asserted here. Do not read a
+pass as "therefore UVA".
+
+STATUS: NEVER EXECUTED. This module was written offline with no CUDA
+compiler; it is compile-checked only. It is a reference-host test.
+
+WHAT IT COVERS AND WHAT IT DOES NOT. It covers the standalone layer under a
+real ``torch.cuda.CUDAGraph`` capture, which is the part that can be asserted
+without booting an engine. It does NOT cover the layer inside the model's
+compiled piecewise forward on Flash-Next; that needs a real boot with
+``VLLM_HOST_EMBED_TABLE`` on, reading the sentinel through a served completion,
+and is the part of the host-embedding check not covered here. Do not treat it as closed on this file
+alone.
+"""
+
+import os
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.vocab_parallel_embedding as vpe
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+ PinnedHostVocabParallelEmbedding,
+)
+
+VOCAB = 4096
+HIDDEN = 256
+DTYPE = torch.bfloat16
+SENTINEL_ID = 1234
+
+pytestmark = [
+ pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA UVA"),
+ 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)",
+ ),
+]
+
+
+@pytest.fixture
+def layer(monkeypatch: pytest.MonkeyPatch) -> PinnedHostVocabParallelEmbedding:
+ """A real host-table embedding at TP=1, so no collective is involved.
+
+ TP=1 is deliberate: the host-embedding work already settled the TP2 collective bytewise. The
+ open question is capture, and a second rank only adds a way for this test
+ to fail for an unrelated reason.
+ """
+ if not vpe.PIN_MEMORY:
+ pytest.skip("pinned memory is unavailable")
+ monkeypatch.setattr(vpe, "get_tensor_model_parallel_rank", lambda: 0)
+ monkeypatch.setattr(vpe, "get_tensor_model_parallel_world_size", lambda: 1)
+ monkeypatch.setattr(vpe, "tensor_model_parallel_all_reduce", lambda value: value)
+
+ embedding = PinnedHostVocabParallelEmbedding(
+ VOCAB, HIDDEN, params_dtype=DTYPE
+ )
+ assert embedding.weight.is_pinned(), "the table must be pinned to be mapped"
+ assert embedding.accelerator_weight_view is not None, "no UVA view"
+ values = torch.arange(VOCAB * HIDDEN, dtype=torch.int32) % 251
+ embedding.weight.data.copy_(values.reshape(VOCAB, HIDDEN).to(DTYPE))
+ return embedding
+
+
+def _ids() -> torch.Tensor:
+ return torch.tensor(
+ [0, 1, SENTINEL_ID, VOCAB - 1, SENTINEL_ID],
+ dtype=torch.int32,
+ device="cuda",
+ )
+
+
+def _bytes(tensor: torch.Tensor) -> bytes:
+ return tensor.contiguous().view(torch.uint8).cpu().numpy().tobytes()
+
+
+def test_captured_replay_is_byte_identical_to_the_eager_forward(layer) -> None:
+ """Arm 1: capture must not change a single output byte."""
+ ids = _ids()
+ eager = _bytes(layer(ids))
+
+ static_ids = ids.clone()
+ graph = torch.cuda.CUDAGraph()
+ # A warm-up on a side stream: capturing a first-ever launch also captures
+ # its lazy initialisation (a lazy load that blocks behind a stream wait).
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ static_out = layer(static_ids)
+ torch.cuda.current_stream().wait_stream(side)
+ with torch.cuda.graph(graph):
+ static_out = layer(static_ids)
+
+ graph.replay()
+ torch.cuda.synchronize()
+
+ assert _bytes(static_out) == eager, (
+ "the captured replay's output differs from the eager forward on the "
+ "same ids and the same table"
+ )
+
+
+def test_a_host_sentinel_written_after_capture_is_visible_to_the_replay(
+ layer,
+) -> None:
+ """Arm 2, THE FAILING CONTROL: does the replay still read host memory?
+
+ Arm 1 passes whether the replay reads the UVA view or a device copy baked
+ in at capture time, because both produce the same bytes for the same ids.
+ Writing the host table AFTER capture separates them: only a real UVA read
+ sees the new row.
+
+ A failure here does not mean the output is wrong today -- it means every
+ byte-equality result taken under capture is evidence about a copy, and the
+ host-table saving (the reason the layer exists) is not what it claims on
+ the serving path.
+ """
+ ids = _ids()
+ static_ids = ids.clone()
+
+ side = torch.cuda.Stream()
+ side.wait_stream(torch.cuda.current_stream())
+ with torch.cuda.stream(side):
+ static_out = layer(static_ids)
+ torch.cuda.current_stream().wait_stream(side)
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ static_out = layer(static_ids)
+
+ graph.replay()
+ torch.cuda.synchronize()
+ before = _bytes(static_out)
+
+ # The sentinel: one row of the PINNED HOST table, written from the CPU
+ # with no device copy anywhere.
+ row = layer.weight.data[SENTINEL_ID]
+ original = row.clone()
+ row.fill_(torch.finfo(DTYPE).max)
+
+ # `after != before` is too weak -- ANY change passes it,
+ # including a wrong one. Build the exact bytes the replay must now produce
+ # and require equality with them. A partially updated output, a torn row
+ # or a different row moving all satisfy inequality and fail this.
+ expected_rows = []
+ for token in ids.tolist():
+ source = row if token == SENTINEL_ID else layer.weight.data[token]
+ expected_rows.append(source.clone())
+ expected = _bytes(torch.stack(expected_rows).to("cuda"))
+
+ graph.replay()
+ torch.cuda.synchronize()
+ after = _bytes(static_out)
+
+ assert after != before, (
+ "the captured replay did not observe a row written directly into the "
+ "pinned host table: it is reading something other than the live host "
+ "storage"
+ )
+ assert after == expected, (
+ "the captured replay changed, but not to the expected post-mutation "
+ "bytes -- it is not a clean read of the mutated host row"
+ )
+
+ # Restore and prove the observation is two-way, not a one-off artefact.
+ row.copy_(original)
+ graph.replay()
+ torch.cuda.synchronize()
+ assert _bytes(static_out) == before
diff --git a/tests/v1/worker/gpu/test_guard_telemetry_limits.py b/tests/v1/worker/gpu/test_guard_telemetry_limits.py
new file mode 100644
index 0000000..9fa7536
--- /dev/null
+++ b/tests/v1/worker/gpu/test_guard_telemetry_limits.py
@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guard telemetry: the documented limits, pinned by test rather than by change.
+
+Three limits of the bounds-guard telemetry are documented rather than fixed in code.
+This module pins the ones that are testable offline. The coverage table and the
+reader/producer stream-ordering argument that go with them are design notes not
+shipped with this tree.
+The three limits:
+
+ * Independent maxima, remaining half -- the (row, value) fields are INDEPENDENT
+ lifetime maxima. The labelling half shipped earlier (the WARNING no
+ longer calls them an "example"); what was missing is a worked example
+ proving the reported pair can be a pair that never occurred. That is
+ ``test_the_reported_pair_can_be_one_that_never_occurred``.
+ * Encoding limits -- the int64 bias encoding loses or falsifies extreme values.
+ Closed by test, not by change: the encoding is correct for every value a
+ correct kernel can produce, and the tests below pin exactly where it stops
+ being able to carry one, plus the fact that the COUNT always survives.
+ * Zero hits -- "zero hits implies no output bit changed" is false in
+ general. See ``test_zero_hits_is_not_bitwise_equivalence`` below; it is a
+ documentation test, and it says so.
+
+No GPU: nothing here allocates a counter buffer or launches a kernel. The
+producer arithmetic is exercised through ``encode_slot``, which is the same
+expression the Triton bodies inline, and the consumer through the real
+``_report_rows``.
+"""
+
+import logging
+
+import pytest
+import torch
+
+from vllm.v1.worker import guard_telemetry as gt
+
+_CPU = torch.device("cpu")
+
+
+@pytest.fixture(autouse=True)
+def _clean_high_water_marks():
+ """``_report_rows`` only reports a count ABOVE what it last saw.
+
+ A test that leaves ``_seen`` behind makes the NEXT test silently measure
+ nothing -- the trap the module's own comment on ``_seen`` describes.
+ """
+ gt.reset_for_test()
+ yield
+ gt.reset_for_test()
+
+
+class _Collector(logging.Handler):
+ def __init__(self) -> None:
+ super().__init__(level=logging.WARNING)
+ self.messages: list[str] = []
+
+ def emit(self, record: logging.LogRecord) -> None:
+ self.messages.append(record.getMessage())
+
+
+@pytest.fixture
+def warnings_from_the_reader():
+ """Capture the reader's WARNINGs off the real logger object.
+
+ Not ``caplog``: vLLM's root "vllm" logger is configured with
+ ``propagate: False``, so records never reach the handler pytest installs
+ on the ROOT logger and the capture would be empty. Attaching to the
+ module's own logger is the only capture that cannot quietly see nothing
+ (a capture that can see nothing measures nothing).
+ """
+ target = getattr(gt.logger, "logger", gt.logger)
+ handler = _Collector()
+ previous_level = target.level
+ target.addHandler(handler)
+ target.setLevel(logging.WARNING)
+ try:
+ yield handler.messages
+ finally:
+ target.removeHandler(handler)
+ target.setLevel(previous_level)
+
+
+# ---------------------------------------------------------------------------
+# Encoding limits -- int64 encoding of extreme values
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ "value",
+ [0, 1, -1, 42, -42, (1 << 31) - 1, -((1 << 31) - 1), 1 << 40],
+)
+def test_ordinary_values_round_trip(value: int) -> None:
+ """Every value a correct kernel can produce survives the encoding.
+
+ A row or index is an int32-range quantity, so this covers the whole domain
+ the guards actually record. The cases below are about metadata that is
+ already corrupt -- which is precisely when a guard fires.
+ """
+ assert gt.decode_slot(gt.encode_slot(value)) == value
+
+
+def test_int32_min_is_falsified_as_nothing_recorded() -> None:
+ """The one value the encoding cannot distinguish from an unwritten slot.
+
+ ``-GUARD_BIAS`` encodes to exactly 0, and 0 is the zeroed-buffer sentinel
+ for "this slot was never written". So a guard that masks an access at row
+ INT32_MIN reports its ROW as absent.
+ """
+ assert gt.encode_slot(-gt.GUARD_BIAS) == 0
+ assert gt.decode_slot(0) is None
+
+
+def test_values_below_int32_min_are_lost_by_atomic_max() -> None:
+ """Below ``-GUARD_BIAS`` the encoded word is negative.
+
+ The slot is only ever written with ``tl.atomic_max`` into a buffer zeroed
+ at allocation, and 0 > any negative, so the store is a no-op and the slot
+ still reads 0. Modelled here as ``max(0, encoded)`` because that is
+ exactly what the atomic does; there is no Triton in this test.
+ """
+ for value in (-gt.GUARD_BIAS - 1, -(1 << 40), -(1 << 62)):
+ encoded = gt.encode_slot(value)
+ assert encoded < 0, value
+ assert gt.decode_slot(max(0, encoded)) is None, value
+
+
+def test_the_encodable_window_is_where_the_constants_say_it_is() -> None:
+ """Pin the published bounds so a later change to GUARD_BIAS moves them."""
+ assert gt.decode_slot(gt.encode_slot(gt.GUARD_MIN_ENCODABLE)) == (
+ gt.GUARD_MIN_ENCODABLE
+ )
+ assert gt.decode_slot(gt.encode_slot(gt.GUARD_MAX_ENCODABLE)) == (
+ gt.GUARD_MAX_ENCODABLE
+ )
+ # One step below the window is the falsified case, not a round trip.
+ assert gt.encode_slot(gt.GUARD_MIN_ENCODABLE - 1) == 0
+ # One step above it does not fit in the int64 slot at all.
+ assert gt.encode_slot(gt.GUARD_MAX_ENCODABLE + 1) > (1 << 63) - 1
+
+
+def test_the_count_survives_an_unencodable_value_and_the_report_says_so(
+ warnings_from_the_reader: list[str],
+) -> None:
+ """The operational half of the encoding limits.
+
+ ``n_masked`` carries no bias and is written with ``atomic_add``, so it
+ advances whatever the offending value was. A report of "3 masked accesses,
+ row None, index None" is therefore a report of three REAL masked faults --
+ not a report of nothing. Before this change the message said only "row-field
+ None", which reads as "nothing recorded" to anyone who has not read the
+ encoding.
+
+ What the message must NOT do is diagnose WHY the field is None. The reader
+ has no ordering dependency on the producer, so an unobserved maximum and an
+ unencodable one read identically; the first draft of this clause claimed
+ "unencodable" outright and was wrong to.
+ """
+ rows = [[0, 0, 0] for _ in range(gt.NUM_SITES)]
+ rows[gt.SITE_MAMBA_COPY] = [3, 0, 0] # count 3, both maxima unencodable
+
+ gt._report_rows(_CPU, rows)
+
+ assert len(warnings_from_the_reader) == 1, warnings_from_the_reader
+ text = warnings_from_the_reader[0]
+ assert "masked 3 out-of-range access(es)" in text
+ assert "row-field None" in text
+ assert "UNOBSERVED OR UNENCODABLE, not absent" in text
+ # It must not claim to know which of the two it is.
+ assert "no ordering dependency on the producer" in text
+
+
+# ---------------------------------------------------------------------------
+# Independent maxima -- the fields are independent; the worked example
+# ---------------------------------------------------------------------------
+
+
+def test_the_reported_pair_can_be_one_that_never_occurred(
+ warnings_from_the_reader: list[str],
+) -> None:
+ """THE EXAMPLE PAIR.
+
+ Two masked accesses at one site, in this order:
+
+ hit A: row 3, offending index 100
+ hit B: row 7, offending index 5
+
+ Each slot independently keeps its lifetime maximum, so the slots end at
+ row 7 and index 100 and the reader reports the pair **(7, 100)** -- a
+ combination that never happened on the device. Row 7's access was at
+ index 5; index 100's access was at row 3.
+
+ This is why the WARNING labels them independent lifetime maxima and why no
+ attribution may be built on the pair: to chase a single access you need
+ the count plus a repro, not this triple. Anyone wanting a coherent
+ (row, value) sample needs the other half -- a capture that writes
+ both slots under one atomic decision -- which this tree does NOT implement.
+ """
+ slots = [0, 0, 0]
+ for row, value in ((3, 100), (7, 5)):
+ slots[0] += 1
+ slots[1] = max(slots[1], gt.encode_slot(row))
+ slots[2] = max(slots[2], gt.encode_slot(value))
+
+ assert gt.decode_slot(slots[1]) == 7
+ assert gt.decode_slot(slots[2]) == 100
+ assert (7, 100) not in {(3, 100), (7, 5)}
+
+ rows = [[0, 0, 0] for _ in range(gt.NUM_SITES)]
+ rows[gt.SITE_GDN_DECODE_STATE] = slots
+
+ gt._report_rows(_CPU, rows)
+
+ assert len(warnings_from_the_reader) == 1, warnings_from_the_reader
+ text = warnings_from_the_reader[0]
+ assert "row-field 7, index-field 100" in text
+ assert "independent lifetime maxima, not one observed pair" in text
+ # ...and no suffix about unobservable fields: both decoded fine.
+ assert "UNOBSERVED OR UNENCODABLE" not in text
+
+
+# ---------------------------------------------------------------------------
+# Zero hits is not bitwise equivalence
+# ---------------------------------------------------------------------------
+
+
+def test_zero_hits_is_not_bitwise_equivalence() -> None:
+ """A DOCUMENTATION TEST. It asserts the guard set's shape, not a behaviour.
+
+ The tempting claim is "zero guard hits implies no output bit
+ changed". It holds only for live, valid inputs, and there are two ways out
+ of it that a counter at zero cannot exclude:
+
+ 1. **The guards are bounds checks.** Every site in ``SITE_NAMES`` masks an
+ OUT-OF-RANGE access and counts it. A wrong-but-in-range row is not an
+ out-of-range access, so it is invisible here by construction -- stated
+ for the mamba copy site in the guard design notes ("Blind to: a column or row that
+ is in range but wrong -- the common case for a recycled slot,
+ producing a silent wrong state copy"). That is the whole of the
+ vllm#55506 / llama.cpp#29092 class, which is why it needed its own
+ instrument (``test_mamba_slot_ownership``) rather than a counter.
+
+ 2. **``NULL_BLOCK_ID`` rows change bytes uncounted.** Block id 0 is the
+ pool's placeholder (``core/block_pool.py:188-191``,
+ ``v1/attention/backends/utils.py:46``), which no request owns. The
+ copy guard's ``<= 0`` branch is DELIBERATELY not counted -- the guard design notes:
+ "Also blind to the ``<= 0`` path, deliberately: NULL_BLOCK_ID is 0 and
+ is the ordinary unallocated marker". Where such a row is zero-filled
+ rather than skipped, inactive bytes are written and no counter moves.
+ Those bytes belong to no live request, so no OUTPUT bit changes -- but
+ "zero hits" is evidence about live rows only, and a claim of bitwise
+ equivalence has to be made about the output tensor, not inferred from
+ the counters.
+
+ So a v2.2 result may say "no guard fired", never "no output bit changed"
+ (measure the property, not the proxy). The check below is that the
+ guard set is still a bounds-check set of the size this argument was
+ written against; if a site is added, this argument is re-read.
+ """
+ assert gt.NUM_SITES == 11
+ assert set(gt.SITE_NAMES) == set(range(gt.NUM_SITES))
+ # The one non-bounds site, which is telemetry and not a masked fault.
+ assert "SKIPPED" in gt.SITE_NAMES[gt.SITE_QSA_RING_SKIP]
+
+
+# ---------------------------------------------------------------------------
+# The diagnostic gates are registered knobs
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+ ("raw", "expected"),
+ [
+ (None, "off"),
+ ("", "off"), # `VAR=` is how scripts and fixtures clear a variable
+ (" ", "off"),
+ ("off", "off"),
+ ("1", "warn"), # continuity with the earlier boolean gates
+ ("warn", "warn"),
+ ("RAISE", "raise"),
+ (" Warn ", "warn"),
+ ],
+)
+def test_e2_diagnostic_gates_normalize_identically(
+ monkeypatch: pytest.MonkeyPatch, raw: str | None, expected: str
+) -> None:
+ """One normalizer for the whole family, so the registry cannot lie.
+
+ Both PLE diagnostic gates are registered in ``vllm.envs`` and read from there by the
+ modules that act on them. If a module re-parsed the variable itself, the
+ value ``vllm.envs`` reports (and hashes, and logs) could differ from the
+ value the check applies.
+ """
+ import vllm.envs as envs
+
+ for name in ("VLLM_E2_PLE_FLAG_CHECK", "VLLM_E2_MAMBA_SLOT_OWNERSHIP"):
+ monkeypatch.delenv(name, raising=False)
+ if raw is not None:
+ monkeypatch.setenv(name, raw)
+ assert envs.environment_variables[name]() == expected, (name, raw)
+
+
+def test_a_set_but_unrecognised_gate_is_a_boot_error(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """Fail closed: a broken gate must not silently disarm its instrument."""
+ import vllm.envs as envs
+
+ monkeypatch.setenv("VLLM_E2_MAMBA_SLOT_OWNERSHIP", "yes")
+ with pytest.raises(ValueError, match="off|warn|raise"):
+ envs.environment_variables["VLLM_E2_MAMBA_SLOT_OWNERSHIP"]()
+
+
+def test_the_host_side_gates_are_not_compile_factors() -> None:
+ """They gate host Python outside any capture and add no graph node.
+
+ Hashing them would invalidate every compile cache the moment a diagnostic
+ boot is taken -- which is precisely when the cache is wanted. If either
+ ever grows a device-side component this test should fail and be revisited,
+ not deleted.
+ """
+ import vllm.envs as envs
+
+ monkeypatch_free = envs.compile_factors()
+ assert "VLLM_E2_PLE_FLAG_CHECK" not in monkeypatch_free
+ assert "VLLM_E2_MAMBA_SLOT_OWNERSHIP" not in monkeypatch_free
+ # ...but they ARE registered, so `vllm.envs` can report them.
+ assert "VLLM_E2_PLE_FLAG_CHECK" in envs.environment_variables
+ assert "VLLM_E2_MAMBA_SLOT_OWNERSHIP" in envs.environment_variables
diff --git a/tests/v1/worker/gpu/test_mamba_slot_ownership.py b/tests/v1/worker/gpu/test_mamba_slot_ownership.py
new file mode 100644
index 0000000..89131ef
--- /dev/null
+++ b/tests/v1/worker/gpu/test_mamba_slot_ownership.py
@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ownership of the mamba state slot a request resumes into.
+
+THE HAZARD (llama.cpp#29092 class; the vllm#55506 failure mode is one
+instance). Mamba state blocks are the one KV class in this tree with neither
+zeroing nor an ownership instrument:
+
+* ``single_type_kv_cache_manager.py:96-99`` gates the recording of newly
+ allocated blocks for zeroing on ``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``) only inspects
+ ``new_block_ids_to_zero`` -- a list mamba blocks are never on.
+* a genuinely fresh request is safe anyway: ``gdn_attn.py:401`` sets
+ ``has_initial_state = context_lens_tensor > 0``, so at ``context_len == 0``
+ the GDN kernel ignores whatever 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 the align copy. The bounds guards bound the
+ indices; a bounds check cannot catch a wrong-but-in-range row, which is the
+ whole of this class.
+
+So the safety of a resuming request rested entirely on the correctness of the
+align copy's row selection -- which is the row-attribution fix, itself the fix for
+having got that row wrong once. These tests give that argument an instrument.
+
+GPU-FREE, INCLUDING THE FAILING CONTROL. The check is host side by
+construction, and ``test_pp_deferred_mamba_block_table`` already drives the
+real production plumbing (``MambaHybridModelState._ensure_align_ctx`` ->
+``MambaSpecDecodeGPUContext.initialize_from_forward_context`` ->
+``run_fused_postprocess_align``) on CPU tensors with the Triton launch
+captured instead of run. The control reuses that scenario builder rather than a
+stand-in (a hook contract needs the real classes), flips the production
+row rule ``ALIGN_BT_ROW_IS_REQ`` off, and requires the check to fire. No GPU
+time is needed; what a reference-host run would add is the end-to-end consequence
+(acceptance-length collapse, constant-token loops), not the assertion.
+
+SCOPE LIMIT, stated because it is a real one: ``run_fused_postprocess_align``
+is engine-thread Python, so a CAPTURED replay does not re-enter it and the
+check does not run for one. It covers the launch path, which is where the row
+is selected.
+"""
+
+import pytest
+
+from tests.v1.worker.gpu.test_pp_deferred_mamba_block_table import (
+ _build_scenario,
+ _run_deferred,
+)
+from vllm.v1.worker import mamba_utils as mamba_utils_mod
+from vllm.v1.worker.mamba_utils import mamba_slot_ownership_violations
+
+# vllm#55506's own reproduction permutation. A NON-IDENTITY batch-row ->
+# req-slot mapping is load-bearing for the control: with an identity mapping
+# ``batch_idx == req_idx`` and the two row rules are indistinguishable
+# (the needle shape is a variable). Four entries is also >= 3 concurrent
+# requests, the concurrency at which tomasruizt reproduced #55506 under
+# PP2 + MTP + prefix caching.
+_PERM = (2, 0, 3, 1)
+
+
+def _arm(monkeypatch: pytest.MonkeyPatch, mode: str = "raise") -> None:
+ monkeypatch.setattr(mamba_utils_mod, "_SLOT_OWNERSHIP_MODE", mode)
+
+
+# ---------------------------------------------------------------------------
+# End-to-end through the real align plumbing
+# ---------------------------------------------------------------------------
+
+
+def test_production_row_rule_reports_no_ownership_violation(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """The green world: with ALIGN_BT_ROW_IS_REQ the check must stay silent.
+
+ A check that fires on correct metadata is worse than no check (a
+ guard rule), so this arm is not decoration: it is what makes the control
+ below evidence about the row rule rather than about the fixture.
+ """
+ _arm(monkeypatch)
+ scenario = _build_scenario(slot_mapping=_PERM)
+ _run_deferred(scenario, monkeypatch)
+ assert scenario.ctx.last_slot_ownership_violations == []
+
+
+def test_bt_row_is_req_false_makes_the_check_fire(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """THE FAILING CONTROL: flip the row rule off, the check must fire.
+
+ This is the vllm#55506 world reproduced offline. The deferred align
+ post-process carries an ``idx_mapping`` captured ``pp_size`` steps earlier,
+ so with ``BT_ROW_IS_REQ`` off the kernel resolves batch row 0 -- which the
+ mapping says is request slot 2 -- to block-table row 0, another request's
+ slot. tomasruizt measured the consequence as acceptance length collapsing
+ to 1.00 for stretches; on this fork it showed as all-NaN logits and
+ constant-token loops.
+
+ If this test passes without the assertion firing, the assertion measures
+ nothing (an instrument must discriminate both worlds).
+ """
+ _arm(monkeypatch)
+ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False)
+ scenario = _build_scenario(slot_mapping=_PERM)
+
+ with pytest.raises(RuntimeError, match="state-slot ownership violation"):
+ _run_deferred(scenario, monkeypatch)
+
+ findings = scenario.ctx.last_slot_ownership_violations
+ assert any(f.startswith("row attribution:") for f in findings), findings
+ # ...and it is not only an index complaint: the block ids differ, so the
+ # copy really would land on another request's recurrent state.
+ assert any(f.startswith("destination disagreement:") for f in findings), findings
+
+
+def test_warn_mode_records_without_stopping_the_step(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """warn is the mode for a long stress boot; raise is for a repro."""
+ _arm(monkeypatch, "warn")
+ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False)
+ scenario = _build_scenario(slot_mapping=_PERM)
+
+ _run_deferred(scenario, monkeypatch)
+
+ assert scenario.ctx.last_slot_ownership_violations
+
+
+def test_off_is_off(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Default: nothing is copied to the host and nothing is recorded."""
+ _arm(monkeypatch, "off")
+ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False)
+ scenario = _build_scenario(slot_mapping=_PERM)
+
+ _run_deferred(scenario, monkeypatch)
+
+ assert scenario.ctx.last_slot_ownership_violations == []
+
+
+def test_an_identity_mapping_hides_the_defect_from_this_check(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """Why the control uses a permutation, asserted rather than asserted-about.
+
+ With an identity mapping ``batch_idx == req_idx``, so the wrong row rule
+ selects the right row anyway and the check is silent although the rule is
+ wrong. That is not a defect in the check -- there is nothing wrong with
+ the step -- but it IS a constraint on any future reproduction attempt:
+ a single-stream or in-order batch cannot exhibit #55506, which is why
+ upstream needed >= 3 concurrent requests to see it at all.
+ """
+ _arm(monkeypatch)
+ monkeypatch.setattr(mamba_utils_mod, "ALIGN_BT_ROW_IS_REQ", False)
+ scenario = _build_scenario(slot_mapping=(0, 1, 2))
+
+ _run_deferred(scenario, monkeypatch)
+
+ assert scenario.ctx.last_slot_ownership_violations == []
+
+
+# ---------------------------------------------------------------------------
+# The three properties, directly on the pure checker
+# ---------------------------------------------------------------------------
+
+_BLOCK_SIZE = 16
+
+
+def _clean_kwargs(**overrides):
+ """One copying request (slot 0) and one bystander, tables disjoint."""
+ kwargs = dict(
+ num_reqs=2,
+ idx_mapping=[0, 1],
+ bt_row_is_req=True,
+ block_tables=[[[10, 11, 12], [20, 21, 22], [30, 31, 32]]],
+ num_rows=3,
+ num_accepted_tokens=[2, 2, 2],
+ state_idx=[0, 0, 0],
+ new_num_computed_tokens=[32, 40, 40],
+ block_size=_BLOCK_SIZE,
+ )
+ kwargs.update(overrides)
+ return kwargs
+
+
+def test_clean_case_is_clean() -> None:
+ assert mamba_slot_ownership_violations(**_clean_kwargs()) == []
+
+
+def test_a_row_that_does_not_reach_the_block_table_is_not_judged() -> None:
+ """``needs_copy`` false means the kernel returns before any dereference.
+
+ Judging such a row would make the check fire on rows the kernel never
+ touches, i.e. on correct metadata.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ # One batch row, mapped to slot 1, whose 40 computed tokens leave
+ # the accepted window inside the running block: needs_copy false.
+ # The row rule is WRONG here (batch row 0 != slot 1) and the check
+ # still says nothing, because the kernel never reaches the table.
+ num_reqs=1,
+ bt_row_is_req=False,
+ idx_mapping=[1, 0],
+ new_num_computed_tokens=[32, 40, 40],
+ )
+ )
+ assert findings == []
+
+
+def test_two_requests_resuming_into_one_state_block_is_a_collision() -> None:
+ """The llama.cpp#29092 shape with the row rule CORRECT.
+
+ Row attribution and destination agreement both pass -- each request reads
+ its own row -- and the tables still hand both of them block 11. Only the
+ collision property catches this, which is why it is separate and not a
+ corollary of the row rule.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ block_tables=[[[10, 11, 12], [20, 11, 22], [30, 31, 32]]],
+ new_num_computed_tokens=[32, 32, 40],
+ )
+ )
+ assert len(findings) == 1, findings
+ assert findings[0].startswith("block collision:")
+ assert "slot 0" in findings[0] and "slot 1" in findings[0]
+
+
+def test_null_block_id_rows_are_not_a_collision() -> None:
+ """Block id 0 is the pool's placeholder; no request owns it.
+
+ ``NULL_BLOCK_ID`` (``v1/attention/backends/utils.py:46``) is popped out of
+ the free queue at construction (``core/block_pool.py:188-191``), so two
+ rows carrying 0 are two unallocated columns, not a shared state block.
+ Treating that as a collision would fire on every partially allocated
+ batch.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ block_tables=[[[10, 0, 12], [20, 0, 22], [30, 31, 32]]],
+ new_num_computed_tokens=[32, 32, 40],
+ )
+ )
+ assert findings == []
+
+
+def test_a_row_outside_the_persistent_table_is_reported_and_not_indexed() -> None:
+ """Out-of-range rows are the bounds guards' business, but the check must
+ not itself index past the table while reporting one."""
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ num_reqs=3,
+ bt_row_is_req=False,
+ idx_mapping=[2, 1, 0],
+ num_rows=2,
+ new_num_computed_tokens=[32, 40, 32],
+ )
+ )
+ # Batch row 2 resolves to table row 2 with only 2 rows registered.
+ assert any(f.startswith("row bound:") for f in findings), findings
+ # Batch row 0 is in range but wrong, and is reported as such -- the two
+ # findings are different properties, not one restated.
+ assert any(f.startswith("row attribution:") for f in findings), findings
+ assert any(f.startswith("destination disagreement:") for f in findings), findings
+
+
+def test_a_request_slot_past_the_decision_arrays_is_reported() -> None:
+ """The per-slot arrays are read at ``req_idx``; a slot past them is
+ malformed metadata, and the checker must say so rather than raise."""
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(num_reqs=1, idx_mapping=[7])
+ )
+ assert len(findings) == 1, findings
+ assert findings[0].startswith("slot bound:")
+
+
+def test_a_negative_request_slot_is_skipped() -> None:
+ """``req_idx < 0`` is the kernel's own inactive-row sentinel."""
+ assert (
+ mamba_slot_ownership_violations(
+ **_clean_kwargs(idx_mapping=[-1, 1], bt_row_is_req=False)
+ )
+ == []
+ )
+
+
+def test_a_short_idx_mapping_is_reported_rather_than_raising() -> None:
+ """An instrument that dies is an instrument that is not there."""
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(num_reqs=3, idx_mapping=[0, 1])
+ )
+ assert any("idx_mapping is shorter" in f for f in findings), findings
+
+
+# ---------------------------------------------------------------------------
+# What the narrowed instrument gained
+# ---------------------------------------------------------------------------
+
+
+def test_a_writer_colliding_with_a_live_non_copying_owner_is_caught() -> None:
+ """A destination-only collision map missed this case.
+
+ Slot 1 is live this step and does NOT copy (40 computed tokens leave its
+ accepted window inside the running block), so it never appeared in a map
+ built from copying requests' destinations. Its current state block is 20.
+ Slot 0 copies into column 1, which also holds block 20. That is a writer
+ landing on a live owner's recurrent state -- the llama.cpp#29092 shape --
+ and it must be reported.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ block_tables=[[[10, 20, 12], [20, 21, 22], [30, 31, 32]]],
+ new_num_computed_tokens=[32, 40, 40],
+ )
+ )
+ assert len(findings) == 1, findings
+ assert findings[0].startswith("block collision:")
+ assert "live source" in findings[0], findings[0]
+
+
+def test_a_negative_destination_column_is_reported_not_wrapped() -> None:
+ """A negative column indexes a Python list FROM THE END.
+
+ Accepted count 1 with computed count 0 gives ``dst_col == -1``. The device
+ guard masks that access. An unbounded reader here would instead fetch the
+ LAST column and could report a collision on a block the kernel never
+ touches -- an instrument inventing a finding is worse than a silent one.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(
+ num_reqs=1,
+ idx_mapping=[0],
+ num_accepted_tokens=[1, 2, 2],
+ new_num_computed_tokens=[0, 40, 40],
+ )
+ )
+ assert findings, "a -1 destination column must be reported"
+ assert all(f.startswith("malformed metadata:") for f in findings), findings
+ # The last column (block 12) must not have been read as a real claim.
+ assert not any("12" in f and "collision" in f for f in findings), findings
+
+
+def test_pointer_registration_disagreement_is_reported() -> None:
+ """The one INDEPENDENT source of truth this layer has.
+
+ Everything else in this checker is derived from the same table it is
+ validating. If the kernel was handed a different address, every other
+ finding in the run is about the wrong tensor, and zero findings would be
+ meaningless rather than reassuring.
+ """
+ findings = mamba_slot_ownership_violations(
+ **_clean_kwargs(registered_ptrs=[0xDEAD0000], view_ptrs=[0xBEEF0000])
+ )
+ assert len(findings) == 1, findings
+ assert findings[0].startswith("pointer registration:")
+
+
+def test_matching_pointers_are_silent() -> None:
+ """A check that fires on correct metadata is worse than no check."""
+ assert (
+ mamba_slot_ownership_violations(
+ **_clean_kwargs(registered_ptrs=[0xCAFE0000], view_ptrs=[0xCAFE0000])
+ )
+ == []
+ )
diff --git a/vllm/envs.py b/vllm/envs.py
index 554c025..157dfaf 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -603,6 +603,30 @@ def _resolve_rust_cli_path() -> str | None:
return raw
+def _e2_diagnostic_mode(name: str) -> str:
+ """Normalize an E2 diagnostic gate to off|warn|raise.
+
+ One normalizer for every knob in the family, so a registry entry and the
+ module that acts on it cannot disagree about what the environment said.
+ ``1`` is accepted as ``warn`` for continuity with the earlier boolean
+ gates. A set-but-unrecognised value is a boot error, not a silent ``off``:
+ a broken gate that disarms its own instrument is the failure mode these
+ diagnostics exist to avoid.
+ """
+ # An EMPTY value means off, not a boot error. Shell scripts and pytest
+ # fixtures both use `VAR=` to mean "unset" (see tests/v1/worker/
+ # test_e2_memstats.py), and rejecting it would turn a routine way of
+ # clearing the variable into a crash.
+ value = os.environ.get(name, "off").strip().lower() or "off"
+ if value == "1":
+ return "warn"
+ if value not in ("off", "warn", "raise"):
+ raise ValueError(
+ f"{name} must be off|warn|raise (or 1 for warn), got {value!r}"
+ )
+ return value
+
+
environment_variables: dict[str, Callable[[], Any]] = {
# ================== Installation Time Env Vars ==================
# Target device of vLLM, supporting [cuda (by default),
@@ -2181,6 +2205,23 @@ environment_variables: dict[str, Callable[[], Any]] = {
# Each op additionally checks its own shape / dtype constraints and falls
# back to the eager path when they do not hold.
"VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))),
+ # PLE diagnostics. Both are off|warn|raise (1 is accepted as warn) and both
+ # default to off. They are registered here so the registry is the single
+ # source of truth for their normalization -- the modules that use them read
+ # these getters rather than calling os.getenv a second time, so the value
+ # the registry reports cannot drift from the value the check applies.
+ # Both are excluded from compile factors in `compile_factors` below: they
+ # gate HOST-side Python outside any capture and add no node to a compiled
+ # graph. If either ever grows a device-side component, remove it from
+ # `ignored_factors` in the same change.
+ #
+ # Read the pull-transport PLE semaphore back before a real launch, on the
+ # publishing rank only (vllm#53899 one-step-lag class).
+ "VLLM_E2_PLE_FLAG_CHECK": lambda: _e2_diagnostic_mode("VLLM_E2_PLE_FLAG_CHECK"),
+ # Host-side align row-rule and state-block collision check.
+ "VLLM_E2_MAMBA_SLOT_OWNERSHIP": lambda: _e2_diagnostic_mode(
+ "VLLM_E2_MAMBA_SLOT_OWNERSHIP"
+ ),
}
@@ -2273,6 +2314,13 @@ def compile_factors() -> dict[str, object]:
"VLLM_DEBUG_DUMP_PATH",
"VLLM_PORT",
"VLLM_CACHE_ROOT",
+ # PLE host-side diagnostics: both gate Python that runs outside any
+ # capture and neither adds a node to a compiled graph, so hashing them
+ # would invalidate every compile cache the moment a diagnostic boot is
+ # taken -- which is exactly when the cache is wanted. Remove from this
+ # set if either ever grows a device-side component.
+ "VLLM_E2_PLE_FLAG_CHECK",
+ "VLLM_E2_MAMBA_SLOT_OWNERSHIP",
# Runtime memory-plan persistence; does not affect compiled graphs.
"VLLM_ENABLE_STARTUP_PLAN",
# Location-only derived paths: where a cache/config directory lives
diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py
index 60fb094..e95653e 100644
--- a/vllm/v1/worker/guard_telemetry.py
+++ b/vllm/v1/worker/guard_telemetry.py
@@ -78,6 +78,59 @@ GUARD_BIAS = 1 << 31
# One (n_masked, row, value) triple per guarded kernel.
SLOTS_PER_SITE = 3
+# Largest value the bias encoding can carry without overflowing the int64 slot.
+GUARD_MAX_ENCODABLE = (1 << 63) - 1 - GUARD_BIAS
+# Smallest one. At exactly -GUARD_BIAS the encoded word is 0, which the decode
+# cannot tell from "nothing recorded"; below it the encoded word is negative
+# and tl.atomic_max against a zeroed buffer never takes it. Both cases lose the
+# VALUE and neither loses the COUNT -- see decode_slot.
+GUARD_MIN_ENCODABLE = -GUARD_BIAS + 1
+
+
+def encode_slot(value: int) -> int:
+ """Producer side of the bias encoding, in Python (for the limits tests).
+
+ The guarded kernels inline `(value + _GUARD_BIAS).to(tl.int64)` and store
+ it with `tl.atomic_max`; this is the same arithmetic, exposed so the
+ encoding's limits can be tested without Triton and without a GPU. It is
+ NOT called by the kernels -- a Python function is not readable from a
+ @triton.jit body, which is why _GUARD_BIAS is a tl.constexpr there.
+ """
+ return int(value) + GUARD_BIAS
+
+
+def decode_slot(biased: int) -> int | None:
+ """Reader side. ``None`` means "this slot recorded nothing".
+
+ THE ENCODING'S LIMITS (closed here by test
+ rather than by change). A slot is an int64 zeroed at allocation and only
+ ever written with `tl.atomic_max`, so the recorded value must be biased to
+ survive negatives, and 0 has to mean "unwritten". Three classes of extreme
+ value are therefore not recoverable from the slot:
+
+ * ``value == -GUARD_BIAS`` (INT32_MIN) encodes to exactly 0 and is
+ FALSIFIED: it reads back as "nothing recorded".
+ * ``value < -GUARD_BIAS`` encodes negative, `atomic_max` against 0 keeps
+ 0, and the value is LOST the same way.
+ * ``value > GUARD_MAX_ENCODABLE`` overflows the int64 slot.
+
+ In every one of those cases the COUNT slot, which carries no bias and is
+ written with `atomic_add`, still advances.
+
+ BUT `None` IS NOT A DIAGNOSIS OF "UNENCODABLE". The reader has
+ NO ordering dependency on the producer -- it copies through a pinned mirror
+ on a side stream with no dependency on the compute stream -- so a count
+ that has become visible with its maxima not yet visible produces exactly
+ the same reading. `_poll_once`'s one-tick confirm delay makes that less
+ likely; it does not make it impossible, and a delay is not a happens-before
+ relationship. So `None` means UNOBSERVED OR UNENCODABLE.
+
+ What it never means is "this field recorded nothing while the count rose".
+ Anything reading these slots must say so, and must not build a fatal
+ decision on the maxima -- those decisions are count-based.
+ """
+ return None if biased == 0 else int(biased) - GUARD_BIAS
+
SITE_MAMBA_COPY = 0
SITE_GDN_DECODE_STATE = 1
SITE_QSA_PRE_INDEXER = 2
@@ -424,8 +477,8 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None:
if n_masked <= _seen.get((device, site), 0):
continue
_seen[(device, site)] = n_masked
- row = biased_row - GUARD_BIAS if biased_row else None
- val = biased_val - GUARD_BIAS if biased_val else None
+ row = decode_slot(biased_row)
+ val = decode_slot(biased_val)
msg = (
f"bounds guard masked {n_masked} out-of-range access(es) "
f"in {SITE_NAMES.get(site, site)} on {device} "
@@ -434,6 +487,25 @@ def _report_rows(device: torch.device, rows: list[list[int]]) -> None:
f"access was suppressed, so this is a MASKED FAULT and not a "
f"correct result -- the metadata is wrong upstream."
)
+ if row is None or val is None:
+ # Corrected wording. A None field has TWO
+ # possible causes and this reader cannot distinguish them, because
+ # it has no happens-before relationship with the producer:
+ # * UNOBSERVED -- the kernel incremented the count and its
+ # atomic_max on this slot is not yet visible to this copy;
+ # * UNENCODABLE -- the value was == INT32_MIN, below it, or past
+ # GUARD_MAX_ENCODABLE, so the bias encoding could not carry it.
+ # Either way it does NOT mean the field recorded nothing. The
+ # earlier wording claimed "unencodable" outright, which the
+ # one-tick confirm delay does not establish.
+ msg += (
+ " A None field is UNOBSERVED OR UNENCODABLE, not absent: "
+ "either this copy predates the slot's atomic_max becoming "
+ "visible, or the value was == INT32_MIN, below it, or past "
+ f"{GUARD_MAX_ENCODABLE}. The reader has no ordering "
+ "dependency on the producer, so it cannot tell which; see "
+ "decode_slot."
+ )
logger.warning("%s", msg)
if _MODE == "raise" or site in _ALWAYS_FATAL_SITES:
_arm(msg)
diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py
index 39d71b3..a6972f1 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 (
@@ -37,6 +38,276 @@ logger = init_logger(__name__)
# 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
+
+_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
@@ -199,6 +470,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,
@@ -923,6 +1199,13 @@ class MambaSpecDecodeGPUContext:
# 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
@@ -941,6 +1224,13 @@ 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
+ )
+
@classmethod
def create(
cls,
@@ -1279,6 +1569,7 @@ class MambaSpecDecodeGPUContext:
# 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,
@@ -1497,6 +1788,15 @@ class MambaSpecDecodeGPUContext:
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](
@@ -1527,10 +1827,124 @@ class MambaSpecDecodeGPUContext:
CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
HAS_IDX_MAPPING=True,
PRECOMPUTED_NEW_COMPUTED=True,
- BT_ROW_IS_REQ=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
+ 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:
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink