From 2ae218c5e7b045e42c55fb26e929d8698728e11c Mon Sep 17 00:00:00 2001 From: halt95 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