From 55b60a9ae9f9518a0ea66f90f3441eb97c23f488 Mon Sep 17 00:00:00 2001 From: TAJh Date: Tue, 20 Oct 2026 07:39:10 +0100 Subject: [PATCH 1/6] feat(debug): add opt-in KV cache address legality diagnostics Validate scheduled token destinations, allocation generations, ownership, logical block order and initialized prefixes for adapted dense PA/MLA paths. Reuse the scheduler output channel and common V1/V2 checker runtime. Audit ordered eager arguments and current FULL metadata without redundant graph callbacks. Track completed physical zero/copy and Mooncake receive lifecycle boundaries, with bounded diagnostics and optional failure latching. Integrate with current upstream main d606456c4, retaining its strided MLA, dummy replay and background receive contracts. Include focused CPU cases, a service regression entry point, user documentation and a release fragment. Present the already-reviewed complete patch as one linear commit so CI's plain git rebase does not replay obsolete intermediate implementations. Validation: 98 focused CPU tests and repository-pinned Ruff check/format pass. NPU service and performance validation remain deferred at user request. Signed-off-by: TAJh --- .../feature_guide/kv_cache_guard.md | 168 ++++ mkdocs.yml | 1 + release_notes/fragments/17833.yaml | 8 + .../one_card/test_kv_cache_guard.py | 95 ++ tests/ut/debug/__init__.py | 2 + tests/ut/debug/test_kv_guard.py | 230 +++++ tests/ut/debug/test_kv_guard_runtime.py | 836 ++++++++++++++++++ tests/ut/debug/test_kv_guard_v2.py | 324 +++++++ .../kv_p2p/test_mooncake_connector.py | 41 + vllm_ascend/attention/attention_v1.py | 15 + vllm_ascend/attention/mla_v1.py | 26 + vllm_ascend/debug/__init__.py | 3 + vllm_ascend/debug/kv_guard.py | 337 +++++++ vllm_ascend/debug/kv_guard_runtime.py | 586 ++++++++++++ vllm_ascend/debug/kv_guard_v2.py | 134 +++ .../kv_transfer/kv_p2p/mooncake_connector.py | 8 +- vllm_ascend/envs.py | 8 + vllm_ascend/patch/__init__.py | 15 + vllm_ascend/patch/platform/__init__.py | 6 + vllm_ascend/patch/platform/patch_kv_guard.py | 85 ++ vllm_ascend/worker/model_runner_v1.py | 20 + vllm_ascend/worker/v2/model_runner.py | 5 + 22 files changed, 2952 insertions(+), 1 deletion(-) create mode 100644 docs/source/user_guide/feature_guide/kv_cache_guard.md create mode 100644 release_notes/fragments/17833.yaml create mode 100644 tests/e2e/pull_request/one_card/test_kv_cache_guard.py create mode 100644 tests/ut/debug/__init__.py create mode 100644 tests/ut/debug/test_kv_guard.py create mode 100644 tests/ut/debug/test_kv_guard_runtime.py create mode 100644 tests/ut/debug/test_kv_guard_v2.py create mode 100644 vllm_ascend/debug/__init__.py create mode 100644 vllm_ascend/debug/kv_guard.py create mode 100644 vllm_ascend/debug/kv_guard_runtime.py create mode 100644 vllm_ascend/debug/kv_guard_v2.py create mode 100644 vllm_ascend/patch/platform/patch_kv_guard.py diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md new file mode 100644 index 000000000000..9cc551fcf0ce --- /dev/null +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -0,0 +1,168 @@ +# KV cache address guard + +KV Guard diagnoses address mappings that can cause cross-request writes, +prefix overwrites, or reads from invalid KV ranges. It compares scheduled +destinations with execution metadata and observed operator arguments. +The feature is disabled by default and supports adapted dense PA and MLA +attention paths in Model Runner V1 and V2. + +## Enable the guard + +```bash +# Model Runner V2 +VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ASCEND_KV_GUARD=1 vllm serve MODEL + +# Model Runner V1 +VLLM_USE_V2_MODEL_RUNNER=0 VLLM_ASCEND_KV_GUARD=1 vllm serve MODEL +``` + +`VLLM_ASCEND_KV_GUARD_ACTION=log` is the default: it reports failures while +preserving diagnostic execution. To stop returning worker results after a +violation or incomplete audit, select `raise`: + +```bash +VLLM_USE_V2_MODEL_RUNNER=1 \ +VLLM_ASCEND_KV_GUARD=1 \ +VLLM_ASCEND_KV_GUARD_ACTION=raise \ +vllm serve MODEL +``` + +The first failure is latched. Catching the exception does not reset the worker; +restart it before reuse. This is a worker-level stop, not cancellation of one +request, and it adds no collective failure coordination across ranks. + +## What is checked + +An allocated block is not necessarily the correct destination for a token. +For example, if requests A and B own blocks 1 and 2, swapping their write slots +still points to allocated memory. Ordinary allocation logs cannot establish +that each token used its intended destination. + +The scheduler supplies changed request-to-block bindings, allocation generations, +released requests, and each request's computed and scheduled token counts through +the existing scheduler output. The worker independently derives the expected +slot of every scheduled token from those bindings and the actual batch order. + +| Constraint | Failure example | +| --- | --- | +| Exact token destination | Slots are swapped between requests or between tokens in one request. | +| Allocation and ownership | A real request accesses another request's block or a recycled allocation. | +| Logical block order | A request reads its own blocks in the wrong logical order. | +| Initialized prefix | A read extends into unwritten data, or a write overwrites a valid prefix. | +| Reserved storage | A write targets the null block instead of negative-slot padding. | +| Collection completeness | A scheduled write, binding, layer, or current graph plan is missing. | + +Each tracked block stores an allocation generation, a valid prefix length, and +request owners. A zero prefix is unwritten; a positive prefix N makes `[0, N)` +readable; an invalid prefix prevents subsequent reads from being accepted. +Completed loads and contiguous appends establish readable prefixes. Completion +and preemption release ownership while preserving reusable cached prefixes. +Physical clearing, block copies, and completed Mooncake receives update validity +at their existing lifecycle boundaries. Failed receives invalidate their +destinations without changing the connector's completion and drain protocol. + +For example, with block size 128 and a loaded 19-token prefix in block 1, the +next token should write slot 147. A dummy write to slot 142 overwrites offset 14; +the guard diagnoses the write and marks the block invalid for subsequent reads. +This example illustrates the metadata checks, not reproduction of a particular +full-model incident. + +## Execution paths and supported configurations + +| Path | Audit evidence and boundary | +| --- | --- | +| Eager | Ordered per-layer operator arguments, with adapted metadata plans where applicable. | +| FULL replay | The current execution's final attention metadata, inspected after forward and graph parameter update calls. Captured or previous-step plans are not reused as evidence. | +| Runner V1 | Bind before `execute_model`; collect through the existing forward context and attention hooks. | +| Runner V2 | Bind before execute; obtain actual batch IDs and metadata after `model_state.prepare_attn`; audit at the end of execute. Persistent request-state indices are not batch positions. | +| Dummy/profile execution | Check writes against existing KV state; ignore discarded scratch reads. | +| Runner V2 startup capture | Require empty request/cache state; do not audit CPU metadata inside device graph capture. | + +Live or cache-preserving Runner V2 recapture is rejected, including pending +receives. Normal FULL replay remains audited. Allocator and kernel block +coordinates must match. + +The guard explicitly rejects CP, DBO, speculation, encoder/decoder execution, +sliding or circular caches, other compressed or state-cache semantics, unknown +attention implementations, and connectors other than Mooncake. These paths +require their own validity rules before they can be supported. + +The guard validates metadata and observed arguments. It does not inspect kernel +memory instructions, verify KV numerical values, or detect a kernel writing to +another address despite correct arguments. It cannot prevent or undo submitted +writes, prove transport byte integrity, or guarantee model output correctness. + +## Read the logs + +Logs use the `[KV_GUARD]` prefix: + +- `guard.enabled`: the configured failure action. +- `guard.violation`: an observed constraint violation, with an operation and + reason such as `write_target_mismatch`, `overwrite_valid_prefix`, or + `read_unwritten_range`. +- `guard.coverage_gap`: required evidence was missing, so the audit was + incomplete. This is not a successful check or an exemption from validation. + +Fields identify the layer, request, pool, cache group, block, allocation +generation, and relevant slot or range when available. Repeated failures are +deduplicated with bounded state. Routine validity transitions are not logged. + +## Validation and performance + +The focused CPU suite exercises allocation generations, ownership, shared +prefixes, exact destinations, read/write order, graph metadata freshness, +failure latching, and the V1/V2 adapters: + +```bash +pytest -q tests/ut/debug/test_kv_guard.py \ + tests/ut/debug/test_kv_guard_runtime.py tests/ut/debug/test_kv_guard_v2.py +``` + +The single-card service regression starts separate disabled and enabled servers, +uses `raise`, repeats a shared prompt, submits concurrent requests, and compares +greedy outputs for both runners in eager and FULL modes: + +```bash +pytest -q tests/e2e/pull_request/one_card/test_kv_cache_guard.py +``` + +CPU mocks establish checker behavior, not NPU execution correctness. Before +merge, record current-revision hardware runs for V1/V2 eager and FULL, prefix +reuse, observed preemption, Mooncake where used, and controlled fault injection. +The service regression does not itself prove that a preemption or prefix-cache +hit occurred; collect those metrics separately. Historical V1 results do not +validate a rebased revision or Runner V2. + +Benchmark disabled and enabled configurations on the same hardware, model, +request trace, sequence lengths, concurrency, seed, and graph mode. Record +successful requests, throughput, TTFT, ITL, CPU usage, and run-to-run variation. +Fix the source commits, actual module load paths, software versions, exact +commands, and logs in the validation report. Do not infer device synchronization +cost from CPU checker timings. + +The scheduler sends changed block/generation tables and compact request ranges. +Expected token slots are temporary and shared across layers in a cache group; +the guard stores no persistent token history, KV snapshots, or content hashes. +Metadata transfers and access-pattern checks are deduplicated, but enabled mode +can still synchronize device metadata to CPU. Measure deployment overhead. + +## Implementation guide and patch lifecycle + +| Component | Responsibility | +| --- | --- | +| `vllm_ascend/debug/kv_guard.py` | Allocation, ownership, prefix validity, and read/write rules. | +| `vllm_ascend/debug/kv_guard_runtime.py` | Independent expected addresses, evidence collection, per-layer audit, and conservative state commit. | +| `vllm_ascend/debug/kv_guard_v2.py` | Sorted batch IDs, execute/update boundaries, and capture restrictions for Runner V2. | +| `vllm_ascend/patch/platform/patch_kv_guard.py` | Conditional scheduler allocation and request-binding instrumentation. | +| Runner and attention implementations | Existing execution boundaries and ordered eager access arguments. | +| Mooncake receiver | Success/failure validity callbacks after receive completion and reformatting. | + +The scheduler patch is conditional because current public hooks do not provide +both allocation generations and the independent scheduler-to-worker request +contract. Its removal path is an upstream allocation-lifecycle and scheduler +output extension that supplies the same evidence without wrapping instance +methods. Changes to those upstream contracts require adapter and lifecycle +regressions, alongside NPU tests for affected execution paths. + +This feature's current scope is address-legality diagnostics. Content tracing, +KV dumps, automatic recovery, and additional cache semantics are separate work. diff --git a/mkdocs.yml b/mkdocs.yml index 094553cd78c1..9abaa25284ce 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -292,6 +292,7 @@ nav: - user_guide/feature_guide/quantization.md - user_guide/feature_guide/sleep_mode.md - user_guide/feature_guide/runtime_guard.md + - user_guide/feature_guide/kv_cache_guard.md - user_guide/feature_guide/structured_output.md - user_guide/feature_guide/lora.md - user_guide/feature_guide/expert_parallelism_load_balancer.md diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml new file mode 100644 index 000000000000..998b026f0464 --- /dev/null +++ b/release_notes/fragments/17833.yaml @@ -0,0 +1,8 @@ +pr: 17833 +entries: + - category: feature + text: >- + Added opt-in KV cache address-legality diagnostics for adapted dense PA + and MLA attention paths in Model Runner V1 and V2. The guard is disabled + by default and can log failures or stop returning worker results after + a violation or incomplete audit. diff --git a/tests/e2e/pull_request/one_card/test_kv_cache_guard.py b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py new file mode 100644 index 000000000000..290f4c3a7c43 --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py @@ -0,0 +1,95 @@ +import json +from concurrent.futures import ThreadPoolExecutor + +import pytest +from vllm.utils.network_utils import get_open_port + +from tests.e2e.conftest import RemoteOpenAIServer, wait_until_npu_memory_free + + +@pytest.mark.e2e_model("Qwen/Qwen3-0.6B") +@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) +@pytest.mark.parametrize( + "graph_mode", + [ + pytest.param( + "NONE", + marks=pytest.mark.e2e_coverage( + arch="dense", + feature="prefix_caching,chunked_prefill", + parallel="TP", + deploy="pd_mix", + hardware="A2", + quantization="BF16", + graph_mode="eager", + ), + ), + pytest.param( + "FULL", + marks=pytest.mark.e2e_coverage( + arch="dense", + feature="prefix_caching,chunked_prefill", + parallel="TP", + deploy="pd_mix", + hardware="A2", + quantization="BF16", + graph_mode="full_graph", + ), + ), + ], +) +@wait_until_npu_memory_free() +def test_guard_preserves_sequential_and_concurrent_generation(use_v2, graph_mode): + model = "Qwen/Qwen3-0.6B" + prompt = "Explain how a key value cache works. " * 40 + arguments = [ + "--max-model-len", + "1024", + "--max-num-seqs", + "4", + "--max-num-batched-tokens", + "256", + "--block-size", + "128", + "--dtype", + "bfloat16", + "--gpu-memory-utilization", + "0.3", + "--enable-prefix-caching", + "--enable-chunked-prefill", + "--compilation-config", + json.dumps({"cudagraph_mode": graph_mode, "cudagraph_capture_sizes": [1, 2, 4]}), + ] + if graph_mode == "NONE": + arguments.append("--enforce-eager") + results = [] + for enabled in (False, True): + port = get_open_port() + environment = { + "VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2)), + "VLLM_ASCEND_KV_GUARD": str(int(enabled)), + "VLLM_ASCEND_KV_GUARD_ACTION": "raise", + } + with ( + RemoteOpenAIServer( + model, + [*arguments, "--port", str(port)], + env_dict=environment, + seed=1234, + server_host="127.0.0.1", + server_port=port, + auto_port=False, + ) as server, + server.get_client() as client, + ): + + def generate(request_index): + response = client.completions.create(model=model, prompt=prompt, temperature=0, max_tokens=8, seed=1234) + assert response.choices[0].text + return response.choices[0].text + + sequential = [generate(request_index) for request_index in range(2)] + with ThreadPoolExecutor(max_workers=4) as executor: + concurrent = list(executor.map(generate, range(4))) + results.append(sequential + concurrent) + assert results[0] == results[1] diff --git a/tests/ut/debug/__init__.py b/tests/ut/debug/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/ut/debug/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/ut/debug/test_kv_guard.py b/tests/ut/debug/test_kv_guard.py new file mode 100644 index 000000000000..ad1e259abb9d --- /dev/null +++ b/tests/ut/debug/test_kv_guard.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""L1 transitions, invalid accesses and exact Case4 address regression.""" + +import importlib.util +import sys +from pathlib import Path + +import pytest + +spec = importlib.util.spec_from_file_location( + "kv_guard_subject", Path(__file__).parents[3] / "vllm_ascend/debug/kv_guard.py" +) +assert spec is not None and spec.loader is not None +subject = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = subject +spec.loader.exec_module(subject) + + +@pytest.fixture +def guard(): + events = [] + guard = subject.KVGuard(lambda event, **fields: events.append({"event": event, **fields})) + guard.configure({0: 128}) + guard.events = events + return guard + + +def bind(guard, request="r", computed=0, epoch=1, pool="pool", ids=(1,)): + guard.bind( + { + "pool_id": pool, + "requests": {request: [[(block, epoch) for block in ids]]}, + "scheduled_requests": {request: (computed, 0)}, + } + ) + + +def access(guard, slots=(), length=0, ids=(1,), request_ids=("r",), phase="forward"): + guard.observe( + [{"group_id": 0, "slots": list(slots), "block_table": [list(ids)], "seq_lens": [length]}], + list(request_ids), + phase, + ) + + +def violations(guard): + return [e for e in guard.events if e["event"] == "guard.violation"] + + +def test_partial_fill_seal_and_read(guard): + bind(guard) + access(guard, range(128, 147), 19) + assert guard.blocks[0, 1].valid_end == 19 + access(guard, range(147, 256), 128) + assert guard.blocks[0, 1].valid_end == 128 + assert not violations(guard) + + +def test_unknown_read_and_skipped_write(guard): + bind(guard) + access(guard, (), 1) + assert violations(guard)[-1]["reason"] == "read_unwritten_range" + access(guard, [130]) + assert violations(guard)[-1]["reason"] == "noncontiguous_write" + assert guard.blocks[0, 1].valid_end == subject.INVALID_PREFIX + + +def test_case4_loaded_prefix_dummy_write_then_dirty_read_without_tensors(guard): + bind(guard) + guard.load({"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) + assert guard.blocks[0, 1].valid_end == 19 + access(guard, [142], 1, request_ids=(), phase="dummy") + dirty_write = violations(guard)[0] + assert dirty_write["reason"] == "dummy_write_to_live_block" + assert dirty_write["block_id"] == 1 and dirty_write["offsets"] == [14] + assert dirty_write["valid_end"] == 19 and dirty_write["owner_request_ids"] == ["r"] + access(guard, (), 19) + assert violations(guard)[-1]["reason"] == "read_invalid_block" + + +def test_loaded_tail_allows_append_but_not_overwrite(guard): + bind(guard, computed=19) + access(guard, [147], 20) + assert not violations(guard) + access(guard, [142], 20) + assert violations(guard)[0]["reason"] == "overwrite_valid_prefix" + + +def test_reallocation_epoch_clears_old_validity_and_taint(guard): + bind(guard, computed=19) + access(guard, [142], phase="dummy", request_ids=()) + bind(guard, request="new", epoch=2) + assert guard.blocks[0, 1].valid_end == 0 + assert guard.blocks[0, 1].owners == {"new"} + guard.events.clear() + access(guard, [128], 1, request_ids=("new",)) + assert not violations(guard) + + +def test_shared_prefix_and_finished_request(guard): + bind(guard, computed=128) + bind(guard, request="shared", computed=128) + access(guard, (), 128, request_ids=("shared",)) + assert not violations(guard) + guard.bind({"released_request_ids": ["r"], "requests": {}}) + assert guard.blocks[0, 1].owners == {"shared"} + access(guard, (), 128) + assert violations(guard)[-1]["reason"] == "unowned_read" + + +def test_cache_lookup_does_not_revive_tainted_same_epoch(guard): + bind(guard, computed=19) + access(guard, [142], phase="dummy", request_ids=()) + bind(guard, request="prefix-hit", computed=19) + assert guard.blocks[0, 1].valid_end == subject.INVALID_PREFIX + + +def test_load_before_scheduler_binding_and_failed_load(guard): + guard.load({"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) + bind(guard) + assert not guard.pending_loads + assert guard.blocks[0, 1].valid_end == 19 + guard.fail_load("r") + access(guard, (), 19) + assert violations(guard)[-1]["reason"] == "read_invalid_block" + + +def test_pending_load_never_applied_to_different_request(guard): + guard.load({"request_id": "old", "local_block_ids": [[1]], "num_computed_tokens": 19}) + bind(guard, request="new", epoch=2) + assert guard.blocks[0, 1].valid_end == 0 + guard.bind({"released_request_ids": ["old"]}) + assert not guard.pending_loads + + +def test_no_per_token_state_log_and_negative_padding(guard): + bind(guard) + for offset in range(64): + access(guard, [128 + offset, -1], offset + 1) + assert not violations(guard) + assert len([e for e in guard.events if e["event"] == "guard.state"]) == 0 + + +def test_null_block_write_is_not_padding_and_taints_later_read(guard): + access(guard, [-1, 0, 14], ids=(0,), request_ids=(), phase="dummy") + assert violations(guard)[0]["reason"] == "write_null_block" + assert violations(guard)[0]["offsets"] == [0, 14] + access(guard, (), 1, ids=(0,)) + assert violations(guard)[-1]["reason"] == "read_null_block" + guard.clear(0, 0) + access(guard, (), 1, ids=(0,)) + assert violations(guard)[-1]["reason"] == "read_null_block" + + +def test_group_pool_and_allocation_isolation(guard): + guard.configure({0: 128, 1: 128}) + bind(guard, computed=19) + guard.bind({"pool_id": "pool", "requests": {"r2": [[], [(1, 1)]]}, "scheduled_requests": {"r2": (7, 0)}}) + assert guard.blocks[0, 1].valid_end == 19 + assert guard.blocks[1, 1].valid_end == 7 + bind(guard, pool="new-pool") + assert guard.blocks[0, 1].valid_end == 0 + assert (1, 1) not in guard.blocks + + +def test_copy_clear_and_read_invalid_destination(guard): + bind(guard, computed=19) + bind(guard, request="dest", ids=(2,)) + guard.copy_blocks(0, [(1, 2)]) + assert guard.blocks[0, 2].valid_end == 19 + access(guard, (), 19, ids=(2,), request_ids=("dest",)) + assert not violations(guard) + guard.clear(0, 1) + guard.copy_blocks(0, [(1, 2)]) + assert violations(guard)[-1]["reason"] == "invalid_copy_source_or_destination" + access(guard, (), 19, ids=(2,), request_ids=("dest",)) + assert violations(guard)[-1]["reason"] == "read_invalid_block" + + +def test_recycling_a_live_block_is_a_lifecycle_violation(guard): + bind(guard, computed=19) + bind(guard, request="new", epoch=2) + assert violations(guard)[0]["reason"] == "recycle_live_block" + + +def test_preempted_request_release_allows_reallocation_without_false_violation(guard): + bind(guard, computed=19) + guard.bind({"released_request_ids": ["r"]}) + access(guard, (), 19) + assert violations(guard)[-1]["reason"] == "unowned_read" + guard.events.clear() + bind(guard, request="new", epoch=2) + access(guard, [128], 1, request_ids=("new",)) + assert not violations(guard) + + +def test_block_copy_uses_original_sources_for_overlapping_pairs(guard): + bind(guard, computed=19) + bind(guard, request="b", computed=7, ids=(2,)) + bind(guard, request="c", ids=(3,)) + guard.copy_blocks(0, [(1, 2), (2, 3)]) + assert guard.blocks[0, 2].valid_end == 19 + assert guard.blocks[0, 3].valid_end == 7 + assert not violations(guard) + + +def test_missing_allocation_generation_does_not_make_a_block_writable(guard): + bind(guard, epoch=0) + assert (0, 1) not in guard.blocks + assert any(e.get("reason") == "allocation_generation_missing" for e in guard.events) + access(guard, [128]) + assert violations(guard)[-1]["reason"] == "unallocated_block" + + +def test_copying_an_empty_prefix_does_not_make_bytes_readable(guard): + bind(guard) + bind(guard, request="dest", ids=(2,)) + guard.copy_blocks(0, [(1, 2)]) + assert not violations(guard) + access(guard, (), 1, ids=(2,), request_ids=("dest",)) + assert violations(guard)[-1]["reason"] == "read_unwritten_range" + + +def test_replaced_pool_drops_old_pending_loads(guard): + bind(guard) + guard.load({"request_id": "next", "local_block_ids": [[1]], "num_computed_tokens": 19}) + bind(guard, request="next", pool="replacement") + assert not guard.pending_loads + assert guard.blocks[0, 1].valid_end == 0 diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py new file mode 100644 index 000000000000..f9c4a02cfeca --- /dev/null +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -0,0 +1,836 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runner plans, ordered operator arguments, graph replay and allocation transport.""" + +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType +from types import SimpleNamespace as NS +from typing import Any, cast + +import pytest +import torch + +ROOT = Path(__file__).parents[3] + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, ROOT / path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture(name="runtime") +def runtime_fixture(monkeypatch): + package = ModuleType("vllm_ascend") + package.__path__ = [] + cast(Any, package).envs = NS(VLLM_ASCEND_KV_GUARD=True, VLLM_ASCEND_KV_GUARD_ACTION="log") + debug = ModuleType("vllm_ascend.debug") + debug.__path__ = [] + monkeypatch.setitem(sys.modules, "vllm_ascend", package) + monkeypatch.setitem(sys.modules, "vllm_ascend.debug", debug) + state = load("guard_runtime_state", "vllm_ascend/debug/kv_guard.py") + monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard", state) + subject = load("guard_runtime_subject", "vllm_ascend/debug/kv_guard_runtime.py") + logger = ModuleType("vllm.logger") + cast(Any, logger).logger = logging.Logger("guard-test") + vllm = ModuleType("vllm") + vllm.__path__ = [] + monkeypatch.setitem(sys.modules, "vllm", vllm) + monkeypatch.setitem(sys.modules, "vllm.logger", logger) + config = NS() + runtime = subject.get_guard(config) + assert subject.get_guard(config) is runtime + runtime.subject = subject + runtime.state.configure({0: 128}) + runtime.layer_groups = {"layer0": 0, "layer1": 0} + runtime.events = [] + runtime.state.emit = lambda event, **fields: runtime.events.append({"event": event, **fields}) + runtime.gap = lambda reason, **fields: runtime.events.append( + {"event": "guard.coverage_gap", "reason": reason, **fields} + ) + return runtime + + +def bind(runtime, request="r", block=1, computed=0): + runtime.state.bind( + {"pool_id": "pool", "requests": {request: [[(block, 1)]]}, "scheduled_requests": {request: (computed, 0)}} + ) + + +def test_access_hooks_validate_case4_without_kv_copy_or_hash(runtime): + bind(runtime, computed=19) + slots, table = torch.tensor([142]), torch.tensor([[1]]) + runner, context = NS(input_batch=NS(req_ids=[])), NS(cudagraph_runtime_mode="NONE") + + def dummy(n): + runtime.write("layer0", slots) + runtime.read("layer0", table, [1]) + return "unchanged-model-result" + + model = runtime.wrap_forward(runner, dummy, lambda: context) + assert model(1, _kv_guard_phase="dummy") == "unchanged-model-result" + assert any(e.get("reason") == "dummy_write_to_live_block" and e["offsets"] == [14] for e in runtime.events) + runner.input_batch.req_ids = ["r"] + runtime.wrap_forward(runner, lambda n: runtime.read("layer0", table, [19]), lambda: context)(1) + assert any(e.get("reason") == "read_invalid_block" for e in runtime.events) + + +def test_full_graph_without_current_metadata_reports_a_gap(runtime): + bind(runtime, computed=19) + runner, context = NS(input_batch=NS(req_ids=[])), NS(cudagraph_runtime_mode="FULL") + + def replay(n): + runtime.write("layer0", torch.tensor([142])) + runtime.read("layer0", torch.tensor([[1]]), [19]) + + runtime.wrap_forward(runner, replay, lambda: context)(1, _kv_guard_phase="dummy") + assert any(e.get("reason") == "full_graph_metadata_missing" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 19 + + +def test_layer_hook_duplicates_do_not_look_like_repeated_writes(runtime): + bind(runtime) + slots, table = torch.arange(128, 147), torch.tensor([[1]]) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + for layer in ("layer0", "layer1"): + runtime.write(layer, slots[:19]) + runtime.read(layer, table, [19]) + + runtime.wrap_forward(runner, model, lambda: context)(19) + assert not any(e["event"] == "guard.violation" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 19 + + +def test_split_decode_prefill_write_ranges_and_read_request_offsets(runtime): + bind(runtime, "decode", 1, 19) + bind(runtime, "prefill", 2, 10) + slots = torch.tensor([147, 266, 267]) + runner, context = NS(input_batch=NS(req_ids=["decode", "prefill"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", slots[:1]) + runtime.write("layer0", slots[1:]) + runtime.read("layer0", torch.tensor([[1]]), [20]) + runtime.read("layer0", torch.tensor([[2]]), [12], request_offset=1) + + runtime.wrap_forward(runner, model, lambda: context)(3) + assert not any(e["event"] == "guard.violation" for e in runtime.events) + + +def test_receive_only_step_binds_before_execute_and_cleans_finished(runtime): + output = NS( + _ascend_kv_guard_context={"pool_id": "pool", "requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (0, 0)}} + ) + runtime.wrap_execute(lambda message: None)(output) + runtime.transfer("load", {"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) + assert runtime.state.blocks[0, 1].valid_end == 19 + runtime.wrap_execute(lambda message: None)( + NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": ["r"], "scheduled_requests": {}}) + ) + assert runtime.state.blocks[0, 1].owners == set() + + +def test_disabled_guard_does_not_create_runtime(runtime, monkeypatch): + sys.modules["vllm_ascend"].envs.VLLM_ASCEND_KV_GUARD = False + config = NS() + assert runtime.subject.get_guard(config) is None + assert not hasattr(config, "_ascend_kv_guard") + + +@pytest.mark.parametrize("kernel_size, accepted", [(128, True), (16, False)]) +def test_unadapted_kernel_block_coordinates_are_rejected(runtime, kernel_size, accepted): + impl = type("AscendAttentionBackendImpl", (), {})() + runtime.config = NS( + parallel_config=NS(), + kv_transfer_config=None, + model_config=NS(), + speculative_config=None, + compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), + ) + runner = NS( + input_batch=NS(block_table=NS(block_tables=[NS(block_size=kernel_size)])), + kv_cache_config=NS( + kv_cache_groups=[NS(kv_cache_spec=NS(block_size=128, head_size=64), layer_names=["layer0"])] + ), + ) + if accepted: + runtime.configure(runner) + assert impl._kv_guard is runtime + else: + with pytest.raises(ValueError, match="unadapted"): + runtime.configure(runner) + + +@pytest.mark.parametrize("use_v2", [False, True]) +def test_scheduler_preserves_output_and_sends_only_changed_generations(runtime, use_v2): + patch = load("guard_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") + + class Pool: + blocks = [NS(block_id=i) for i in range(3)] + null_block = blocks[0] + + def get_new_blocks(self, count): + return [self.blocks[1]] + + class Scheduler: + def __init__(self): + self.vllm_config = NS(use_v2_model_runner=use_v2) + self.pool = Pool() + self.mapping = {"r": [[1]]} + self.kv_cache_manager = NS( + block_pool=self.pool, get_block_ids=lambda request: self.mapping[request], free=lambda request: None + ) + self.output = NS( + scheduled_new_reqs=[NS(req_id="r", num_computed_tokens=0)], + scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), + num_scheduled_tokens={"r": 1}, + finished_req_ids=set(), + ) + + def schedule(self): + return self.output + + patch.install_scheduler_guard(Scheduler) + patch.install_scheduler_guard(Scheduler) + scheduler = Scheduler() + scheduler.pool.get_new_blocks(1) + assert scheduler.schedule() is scheduler.output + binding = scheduler.output._ascend_kv_guard_context + assert binding["requests"]["r"] == [[(1, 1)]] + assert binding["scheduled_requests"] == {"r": (0, 1)} + scheduler.schedule() + assert scheduler.output._ascend_kv_guard_context["requests"] == {} + scheduler.pool.get_new_blocks(1) + scheduler.schedule() + assert scheduler.output._ascend_kv_guard_context["requests"]["r"] == [[(1, 2)]] + scheduler.kv_cache_manager.free(NS(request_id="r")) + scheduler.schedule() + assert scheduler.output._ascend_kv_guard_context["released_request_ids"] == ["r"] + assert "r" in scheduler.output._ascend_kv_guard_context["requests"] + scheduler.output.finished_req_ids = {"r"} + scheduler.output.num_scheduled_tokens = {} + scheduler.schedule() + assert scheduler.output._ascend_kv_guard_context == { + "pool_id": binding["pool_id"], + "requests": {}, + "released_request_ids": ["r"], + "scheduled_requests": {}, + } + + +@pytest.mark.parametrize("slots, mismatch", [([128, 129], False), ([129, 128], True)]) +def test_full_graph_update_audits_final_metadata_without_guard_callbacks(runtime, monkeypatch, slots, mismatch): + order: list[Any] = [] + monkeypatch.setattr( + torch, + "npu", + NS( + NPUGraph=object, + graph_task_update_begin=lambda *args: order.append("begin"), + graph_task_update_end=lambda *args: order.append("end"), + ), + raising=False, + ) + monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) + utils = ModuleType("vllm_ascend.utils") + cast(Any, utils).weak_ref_tensors = lambda value: value + monkeypatch.setitem(sys.modules, "vllm_ascend.utils", utils) + graph = load("guard_graph_update_subject", "vllm_ascend/compilation/updatable_graph.py") + bind(runtime) + runtime.layer_groups = {"layer0": 0} + runtime._scheduled = {"r": (0, 2)} + item = metadata(torch.tensor([-1, -1]), torch.tensor([[1]]), torch.tensor([0])) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) + + def operation(**kwargs): + order.append(("operation", kwargs["context_lens"].tolist())) + item.slot_mapping[:] = torch.tensor(slots) + item.seq_lens[:] = kwargs["context_lens"] + + task = graph.GraphUpdateTask( + operation=operation, + kwargs={"block_table": item.block_tables, "context_lens": torch.tensor([0])}, + provider=NS(layer_name="layer0"), + provider_index=0, + handle="handle", + event=NS(record=lambda stream: order.append("event")), + ) + bound = task.bind({"context_lens": torch.tensor([2])}) + + def replay(count): + bound.apply("stream") + assert not runtime._accesses + return "result" + + assert runtime.wrap_forward(runner, replay, lambda: context)(2) == "result" + assert order == ["begin", ("operation", [2]), "end", "event"] + assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == mismatch + assert not any(event["event"] == "guard.coverage_gap" for event in runtime.events) + if not mismatch: + assert runtime.state.blocks[0, 1].valid_end == 2 + + +def test_v1_dummy_wrapper_preserves_decoder_replay_argument(runtime): + calls = [] + runner = NS(input_batch=NS(req_ids=[])) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata=None) + + def forward(count, *, decoder_replay_dummy): + calls.append((count, decoder_replay_dummy)) + return "unchanged" + + wrapped = runtime.wrap_forward(runner, forward, lambda: context) + assert wrapped(8, decoder_replay_dummy=True, _kv_guard_phase="warmup") == "unchanged" + assert calls == [(8, True)] + assert not runtime.events + + +def test_repeated_violations_are_deduplicated_without_hiding_other_blocks(runtime): + messages = [] + runtime.logger = NS( + info=lambda *args: messages.append(("info", args)), + warning=lambda *args: messages.append(("warning", args)), + error=lambda *args: messages.append(("error", args)), + ) + runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=1) + runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=1) + runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=2) + assert [level for level, args in messages].count("error") == 2 + + +def test_eager_raw_key_read_does_not_require_paged_cache_read_metadata(runtime): + bind(runtime) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", torch.tensor([128])) + runtime.read("layer0", None, [1]) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not any(e["event"] == "guard.coverage_gap" for e in runtime.events) + + +def test_same_layer_double_write_is_not_deduplicated(runtime): + bind(runtime) + slots = torch.tensor([128]) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", slots) + runtime.write("layer0", slots) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert any(e.get("reason") == "overwrite_valid_prefix" for e in runtime.events) + + +def test_equivalent_cloned_slots_across_layers_remain_one_logical_write(runtime): + bind(runtime) + slots = torch.tensor([128]) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", slots) + runtime.write("layer1", slots.clone()) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert not any(e["event"] == "guard.violation" for e in runtime.events) + + +def metadata(slots, table=None, lengths=None): + return NS(slot_mapping=slots, block_tables=table, seq_lens=lengths, num_actual_tokens=len(slots)) + + +def test_runner_plan_detects_case4_without_any_operator_hook(runtime): + bind(runtime, computed=19) + item = metadata(torch.tensor([142]), torch.tensor([[1]]), torch.tensor([19])) + runner = NS(input_batch=NS(req_ids=[])) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runtime.wrap_forward(runner, lambda n: "result", lambda: context)(1, _kv_guard_phase="dummy") + violation = next(e for e in runtime.events if e.get("reason") == "dummy_write_to_live_block") + assert violation["evidence"] == "execution_plan" and violation["offsets"] == [14] + assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX + + +def test_final_metadata_is_read_after_forward_parameter_updates_and_copied_once(runtime, monkeypatch): + bind(runtime) + slots, table, lengths = torch.tensor([-1]), torch.tensor([[1]]), torch.tensor([0]) + item = metadata(slots, table, lengths) + runner = NS(input_batch=NS(req_ids=["r"]), kv_caches=object()) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) + copies = [] + original = runtime._list + + def record_copy(tensor): + copies.append(tensor.data_ptr()) + return original(tensor) + + monkeypatch.setattr(runtime, "_list", record_copy) + + def replay(n): + assert not copies # No D2H before asynchronous graph parameter updates. + assert runtime.state.blocks[0, 1].valid_end == 0 + slots[0], lengths[0] = 128, 1 + return "result" + + assert runtime.wrap_forward(runner, replay, lambda: context)(1) == "result" + assert runtime.state.blocks[0, 1].valid_end == 1 + assert len(copies) == 3 and len(set(copies)) == 3 + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_current_full_graph_metadata_ignores_operator_callbacks(runtime): + bind(runtime, computed=19) + item = metadata(torch.tensor([147]), torch.tensor([[1]]), [20]) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor([142])), lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 20 + assert not any(e["event"] == "guard.violation" for e in runtime.events) + + +def test_unadapted_ring_cache_is_rejected_instead_of_silently_reducing_checks(runtime): + spec = type("CircularBufferSpec", (), {"block_size": 128, "head_size": 64})() + runtime.config = NS( + parallel_config=NS(), + kv_transfer_config=None, + model_config=NS(is_hybrid=True), + speculative_config=None, + compilation_config=NS(static_forward_context={"ring": NS(kv_cache=object())}), + ) + runner = NS( + input_batch=NS(req_ids=[]), kv_cache_config=NS(kv_cache_groups=[NS(kv_cache_spec=spec, layer_names=["ring"])]) + ) + with pytest.raises(ValueError, match="unadapted cache validity"): + runtime.configure(runner) + + +def test_missing_layer_metadata_does_not_advance_group_validity(runtime): + bind(runtime) + item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) + assert any(e.get("reason") == "metadata_write_mapping_missing:layer1" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_operator_read_before_write_cannot_be_legalized_by_later_write(runtime): + bind(runtime) + table = torch.tensor([[1]]) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.read("layer0", table, [1]) + runtime.write("layer0", torch.tensor([128])) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert any(e.get("reason") == "read_unwritten_range" for e in runtime.events) + + +def test_later_illegal_write_does_not_retroactively_invalidate_earlier_read(runtime): + bind(runtime, computed=19) + runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.read("layer0", torch.tensor([[1]]), [19]) + runtime.write("layer0", torch.tensor([142])) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert any(e.get("reason") == "overwrite_valid_prefix" for e in runtime.events) + assert not any(e.get("reason") == "read_invalid_block" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX + + +def test_one_layers_write_cannot_hide_another_layers_missing_write(runtime): + bind(runtime) + slots, table = torch.tensor([128]), torch.tensor([[1]]) + item = metadata(slots, table, [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(n): + runtime.write("layer0", slots) + runtime.read("layer0", table, [1]) + runtime.read("layer1", table, [1]) + + runtime.wrap_forward(runner, model, lambda: context)(1) + violation = next(e for e in runtime.events if e.get("reason") == "read_unwritten_range") + assert violation["layer"] == "layer1" and violation["evidence"] == "operator_arguments" + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_missing_known_operator_layer_is_reported_and_not_committed(runtime): + bind(runtime) + item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + runtime.wrap_forward(runner, lambda n: runtime.write("layer0", item.slot_mapping), lambda: context)(1) + assert any(e.get("reason") == "operator_layer_not_observed:layer1" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_failed_forward_preserves_exception_and_does_not_commit_planned_writes(runtime): + bind(runtime) + item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(n): + raise RuntimeError("model failure") + + with pytest.raises(RuntimeError, match="model failure"): + runtime.wrap_forward(runner, model, lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 0 and runtime._phase is None + + +def test_only_needed_block_table_rows_and_columns_are_copied(runtime, monkeypatch): + bind(runtime) + table = torch.full((32, 1024), 99) + table[0, 0] = 1 + item = metadata(torch.tensor([128]), table, [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + shapes = [] + original = runtime._list + + def record_shape(value): + shapes.append(tuple(getattr(value, "shape", ()))) + return original(value) + + monkeypatch.setattr(runtime, "_list", record_shape) + runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) + assert (1, 1) in shapes and (32, 1024) not in shapes + assert not any(e["event"] == "guard.violation" for e in runtime.events) + + +def test_unknown_backend_is_rejected_without_mutating_it(runtime): + class OtherAttentionImpl: + __slots__ = () + + impl = OtherAttentionImpl() + runtime.layer_groups.clear() + runtime.config = NS( + parallel_config=NS(), + kv_transfer_config=None, + model_config=NS(), + speculative_config=None, + compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), + ) + runner = NS( + input_batch=NS(req_ids=["r"]), + kv_cache_config=NS( + kv_cache_groups=[NS(kv_cache_spec=NS(block_size=128, head_size=64), layer_names=["layer0"])] + ), + ) + with pytest.raises(ValueError, match="unadapted attention implementation"): + runtime.configure(runner) + assert not hasattr(impl, "_kv_guard") + + +def test_eager_uses_operator_arguments_without_building_unused_plans(runtime, monkeypatch): + bind(runtime) + slots, table = torch.tensor([128]), torch.tensor([[1]]) + item = metadata(slots, table, [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + + def unused_plan(*args): + raise AssertionError("eager must not materialize an unused metadata plan") + + monkeypatch.setattr(runtime, "_metadata_accesses", unused_plan) + + def model(n): + for layer in ("layer0", "layer1"): + runtime.write(layer, slots) + runtime.read(layer, table, [1]) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_full_metadata_skips_redundant_hooks_and_builds_one_shared_plan(runtime, monkeypatch): + bind(runtime) + slots, table = torch.tensor([128]), torch.tensor([[1]]) + item = metadata(slots, table, [1]) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["r"])) + parsed = [] + original = runtime._metadata_accesses + + def parse(layer, value): + parsed.append(layer) + return original(layer, value) + + monkeypatch.setattr(runtime, "_metadata_accesses", parse) + + def replay(n): + for layer in ("layer0", "layer1"): + runtime.write(layer, slots) + runtime.read(layer, table, [1]) + assert not runtime._accesses + + runtime.wrap_forward(runner, replay, lambda: context)(1) + assert len(parsed) == 1 + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_unique_layer_patterns_share_owner_snapshot_and_reuse_checker(runtime, monkeypatch): + bind(runtime) + slots, table = torch.tensor([128]), torch.tensor([[1]]) + context = NS(cudagraph_runtime_mode="NONE") + runner = NS(input_batch=NS(req_ids=["r"])) + owner_copies = [] + original = runtime.subject.replace + scratch = runtime._scratch + + def copied(block, **fields): + if "owners" in fields: + owner_copies.append(block) + return original(block, **fields) + + monkeypatch.setattr(runtime.subject, "replace", copied) + + def unexpected_checker(*args): + raise AssertionError("the scratch checker must be reused") + + monkeypatch.setattr(runtime.subject, "KVGuard", unexpected_checker) + + def model(n): + for layer in ("layer0", "layer1"): + runtime.write(layer, slots.clone()) + runtime.read(layer, table, [1]) + + runtime.wrap_forward(runner, model, lambda: context)(1) + assert runtime._scratch is scratch and len(owner_copies) == 1 + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_execute_and_update_bind_once_without_skipping_receive_only_or_next_step(runtime, monkeypatch): + output = NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": []}) + runner = NS(requests={}) + bound = [] + original = runtime.state.bind + + def record(binding): + bound.append(binding) + return original(binding) + + monkeypatch.setattr(runtime.state, "bind", record) + update = runtime.wrap_update(runner, lambda value: None) + execute = runtime.wrap_execute(update) + execute(output) + assert len(bound) == 1 and runtime._scheduled is None + execute(output) + assert len(bound) == 2 and runtime._scheduled is None + runtime.wrap_execute(lambda value: None)(output) + assert len(bound) == 3 and runtime._scheduled is None + + +def test_shared_unsupported_metadata_reports_missing_read_coverage_for_each_group(runtime): + runtime.layer_groups = {"layer0": 0, "layer1": 1} + runtime.state.configure({0: 128, 1: 128}) + item = metadata(torch.tensor([-1])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=[])) + runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) + assert runtime._missing_groups == {0, 1} + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +@pytest.mark.parametrize("swapped", [False, True]) +def test_scheduler_plan_detects_cross_request_slot_swaps(runtime, mode, swapped): + bind(runtime, "A", 1) + bind(runtime, "B", 2) + slots = torch.tensor([256, 128] if swapped else [128, 256]) + table = torch.tensor([[1], [2]]) + item = metadata(slots, table, [1, 1]) + context = NS(cudagraph_runtime_mode=mode, attn_metadata={"layer0": item, "layer1": item}) + runner = NS(input_batch=NS(req_ids=["A", "B"])) + output = NS(_ascend_kv_guard_context={"requests": {}, "scheduled_requests": {"A": (0, 1), "B": (0, 1)}}) + + def model(n): + for layer in runtime.layer_groups: + runtime.write(layer, slots) + runtime.read(layer, table, [1, 1]) + + forward = runtime.wrap_forward(runner, model, lambda: context) + runtime.wrap_execute(lambda message: forward(2))(output) + violations = [e for e in runtime.events if e["event"] == "guard.violation"] + if swapped: + mismatch = next(e for e in violations if e["reason"] == "write_target_mismatch") + assert mismatch["request_id"] == "A" and mismatch["slot"] == 256 and mismatch["expected_slot"] == 128 + assert all(b.valid_end == runtime.subject.INVALID_PREFIX for b in runtime.state.blocks.values()) + else: + assert not violations + assert runtime._scheduled is None + + +def test_split_mixed_batch_uses_explicit_token_offset_and_runner_request_order(runtime): + bind(runtime, "prefill", 2, 10) + bind(runtime, "decode", 1, 19) + runtime._scheduled = {"prefill": (10, 2), "decode": (19, 1)} + runner = NS(input_batch=NS(req_ids=["decode", "prefill"])) + context = NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", torch.tensor([266, 267]), token_offset=1) + runtime.write("layer0", torch.tensor([147])) + runtime.read("layer0", torch.tensor([[1], [2]]), [20, 12]) + + runtime.wrap_forward(runner, model, lambda: context)(3) + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +@pytest.mark.parametrize( + "slots, mismatch", [([128, 129, -1], False), ([129, 128, -1], True), ([128, -1, -1], True), ([128, 129, 256], True)] +) +def test_exact_token_positions_distinguish_permutation_missing_write_and_padding(runtime, slots, mismatch): + bind(runtime) + runtime._scheduled = {"r": (0, 2)} + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="NONE") + runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor(slots)), lambda: context)(3) + assert any(e.get("reason") == "write_target_mismatch" for e in runtime.events) == mismatch + + +def test_missing_scheduler_range_does_not_advance_validity(runtime): + bind(runtime) + runtime._scheduled = {"r": (None, 1)} + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="NONE") + runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor([128])), lambda: context)(1) + assert any(e.get("reason") == "scheduler_write_range_missing" and e["request_id"] == "r" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_raise_mode_stops_output_and_remains_latched_after_caller_catches(runtime): + bind(runtime, computed=19) + runtime.action = "raise" + runtime.state.emit = runtime.report + runner = NS(input_batch=NS(req_ids=[])) + context = NS(cudagraph_runtime_mode="NONE") + completed = [] + + def model(n): + runtime.write("layer0", torch.tensor([142])) + completed.append("kernel submitted") + return "must not escape" + + forward = runtime.wrap_forward(runner, model, lambda: context) + with pytest.raises(runtime.subject.KVGuardError, match="dummy_write_to_live_block"): + forward(1, _kv_guard_phase="dummy") + assert completed == ["kernel submitted"] # Auditing cannot undo submitted operations. + with pytest.raises(runtime.subject.KVGuardError): + forward(1, _kv_guard_phase="dummy") + assert completed == ["kernel submitted"] and runtime._phase is None + + +def test_raise_mode_stops_on_coverage_gap_before_execution(runtime): + runtime.action = "raise" + runtime.gap = runtime.subject.GuardRuntime.gap.__get__(runtime) + called = [] + with pytest.raises(runtime.subject.KVGuardError, match="scheduler_write_plan_missing"): + runtime.wrap_execute(lambda message: called.append(True))(NS(_ascend_kv_guard_context={"requests": {}})) + assert not called + + +def test_physical_zero_and_copy_preserve_only_proven_validity(runtime): + bind(runtime, "source", 1, 19) + output = NS( + _ascend_kv_guard_context={ + "pool_id": "pool", + "requests": {"dest": [[(2, 1)]]}, + "scheduled_requests": {"dest": (19, 0)}, + }, + new_block_ids_to_zero=[2], + kv_cache_block_copies=[NS(src_block_id=1, dst_block_id=2)], + ) + + def execute(message): + assert runtime.state.blocks[0, 2].valid_end == 0 + runtime.zero_blocks(message.new_block_ids_to_zero) + assert runtime.state.blocks[0, 2].valid_end == 0 + runtime.copy_blocks(message.kv_cache_block_copies) + assert runtime.state.blocks[0, 2].valid_end == 19 + + runtime.wrap_execute(execute)(output) + assert not any(e["event"] == "guard.violation" for e in runtime.events) + runtime.zero_blocks([2]) + assert runtime.state.blocks[0, 2].valid_end == runtime.subject.INVALID_PREFIX + assert any(e.get("reason") == "clear_live_block" for e in runtime.events) + + +def test_read_block_order_within_one_request_is_checked_against_scheduler(runtime): + runtime.state.bind( + {"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}, "scheduled_requests": {"r": (256, 1)}} + ) + runtime._scheduled = {"r": (256, 1)} + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="NONE") + runtime.wrap_forward(runner, lambda n: runtime.read("layer0", torch.tensor([[2, 1]]), [256]), lambda: context)(1) + assert any(e.get("reason") == "read_target_mismatch" for e in runtime.events) + + +def test_shared_prefix_read_and_tail_append_remain_legal(runtime): + runtime.state.bind( + { + "pool_id": "pool", + "requests": {"A": [[(1, 1), (2, 1)]], "B": [[(1, 1), (3, 1)]]}, + "scheduled_requests": {"A": (128, 1), "B": (128, 1)}, + } + ) + runtime._scheduled = {"A": (128, 1), "B": (128, 1)} + runner = NS(input_batch=NS(req_ids=["B", "A"])) + context = NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", torch.tensor([384, 256])) + runtime.read("layer0", torch.tensor([[1, 3], [1, 2]]), [129, 129]) + + runtime.wrap_forward(runner, model, lambda: context)(2) + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_truncated_write_observation_is_incomplete_even_if_read_only_uses_old_prefix(runtime): + bind(runtime, computed=19) + runtime._scheduled = {"r": (19, 2)} + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="NONE") + + def model(n): + runtime.write("layer0", torch.tensor([147])) + runtime.read("layer0", torch.tensor([[1]]), [19]) + + runtime.wrap_forward(runner, model, lambda: context)(2) + assert any(e.get("reason") == "write_tokens_not_observed" for e in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 19 + + +def test_copy_to_null_or_unbound_block_cannot_silently_pass(runtime): + bind(runtime, computed=19) + runtime.copy_blocks([NS(src_block_id=1, dst_block_id=0), NS(src_block_id=1, dst_block_id=99)]) + assert any(e.get("reason") == "write_null_block" and e["operation"] == "copy" for e in runtime.events) + assert any(e.get("reason") == "copy_destination_unbound" and e["block_id"] == 99 for e in runtime.events) + + +def test_scheduler_optional_zero_and_copy_fields_accept_none(runtime): + output = NS( + _ascend_kv_guard_context={"requests": {}, "scheduled_requests": {}}, + new_block_ids_to_zero=None, + kv_cache_block_copies=None, + ) + assert runtime.wrap_execute(lambda message: "executed")(output) == "executed" + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) diff --git a/tests/ut/debug/test_kv_guard_v2.py b/tests/ut/debug/test_kv_guard_v2.py new file mode 100644 index 000000000000..2e8581d865ab --- /dev/null +++ b/tests/ut/debug/test_kv_guard_v2.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MRV2 batch ordering, graph updates, dummy writes and lifecycle integration.""" + +import sys +from types import SimpleNamespace as NS +from typing import Any + +import pytest +import torch + +from .test_kv_guard_runtime import bind, load, metadata +from .test_kv_guard_runtime import runtime_fixture as runtime_fixture + + +@pytest.fixture +def adapter(runtime, monkeypatch): + monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_runtime", runtime.subject) + subject = load("guard_v2_subject", "vllm_ascend/debug/kv_guard_v2.py") + impl = type("AscendAttentionBackendImpl", (), {})() + runtime.config = NS( + parallel_config=NS(), + model_config=NS(), + speculative_config=None, + kv_transfer_config=None, + compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), + _ascend_kv_guard=runtime, + ) + runner: Any = NS( + vllm_config=runtime.config, + block_tables=NS(block_sizes=[128], kernel_block_sizes=[128]), + kv_cache_config=NS( + kv_cache_groups=[NS(layer_names=["layer0"], kv_cache_spec=NS(block_size=128, head_size=64))] + ), + model_state=NS(prepare_attn=lambda batch, mode: runner.next_metadata), + next_metadata={}, + ) + runtime.layer_groups.clear() + result = subject.KVGuardV2(runner, runtime) + result.configure() + result.subject = subject + return result + + +def message(requests=None, scheduled=None, released=()): + scheduled = {"r": (0, 1)} if scheduled is None else scheduled + return NS( + total_num_scheduled_tokens=sum(count for _, count in scheduled.values()), + _ascend_kv_guard_context={ + "pool_id": "pool", + "requests": {"r": [[(1, 1)]]} if requests is None else requests, + "scheduled_requests": scheduled, + "released_request_ids": list(released), + }, + ) + + +def prepare(adapter, ids, slots, tables, lengths, mode="NONE"): + adapter.runner.next_metadata = {"layer0": metadata(torch.tensor(slots), torch.tensor(tables), lengths)} + return adapter.runner.model_state.prepare_attn(NS(req_ids=ids, num_reqs=len(ids)), mode) + + +def reasons(runtime): + return {event.get("reason") for event in runtime.events} + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_sorted_v2_batch_uses_batch_ids_instead_of_request_state_indices(adapter, mode): + runtime = adapter.runtime + packet = message(requests={"A": [[(1, 1)]], "B": [[(2, 1)]]}, scheduled={"A": (0, 1), "B": (0, 2)}) + adapter.runner.req_states = NS(req_id_to_index={"A": 9, "B": 3}) + + def execute(output): + prepare(adapter, ["B", "A"], [256, 257, 128], [[2], [1]], [2, 1], mode) + if mode == "NONE": + runtime.write("layer0", torch.tensor([256, 257, 128])) + runtime.read("layer0", torch.tensor([[2], [1]]), [2, 1]) + return "unchanged" + + assert adapter.wrap_execute(execute)(packet) == "unchanged" + assert not runtime.events + assert runtime.state.blocks[0, 2].valid_end == 2 + assert runtime.state.blocks[0, 1].valid_end == 1 + assert runtime._phase is None and runtime._scheduled is None + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_v2_cross_request_swaps_are_not_legalized_by_allocated_blocks(adapter, mode): + packet = message(requests={"A": [[(1, 1)]], "B": [[(2, 1)]]}, scheduled={"A": (0, 1), "B": (0, 1)}) + adapter.wrap_execute(lambda output: prepare(adapter, ["A", "B"], [256, 128], [[1], [2]], [1, 1], mode))(packet) + assert "write_target_mismatch" in reasons(adapter.runtime) + + +def test_v2_full_replay_observes_current_metadata_after_parameter_update(adapter): + runtime = adapter.runtime + packet = message(requests={"r": [[(2, 1)]]}) + + def execute(output): + item = prepare(adapter, ["r"], [128, -1], [[1], [0]], [1, 0], "FULL")["layer0"] + # Replay/parameter updates mutate the same buffers after collection. + item.slot_mapping[0] = 256 + item.block_tables[0, 0] = 2 + runtime.write("layer0", torch.tensor([128])) + assert runtime._accesses == [] + return "replayed" + + assert adapter.wrap_execute(execute)(packet) == "replayed" + assert not runtime.events + assert runtime.state.blocks[0, 2].valid_end == 1 + + +def test_v2_full_cannot_reuse_last_steps_metadata(adapter): + forward = adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [128], [[1]], [1], "FULL")) + forward(message()) + adapter.runner.model_state.prepare_attn = lambda batch, mode: None + adapter.wrap_execute(lambda output: None)(message(requests={}, scheduled={"r": (1, 1)})) + assert "v2_attention_metadata_missing" in reasons(adapter.runtime) + assert adapter.runtime.state.blocks[0, 1].valid_end == 1 + + +@pytest.mark.parametrize("ids", [["r", "r"], [], ["other"]]) +def test_v2_missing_duplicate_or_unscheduled_batch_ids_are_coverage_gaps(adapter, ids): + adapter.wrap_execute(lambda output: prepare(adapter, ids, [128], [[1]], [1]))(message()) + assert "v2_batch_request_mismatch" in reasons(adapter.runtime) + assert adapter.runtime.state.blocks[0, 1].valid_end <= 0 + + +def test_v2_receive_only_and_finished_steps_do_not_require_attention_metadata(adapter): + sentinel = object() + execute = adapter.wrap_execute(lambda output: sentinel) + assert execute(message(scheduled={})) is sentinel + assert "r" in adapter.runtime.state.requests + assert execute(message(requests={}, scheduled={}, released=["r"])) is sentinel + assert adapter.runtime.state.requests == {} + assert not adapter.runtime.events + + +@pytest.mark.parametrize("profile", [False, True]) +def test_v2_dummy_or_profile_writes_still_detect_case4_without_a_scheduler_binding(adapter, profile): + runtime = adapter.runtime + bind(runtime, computed=19) + packet = NS(total_num_scheduled_tokens=1) + adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [142], [[1]], [19]))( + packet, dummy_run=True, is_profile=profile + ) + assert "dummy_write_to_live_block" in reasons(runtime) + assert "scheduler_binding_missing" not in reasons(runtime) + assert runtime.state.blocks[0, 1].valid_end < 0 + adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [147], [[1]], [20]))( + message(requests={}, scheduled={"r": (19, 1)}) + ) + assert "read_invalid_block" in reasons(runtime) + + +def test_v2_dummy_scratch_reads_are_ignored_but_null_writes_are_reported(adapter): + packet = NS(total_num_scheduled_tokens=1) + adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [-1], [[0]], [1]))(packet, dummy_run=True) + assert not adapter.runtime.events + adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [0], [[0]], [1]))(packet, dummy_run=True) + assert "write_null_block" in reasons(adapter.runtime) + + +def test_v2_profile_without_attention_and_positional_flags_preserve_arguments(adapter): + token = object() + calls = [] + + def execute(output, *args): + calls.append((output, args)) + return token + + packet = NS(total_num_scheduled_tokens=16) + assert adapter.wrap_execute(execute)(packet, token, True, True, True) is token + assert calls == [(packet, (token, True, True, True))] + assert not adapter.runtime.events + + +def test_v2_failed_forward_does_not_commit_writes_or_leak_scope(adapter): + def execute(output): + prepare(adapter, ["r"], [128], [[1]], [1]) + adapter.runtime.write("layer0", torch.tensor([128])) + raise RuntimeError("original failure") + + with pytest.raises(RuntimeError, match="original failure"): + adapter.wrap_execute(execute)(message()) + assert adapter.runtime.state.blocks[0, 1].valid_end == 0 + assert adapter.runtime._phase is None and adapter.runtime._scheduled is None + + +def test_v2_clears_then_copies_after_physical_update(adapter): + runtime = adapter.runtime + runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}, "scheduled_requests": {"r": (4, 0)}}) + packet = NS(new_block_ids_to_zero=[2], kv_cache_block_copies=[NS(src_block_id=1, dst_block_id=2)]) + order = [] + clear, copy = runtime.zero_blocks, runtime.copy_blocks + + def cleared(blocks): + order.append("clear") + clear(blocks) + + def copied(blocks): + order.append("copy") + copy(blocks) + + runtime.zero_blocks = cleared + runtime.copy_blocks = copied + adapter.wrap_update(lambda output: order.append("physical"))(packet) + assert order == ["physical", "clear", "copy"] + assert runtime.state.blocks[0, 2].valid_end == 4 + adapter.wrap_update(lambda output: None)(NS(new_block_ids_to_zero=None, kv_cache_block_copies=None)) + + +def test_v2_startup_capture_does_not_copy_metadata_inside_capture(adapter): + calls = [] + adapter.runtime.observe_metadata = lambda metadata: calls.append(metadata) + capture = adapter.wrap_capture(lambda: prepare(adapter, ["dummy"], [-1], [[0]], [1])) + capture() + assert calls == [] and not adapter.runtime.events + + +@pytest.mark.parametrize("released", [False, True]) +def test_v2_live_or_prefix_preserving_recapture_is_rejected(adapter, released): + bind(adapter.runtime, computed=19) + if released: + adapter.runtime.state.bind({"released_request_ids": ["r"]}) + with pytest.raises(ValueError, match="recapture is unsupported"): + adapter.wrap_capture(lambda: pytest.fail("capture must not run"))() + + +def test_v2_kernel_allocator_coordinate_mismatch_is_rejected(adapter): + adapter.runner.block_tables.kernel_block_sizes = [64] + with pytest.raises(ValueError, match="allocator/kernel"): + adapter.configure() + + +def test_v2_install_disabled_does_not_wrap_methods(adapter, monkeypatch): + from vllm_ascend import envs + + monkeypatch.setattr(envs, "VLLM_ASCEND_KV_GUARD", False) + assert adapter.subject.install_guard(NS(vllm_config=NS())) is None + + +def enable_raise(runtime): + runtime.action = "raise" + append = runtime.state.emit + + def report(event, **fields): + append(event, **fields) + runtime.report(event, **fields) + + runtime.state.emit = report + runtime.gap = lambda reason, **fields: report("guard.coverage_gap", reason=reason, **fields) + + +def test_v2_raise_stops_output_and_latches_before_the_next_model_call(adapter): + enable_raise(adapter.runtime) + calls = [] + + def execute(output): + prepare(adapter, ["r"], [129], [[1]], [1]) + calls.append(True) + return "must not escape" + + forward = adapter.wrap_execute(execute) + for _ in range(2): + with pytest.raises(adapter.runtime.subject.KVGuardError): + forward(message()) + assert calls == [True] + + +def test_v2_missing_binding_and_batch_gap_stop_before_kernel_submission(adapter): + enable_raise(adapter.runtime) + calls = [] + with pytest.raises(adapter.runtime.subject.KVGuardError, match="scheduler_binding_missing"): + adapter.wrap_execute(lambda output: calls.append(True))(NS(total_num_scheduled_tokens=1)) + assert calls == [] + + +def test_v2_wrong_batch_ids_raise_before_the_model_runs(adapter): + enable_raise(adapter.runtime) + calls = [] + + def execute(output): + prepare(adapter, ["unknown"], [128], [[1]], [1]) + calls.append(True) + + with pytest.raises(adapter.runtime.subject.KVGuardError, match="v2_batch_request_mismatch"): + adapter.wrap_execute(execute)(message()) + assert calls == [] + + +def test_v2_install_wraps_real_runner_boundaries_and_reinitializes_model_state(adapter): + runner = adapter.runner + runner.execute_model = lambda output: prepare(adapter, ["r"], [128], [[1]], [1]) + runner.update_requests = lambda output: None + runner.capture_model = lambda: "capture" + installed = adapter.subject.install_guard(runner) + assert installed is not None + # Profiling can replace the model-state object before real cache setup. + runner.model_state = NS(prepare_attn=lambda batch, mode: runner.next_metadata) + installed.configure() + wrapped = runner.model_state.prepare_attn + installed.configure() + assert runner.model_state.prepare_attn is wrapped + assert runner.capture_model() == "capture" + runner.execute_model(message()) + assert adapter.runtime.state.blocks[0, 1].valid_end == 1 + assert not adapter.runtime.events + + +def test_v2_keyword_calls_preserve_upstream_parameter_names(adapter): + def execute(output): + adapter.runner.next_metadata = {"layer0": metadata(torch.tensor([128]), torch.tensor([[1]]), [1])} + return adapter.runner.model_state.prepare_attn(input_batch=NS(req_ids=["r"], num_reqs=1), cudagraph_mode="NONE") + + adapter.wrap_execute(execute)(scheduler_output=message()) + adapter.wrap_update(lambda output: "unchanged")(scheduler_output=NS()) + assert adapter.runtime.state.blocks[0, 1].valid_end == 1 + assert not adapter.runtime.events + + +def test_v2_pending_receive_also_blocks_graph_recapture(adapter): + adapter.runtime.state.load({"request_id": "incoming", "local_block_ids": [[1]], "num_computed_tokens": 19}) + with pytest.raises(ValueError, match="recapture is unsupported"): + adapter.wrap_capture(lambda: pytest.fail("capture must not run"))() diff --git a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py index 6ba14a836eea..9a9f29a3cdc1 100644 --- a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py +++ b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py @@ -5260,6 +5260,47 @@ def test_unused_remote_ports_are_notified_once_by_primary_receiver(primary, has_ assert sent == ([("r", "idle-peer", 7000, ports)] if primary and has_ports else []) +@pytest.mark.parametrize("awaiting_kvs", [False, True]) +@pytest.mark.parametrize("transfer_failed", [False, True]) +def test_guard_receive_callback_preserves_completion_contract(awaiting_kvs, transfer_failed): + receiver = KVCacheRecvingThread.__new__(KVCacheRecvingThread) + receiver._is_failed_recv_request = MagicMock(side_effect=[False, transfer_failed]) + receiver._transfer_kv_cache_all_groups = MagicMock() + receiver._mark_request_task_done = MagicMock(return_value=True) + receiver.pending_reformat_lock = threading.Lock() + receiver.pending_reformat = {"r": object()} + receiver.proc_not_transfer_request_lock = threading.Lock() + receiver.proc_not_transfer_request = {} + receiver._reformat_pending_kv_caches = MagicMock() + receiver._clear_failed_recv_request = MagicMock() + receiver.task_tracker = MagicMock() + receiver.request_queue = MagicMock() + receiver._send_done_signal_to_free_remote_port = MagicMock() + receiver._send_done_recv_signal = MagicMock() + receiver._kv_guard = MagicMock() + meta = { + "request_id": "r", + "remote_request_id": "remote-r", + "remote_host": "host", + "remote_handshake_port": 1, + "remote_port_send_num": {}, + "all_task_done": True, + "awaiting_kvs": awaiting_kvs, + } + receiver._handle_request(meta) + receiver._kv_guard.transfer.assert_called_once_with("error" if transfer_failed else "load", meta) + if transfer_failed: + receiver._reformat_pending_kv_caches.assert_not_called() + else: + receiver._reformat_pending_kv_caches.assert_called_once_with("r") + if awaiting_kvs: + receiver.task_tracker.update_done_task_count.assert_called_once_with("r") + else: + receiver.task_tracker.update_done_task_count.assert_called_once_with("r", report_completion=False) + receiver.request_queue.task_done.assert_called_once() + receiver._send_done_recv_signal.assert_called_once() + + @pytest.mark.parametrize("pulls", [1, 2, 4]) def test_head_reformat_preserves_token_head_order_at_scatter_boundary(monkeypatch, pulls): """Validate PD's CPU reshape arithmetic, not the external NPU scatter op.""" diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py index 8e80bb018c37..868bca368e1d 100644 --- a/vllm_ascend/attention/attention_v1.py +++ b/vllm_ascend/attention/attention_v1.py @@ -934,6 +934,9 @@ def forward_fused_infer_attention( key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params( key, value, attn_metadata, kv_cache ) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.read(self._layer_name, block_table, actual_seq_lengths_kv) num_tokens = attn_metadata.actual_seq_lengths_q[-1] query = query[:num_tokens] if ( @@ -1077,6 +1080,9 @@ def _forward_fia_chunked_prefill_split( seq_lens_list = attn_metadata.seq_lens_list num_tokens = int(actual_seq_qlen[-1]) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.read(self._layer_name, block_table, seq_lens_list) # decode part if num_decode_tokens > 0: decode_out, _ = DeviceOperator.npu_fused_infer_attention_score( @@ -1142,6 +1148,9 @@ def forward_paged_attention( attn_metadata: AscendMetadata, output: torch.Tensor | None = None, ) -> torch.Tensor: + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.read(self._layer_name, attn_metadata.block_tables, attn_metadata.seq_lens) if _EXTRA_CTX.capturing: return self.full_graph_pa(query, attn_metadata, output) torch_npu._npu_paged_attention( @@ -1195,6 +1204,9 @@ def do_kv_cache_update( if self.key_cache is None: self.key_cache, self.value_cache = kv_cache[0], kv_cache[1] + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.write(layer.layer_name, slot_mapping) DeviceOperator.reshape_and_cache( key=key, value=value, @@ -1229,6 +1241,9 @@ def reshape_and_cache( # quick fix to make sure slots is int32 for cross attention case. # see: https://github.com/vllm-project/vllm/blob/ce88756b967c2c5006746a424c15dd59a284ed8c/vllm/model_executor/layers/attention/cross_attention.py#L117 slots_to_cache = slots[: attn_metadata.num_actual_tokens] if not encoder_decoder else slots.to(torch.int32) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.write(self._layer_name, slots_to_cache) DeviceOperator.reshape_and_cache( key=key_to_cache, value=value_to_cache, diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index 189c0a92cfd4..73a891b7d4bc 100644 --- a/vllm_ascend/attention/mla_v1.py +++ b/vllm_ascend/attention/mla_v1.py @@ -1410,6 +1410,15 @@ def _compute_prefill_context( pe_cache = cache_k_pe k_pe = torch.empty(toks, num_heads, rope_dim, dtype=q_pe.dtype, device=q_pe.device) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.read( + self._kv_guard_layer_name, + prefill_metadata.block_table, + context_seq_len_npu, + prefill_metadata.chunked_context.starts[i], + request_offset=attn_metadata.num_decodes, + ) DeviceOperator.kv_cache_load( cache_kv_c, pe_cache, @@ -1617,6 +1626,9 @@ def exec_kv_decode( slots: torch.Tensor, return_current_kv: bool = False, ): + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.write(self._kv_guard_layer_name, slots) if self._uses_strided_kv_adapter() and any(not cache.is_contiguous() for cache in kv_cache): current_k_pe, current_k_nope = self._exec_strided_kv_writer(kv_no_split, cos, sin, kv_cache, slots) if return_current_kv: @@ -2423,6 +2435,13 @@ def exec_kv_prefill( *, attn_metadata: AscendMLAMetadata | None = None, ): + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.write( + self._kv_guard_layer_name, + slots, + token_offset=attn_metadata.num_decode_tokens if attn_metadata else 0, + ) pcp_local_range = None if self.pcp_enabled: assert attn_metadata is not None @@ -2532,6 +2551,9 @@ def _forward_decode( dequant_scale_q_nope = decode_preprocess_res.dequant_scale_q_nope decode_meta = attn_metadata.decode assert decode_meta is not None + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.read(self._kv_guard_layer_name, decode_meta.block_table, decode_meta.seq_lens_list) # TODO: The CANN package is expected to support num_heads that are not # powers of 2 in 2026 Q2. Once supported, all padding operations under # `if self.head_padding > 0` in this function can be removed. @@ -2780,6 +2802,9 @@ def mla_preprocess_only_decode(self, hidden_states, kv_cache, attn_metadata): from cann_ops_transformer import mla_prolog # type: ignore[import-not-found,import-untyped] # noqa: PLC0415 bsz = attn_metadata.num_decode_tokens + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.write(self._kv_guard_layer_name, attn_metadata.slot_mapping[:bsz]) history_slots = attn_metadata.slot_mapping[:bsz] return_current_kv = self._decode_requires_current_kv(attn_metadata) if return_current_kv: @@ -3097,6 +3122,7 @@ def forward( # Profiling run. return output.fill_(0) + self._kv_guard_layer_name = layer_name num_actual_tokens = self.get_num_actual_tokens(attn_metadata) assert ( attn_metadata.num_decodes is not None diff --git a/vllm_ascend/debug/__init__.py b/vllm_ascend/debug/__init__.py new file mode 100644 index 000000000000..fcf134c0cd01 --- /dev/null +++ b/vllm_ascend/debug/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in diagnostics. Importing this package does not enable instrumentation.""" diff --git a/vllm_ascend/debug/kv_guard.py b/vllm_ascend/debug/kv_guard.py new file mode 100644 index 000000000000..fb805acbc5a9 --- /dev/null +++ b/vllm_ascend/debug/kv_guard.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""KV address legality: allocation generation, valid prefix and request owners. + +A prefix of zero is unwritten; INVALID_PREFIX is tainted. Loading and appending +produce the same readable prefix, so their provenance is not persistent state. +""" + +from __future__ import annotations + +import threading +from bisect import bisect_right +from dataclasses import dataclass, field +from typing import Any + +INVALID_PREFIX = -1 +# Group and logical block index are the positions in this nested table. +BlockTable = list[list[tuple[int, int]]] + + +@dataclass(slots=True) +class Block: + epoch: int + valid_end: int = 0 + owners: set[str] = field(default_factory=set) + + +class KVGuard: + """One worker's pool, with request bindings and a reverse ownership index.""" + + def __init__(self, emit: Any): + self.emit = emit + self.pool_id: str | None = None + self.blocks: dict[tuple[int, int], Block] = {} + self.requests: dict[str, BlockTable] = {} + self.pending_loads: dict[str, dict[str, Any]] = {} + self.group_sizes: dict[int, int] = {} + self._lock = threading.RLock() + + def configure(self, group_sizes: dict[int, int]) -> None: + self.group_sizes = group_sizes + + def _violation(self, key: tuple[int, int], block: Block | None, operation: str, reason: str, **fields: Any) -> None: + self.emit( + "guard.violation", + operation=operation, + reason=reason, + group_id=key[0], + block_id=key[1], + pool_id=self.pool_id, + alloc_epoch=block.epoch if block else None, + valid_end=block.valid_end if block else 0, + owner_request_ids=sorted(block.owners) if block else [], + **fields, + ) + + def bind(self, context: dict[str, Any]) -> None: + with self._lock: + pool = context.get("pool_id", self.pool_id) + if pool != self.pool_id: + self.blocks.clear() + self.requests.clear() + if self.pool_id is not None: + self.pending_loads.clear() + self.pool_id = pool + for request_id in context.get("released_request_ids", ()): + for group, refs in enumerate(self.requests.pop(request_id, ())): + for block_id, _ in refs: + block = self.blocks.get((group, block_id)) + if block is not None: + block.owners.discard(request_id) + self.pending_loads.pop(request_id, None) + for request_id, groups in context.get("requests", {}).items(): + keys = {(group, block) for group, refs in enumerate(groups) for block, _ in refs} + for group, refs in enumerate(self.requests.get(request_id, ())): + for block_id, _ in refs: + key = group, block_id + if key not in keys and key in self.blocks: + self.blocks[key].owners.discard(request_id) + self.requests[request_id] = groups + computed = context.get("scheduled_requests", {}).get(request_id, (0, 0))[0] or 0 + for group, refs in enumerate(groups): + size = self.group_sizes[group] + for logical, (block_id, epoch) in enumerate(refs): + if block_id == 0: + continue + if epoch <= 0: + self.emit( + "guard.coverage_gap", + reason="allocation_generation_missing", + group_id=group, + block_id=block_id, + request_id=request_id, + ) + continue + key = group, block_id + block = self.blocks.get(key) + if block is None or block.epoch != epoch: + if block is not None and block.owners: + self._violation(key, block, "recycle", "recycle_live_block", request_id=request_id) + block = Block(epoch) + self.blocks[key] = block + block.owners.add(request_id) + valid = min(size, max(0, computed - logical * size)) + # Reuse may establish validity, but never revive a tainted + # allocation or a destination awaiting a physical copy. + if ( + valid + and block.valid_end == 0 + and block_id not in context.get("uninitialized_block_ids", ()) + ): + block.valid_end = valid + pending = self.pending_loads.pop(request_id, None) + if pending is not None: + self._load(pending) + + def load(self, meta: dict[str, Any]) -> None: + with self._lock: + if meta["request_id"] not in self.requests or not self.group_sizes: + self.pending_loads[meta["request_id"]] = { + key: meta[key] for key in ("request_id", "local_block_ids", "num_computed_tokens") if key in meta + } + else: + self._load(meta) + + def _load(self, meta: dict[str, Any]) -> None: + request_id = meta["request_id"] + expected = { + (group, block): (logical, epoch) + for group, refs in enumerate(self.requests.get(request_id, ())) + for logical, (block, epoch) in enumerate(refs) + } + for group, ids in enumerate(meta.get("local_block_ids", ())): + for block_id in ids: + key = group, block_id + ref, block = expected.get(key), self.blocks.get(key) + if ref is None or block is None or request_id not in block.owners: + self._violation(key, block, "load", "unowned_destination", request_id=request_id) + continue + logical, epoch = ref + if block.epoch != epoch: + self._violation(key, block, "load", "stale_allocation_epoch", request_id=request_id) + continue + computed = meta.get("num_computed_tokens") + if computed is None: + self.emit( + "guard.coverage_gap", reason="load_valid_length_missing", group_id=group, block_id=block_id + ) + continue + size = self.group_sizes[group] + block.valid_end = min(size, max(0, computed - logical * size)) + + def fail_load(self, request_id: str) -> None: + with self._lock: + self.pending_loads.pop(request_id, None) + for group, refs in enumerate(self.requests.get(request_id, ())): + for block_id, epoch in refs: + block = self.blocks.get((group, block_id)) + if block is not None and block.epoch == epoch: + block.valid_end = INVALID_PREFIX + + def clear(self, group_id: int, block_id: int) -> None: + with self._lock: + block = self.blocks.get((group_id, block_id)) + if block is not None: + block.valid_end = INVALID_PREFIX + + def copy_blocks(self, group_id: int, copies: list[tuple[int, int]]) -> None: + """Match the operator's snapshot-all-sources-before-writing semantics.""" + with self._lock: + source_ends = { + source: block.valid_end + for source, _ in copies + if (block := self.blocks.get((group_id, source))) is not None + } + for source_id, destination_id in copies: + valid = source_ends.get(source_id, INVALID_PREFIX) + key = group_id, destination_id + target = self.blocks.get(key) + if valid == INVALID_PREFIX or target is None or not target.owners: + self._violation( + key, target, "copy", "invalid_copy_source_or_destination", source_block_id=source_id + ) + if target is not None: + target.valid_end = INVALID_PREFIX + else: + target.valid_end = valid + + def observe(self, groups: list[dict[str, Any]], request_ids: list[str], phase: str) -> None: + """Audit WRITE_SLOT then READ_RANGE using attention's consumed metadata. + + An illegal write taints its destination so subsequent reads diagnose + an invalid prefix, even when the block address itself remains legal. + """ + with self._lock: + for group in groups: + group_id = group["group_id"] + size = self.group_sizes[group_id] + ranges: dict[int, list[int]] = {} + expected = group.get("expected_slots") + query_ends = group.get("query_ends", ()) + token_offset = group.get("token_offset", 0) + mismatched: set[int] = set() + slots = group["slots"] + # List equality handles the common valid path in C. Only an + # actual mismatch needs per-token Python diagnostics. + count = min(len(slots), max(0, len(expected) - token_offset)) if expected is not None else 0 + if expected is not None and ( + slots[:count] != expected[token_offset : token_offset + count] + or any(slot >= 0 for slot in slots[count:]) + ): + for index, slot in enumerate(slots): + token_index = token_offset + index + target = expected[token_index] if token_index < len(expected) else None + if (target is not None and slot == target) or (target is None and slot < 0): + continue + destination = slot if slot >= 0 else target + assert destination is not None + block_id = destination // size + key = group_id, block_id + block = self.blocks.get(key) + if block_id in mismatched: + continue + self._violation( + key, + block, + "write", + "write_target_mismatch", + token_index=token_index, + slot=slot, + expected_slot=target, + request_id=request_ids[bisect_right(query_ends, token_index)] + if query_ends and token_index < query_ends[-1] + else None, + phase=phase, + ) + mismatched.add(block_id) + if block is not None: + block.valid_end = INVALID_PREFIX + for slot in slots: + if slot < 0: + continue + ranges.setdefault(slot // size, []).append(slot % size) + for block_id, offsets in ranges.items(): + key, offsets = (group_id, block_id), sorted(offsets) + block = self.blocks.get(key) + if block_id in mismatched: + continue + if block_id == 0: + self._violation(key, None, "write", "write_null_block", offsets=offsets, phase=phase) + continue + if phase == "warmup" and block is None: + continue # Capture may use explicitly allocated scratch. + reason = None + if block is None: + reason = "unallocated_block" + elif phase in ("dummy", "warmup"): + reason = "dummy_write_to_live_block" + elif not block.owners.intersection(request_ids): + reason = "unowned_write" + elif block.valid_end == INVALID_PREFIX: + reason = "write_invalid_block" + elif offsets[0] < block.valid_end: + reason = "overwrite_valid_prefix" + elif offsets != list(range(block.valid_end, block.valid_end + len(offsets))): + reason = "noncontiguous_write" + if reason is not None: + self._violation( + key, block, "write", reason, offsets=offsets, phase=phase, request_ids=request_ids + ) + if block is not None: + block.valid_end = INVALID_PREFIX + elif block is not None: + block.valid_end = offsets[-1] + 1 + for row, length in enumerate(group.get("seq_lens", ())): + if length <= 0: + continue + request_id = request_ids[row] if row < len(request_ids) else None + table = group["block_table"][row] + expected_tables = group.get("expected_block_table") + expected_lengths = group.get("expected_seq_lens") + if expected_lengths is not None and ( + row >= len(expected_lengths) or length > expected_lengths[row] + ): + self.emit( + "guard.violation", + operation="read", + reason="read_length_mismatch", + group_id=group_id, + request_id=request_id, + length=length, + phase=phase, + ) + continue + needed = (length + size - 1) // size + if needed > len(table): + self.emit( + "guard.violation", + operation="read", + reason="short_block_table", + group_id=group_id, + phase=phase, + ) + continue + for logical, block_id in enumerate(table[:needed]): + if expected_tables is not None and block_id != expected_tables[row][logical]: + self._violation( + (group_id, block_id), + self.blocks.get((group_id, block_id)), + "read", + "read_target_mismatch", + request_id=request_id, + logical_block_index=logical, + expected_block_id=expected_tables[row][logical], + phase=phase, + ) + continue + if block_id == 0: + if request_id is not None: + reason = "read_null_block" + self._violation( + (group_id, block_id), None, "read", reason, phase=phase, request_id=request_id + ) + continue + key, end = (group_id, block_id), min(size, length - logical * size) + block = self.blocks.get(key) + if block is None: + reason = "read_unallocated_block" + elif block.valid_end == INVALID_PREFIX: + reason = "read_invalid_block" + elif request_id not in block.owners: + reason = "unowned_read" + elif end > block.valid_end: + reason = "read_unwritten_range" + else: + continue + self._violation(key, block, "read", reason, end_offset=end, phase=phase, request_id=request_id) diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py new file mode 100644 index 000000000000..afe57f4d6bfa --- /dev/null +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -0,0 +1,586 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker L1 guard with one runner execution-plan collection point. + +No KV snapshots, hashes, general-purpose trace or offline reader. +Adapted per-layer metadata supplies declared access plans. Operator hooks +preserve argument order; scheduler ranges independently identify each token. +Only small metadata is copied, after the forward and graph parameter updates; +the checker neither inspects KV values nor proves kernel memory accesses. +""" + +from __future__ import annotations + +import functools +from collections import OrderedDict +from contextlib import contextmanager +from dataclasses import dataclass, replace +from typing import Any + +from vllm_ascend.debug.kv_guard import INVALID_PREFIX, KVGuard + +_MAX_REPORT_KEYS = 1024 + + +@dataclass(slots=True) +class WriteAccess: + layer: str + slots: Any + token_offset: int = 0 + + +@dataclass(slots=True) +class ReadAccess: + layer: str + table: Any + lengths: Any + starts: Any = None + request_offset: int = 0 + + +@dataclass(slots=True) +class ExpectedAccess: + slots: list[int] + query_ends: list[int] + tables: list[list[int]] + lengths: list[int] + + +class KVGuardError(RuntimeError): + """Stop returning model results after a violation or incomplete audit.""" + + +class GuardRuntime: + def __init__(self, vllm_config: Any): + from vllm.logger import logger + + self.logger = logger + self.config = vllm_config + self.state = KVGuard(self.report) + self._scratch = KVGuard(self._emit_scratch) + self._audit_layer = "" + self._audit_evidence = "" + self.layer_groups: dict[str, int] = {} + self._accesses: list[WriteAccess | ReadAccess] = [] + self._metadata: dict[str, Any] = {} + self._missing_groups: set[int] = set() + self._phase: str | None = None + self._full_graph = False + self._reported: OrderedDict[tuple, None] = OrderedDict() + self._scheduled: dict[str, tuple[int | None, int]] | None = None + self._failure: str | None = None + self.action = "log" + + def _check_failure(self) -> None: + if self.action == "raise" and self._failure is not None: + raise KVGuardError(f"KV guard stopped model output: {self._failure}; restart the worker before reuse") + + def _emit_scratch(self, event: str, **fields: Any) -> None: + if event == "guard.violation": + self.state.emit(event, layer=self._audit_layer, evidence=self._audit_evidence, **fields) + + def report(self, event: str, **fields: Any) -> None: + failed = event in ("guard.violation", "guard.coverage_gap") + if failed: + if self._failure is None: + self._failure = str(fields.get("reason", event)) + key = ( + event, + *( + fields.get(name) + for name in ( + "operation", + "reason", + "pool_id", + "group_id", + "block_id", + "alloc_epoch", + "layer", + "request_id", + ) + ), + tuple(fields.get("request_ids", ())), + ) + if key in self._reported: + self._reported.move_to_end(key) + return + self._reported[key] = None + if len(self._reported) > _MAX_REPORT_KEYS: + self._reported.popitem(last=False) + payload = " ".join(f"{key}={value}" for key, value in fields.items()) + log = self.logger.error if failed else self.logger.info + log("[KV_GUARD] %s %s", event, payload) + + def gap(self, reason: str, **fields: Any) -> None: + self.report("guard.coverage_gap", reason=reason, **fields) + + def configure(self, runner: Any) -> None: + from vllm_ascend import envs + + self.action = envs.VLLM_ASCEND_KV_GUARD_ACTION + if self.action not in ("log", "raise"): + raise ValueError("VLLM_ASCEND_KV_GUARD_ACTION must be 'log' or 'raise'") + config = self.config + parallel = config.parallel_config + transfer = config.kv_transfer_config + unsupported = any( + ( + getattr(config.model_config, "is_encoder_decoder", False), + config.speculative_config is not None, + getattr(parallel, "prefill_context_parallel_size", 1) != 1, + getattr(parallel, "decode_context_parallel_size", 1) != 1, + getattr(parallel, "enable_dbo", False), + transfer is not None and transfer.kv_connector != "MooncakeConnector", + ) + ) + sizes = {} + tables = getattr(getattr(getattr(runner, "input_batch", None), "block_table", None), "block_tables", ()) + for group_id, group in enumerate(runner.kv_cache_config.kv_cache_groups): + spec = group.kv_cache_spec + sizes[group_id] = spec.block_size + if tables and tables[group_id].block_size != spec.block_size: + unsupported = True + for name in group.layer_names: + self.layer_groups[name] = group_id + layer_spec = getattr(spec, "kv_cache_specs", {}).get(name, spec) + append = ( + hasattr(layer_spec, "head_size") + and getattr(layer_spec, "sliding_window", None) is None + and getattr(layer_spec, "compress_ratio", 1) == 1 + and type(layer_spec).__name__ != "CircularBufferSpec" + ) + if not append: + raise ValueError(f"KV L1 guard: unadapted cache validity rules: {type(layer_spec).__name__}") + layer = config.compilation_config.static_forward_context.get(name) + impl = getattr(layer, "impl", None) + if impl is not None and type(impl).__name__ in ("AscendAttentionBackendImpl", "AscendMLAImpl"): + impl._kv_guard = self + else: + raise ValueError(f"KV L1 guard: unadapted attention implementation: {type(impl).__name__}") + if unsupported: + raise ValueError( + "KV L1 guard: unadapted execution mode or block coordinates " + "(CP, DBO, speculation, encoder/decoder, connector)" + ) + self.state.configure(sizes) + self.report("guard.enabled", failure_action=self.action) + + def write(self, layer: str, slots: Any, token_offset: int = 0) -> None: + if self._phase is not None and not self._full_graph: + self._accesses.append(WriteAccess(layer, slots, token_offset)) + + def read(self, layer: str, table: Any, lengths: Any, starts: Any = None, request_offset: int = 0) -> None: + if self._phase is not None and not self._full_graph and table is not None: + self._accesses.append(ReadAccess(layer, table, lengths, starts, request_offset)) + + @staticmethod + def _list(value: Any) -> list: + if isinstance(value, (list, tuple)): + return list(value) + return value.detach().cpu().tolist() + + @staticmethod + def _field(metadata: Any, *names: str) -> Any: + return next((value for name in names if (value := getattr(metadata, name, None)) is not None), None) + + def collect_metadata(self, metadata: Any) -> None: + """Keep references/coverage only; materialize the selected source later.""" + self._metadata, self._missing_groups = {}, set() + if metadata is None: + return # Profiling and eager operator-only tests may have no attention plan. + if not isinstance(metadata, dict): + self.gap("microbatch_metadata_unadapted") + self._missing_groups.update(self.layer_groups.values()) + return + for layer, group in self.layer_groups.items(): + item = metadata.get(layer) + if item is None: + self.gap(f"metadata_write_mapping_missing:{layer}") + self._missing_groups.add(group) + continue + self._metadata[layer] = item + + def _metadata_accesses(self, layer: str, item: Any) -> list[WriteAccess | ReadAccess]: + slots = self._field(item, "slot_mapping") + if slots is None: + self.gap(f"metadata_write_mapping_missing:{layer}") + self._missing_groups.add(self.layer_groups[layer]) + return [] + count = getattr(item, "num_actual_tokens", None) + if count is not None: + slots = slots[:count] + accesses: list[WriteAccess | ReadAccess] = [WriteAccess(layer, slots)] + table = self._field(item, "block_tables", "block_table", "block_table_tensor") + lengths = self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") + if table is None or lengths is None: + self.gap(f"metadata_read_range_missing:{layer}") + self._missing_groups.add(self.layer_groups[layer]) + else: + accesses.append(ReadAccess(layer, table, lengths)) + return accesses + + @staticmethod + def _identity(value: Any) -> Any: + if value is None: + return None + if isinstance(value, (list, tuple)): + return tuple(GuardRuntime._identity(item) for item in value) + if isinstance(value, (int, float, str)): + return value + return value.data_ptr(), tuple(value.shape), tuple(value.stride()) + + def _expected_write_plan(self, runner: Any, group_id: int) -> ExpectedAccess | None: + """Derive destinations from the scheduler, never from the audited mapping. + + Per-token lists exist only for this forward and are shared by all layers + in a cache group. No persistent token-ID history or KV tensor is copied. + """ + if self._phase != "forward" or self._scheduled is None: + return None + size = self.state.group_sizes[group_id] + slots: list[int] = [] + query_ends: list[int] = [] + tables: list[list[int]] = [] + lengths: list[int] = [] + for request_id in runner.input_batch.req_ids: + computed, count = self._scheduled.get(request_id, (None, 0)) + if computed is None or computed < 0 or count <= 0: + self.gap("scheduler_write_range_missing", request_id=request_id) + self._missing_groups.add(group_id) + return None + groups = self.state.requests.get(request_id, ()) + refs = groups[group_id] if group_id < len(groups) else () + needed = (computed + count + size - 1) // size + if len(refs) < needed: + self.gap("scheduler_read_blocks_missing", request_id=request_id, group_id=group_id) + self._missing_groups.add(group_id) + return None + tables.append([block for block, _ in refs[:needed]]) + lengths.append(computed + count) + for logical in range(computed // size, (computed + count - 1) // size + 1): + block_id, _ = refs[logical] + start = max(computed, logical * size) - logical * size + end = min(computed + count, (logical + 1) * size) - logical * size + slots.extend(range(block_id * size + start, block_id * size + end)) + query_ends.append(len(slots)) + return ExpectedAccess(slots, query_ends, tables, lengths) + + def _audit(self, runner: Any) -> None: + phase = self._phase + assert phase is not None + # FULL requires this replay's final metadata. Never reuse a captured + # address plan when the current metadata is missing. + if self._full_graph and not self._metadata: + if phase != "warmup": + self.gap("full_graph_metadata_missing") + return + observed: dict[str, list[WriteAccess | ReadAccess]] = {} + for access in self._accesses: + observed.setdefault(access.layer, []).append(access) + if not self._metadata and not observed: + if phase != "warmup": + self.gap("execution_access_plan_missing") + return + + copies: dict[Any, list] = {} + results: dict[int, list[dict]] = {} + audited: set[Any] = set() + templates: dict[tuple[int, int], list[WriteAccess | ReadAccess]] = {} + baseline: dict = {} + shadow = self._scratch + shadow.configure(self.state.group_sizes) + pool_id = shadow.pool_id = self.state.pool_id + expected: dict[int, ExpectedAccess | None] = {} + for layer in dict.fromkeys((*self._metadata, *observed)): + group_id = self.layer_groups.get(layer) + if group_id is None: + self.gap(f"unbound_access_layer:{layer}") + continue + actual = observed.get(layer, []) + if observed and not actual and not self._full_graph: + self.gap(f"operator_layer_not_observed:{layer}") + self._missing_groups.add(group_id) + if actual and not self._full_graph: + accesses = actual + evidence = "operator_arguments" + elif layer in self._metadata: + item = self._metadata[layer] + template_key = id(item), group_id + if template_key not in templates: + templates[template_key] = self._metadata_accesses(layer, item) + accesses = templates[template_key] + evidence = "execution_plan" + else: + accesses = actual + evidence = "operator_arguments" + signature = ( + group_id, + evidence, + tuple( + ("write", self._identity(a.slots), a.token_offset) + if isinstance(a, WriteAccess) + else ( + "read", + self._identity(a.table), + self._identity(a.lengths), + self._identity(a.starts), + a.request_offset, + ) + for a in accesses + ), + ) + if signature in audited: + continue + + def values(value): + key = self._identity(value) + if key not in copies: + copies[key] = self._list(value) + return copies[key] + + groups: list[tuple[dict[str, Any], int]] = [] + write_ranges: list[tuple[int, int]] = [] + keys: set[tuple[int, int]] = set() + write_keys: set[tuple[int, int]] = set() + size = self.state.group_sizes[group_id] + if group_id not in expected: + expected[group_id] = self._expected_write_plan(runner, group_id) + plan = expected[group_id] + for access in accesses: + entry: dict[str, Any] = { + "group_id": group_id, + "slots": [], + "block_table": [], + "seq_lens": [], + } + if isinstance(access, WriteAccess): + slots = values(access.slots) + entry["slots"] = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] + if plan is not None: + entry["expected_slots"], entry["query_ends"] = plan.slots, plan.query_ends + entry["token_offset"] = access.token_offset + write_ranges.append((access.token_offset, access.token_offset + len(entry["slots"]))) + write_keys.update((group_id, slot // size) for slot in entry["slots"] if slot >= 0) + keys.update(write_keys) + offset = 0 + else: + if self._phase in ("dummy", "warmup"): + # Dummy outputs are discarded. Their writes can corrupt + # live KV, but reading reserved scratch is expected. + continue + lengths = values(access.lengths) + if access.starts is not None: + lengths = [a + b for a, b in zip(lengths, values(access.starts))] + entry["seq_lens"] = lengths + needed = max(((length + size - 1) // size for length in lengths), default=0) + if isinstance(access.table, (list, tuple)): + entry["block_table"] = [row[:needed] for row in access.table[: len(lengths)]] + else: + entry["block_table"] = values(access.table[: len(lengths), :needed]) + keys.update((group_id, block) for row in entry["block_table"] for block in row if block >= 0) + offset = access.request_offset + if plan is not None: + entry["expected_block_table"] = plan.tables[offset:] + entry["expected_seq_lens"] = plan.lengths[offset:] + groups.append((entry, offset)) + + if plan is not None: + covered = 0 + for start, end in sorted(write_ranges): + if start > covered: + break + covered = max(covered, end) + if covered < len(plan.slots): + self.gap("write_tokens_not_observed", layer=layer, first_missing_token=covered) + self._missing_groups.add(group_id) + + with self.state._lock: + # Owners are immutable within the audit. Copy them once per + # touched block; read-only state can also be shared by layers. + for key in keys: + if key not in baseline and key in self.state.blocks: + initial = self.state.blocks[key] + baseline[key] = replace(initial, owners=set(initial.owners)) + shadow.blocks = { + key: replace(baseline[key]) if key in write_keys else baseline[key] + for key in keys + if key in baseline + } + self._audit_layer, self._audit_evidence = layer, evidence + request_ids = list(runner.input_batch.req_ids) if self._phase == "forward" else [] + for entry, offset in groups: + shadow.observe([entry], request_ids[offset:], phase) + audited.add(signature) + results.setdefault(group_id, []).append(shadow.blocks) + + # Each layer starts from the same pre-forward validity. Another layer's + # write cannot legalize an earlier read or hide a missing write. Merge + # conservatively; no validity is advanced on failed/partial collection. + with self.state._lock: + for group_id, versions_by_layer in results.items(): + touched = {key for versions in versions_by_layer for key in versions} + for key in touched: + block = self.state.blocks.get(key) + if block is None: + continue + versions = [versions.get(key, baseline.get(key, block)) for versions in versions_by_layer] + if self.state.pool_id != pool_id or any(b.epoch != block.epoch for b in versions): + self.gap("allocation_changed_during_access_audit") + continue + if any(b.valid_end == INVALID_PREFIX for b in versions): + block.valid_end = INVALID_PREFIX + elif ( + block.valid_end != INVALID_PREFIX + and group_id not in self._missing_groups + and self._phase == "forward" + ): + valid = min(b.valid_end for b in versions) + if valid > block.valid_end: + block.valid_end = valid + + def observe_metadata(self, metadata: Any) -> None: + try: + self.collect_metadata(metadata) + except Exception as exc: + self._metadata = {} + self._missing_groups = set(self.layer_groups.values()) + self.gap(f"metadata_collection_failed:{type(exc).__name__}") + + @contextmanager + def forward_scope( + self, runner: Any, *, phase: str, metadata: Any = None, full_graph: bool = False, required: bool = True + ): + """Audit one completed forward; never advance validity on exceptions.""" + self._check_failure() + self._phase, self._full_graph = phase, full_graph + self._accesses = [] + self.observe_metadata(metadata) + try: + yield + if required or self._metadata or self._accesses: + try: + self._audit(runner) + except Exception as exc: + self.gap(f"access_audit_failed:{type(exc).__name__}") + self._check_failure() + finally: + self._phase = None + + def wrap_forward(self, runner: Any, forward: Any, get_context: Any): + @functools.wraps(forward) + def observed(num_tokens, *args, **kwargs): + context = get_context() + context._kv_guard = self + with self.forward_scope( + runner, + phase=kwargs.pop("_kv_guard_phase", "forward"), + metadata=getattr(context, "attn_metadata", None), + full_graph="FULL" in str(context.cudagraph_runtime_mode), + ): + return forward(num_tokens, *args, **kwargs) + + return observed + + def wrap_update(self, runner: Any, update: Any): + @functools.wraps(update) + def observed(output, *args, **kwargs): + result = update(output, *args, **kwargs) + binding = getattr(output, "_ascend_kv_guard_context", None) + if binding is None: + self.gap("scheduler_binding_missing") + self._check_failure() + return result + for request_id, groups in binding["requests"].items(): + actual = runner.requests.get(request_id) + expected = [[block for block, _ in refs] for refs in groups] + if actual is None or [list(g) for g in actual.block_ids] != expected: + self.report( + "guard.violation", + operation="bind", + reason="worker_block_mapping_mismatch", + request_id=request_id, + ) + self._check_failure() + return result + + return observed + + def wrap_execute(self, execute: Any): + @functools.wraps(execute) + def observed(output, *args, **kwargs): + self._check_failure() + binding = getattr(output, "_ascend_kv_guard_context", None) + if binding is not None: + # Receive-only and idle DP steps can return before _update_states. + uninitialized = set(getattr(output, "new_block_ids_to_zero", None) or ()) + uninitialized.update( + copy.dst_block_id for copy in (getattr(output, "kv_cache_block_copies", None) or ()) + ) + self.state.bind({**binding, "uninitialized_block_ids": uninitialized}) + self._scheduled = binding.get("scheduled_requests") + if self._scheduled is None: + self.gap("scheduler_write_plan_missing") + else: + self.gap("scheduler_binding_missing") + self._scheduled = None + try: + self._check_failure() + return execute(output, *args, **kwargs) + finally: + self._scheduled = None + + return observed + + def zero_blocks(self, block_ids: list[int]) -> None: + with self.state._lock: + for group_id in self.state.group_sizes: + for block_id in block_ids: + if block_id == 0: + self.state.clear(group_id, block_id) + continue + block = self.state.blocks.get((group_id, block_id)) + # A fresh allocation is already UNKNOWN. Clearing it must + # not turn its first legitimate write into an invalid write. + if block is not None and block.valid_end != 0: + if block.valid_end > 0 and block.owners: + self.state._violation((group_id, block_id), block, "clear", "clear_live_block") + self.state.clear(group_id, block_id) + self._check_failure() + + def copy_blocks(self, copies: Any) -> None: + with self.state._lock: + for copy in copies: + if copy.dst_block_id == 0: + for group_id in self.state.group_sizes: + self.state._violation((group_id, 0), None, "copy", "write_null_block") + elif not any((group_id, copy.dst_block_id) in self.state.blocks for group_id in self.state.group_sizes): + self.gap("copy_destination_unbound", block_id=copy.dst_block_id) + for group_id in self.state.group_sizes: + pairs = [ + (copy.src_block_id, copy.dst_block_id) + for copy in copies + if (group_id, copy.dst_block_id) in self.state.blocks + ] + self.state.copy_blocks(group_id, pairs) + self._check_failure() + + def transfer(self, event: str, meta: dict[str, Any]) -> None: + try: + if event == "load": + self.state.load(meta) + else: + self.state.fail_load(meta["request_id"]) + except Exception as exc: + self.gap(f"load_hook_failed:{type(exc).__name__}") + + +def get_guard(vllm_config: Any) -> GuardRuntime | None: + from vllm_ascend import envs + + if not envs.VLLM_ASCEND_KV_GUARD: + return None + guard = getattr(vllm_config, "_ascend_kv_guard", None) + if guard is None: + guard = GuardRuntime(vllm_config) + vllm_config._ascend_kv_guard = guard + return guard diff --git a/vllm_ascend/debug/kv_guard_v2.py b/vllm_ascend/debug/kv_guard_v2.py new file mode 100644 index 000000000000..54b103ba4ba3 --- /dev/null +++ b/vllm_ascend/debug/kv_guard_v2.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MRV2 batch adapter for the shared KV address checker. + +Batch IDs come from prepare_attn after sorting, not RequestState row order. +Metadata is retained by reference and copied only after forward/graph updates. +""" + +import functools +from types import SimpleNamespace +from typing import Any, cast + +from vllm_ascend.debug.kv_guard_runtime import GuardRuntime, get_guard + + +class KVGuardV2: + def __init__(self, runner: Any, runtime: GuardRuntime): + self.runner = runner + self.runtime = runtime + self.input_batch = SimpleNamespace(req_ids=[]) + self._prepared = False + + def configure(self) -> None: + tables = self.runner.block_tables + sizes = [group.kv_cache_spec.block_size for group in self.runner.kv_cache_config.kv_cache_groups] + if list(tables.block_sizes) != sizes or list(tables.kernel_block_sizes) != sizes: + raise ValueError("KV guard: unadapted MRV2 allocator/kernel block coordinates") + self.runtime.configure(self.runner) + prepare = self.runner.model_state.prepare_attn + if getattr(prepare, "_ascend_kv_guard", False): + return + + @functools.wraps(prepare) + def observed(input_batch, cudagraph_mode, *args, **kwargs): + metadata = prepare(input_batch, cudagraph_mode, *args, **kwargs) + if self.runtime._phase is not None: + self.input_batch = SimpleNamespace(req_ids=list(input_batch.req_ids[: input_batch.num_reqs])) + self._prepared = True + self.runtime._full_graph = "FULL" in str(cudagraph_mode) + self.runtime.observe_metadata(metadata) + scheduled = self.runtime._scheduled + if self.runtime._phase == "forward" and scheduled is not None: + ids = self.input_batch.req_ids + if len(ids) != len(set(ids)) or set(ids) != set(scheduled): + self.runtime.gap("v2_batch_request_mismatch") + self.runtime._missing_groups.update(self.runtime.layer_groups.values()) + self.runtime._check_failure() + return metadata + + cast(Any, observed)._ascend_kv_guard = True + self.runner.model_state.prepare_attn = observed + + def wrap_execute(self, execute: Any): + def forward(output, *args, **kwargs): + # Preserve positional and keyword calls across the supported vLLM + # lanes; forward every argument unchanged to the original method. + flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) + flags.update(kwargs) + dummy = bool(flags.get("dummy_run", False)) + profile = bool(flags.get("is_profile", False)) + phase = "warmup" if profile else "dummy" if dummy else "forward" + required = bool(output.total_num_scheduled_tokens) and not ( + profile and flags.get("skip_attn_for_dummy_run") + ) + self.input_batch = SimpleNamespace(req_ids=[]) + self._prepared = False + + with self.runtime.forward_scope(self, phase=phase, required=required): + result = execute(output, *args, **kwargs) + if phase == "forward" and required and not self._prepared: + self.runtime.gap("v2_attention_metadata_missing") + self.runtime._missing_groups.update(self.runtime.layer_groups.values()) + return result + + execute_bound = self.runtime.wrap_execute(forward) + + @functools.wraps(execute) + def observed(scheduler_output, *args, **kwargs): + flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) + flags.update(kwargs) + # Dummy/profile outputs are synthesized without a scheduler binding. + # They must still audit writes against any existing live KV state. + if flags.get("dummy_run", False) or flags.get("is_profile", False): + self.runtime._scheduled = None + return forward(scheduler_output, *args, **kwargs) + return execute_bound(scheduler_output, *args, **kwargs) + + return observed + + def wrap_update(self, update: Any): + @functools.wraps(update) + def observed(scheduler_output, *args, **kwargs): + result = update(scheduler_output, *args, **kwargs) + # The parent has completed the physical operations at this point. + zeroed = getattr(scheduler_output, "new_block_ids_to_zero", None) + if zeroed: + self.runtime.zero_blocks(zeroed) + copies = getattr(scheduler_output, "kv_cache_block_copies", None) + if copies: + self.runtime.copy_blocks(copies) + return result + + return observed + + def wrap_capture(self, capture: Any): + @functools.wraps(capture) + def observed(*args, **kwargs): + self.runtime._check_failure() + with self.runtime.state._lock: + if ( + self.runtime.state.requests + or self.runtime.state.pending_loads + or any(block.valid_end != 0 for block in self.runtime.state.blocks.values()) + ): + raise ValueError( + "KV guard: capture requires empty KV state; live/cache-preserving recapture is unsupported" + ) + # Startup capture has no request KV to protect. Do not synchronize + # metadata to the CPU inside a device graph capture. Real replays + # are audited by wrap_execute using their current metadata. + return capture(*args, **kwargs) + + return observed + + +def install_guard(runner: Any) -> KVGuardV2 | None: + runtime = get_guard(runner.vllm_config) + if runtime is None: + return None + adapter = KVGuardV2(runner, runtime) + cast(Any, runner).execute_model = adapter.wrap_execute(runner.execute_model) + cast(Any, runner).update_requests = adapter.wrap_update(runner.update_requests) + cast(Any, runner).capture_model = adapter.wrap_capture(runner.capture_model) + return adapter diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py index 3f2ccbe4c775..6e6c8063723c 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py @@ -59,6 +59,7 @@ from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec, AscendSlidingWindowMLASpec +from vllm_ascend.debug.kv_guard_runtime import get_guard from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import global_te from vllm_ascend.distributed.kv_transfer.utils.utils import ( PD_QOS_DEFAULT, @@ -509,6 +510,7 @@ def __init__( self.request_task_counts_lock = threading.Lock() self.task_tracker = KVCacheTaskTracker() + self._kv_guard = get_guard(vllm_config) self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata) @@ -757,7 +759,8 @@ def _handle_request(self, req_meta: dict[str, Any]): finally: all_tasks_done = self._mark_request_task_done(request_id, all_task_done) if all_tasks_done: - if transfer_failed or self._is_failed_recv_request(request_id): + transfer_failed = transfer_failed or self._is_failed_recv_request(request_id) + if transfer_failed: with self.pending_reformat_lock: self.pending_reformat.pop(request_id, None) else: @@ -774,6 +777,9 @@ def _handle_request(self, req_meta: dict[str, Any]): e, ) # Drain the failure with completion, after all local writes end. + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.transfer("error" if transfer_failed else "load", req_meta) if req_meta.get("awaiting_kvs", True): self.task_tracker.update_done_task_count(request_id) else: diff --git a/vllm_ascend/envs.py b/vllm_ascend/envs.py index ef496f301baa..442624107dc7 100644 --- a/vllm_ascend/envs.py +++ b/vllm_ascend/envs.py @@ -29,6 +29,14 @@ env_variables: dict[str, Callable[[], Any]] = { + # Worker V1/V2 KV read/write legality. 0 (default) disables all hooks, 1 enables + # address/valid-prefix checks and violation logs. Dense PA/MLA only; unadapted + # features fail initialization. No KV content dumps or checksums. Not sensitive. + "VLLM_ASCEND_KV_GUARD": lambda: bool(int(os.getenv("VLLM_ASCEND_KV_GUARD", "0"))), + # Enabled guard only: 'log' (default) reports failures; 'raise' stops worker + # output after a violation/coverage gap. Does not undo submitted kernels. + # No credentials or sensitive configuration. + "VLLM_ASCEND_KV_GUARD_ACTION": lambda: os.getenv("VLLM_ASCEND_KV_GUARD_ACTION", "log"), # max compile thread number for package building. Usually, it is set to # the number of CPU cores. If not set, the default value is None, which # means all number of CPU cores will be used. diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 76894c1c4304..182396f025ea 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1676,3 +1676,18 @@ # Remove this patch when upstream provides a backend capability hook for # enabling MoE sequence parallelism with DP=1. # +# ** File: platform/patch_kv_guard.py ** +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# 1. `vllm.v1.core.sched.scheduler.Scheduler` +# Why: +# Optional KV address diagnostics need allocation generations and request +# bindings independently of worker attention metadata. +# How: +# When KV Guard is enabled, wrap allocation, release, and schedule calls +# and attach changed bindings and token ranges to SchedulerOutput. +# Related PR (if no, explain why): +# No upstream PR yet; this is the opt-in Ascend KV Guard integration. +# Future Plan: +# Replace the wrappers with upstream allocation lifecycle and scheduler +# output extension points that provide the same independent evidence. +# diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index 13c53e5a0c33..eb3d0e3ebb7c 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -24,8 +24,14 @@ import vllm_ascend.patch.platform.patch_parallel_config # noqa import vllm_ascend.patch.platform.patch_pp_mtp # noqa import vllm_ascend.patch.platform.patch_use_v2_model_runner # noqa +from vllm_ascend import envs as _kv_guard_envs from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile +if _kv_guard_envs.VLLM_ASCEND_KV_GUARD: + from vllm_ascend.patch.platform.patch_kv_guard import apply_patch as _apply_kv_guard + + _apply_kv_guard() + if get_current_hardware_profile().supports(HardwareCapability.STANDARD_MAMBA_PATCH): import vllm_ascend.patch.platform.patch_mamba_config # noqa else: diff --git a/vllm_ascend/patch/platform/patch_kv_guard.py b/vllm_ascend/patch/platform/patch_kv_guard.py new file mode 100644 index 000000000000..92759bb375e2 --- /dev/null +++ b/vllm_ascend/patch/platform/patch_kv_guard.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Transport allocation generations to the Worker L1 state machine. + +No lifecycle trace, extra publisher, token hashes or per-token log messages. +Changed bindings, releases and compact per-request write ranges use the existing channel. +""" + +import functools +import uuid +from typing import Any, cast + + +def install_scheduler_guard(scheduler_cls): + if getattr(scheduler_cls.__init__, "_ascend_kv_guard", False): + return + original_init = scheduler_cls.__init__ + + @functools.wraps(original_init) + def initialize(self, *args, **kwargs): + original_init(self, *args, **kwargs) + pool = self.kv_cache_manager.block_pool + counters = [0] * len(pool.blocks) + pool_id = uuid.uuid4().hex + allocate = pool.get_new_blocks + + @functools.wraps(allocate) + def allocated(*args, **kwargs): + blocks = allocate(*args, **kwargs) + for block in blocks: + counters[block.block_id] += 1 + return blocks + + pool.get_new_blocks = allocated + schedule = self.schedule + mappings: dict[str, tuple] = {} + released: set[str] = set() + free = self.kv_cache_manager.free + + @functools.wraps(free) + def freed(request, *args, **kwargs): + result = free(request, *args, **kwargs) + released.add(request.request_id) + return result + + self.kv_cache_manager.free = freed + + @functools.wraps(schedule) + def scheduled(*args, **kwargs): + output = schedule(*args, **kwargs) + computed = {r.req_id: r.num_computed_tokens for r in output.scheduled_new_reqs} + cached = output.scheduled_cached_reqs + computed.update(zip(cached.req_ids, cached.num_computed_tokens)) + changed = {} + for request_id in output.finished_req_ids | released: + mappings.pop(request_id, None) + for request_id in output.num_scheduled_tokens: + blocks = [list(group) for group in self.kv_cache_manager.get_block_ids(request_id)] + refs = [[(block, counters[block]) for block in ids] for ids in blocks] + signature = tuple(tuple(group) for group in refs) + if mappings.get(request_id) != signature: + changed[request_id] = refs + mappings[request_id] = signature + output._ascend_kv_guard_context = { + "pool_id": pool_id, + "requests": changed, + "released_request_ids": sorted(output.finished_req_ids | released), + "scheduled_requests": { + request_id: (computed.get(request_id), count) + for request_id, count in output.num_scheduled_tokens.items() + }, + } + released.clear() + return output + + self.schedule = scheduled + + cast(Any, initialize)._ascend_kv_guard = True + scheduler_cls.__init__ = initialize + + +def apply_patch(): + from vllm.v1.core.sched.scheduler import Scheduler + + install_scheduler_guard(Scheduler) diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 144bc1aa8c22..fe766dd70fb5 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -155,6 +155,7 @@ ) from vllm_ascend.compilation.breakable_aclgraph import BreakableACLGraphWrapper from vllm_ascend.core.kv_cache_interface import is_circular_kv_cache_spec +from vllm_ascend.debug.kv_guard_runtime import get_guard from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h import get_prebound_copy_sfa_slots from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise_cache_layout import ( @@ -395,6 +396,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) + + self._kv_guard = get_guard(vllm_config) + if self._kv_guard is not None: + cast(Any, self)._model_forward = self._kv_guard.wrap_forward(self, self._model_forward, get_forward_context) + self.execute_model = self._kv_guard.wrap_execute(self.execute_model) + cast(Any, self)._update_states = self._kv_guard.wrap_update(self, self._update_states) self.lookback_token_ids: CpuGpuBuffer | None = None self.device_metadata_executor: DeviceMetadataExecutor | None = None @@ -1016,6 +1023,9 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.copy_blocks(scheduler_output.kv_cache_block_copies) scheduler_output = replace( scheduler_output, new_block_ids_to_zero=[], @@ -4587,6 +4597,9 @@ def dummy_drafter_compute_logits(hidden_states): if not is_graph_capturing and self.ascend_config.enable_force_eplb \ and self.vllm_config.model_config.is_moe: build_force_eplb_topk(self.device, self.max_num_tokens) + guard_kwargs: dict[str, Any] = {} + if getattr(self, "_kv_guard", None) is not None: + guard_kwargs["_kv_guard_phase"] = "warmup" if is_profile or is_graph_capturing else "dummy" outputs = self._model_forward( num_tokens_padded, input_ids, @@ -4594,6 +4607,7 @@ def dummy_drafter_compute_logits(hidden_states): intermediate_tensors, inputs_embeds, decoder_replay_dummy=True, + **guard_kwargs, ) if active_device_metadata_executor is not None and active_device_metadata_executor.submission_in_flight: active_device_metadata_executor.release() @@ -5070,6 +5084,9 @@ def _fill_c8_mxfp_v_scale_caches(self, kv_caches: dict[str, Any]) -> None: def _zero_block_ids(self, block_ids: list[int]) -> None: """Prepare freshly assigned attention pages, including static C8 scales.""" super()._zero_block_ids(block_ids) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.zero_blocks(block_ids) # MRV1 can receive new attention block IDs without a KV zeroer. Scale # restoration must still run: Mamba owns the same physical pages while # those IDs are assigned to its cache groups. @@ -5102,6 +5119,9 @@ def initialize_kv_cache_tensors( kv_caches = self._reshape_kv_cache_tensors( kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes ) + guard = getattr(self, "_kv_guard", None) + if guard is not None: + guard.configure(self) # Set up cross-layer KV cache sharing for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 46252f76c618..10c6bea66ea2 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -72,6 +72,7 @@ _finish_profiling_chunk_timing, _start_profiling_chunk_timing, ) +from vllm_ascend.debug.kv_guard_v2 import install_guard from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h.dspark_context import ( bind_dspark_context_receiver, configure_dspark_kv_transfer, @@ -281,6 +282,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.debugger = AclGraphDumper(dump_cfg) self.runtime_guard = RuntimeGuardProcessor.bind(self) + self._kv_guard_v2 = install_guard(self) @property def pcp_manager_cls(self) -> type[AscendPCPManager]: @@ -528,6 +530,9 @@ def initialize_kv_cache( self.model_state.kvpp_runtime = self.kvpp self._fill_c8_mxfp_v_scale_caches() + guard = getattr(self, "_kv_guard_v2", None) + if guard is not None: + guard.configure() def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" From 51ae23149d487e82a5906896b19293d8fc6bd2c6 Mon Sep 17 00:00:00 2001 From: TAJh Date: Tue, 27 Oct 2026 01:04:22 +0100 Subject: [PATCH 2/6] refactor(debug): configure and share KV cache guard runtime Configure per-engine KV address diagnostics and failure actions through additional_config.kv_guard_config instead of environment variables. Gate scheduler hooks per instance and share runner installation and execution state across V1 and V2, retaining only the V2-specific adapter. Signed-off-by: TAJh --- .../configuration/additional_config.md | 1 + .../feature_guide/kv_cache_guard.md | 51 ++++-- release_notes/fragments/17833.yaml | 5 +- .../one_card/test_kv_cache_guard.py | 14 +- tests/ut/debug/test_kv_guard_runtime.py | 86 ++++++++- tests/ut/debug/test_kv_guard_v2.py | 26 ++- vllm_ascend/ascend_config.py | 3 + vllm_ascend/debug/kv_guard_config.py | 25 +++ vllm_ascend/debug/kv_guard_runtime.py | 166 ++++++++++++++++-- vllm_ascend/debug/kv_guard_v2.py | 134 -------------- vllm_ascend/envs.py | 8 - vllm_ascend/patch/__init__.py | 5 +- vllm_ascend/patch/platform/__init__.py | 7 +- vllm_ascend/patch/platform/patch_kv_guard.py | 4 + vllm_ascend/worker/model_runner_v1.py | 8 +- vllm_ascend/worker/v2/model_runner.py | 8 +- 16 files changed, 328 insertions(+), 223 deletions(-) create mode 100644 vllm_ascend/debug/kv_guard_config.py delete mode 100644 vllm_ascend/debug/kv_guard_v2.py diff --git a/docs/source/user_guide/configuration/additional_config.md b/docs/source/user_guide/configuration/additional_config.md index f1049218254b..8319e780bc71 100644 --- a/docs/source/user_guide/configuration/additional_config.md +++ b/docs/source/user_guide/configuration/additional_config.md @@ -50,6 +50,7 @@ The following table lists additional configuration options available in vLLM Asc |-------------------------------------|------|---------|-----------------------------------------------------------------------------------------------------------| | `xlite_graph_config` | dict | `{}` | Configuration options for Xlite graph mode | | `finegrained_tp_config` | dict | `{}` | Configuration options for fine-grained tensor parallelism. See [Fine-Grained Tensor Parallelism](../feature_guide/Fine_grained_TP.md). | +| `kv_guard_config` | dict | `{}` | Per-engine KV cache address diagnostics for Model Runner V1 and V2. `enabled` defaults to `false`; `action` accepts `"log"` (default) or `"raise"`. See [KV cache address guard](../feature_guide/kv_cache_guard.md). | | `ascend_compilation_config` | dict | `{}` | Configuration options for ascend compilation | | `ascend_warmup_config` | dict | `{}` | Configuration options for startup warmup that overlaps weight loading | | `eplb_config` | dict | `{}` | Runner-specific EPLB extensions. See [Expert Parallelism Load Balancer](../feature_guide/expert_parallelism_load_balancer.md). | diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md index 9cc551fcf0ce..bf889f9c9574 100644 --- a/docs/source/user_guide/feature_guide/kv_cache_guard.md +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -8,25 +8,37 @@ attention paths in Model Runner V1 and V2. ## Enable the guard -```bash -# Model Runner V2 -VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ASCEND_KV_GUARD=1 vllm serve MODEL +Use the same per-engine configuration for Model Runner V1 and V2: -# Model Runner V1 -VLLM_USE_V2_MODEL_RUNNER=0 VLLM_ASCEND_KV_GUARD=1 vllm serve MODEL +```bash +vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true}}' ``` -`VLLM_ASCEND_KV_GUARD_ACTION=log` is the default: it reports failures while -preserving diagnostic execution. To stop returning worker results after a -violation or incomplete audit, select `raise`: +| `kv_guard_config` option | Type | Default | Meaning | +| --- | --- | --- | --- | +| `enabled` | bool | `false` | Enable KV address diagnostics for this engine. | +| `action` | `"log"` or `"raise"` | `"log"` | Report failures or stop returning worker results. | + +`action="log"` reports failures while preserving diagnostic execution. To stop +returning worker results after a violation or incomplete audit, select `raise`: ```bash -VLLM_USE_V2_MODEL_RUNNER=1 \ -VLLM_ASCEND_KV_GUARD=1 \ -VLLM_ASCEND_KV_GUARD_ACTION=raise \ -vllm serve MODEL +vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' ``` +For the Python API, pass the same dictionary through `additional_config`: + +```python +from vllm import LLM + +llm = LLM(model="MODEL", additional_config={"kv_guard_config": {"enabled": True, "action": "raise"}}) +``` + +The scheduler and workers read this configuration independently for each engine. +Omitting `kv_guard_config` or setting `enabled=false` disables instrumentation. +Invalid option names, non-boolean `enabled` values, and unknown actions are rejected. +KV guard environment variables are no longer used. + The first failure is latched. Catching the exception does not reset the worker; restart it before reuse. This is a worker-level stop, not cancellation of one request, and it adds no collective failure coordination across ranks. @@ -151,15 +163,20 @@ can still synchronize device metadata to CPU. Measure deployment overhead. | Component | Responsibility | | --- | --- | | `vllm_ascend/debug/kv_guard.py` | Allocation, ownership, prefix validity, and read/write rules. | -| `vllm_ascend/debug/kv_guard_runtime.py` | Independent expected addresses, evidence collection, per-layer audit, and conservative state commit. | -| `vllm_ascend/debug/kv_guard_v2.py` | Sorted batch IDs, execute/update boundaries, and capture restrictions for Runner V2. | +| `vllm_ascend/debug/kv_guard_config.py` | Validated per-engine options shared by the scheduler and workers. | +| `vllm_ascend/debug/kv_guard_runtime.py` | Shared V1/V2 installation, scheduler binding, evidence collection, audit, and conservative state commit. Includes the V2 batch and capture adapter. | | `vllm_ascend/patch/platform/patch_kv_guard.py` | Conditional scheduler allocation and request-binding instrumentation. | | Runner and attention implementations | Existing execution boundaries and ordered eager access arguments. | | Mooncake receiver | Success/failure validity callbacks after receive completion and reformatting. | -The scheduler patch is conditional because current public hooks do not provide -both allocation generations and the independent scheduler-to-worker request -contract. Its removal path is an upstream allocation-lifecycle and scheduler +Both runners use one `GuardRuntime` and one address-checking state machine. +The internal V2 adapter only handles sorted batch IDs, metadata timing, and +capture boundaries; it does not implement another checker. + +The scheduler constructor is wrapped once; allocation, release, and schedule +methods are instrumented only for engines with `kv_guard_config.enabled=true`. +Current public hooks do not provide both allocation generations and the +independent scheduler-to-worker request contract. Its removal path is an upstream allocation-lifecycle and scheduler output extension that supplies the same evidence without wrapping instance methods. Changes to those upstream contracts require adapter and lifecycle regressions, alongside NPU tests for affected execution paths. diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml index 998b026f0464..7f7092f9af7a 100644 --- a/release_notes/fragments/17833.yaml +++ b/release_notes/fragments/17833.yaml @@ -4,5 +4,6 @@ entries: text: >- Added opt-in KV cache address-legality diagnostics for adapted dense PA and MLA attention paths in Model Runner V1 and V2. The guard is disabled - by default and can log failures or stop returning worker results after - a violation or incomplete audit. + by default and configured per engine through additional_config.kv_guard_config + (enabled and action). Both runners share one checker, which can log failures + or stop returning worker results after a violation or incomplete audit. diff --git a/tests/e2e/pull_request/one_card/test_kv_cache_guard.py b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py index 290f4c3a7c43..0de7114cc7d6 100644 --- a/tests/e2e/pull_request/one_card/test_kv_cache_guard.py +++ b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py @@ -65,15 +65,17 @@ def test_guard_preserves_sequential_and_concurrent_generation(use_v2, graph_mode results = [] for enabled in (False, True): port = get_open_port() - environment = { - "VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2)), - "VLLM_ASCEND_KV_GUARD": str(int(enabled)), - "VLLM_ASCEND_KV_GUARD_ACTION": "raise", - } + environment = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2))} with ( RemoteOpenAIServer( model, - [*arguments, "--port", str(port)], + [ + *arguments, + "--additional-config", + json.dumps({"kv_guard_config": {"enabled": enabled, "action": "raise"}}), + "--port", + str(port), + ], env_dict=environment, seed=1234, server_host="127.0.0.1", diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py index f9c4a02cfeca..1d0e94759072 100644 --- a/tests/ut/debug/test_kv_guard_runtime.py +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -29,11 +29,14 @@ def load(name, path): def runtime_fixture(monkeypatch): package = ModuleType("vllm_ascend") package.__path__ = [] - cast(Any, package).envs = NS(VLLM_ASCEND_KV_GUARD=True, VLLM_ASCEND_KV_GUARD_ACTION="log") debug = ModuleType("vllm_ascend.debug") debug.__path__ = [] monkeypatch.setitem(sys.modules, "vllm_ascend", package) monkeypatch.setitem(sys.modules, "vllm_ascend.debug", debug) + utils = load("guard_config_utils", "vllm_ascend/config_utils.py") + monkeypatch.setitem(sys.modules, "vllm_ascend.config_utils", utils) + options = load("guard_config_subject", "vllm_ascend/debug/kv_guard_config.py") + monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_config", options) state = load("guard_runtime_state", "vllm_ascend/debug/kv_guard.py") monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard", state) subject = load("guard_runtime_subject", "vllm_ascend/debug/kv_guard_runtime.py") @@ -43,7 +46,7 @@ def runtime_fixture(monkeypatch): vllm.__path__ = [] monkeypatch.setitem(sys.modules, "vllm", vllm) monkeypatch.setitem(sys.modules, "vllm.logger", logger) - config = NS() + config = NS(additional_config={"kv_guard_config": {"enabled": True}}) runtime = subject.get_guard(config) assert subject.get_guard(config) is runtime runtime.subject = subject @@ -139,12 +142,77 @@ def test_receive_only_step_binds_before_execute_and_cleans_finished(runtime): def test_disabled_guard_does_not_create_runtime(runtime, monkeypatch): - sys.modules["vllm_ascend"].envs.VLLM_ASCEND_KV_GUARD = False + monkeypatch.setenv("VLLM_ASCEND_KV_GUARD", "1") + monkeypatch.setenv("VLLM_ASCEND_KV_GUARD_ACTION", "raise") config = NS() assert runtime.subject.get_guard(config) is None assert not hasattr(config, "_ascend_kv_guard") +@pytest.mark.parametrize("use_v2", [False, True]) +def test_guard_configuration_is_per_engine(runtime, use_v2, monkeypatch): + monkeypatch.setenv("VLLM_ASCEND_KV_GUARD", "0") + configurations = [ + NS(use_v2_model_runner=use_v2, additional_config={"kv_guard_config": options}) + for options in ({"enabled": True}, {"enabled": False}, {"enabled": True, "action": "raise"}) + ] + log_guard, disabled_guard, raise_guard = [runtime.subject.get_guard(config) for config in configurations] + assert log_guard is not None and log_guard.action == "log" + assert disabled_guard is None + assert raise_guard is not None and raise_guard.action == "raise" + assert log_guard is not raise_guard + assert runtime.subject.get_guard(configurations[0]) is log_guard + assert runtime.subject.get_guard(configurations[2]) is raise_guard + + +@pytest.mark.parametrize( + "options", [None, [], True, {"enabled": "false"}, {"enabled": 1}, {"action": "stop"}, {"enable": True}] +) +def test_guard_rejects_invalid_configuration(runtime, options): + config = NS(additional_config={"kv_guard_config": options}) + with pytest.raises(ValueError): + runtime.subject.get_guard(config) + assert not hasattr(config, "_ascend_kv_guard") + + +def test_v1_installs_the_shared_runtime_at_existing_boundaries(runtime, monkeypatch): + context = NS(cudagraph_runtime_mode="NONE") + forward_context = ModuleType("vllm.forward_context") + cast(Any, forward_context).get_forward_context = lambda: context + monkeypatch.setitem(sys.modules, "vllm.forward_context", forward_context) + calls = [] + result = object() + runner = NS(vllm_config=runtime.config, input_batch=NS(req_ids=["r"]), requests={"r": NS(block_ids=[[1]])}) + + def forward(num_tokens, *, marker): + assert marker is result + assert runtime._phase == "forward" + runtime.write("layer0", torch.tensor([128])) + return marker + + def update(output): + calls.append("update") + + def execute(output, *, marker): + assert runtime._scheduled == {"r": (0, 1)} + runner._update_states(output) + return runner._model_forward(1, marker=marker) + + runner._model_forward = forward + runner._update_states = update + runner.execute_model = execute + assert runtime.subject.install_guard(runner) is runtime + packet = NS( + _ascend_kv_guard_context={"pool_id": "pool", "requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (0, 1)}} + ) + assert runner.execute_model(packet, marker=result) is result + assert context._kv_guard is runtime + assert calls == ["update"] + assert runtime.state.blocks[0, 1].valid_end == 1 + assert runtime._phase is None and runtime._scheduled is None + assert not runtime.events + + @pytest.mark.parametrize("kernel_size, accepted", [(128, True), (16, False)]) def test_unadapted_kernel_block_coordinates_are_rejected(runtime, kernel_size, accepted): impl = type("AscendAttentionBackendImpl", (), {})() @@ -181,8 +249,8 @@ def get_new_blocks(self, count): return [self.blocks[1]] class Scheduler: - def __init__(self): - self.vllm_config = NS(use_v2_model_runner=use_v2) + def __init__(self, vllm_config): + self.vllm_config = vllm_config self.pool = Pool() self.mapping = {"r": [[1]]} self.kv_cache_manager = NS( @@ -200,7 +268,13 @@ def schedule(self): patch.install_scheduler_guard(Scheduler) patch.install_scheduler_guard(Scheduler) - scheduler = Scheduler() + scheduler = Scheduler(NS(use_v2_model_runner=use_v2, additional_config={"kv_guard_config": {"enabled": True}})) + for config in (NS(additional_config=None), NS(additional_config={"kv_guard_config": {"enabled": False}})): + disabled = Scheduler(vllm_config=config) + assert "schedule" not in vars(disabled) + assert "get_new_blocks" not in vars(disabled.pool) + assert disabled.schedule() is disabled.output + assert not hasattr(disabled.output, "_ascend_kv_guard_context") scheduler.pool.get_new_blocks(1) assert scheduler.schedule() is scheduler.output binding = scheduler.output._ascend_kv_guard_context diff --git a/tests/ut/debug/test_kv_guard_v2.py b/tests/ut/debug/test_kv_guard_v2.py index 2e8581d865ab..de400b646133 100644 --- a/tests/ut/debug/test_kv_guard_v2.py +++ b/tests/ut/debug/test_kv_guard_v2.py @@ -2,23 +2,23 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MRV2 batch ordering, graph updates, dummy writes and lifecycle integration.""" -import sys from types import SimpleNamespace as NS from typing import Any import pytest import torch -from .test_kv_guard_runtime import bind, load, metadata +from .test_kv_guard_runtime import bind, metadata from .test_kv_guard_runtime import runtime_fixture as runtime_fixture @pytest.fixture -def adapter(runtime, monkeypatch): - monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_runtime", runtime.subject) - subject = load("guard_v2_subject", "vllm_ascend/debug/kv_guard_v2.py") +def adapter(runtime): + subject = runtime.subject impl = type("AscendAttentionBackendImpl", (), {})() runtime.config = NS( + use_v2_model_runner=True, + additional_config={"kv_guard_config": {"enabled": True}}, parallel_config=NS(), model_config=NS(), speculative_config=None, @@ -36,8 +36,9 @@ def adapter(runtime, monkeypatch): next_metadata={}, ) runtime.layer_groups.clear() - result = subject.KVGuardV2(runner, runtime) - result.configure() + result = subject._V2RunnerAdapter(runner, runtime) + runtime._runner_adapter = result + runtime.configure(runner) result.subject = subject return result @@ -232,11 +233,8 @@ def test_v2_kernel_allocator_coordinate_mismatch_is_rejected(adapter): adapter.configure() -def test_v2_install_disabled_does_not_wrap_methods(adapter, monkeypatch): - from vllm_ascend import envs - - monkeypatch.setattr(envs, "VLLM_ASCEND_KV_GUARD", False) - assert adapter.subject.install_guard(NS(vllm_config=NS())) is None +def test_v2_install_disabled_does_not_wrap_methods(adapter): + assert adapter.subject.install_guard(NS(vllm_config=NS(use_v2_model_runner=True))) is None def enable_raise(runtime): @@ -297,9 +295,9 @@ def test_v2_install_wraps_real_runner_boundaries_and_reinitializes_model_state(a assert installed is not None # Profiling can replace the model-state object before real cache setup. runner.model_state = NS(prepare_attn=lambda batch, mode: runner.next_metadata) - installed.configure() + installed.configure(runner) wrapped = runner.model_state.prepare_attn - installed.configure() + installed.configure(runner) assert runner.model_state.prepare_attn is wrapped assert runner.capture_model() == "capture" runner.execute_model(message()) diff --git a/vllm_ascend/ascend_config.py b/vllm_ascend/ascend_config.py index 55a0e9971b3b..d3d6c0f8b408 100644 --- a/vllm_ascend/ascend_config.py +++ b/vllm_ascend/ascend_config.py @@ -29,6 +29,7 @@ from vllm.utils.math_utils import cdiv from vllm_ascend.config_utils import config +from vllm_ascend.debug.kv_guard_config import KVGuardConfig from vllm_ascend.device.device_config import is_950 from vllm_ascend.device.hardware_profile import DeviceAdaptorFamily, HardwareCapability, get_current_hardware_profile @@ -426,6 +427,7 @@ class AscendConfig: "enable_prefill_mc2": false, "multistream_overlap_shared_expert": false, "enable_kv_nz": false, + "kv_guard_config": {"enabled": false, "action": "log"}, "kv_pool_range_debug": false, "enable_mc2_hierarchy_comm": false, "enable_dsa_cp": false, @@ -575,6 +577,7 @@ class AscendConfig: enable_prefill_mc2: bool = False multistream_overlap_shared_expert: bool = False enable_kv_nz: bool = False + kv_guard_config: KVGuardConfig = dataclasses.field(default_factory=KVGuardConfig) kv_pool_range_debug: bool = False enable_mc2_hierarchy_comm: bool = False # deprecated, will be replaced by mc2_comm_alg = "hierarchy" enable_dsa_cp: bool = False diff --git a/vllm_ascend/debug/kv_guard_config.py b/vllm_ascend/debug/kv_guard_config.py new file mode 100644 index 000000000000..93fecd8a7a29 --- /dev/null +++ b/vllm_ascend/debug/kv_guard_config.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-engine KV guard options, usable before AscendConfig initialization.""" + +from typing import Any, Literal + +from pydantic import StrictBool + +from vllm_ascend.config_utils import config + + +@config +class KVGuardConfig: + """KV address diagnostics configured through additional_config.""" + + enabled: StrictBool = False + action: Literal["log", "raise"] = "log" + + @classmethod + def from_vllm_config(cls, vllm_config: Any) -> "KVGuardConfig": + additional_config = getattr(vllm_config, "additional_config", None) or {} + options = additional_config.get("kv_guard_config", {}) + if not isinstance(options, dict): + raise ValueError("additional_config.kv_guard_config must be a dict") + return cls(**options) diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py index afe57f4d6bfa..79357d5f8c35 100644 --- a/vllm_ascend/debug/kv_guard_runtime.py +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -15,9 +15,11 @@ from collections import OrderedDict from contextlib import contextmanager from dataclasses import dataclass, replace -from typing import Any +from types import SimpleNamespace +from typing import Any, cast from vllm_ascend.debug.kv_guard import INVALID_PREFIX, KVGuard +from vllm_ascend.debug.kv_guard_config import KVGuardConfig _MAX_REPORT_KEYS = 1024 @@ -69,7 +71,8 @@ def __init__(self, vllm_config: Any): self._reported: OrderedDict[tuple, None] = OrderedDict() self._scheduled: dict[str, tuple[int | None, int]] | None = None self._failure: str | None = None - self.action = "log" + self.action = KVGuardConfig.from_vllm_config(vllm_config).action + self._runner_adapter: _V2RunnerAdapter | None = None def _check_failure(self) -> None: if self.action == "raise" and self._failure is not None: @@ -115,11 +118,6 @@ def gap(self, reason: str, **fields: Any) -> None: self.report("guard.coverage_gap", reason=reason, **fields) def configure(self, runner: Any) -> None: - from vllm_ascend import envs - - self.action = envs.VLLM_ASCEND_KV_GUARD_ACTION - if self.action not in ("log", "raise"): - raise ValueError("VLLM_ASCEND_KV_GUARD_ACTION must be 'log' or 'raise'") config = self.config parallel = config.parallel_config transfer = config.kv_transfer_config @@ -163,6 +161,8 @@ def configure(self, runner: Any) -> None: "(CP, DBO, speculation, encoder/decoder, connector)" ) self.state.configure(sizes) + if self._runner_adapter is not None: + self._runner_adapter.configure() self.report("guard.enabled", failure_action=self.action) def write(self, layer: str, slots: Any, token_offset: int = 0) -> None: @@ -505,10 +505,12 @@ def observed(output, *args, **kwargs): return observed - def wrap_execute(self, execute: Any): - @functools.wraps(execute) - def observed(output, *args, **kwargs): - self._check_failure() + @contextmanager + def execute_scope(self, output: Any, *, bind: bool = True): + """Share scheduler binding and cleanup across both runner versions.""" + self._check_failure() + self._scheduled = None + if bind: binding = getattr(output, "_ascend_kv_guard_context", None) if binding is not None: # Receive-only and idle DP steps can return before _update_states. @@ -522,12 +524,17 @@ def observed(output, *args, **kwargs): self.gap("scheduler_write_plan_missing") else: self.gap("scheduler_binding_missing") - self._scheduled = None - try: - self._check_failure() + try: + self._check_failure() + yield + finally: + self._scheduled = None + + def wrap_execute(self, execute: Any): + @functools.wraps(execute) + def observed(output, *args, **kwargs): + with self.execute_scope(output): return execute(output, *args, **kwargs) - finally: - self._scheduled = None return observed @@ -575,12 +582,133 @@ def transfer(self, event: str, meta: dict[str, Any]) -> None: def get_guard(vllm_config: Any) -> GuardRuntime | None: - from vllm_ascend import envs - - if not envs.VLLM_ASCEND_KV_GUARD: + if not KVGuardConfig.from_vllm_config(vllm_config).enabled: return None guard = getattr(vllm_config, "_ascend_kv_guard", None) if guard is None: guard = GuardRuntime(vllm_config) vllm_config._ascend_kv_guard = guard return guard + + +class _V2RunnerAdapter: + """Only V2 batch ordering, metadata timing and capture need adaptation.""" + + def __init__(self, runner: Any, runtime: GuardRuntime): + self.runner = runner + self.runtime = runtime + self.input_batch = SimpleNamespace(req_ids=[]) + self._prepared = False + + def configure(self) -> None: + tables = self.runner.block_tables + sizes = [group.kv_cache_spec.block_size for group in self.runner.kv_cache_config.kv_cache_groups] + if list(tables.block_sizes) != sizes or list(tables.kernel_block_sizes) != sizes: + raise ValueError("KV guard: unadapted MRV2 allocator/kernel block coordinates") + prepare = self.runner.model_state.prepare_attn + if getattr(prepare, "_ascend_kv_guard", False): + return + + @functools.wraps(prepare) + def observed(input_batch, cudagraph_mode, *args, **kwargs): + metadata = prepare(input_batch, cudagraph_mode, *args, **kwargs) + if self.runtime._phase is not None: + self.input_batch = SimpleNamespace(req_ids=list(input_batch.req_ids[: input_batch.num_reqs])) + self._prepared = True + self.runtime._full_graph = "FULL" in str(cudagraph_mode) + self.runtime.observe_metadata(metadata) + scheduled = self.runtime._scheduled + if self.runtime._phase == "forward" and scheduled is not None: + ids = self.input_batch.req_ids + if len(ids) != len(set(ids)) or set(ids) != set(scheduled): + self.runtime.gap("v2_batch_request_mismatch") + self.runtime._missing_groups.update(self.runtime.layer_groups.values()) + self.runtime._check_failure() + return metadata + + cast(Any, observed)._ascend_kv_guard = True + self.runner.model_state.prepare_attn = observed + + def wrap_execute(self, execute: Any): + @functools.wraps(execute) + def observed(scheduler_output, *args, **kwargs): + # Forward positional and keyword arguments unchanged across vLLM lanes. + flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) + flags.update(kwargs) + dummy = bool(flags.get("dummy_run", False)) + profile = bool(flags.get("is_profile", False)) + phase = "warmup" if profile else "dummy" if dummy else "forward" + required = bool(scheduler_output.total_num_scheduled_tokens) and not ( + profile and flags.get("skip_attn_for_dummy_run") + ) + self.input_batch = SimpleNamespace(req_ids=[]) + self._prepared = False + + # Synthetic dummy/profile outputs have no scheduler binding, but + # their writes must still be checked against the shared live KV state. + with ( + self.runtime.execute_scope(scheduler_output, bind=not (dummy or profile)), + self.runtime.forward_scope(self, phase=phase, required=required), + ): + result = execute(scheduler_output, *args, **kwargs) + if phase == "forward" and required and not self._prepared: + self.runtime.gap("v2_attention_metadata_missing") + self.runtime._missing_groups.update(self.runtime.layer_groups.values()) + return result + + return observed + + def wrap_update(self, update: Any): + @functools.wraps(update) + def observed(scheduler_output, *args, **kwargs): + result = update(scheduler_output, *args, **kwargs) + # The parent has completed the physical operations at this point. + zeroed = getattr(scheduler_output, "new_block_ids_to_zero", None) + if zeroed: + self.runtime.zero_blocks(zeroed) + copies = getattr(scheduler_output, "kv_cache_block_copies", None) + if copies: + self.runtime.copy_blocks(copies) + return result + + return observed + + def wrap_capture(self, capture: Any): + @functools.wraps(capture) + def observed(*args, **kwargs): + self.runtime._check_failure() + with self.runtime.state._lock: + if ( + self.runtime.state.requests + or self.runtime.state.pending_loads + or any(block.valid_end != 0 for block in self.runtime.state.blocks.values()) + ): + raise ValueError( + "KV guard: capture requires empty KV state; live/cache-preserving recapture is unsupported" + ) + # Startup capture has no request KV to protect. Do not synchronize + # metadata to the CPU inside a device graph capture. Real replays + # are audited by wrap_execute using their current metadata. + return capture(*args, **kwargs) + + return observed + + +def install_guard(runner: Any) -> GuardRuntime | None: + """Install one per-engine guard with only the runner-specific hooks.""" + runtime = get_guard(runner.vllm_config) + if runtime is None: + return None + if getattr(runner.vllm_config, "use_v2_model_runner", False): + adapter = _V2RunnerAdapter(runner, runtime) + runtime._runner_adapter = adapter + cast(Any, runner).execute_model = adapter.wrap_execute(runner.execute_model) + cast(Any, runner).update_requests = adapter.wrap_update(runner.update_requests) + cast(Any, runner).capture_model = adapter.wrap_capture(runner.capture_model) + else: + from vllm.forward_context import get_forward_context + + cast(Any, runner)._model_forward = runtime.wrap_forward(runner, runner._model_forward, get_forward_context) + cast(Any, runner).execute_model = runtime.wrap_execute(runner.execute_model) + cast(Any, runner)._update_states = runtime.wrap_update(runner, runner._update_states) + return runtime diff --git a/vllm_ascend/debug/kv_guard_v2.py b/vllm_ascend/debug/kv_guard_v2.py deleted file mode 100644 index 54b103ba4ba3..000000000000 --- a/vllm_ascend/debug/kv_guard_v2.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MRV2 batch adapter for the shared KV address checker. - -Batch IDs come from prepare_attn after sorting, not RequestState row order. -Metadata is retained by reference and copied only after forward/graph updates. -""" - -import functools -from types import SimpleNamespace -from typing import Any, cast - -from vllm_ascend.debug.kv_guard_runtime import GuardRuntime, get_guard - - -class KVGuardV2: - def __init__(self, runner: Any, runtime: GuardRuntime): - self.runner = runner - self.runtime = runtime - self.input_batch = SimpleNamespace(req_ids=[]) - self._prepared = False - - def configure(self) -> None: - tables = self.runner.block_tables - sizes = [group.kv_cache_spec.block_size for group in self.runner.kv_cache_config.kv_cache_groups] - if list(tables.block_sizes) != sizes or list(tables.kernel_block_sizes) != sizes: - raise ValueError("KV guard: unadapted MRV2 allocator/kernel block coordinates") - self.runtime.configure(self.runner) - prepare = self.runner.model_state.prepare_attn - if getattr(prepare, "_ascend_kv_guard", False): - return - - @functools.wraps(prepare) - def observed(input_batch, cudagraph_mode, *args, **kwargs): - metadata = prepare(input_batch, cudagraph_mode, *args, **kwargs) - if self.runtime._phase is not None: - self.input_batch = SimpleNamespace(req_ids=list(input_batch.req_ids[: input_batch.num_reqs])) - self._prepared = True - self.runtime._full_graph = "FULL" in str(cudagraph_mode) - self.runtime.observe_metadata(metadata) - scheduled = self.runtime._scheduled - if self.runtime._phase == "forward" and scheduled is not None: - ids = self.input_batch.req_ids - if len(ids) != len(set(ids)) or set(ids) != set(scheduled): - self.runtime.gap("v2_batch_request_mismatch") - self.runtime._missing_groups.update(self.runtime.layer_groups.values()) - self.runtime._check_failure() - return metadata - - cast(Any, observed)._ascend_kv_guard = True - self.runner.model_state.prepare_attn = observed - - def wrap_execute(self, execute: Any): - def forward(output, *args, **kwargs): - # Preserve positional and keyword calls across the supported vLLM - # lanes; forward every argument unchanged to the original method. - flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) - flags.update(kwargs) - dummy = bool(flags.get("dummy_run", False)) - profile = bool(flags.get("is_profile", False)) - phase = "warmup" if profile else "dummy" if dummy else "forward" - required = bool(output.total_num_scheduled_tokens) and not ( - profile and flags.get("skip_attn_for_dummy_run") - ) - self.input_batch = SimpleNamespace(req_ids=[]) - self._prepared = False - - with self.runtime.forward_scope(self, phase=phase, required=required): - result = execute(output, *args, **kwargs) - if phase == "forward" and required and not self._prepared: - self.runtime.gap("v2_attention_metadata_missing") - self.runtime._missing_groups.update(self.runtime.layer_groups.values()) - return result - - execute_bound = self.runtime.wrap_execute(forward) - - @functools.wraps(execute) - def observed(scheduler_output, *args, **kwargs): - flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) - flags.update(kwargs) - # Dummy/profile outputs are synthesized without a scheduler binding. - # They must still audit writes against any existing live KV state. - if flags.get("dummy_run", False) or flags.get("is_profile", False): - self.runtime._scheduled = None - return forward(scheduler_output, *args, **kwargs) - return execute_bound(scheduler_output, *args, **kwargs) - - return observed - - def wrap_update(self, update: Any): - @functools.wraps(update) - def observed(scheduler_output, *args, **kwargs): - result = update(scheduler_output, *args, **kwargs) - # The parent has completed the physical operations at this point. - zeroed = getattr(scheduler_output, "new_block_ids_to_zero", None) - if zeroed: - self.runtime.zero_blocks(zeroed) - copies = getattr(scheduler_output, "kv_cache_block_copies", None) - if copies: - self.runtime.copy_blocks(copies) - return result - - return observed - - def wrap_capture(self, capture: Any): - @functools.wraps(capture) - def observed(*args, **kwargs): - self.runtime._check_failure() - with self.runtime.state._lock: - if ( - self.runtime.state.requests - or self.runtime.state.pending_loads - or any(block.valid_end != 0 for block in self.runtime.state.blocks.values()) - ): - raise ValueError( - "KV guard: capture requires empty KV state; live/cache-preserving recapture is unsupported" - ) - # Startup capture has no request KV to protect. Do not synchronize - # metadata to the CPU inside a device graph capture. Real replays - # are audited by wrap_execute using their current metadata. - return capture(*args, **kwargs) - - return observed - - -def install_guard(runner: Any) -> KVGuardV2 | None: - runtime = get_guard(runner.vllm_config) - if runtime is None: - return None - adapter = KVGuardV2(runner, runtime) - cast(Any, runner).execute_model = adapter.wrap_execute(runner.execute_model) - cast(Any, runner).update_requests = adapter.wrap_update(runner.update_requests) - cast(Any, runner).capture_model = adapter.wrap_capture(runner.capture_model) - return adapter diff --git a/vllm_ascend/envs.py b/vllm_ascend/envs.py index 442624107dc7..ef496f301baa 100644 --- a/vllm_ascend/envs.py +++ b/vllm_ascend/envs.py @@ -29,14 +29,6 @@ env_variables: dict[str, Callable[[], Any]] = { - # Worker V1/V2 KV read/write legality. 0 (default) disables all hooks, 1 enables - # address/valid-prefix checks and violation logs. Dense PA/MLA only; unadapted - # features fail initialization. No KV content dumps or checksums. Not sensitive. - "VLLM_ASCEND_KV_GUARD": lambda: bool(int(os.getenv("VLLM_ASCEND_KV_GUARD", "0"))), - # Enabled guard only: 'log' (default) reports failures; 'raise' stops worker - # output after a violation/coverage gap. Does not undo submitted kernels. - # No credentials or sensitive configuration. - "VLLM_ASCEND_KV_GUARD_ACTION": lambda: os.getenv("VLLM_ASCEND_KV_GUARD_ACTION", "log"), # max compile thread number for package building. Usually, it is set to # the number of CPU cores. If not set, the default value is None, which # means all number of CPU cores will be used. diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 182396f025ea..f7f875a8f83f 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1683,8 +1683,9 @@ # Optional KV address diagnostics need allocation generations and request # bindings independently of worker attention metadata. # How: -# When KV Guard is enabled, wrap allocation, release, and schedule calls -# and attach changed bindings and token ranges to SchedulerOutput. +# Wrap the constructor once, then instrument allocation, release, and +# schedule only for instances with additional_config.kv_guard_config.enabled. +# Attach changed bindings and token ranges to SchedulerOutput. # Related PR (if no, explain why): # No upstream PR yet; this is the opt-in Ascend KV Guard integration. # Future Plan: diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index eb3d0e3ebb7c..6452e446caa3 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -24,13 +24,10 @@ import vllm_ascend.patch.platform.patch_parallel_config # noqa import vllm_ascend.patch.platform.patch_pp_mtp # noqa import vllm_ascend.patch.platform.patch_use_v2_model_runner # noqa -from vllm_ascend import envs as _kv_guard_envs from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile +from vllm_ascend.patch.platform.patch_kv_guard import apply_patch as _apply_kv_guard -if _kv_guard_envs.VLLM_ASCEND_KV_GUARD: - from vllm_ascend.patch.platform.patch_kv_guard import apply_patch as _apply_kv_guard - - _apply_kv_guard() +_apply_kv_guard() if get_current_hardware_profile().supports(HardwareCapability.STANDARD_MAMBA_PATCH): import vllm_ascend.patch.platform.patch_mamba_config # noqa diff --git a/vllm_ascend/patch/platform/patch_kv_guard.py b/vllm_ascend/patch/platform/patch_kv_guard.py index 92759bb375e2..1a64b60878d1 100644 --- a/vllm_ascend/patch/platform/patch_kv_guard.py +++ b/vllm_ascend/patch/platform/patch_kv_guard.py @@ -10,6 +10,8 @@ import uuid from typing import Any, cast +from vllm_ascend.debug.kv_guard_config import KVGuardConfig + def install_scheduler_guard(scheduler_cls): if getattr(scheduler_cls.__init__, "_ascend_kv_guard", False): @@ -19,6 +21,8 @@ def install_scheduler_guard(scheduler_cls): @functools.wraps(original_init) def initialize(self, *args, **kwargs): original_init(self, *args, **kwargs) + if not KVGuardConfig.from_vllm_config(self.vllm_config).enabled: + return pool = self.kv_cache_manager.block_pool counters = [0] * len(pool.blocks) pool_id = uuid.uuid4().hex diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index fe766dd70fb5..1dfd3ee058de 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -155,7 +155,7 @@ ) from vllm_ascend.compilation.breakable_aclgraph import BreakableACLGraphWrapper from vllm_ascend.core.kv_cache_interface import is_circular_kv_cache_spec -from vllm_ascend.debug.kv_guard_runtime import get_guard +from vllm_ascend.debug.kv_guard_runtime import install_guard from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h import get_prebound_copy_sfa_slots from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise_cache_layout import ( @@ -397,11 +397,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - self._kv_guard = get_guard(vllm_config) - if self._kv_guard is not None: - cast(Any, self)._model_forward = self._kv_guard.wrap_forward(self, self._model_forward, get_forward_context) - self.execute_model = self._kv_guard.wrap_execute(self.execute_model) - cast(Any, self)._update_states = self._kv_guard.wrap_update(self, self._update_states) + self._kv_guard = install_guard(self) self.lookback_token_ids: CpuGpuBuffer | None = None self.device_metadata_executor: DeviceMetadataExecutor | None = None diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 10c6bea66ea2..36b2682301f3 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -72,7 +72,7 @@ _finish_profiling_chunk_timing, _start_profiling_chunk_timing, ) -from vllm_ascend.debug.kv_guard_v2 import install_guard +from vllm_ascend.debug.kv_guard_runtime import install_guard from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h.dspark_context import ( bind_dspark_context_receiver, configure_dspark_kv_transfer, @@ -282,7 +282,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.debugger = AclGraphDumper(dump_cfg) self.runtime_guard = RuntimeGuardProcessor.bind(self) - self._kv_guard_v2 = install_guard(self) + self._kv_guard = install_guard(self) @property def pcp_manager_cls(self) -> type[AscendPCPManager]: @@ -530,9 +530,9 @@ def initialize_kv_cache( self.model_state.kvpp_runtime = self.kvpp self._fill_c8_mxfp_v_scale_caches() - guard = getattr(self, "_kv_guard_v2", None) + guard = getattr(self, "_kv_guard", None) if guard is not None: - guard.configure() + guard.configure(self) def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" From c95f8b0608068c5685ff2d6205499936ad0edce4 Mon Sep 17 00:00:00 2001 From: TAJh Date: Fri, 6 Nov 2026 06:25:17 +0000 Subject: [PATCH 3/6] perf(debug): preserve KV load state and overlap guard audits Track external-load logical ranges and generations through receive completion, preserve unaffected prefixes on failure, and safely drain cancelled submissions. Freeze step metadata and overlap copies with ordinary single-worker async generation. Validate before output release or the next forward, retaining every step's checks. Other layouts keep synchronous validation. Add lifecycle, async audit and scope regressions; document the behavior. Validated on 179 with 192 focused Linux tests, four receiver contracts and V1/V2 eager/FULL normal and injected-fault cases. Signed-off-by: TAJh --- .../feature_guide/kv_cache_guard.md | 48 +- release_notes/fragments/17833.yaml | 19 + tests/ut/debug/__init__.py | 22 + tests/ut/debug/test_kv_guard.py | 83 ++- tests/ut/debug/test_kv_guard_async.py | 211 +++++++ tests/ut/debug/test_kv_guard_loads.py | 527 ++++++++++++++++++ tests/ut/debug/test_kv_guard_runtime.py | 118 +++- tests/ut/debug/test_kv_guard_scope.py | 55 ++ tests/ut/debug/test_kv_guard_v2.py | 48 +- .../kv_p2p/test_mooncake_connector.py | 4 +- vllm_ascend/debug/kv_guard.py | 144 +++-- vllm_ascend/debug/kv_guard_connector.py | 186 +++++++ vllm_ascend/debug/kv_guard_runtime.py | 477 ++++++++++++++-- .../kv_transfer/kv_p2p/mooncake_connector.py | 5 - vllm_ascend/patch/__init__.py | 1 + vllm_ascend/patch/platform/patch_kv_guard.py | 99 +++- 16 files changed, 1917 insertions(+), 130 deletions(-) create mode 100644 tests/ut/debug/test_kv_guard_async.py create mode 100644 tests/ut/debug/test_kv_guard_loads.py create mode 100644 tests/ut/debug/test_kv_guard_scope.py create mode 100644 vllm_ascend/debug/kv_guard_connector.py diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md index bf889f9c9574..3ed92851f7dc 100644 --- a/docs/source/user_guide/feature_guide/kv_cache_guard.md +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -69,9 +69,33 @@ request owners. A zero prefix is unwritten; a positive prefix N makes `[0, N)` readable; an invalid prefix prevents subsequent reads from being accepted. Completed loads and contiguous appends establish readable prefixes. Completion and preemption release ownership while preserving reusable cached prefixes. -Physical clearing, block copies, and completed Mooncake receives update validity -at their existing lifecycle boundaries. Failed receives invalidate their -destinations without changing the connector's completion and drain protocol. +The scheduler's computed-token count describes a consumer's read requirement; +it does not establish validity. A same-step prefix producer must first supply +legal write evidence, while an already valid prefix still rejects overwrites. +Physical clearing, block copies, and completed receives update validity at their +existing lifecycle boundaries. The scheduler records external-load destinations, +allocation generations, and logical token ranges when slots are allocated, +including requests that are still waiting and have no scheduled tokens. +Connector transport offsets are not interpreted as initialized-prefix lengths. + +The common connector methods provide completion evidence. Async receive results +authorize later forwards. Sync loads become readable for a cache group only +after its required layer waits; a missing wait reports incomplete coverage. +FULL replay without observable sync-load completion also reports incomplete +coverage. No extra connector polling or device waits are introduced. +Failed loads discard only their destination blocks, preserving untouched local +prefix blocks. A failed transfer can have overwritten the beginning of a partial +destination block, so that entire block must be recomputed. +Block failures remain attached to their submission until its completion is +drained, even if the connector reports failures and completion in different steps. +Completions must match the submission snapshot and current allocation. Reusing +a request ID while its old receive is outstanding reports incomplete coverage, +because the connector's request-ID-only result cannot distinguish both loads. +Known old submissions may drain after cancellation without changing a new +allocation or treating cancellation itself as a guard failure. +Framework-authorized prefix replay truncates the specified validity range before +write auditing; an ordinary prefix overwrite still fails. Transport drain and +completion results are returned unchanged. For example, with block size 128 and a loaded 19-token prefix in block 1, the next token should write slot 147. A dummy write to slot 142 overwrites offset 14; @@ -90,6 +114,13 @@ full-model incident. | Dummy/profile execution | Check writes against existing KV state; ignore discarded scratch reads. | | Runner V2 startup capture | Require empty request/cache state; do not audit CPU metadata inside device graph capture. | +Canonical FULL plans can prepare a CPU validity transaction while the final +metadata copies complete. The transaction is committed only when every actual +plan exactly matches the independent scheduler-derived plan and the tracked +state is unchanged. Other plans and suspected violations use the complete audit. +This optimization preserves the audit and failure boundary, including its +device-completion wait. + Live or cache-preserving Runner V2 recapture is rejected, including pending receives. Normal FULL replay remains audited. Allocator and kernel block coordinates must match. @@ -98,6 +129,10 @@ The guard explicitly rejects CP, DBO, speculation, encoder/decoder execution, sliding or circular caches, other compressed or state-cache semantics, unknown attention implementations, and connectors other than Mooncake. These paths require their own validity rules before they can be supported. +The Ascend block-transfer implementation is registered as `MooncakeConnectorV1`; +this name and the legacy `MooncakeConnector` name pass the guard's connector check. +The shared load adapter is not an expansion of this connector support list; +other backends still require end-to-end validation before enabling them. The guard validates metadata and observed arguments. It does not inspect kernel memory instructions, verify KV numerical values, or detect a kernel writing to @@ -115,6 +150,10 @@ Logs use the `[KV_GUARD]` prefix: - `guard.coverage_gap`: required evidence was missing, so the audit was incomplete. This is not a successful check or an exemption from validation. +Runner V2 may submit async loads during `sample_tokens` post-forward cleanup. +Load submission and completion are checked at the connector cleanup boundary, +after the separately audited `execute_model` forward. + Fields identify the layer, request, pool, cache group, block, allocation generation, and relevant slot or range when available. Repeated failures are deduplicated with bounded state. Routine validity transitions are not logged. @@ -127,7 +166,8 @@ failure latching, and the V1/V2 adapters: ```bash pytest -q tests/ut/debug/test_kv_guard.py \ - tests/ut/debug/test_kv_guard_runtime.py tests/ut/debug/test_kv_guard_v2.py + tests/ut/debug/test_kv_guard_runtime.py tests/ut/debug/test_kv_guard_v2.py \ + tests/ut/debug/test_kv_guard_loads.py ``` The single-card service regression starts separate disabled and enabled servers, diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml index 7f7092f9af7a..e13c920457b4 100644 --- a/release_notes/fragments/17833.yaml +++ b/release_notes/fragments/17833.yaml @@ -7,3 +7,22 @@ entries: by default and configured per engine through additional_config.kv_guard_config (enabled and action). Both runners share one checker, which can log failures or stop returning worker results after a violation or incomplete audit. + - category: bugfix + text: >- + KV guard no longer treats a consumer's computed-token count as proof of + initialized cache bytes. Same-step shared-prefix writes are checked against + recorded validity, which continues to reject overwrites of valid prefixes. + External-load readiness now uses submission-time logical ranges and + allocation generations rather than transport-local computed-token offsets. + Failed loads preserve unaffected local prefixes, and explicit framework + tail replay is distinguished from an illegal prefix overwrite. + Failed-block records are retained across connector polls until receive + completion, and the registered Ascend MooncakeConnectorV1 name is accepted. + Runner V2 async loads are finalized at connector cleanup, including when + submission is deferred from execute_model to sample_tokens. + Cancelled submissions can drain without modifying a later allocation. + - category: performance + text: >- + Ordinary async generation with one model worker defers KV guard validation to output consumption + or the next forward boundary, allowing address metadata copies to overlap + model execution while retaining every step's access checks. diff --git a/tests/ut/debug/__init__.py b/tests/ut/debug/__init__.py index 208f01a7cb5e..44d0f6a6de8e 100644 --- a/tests/ut/debug/__init__.py +++ b/tests/ut/debug/__init__.py @@ -1,2 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def load_plan(guard, request="r", end=19, start=0, *, is_async=True): + guard._test_load_id = getattr(guard, "_test_load_id", 0) + 1 + return { + "pool_id": guard.pool_id, + "load_id": guard._test_load_id, + "request_id": request, + "block_table": [list(refs) for refs in guard.requests[request]], + "token_start": start, + "token_end": end, + "is_async": is_async, + } + + +def complete_load(guard, request="r", end=19, start=0): + meta = load_plan(guard, request, end, start) + guard.begin_load(meta) + guard.load(meta) + return meta diff --git a/tests/ut/debug/test_kv_guard.py b/tests/ut/debug/test_kv_guard.py index ad1e259abb9d..d3d09eafbe1f 100644 --- a/tests/ut/debug/test_kv_guard.py +++ b/tests/ut/debug/test_kv_guard.py @@ -8,6 +8,8 @@ import pytest +from . import complete_load, load_plan + spec = importlib.util.spec_from_file_location( "kv_guard_subject", Path(__file__).parents[3] / "vllm_ascend/debug/kv_guard.py" ) @@ -34,6 +36,8 @@ def bind(guard, request="r", computed=0, epoch=1, pool="pool", ids=(1,)): "scheduled_requests": {request: (computed, 0)}, } ) + if computed: + complete_load(guard, request, computed) def access(guard, slots=(), length=0, ids=(1,), request_ids=("r",), phase="forward"): @@ -57,6 +61,50 @@ def test_partial_fill_seal_and_read(guard): assert not violations(guard) +@pytest.mark.parametrize("reverse", [False, True]) +def test_same_step_prefix_consumer_cannot_prevalidate_producers_write(guard, reverse): + items = [("producer", [[(1, 1), (2, 1)]]), ("consumer", [[(1, 1), (3, 1)]])] + guard.bind( + { + "pool_id": "pool", + "requests": dict(reversed(items) if reverse else items), + "scheduled_requests": {"producer": (0, 256), "consumer": (128, 128)}, + } + ) + assert guard.blocks[0, 1].valid_end == 0 + guard.observe( + [ + { + "group_id": 0, + "slots": list(range(128, 384)) + list(range(384, 512)), + "block_table": [[1, 2], [1, 3]], + "seq_lens": [256, 256], + } + ], + ["producer", "consumer"], + "forward", + ) + assert not violations(guard) + assert all(guard.blocks[0, block].valid_end == 128 for block in (1, 2, 3)) + + +def test_computed_without_write_evidence_does_not_make_unknown_prefix_readable(guard): + guard.bind({"requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (128, 0)}}) + access(guard, (), 128) + assert violations(guard)[-1]["reason"] == "read_unwritten_range" + + +def test_recorded_prefix_survives_release_and_same_epoch_reuse(guard): + bind(guard) + access(guard, range(128, 256), 128) + guard.bind({"released_request_ids": ["r"]}) + guard.bind({"requests": {"next": [[(1, 1)]]}, "scheduled_requests": {"next": (128, 0)}}) + access(guard, (), 128, request_ids=("next",)) + assert not violations(guard) + access(guard, [128], 128, request_ids=("next",)) + assert violations(guard)[0]["reason"] == "overwrite_valid_prefix" + + def test_unknown_read_and_skipped_write(guard): bind(guard) access(guard, (), 1) @@ -68,7 +116,7 @@ def test_unknown_read_and_skipped_write(guard): def test_case4_loaded_prefix_dummy_write_then_dirty_read_without_tensors(guard): bind(guard) - guard.load({"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) + complete_load(guard) assert guard.blocks[0, 1].valid_end == 19 access(guard, [142], 1, request_ids=(), phase="dummy") dirty_write = violations(guard)[0] @@ -112,26 +160,28 @@ def test_shared_prefix_and_finished_request(guard): def test_cache_lookup_does_not_revive_tainted_same_epoch(guard): bind(guard, computed=19) access(guard, [142], phase="dummy", request_ids=()) - bind(guard, request="prefix-hit", computed=19) + guard.bind({"requests": {"prefix-hit": [[(1, 1)]]}, "scheduled_requests": {"prefix-hit": (19, 0)}}) assert guard.blocks[0, 1].valid_end == subject.INVALID_PREFIX -def test_load_before_scheduler_binding_and_failed_load(guard): - guard.load({"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) +def test_unregistered_completion_cannot_establish_validity(guard): bind(guard) - assert not guard.pending_loads - assert guard.blocks[0, 1].valid_end == 19 - guard.fail_load("r") + guard.load(load_plan(guard)) + assert guard.blocks[0, 1].valid_end == 0 access(guard, (), 19) - assert violations(guard)[-1]["reason"] == "read_invalid_block" + assert violations(guard)[-1]["reason"] == "read_unwritten_range" + assert any(e.get("reason") == "stale_load_completion" for e in guard.events) -def test_pending_load_never_applied_to_different_request(guard): - guard.load({"request_id": "old", "local_block_ids": [[1]], "num_computed_tokens": 19}) +def test_old_completion_never_applied_to_different_request(guard): + bind(guard, request="old") + meta = load_plan(guard, "old") + guard.begin_load(meta) + guard.bind({"released_request_ids": ["old"]}) bind(guard, request="new", epoch=2) + guard.load(meta) assert guard.blocks[0, 1].valid_end == 0 - guard.bind({"released_request_ids": ["old"]}) - assert not guard.pending_loads + assert not guard.loads def test_no_per_token_state_log_and_negative_padding(guard): @@ -157,6 +207,7 @@ def test_group_pool_and_allocation_isolation(guard): guard.configure({0: 128, 1: 128}) bind(guard, computed=19) guard.bind({"pool_id": "pool", "requests": {"r2": [[], [(1, 1)]]}, "scheduled_requests": {"r2": (7, 0)}}) + complete_load(guard, "r2", 7) assert guard.blocks[0, 1].valid_end == 19 assert guard.blocks[1, 1].valid_end == 7 bind(guard, pool="new-pool") @@ -222,9 +273,11 @@ def test_copying_an_empty_prefix_does_not_make_bytes_readable(guard): assert violations(guard)[-1]["reason"] == "read_unwritten_range" -def test_replaced_pool_drops_old_pending_loads(guard): +def test_replaced_pool_drops_old_loads(guard): bind(guard) - guard.load({"request_id": "next", "local_block_ids": [[1]], "num_computed_tokens": 19}) + meta = load_plan(guard) + guard.begin_load(meta) bind(guard, request="next", pool="replacement") - assert not guard.pending_loads + guard.load(meta) + assert not guard.loads assert guard.blocks[0, 1].valid_end == 0 diff --git a/tests/ut/debug/test_kv_guard_async.py b/tests/ut/debug/test_kv_guard_async.py new file mode 100644 index 000000000000..0c3a0c8af67f --- /dev/null +++ b/tests/ut/debug/test_kv_guard_async.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deferred evidence must remain exact and precede output/state consumption.""" + +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from types import SimpleNamespace as NS + +import pytest +import torch + +from .test_kv_guard_runtime import bind, metadata +from .test_kv_guard_runtime import runtime_fixture as runtime_fixture + + +def deferred_forward(runtime, mode="FULL", slot=147): + bind(runtime, computed=19) + runtime._defer_audit = True + runtime._scheduled = {"r": (19, 1)} + slots, table, lengths = torch.tensor([slot]), torch.tensor([[1]]), torch.tensor([20]) + item = metadata(slots, table, lengths) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + + def forward(count): + if mode != "FULL": + for layer in runtime.layer_groups: + runtime.write(layer, slots) + runtime.read(layer, table, lengths) + + runtime.wrap_forward(runner, forward, lambda: context)(1) + return slots, table, lengths + + +@pytest.mark.parametrize("mode", ["FULL", "NONE"]) +@pytest.mark.parametrize("slot", [147, 148]) +def test_deferred_audit_freezes_final_values_without_publishing_validity(runtime, mode, slot): + slots, table, lengths = deferred_forward(runtime, mode, slot) + assert runtime.state.blocks[0, 1].valid_end == 19 + slots[0], table[0, 0], lengths[0] = 0, 2, 99 + runtime._finish_pending() + mismatch = any(event.get("reason") == "write_target_mismatch" for event in runtime.events) + assert mismatch == (slot != 147) + if slot == 147: + assert runtime.state.blocks[0, 1].valid_end == 20 + assert not runtime.events + + +@pytest.mark.parametrize("mode", ["FULL", "NONE"]) +def test_invalid_snapshot_cannot_escape_through_async_output(runtime, mode): + runtime.action = "raise" + runtime.state.emit = runtime.report + deferred_forward(runtime, mode, 148) + consumed = [] + output = NS(get_output=lambda: consumed.append(True)) + assert runtime._guard_output(output) is output + with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): + output.get_output() + assert not consumed + with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): + output.get_output() + assert not consumed + + +def test_direct_output_finishes_audit_before_return(runtime): + deferred_forward(runtime) + marker = object() + assert runtime._guard_output(marker) is marker + assert runtime.state.blocks[0, 1].valid_end == 20 + + +def test_completed_output_still_honors_a_later_sticky_failure(runtime): + runtime.action = "raise" + deferred_forward(runtime) + consumed = [] + output = runtime._guard_output(NS(get_output=lambda: consumed.append(True))) + runtime._finish_pending() + runtime._failure = "later_step_failed" + with pytest.raises(runtime.subject.KVGuardError, match="later_step_failed"): + output.get_output() + assert not consumed + + +def test_readonly_output_fallback_cannot_bypass_a_sticky_failure(runtime): + class Output: + __slots__ = () + + def get_output(self): + return "result" + + runtime.action = "raise" + deferred_forward(runtime) + assert runtime._guard_output(Output()).get_output() == "result" + runtime._failure = "later_step_failed" + with pytest.raises(runtime.subject.KVGuardError, match="later_step_failed"): + runtime._guard_output(Output()) + + +def test_next_binding_and_clear_follow_the_predecessor_audit(runtime): + deferred_forward(runtime) + pending = runtime._pending + consumed = [] + output = runtime._guard_output(NS(get_output=lambda: consumed.append(True))) + original_bind = runtime.state.bind + calls = [] + + def checked_bind(context): + assert runtime.state.blocks[0, 1].valid_end == 20 + calls.append("bind") + original_bind(context) + + runtime.state.bind = checked_bind + packet = NS( + _ascend_kv_guard_context={ + "pool_id": "pool", + "requests": {"next": [[(2, 1)]]}, + "released_request_ids": ["r"], + "scheduled_requests": {"next": (0, 1)}, + } + ) + with runtime.execute_scope(packet): + runtime.zero_blocks([2]) + assert not calls and "r" in runtime.state.requests + runtime._finish_pending() + assert calls == ["bind"] and "r" not in runtime.state.requests + assert runtime.state.blocks[0, 2].valid_end == 0 + # A late reader reuses the completed decision, not the next step's buffers/state. + runtime.state.blocks[0, 1].epoch = 9 + output.get_output() + assert consumed == [True] and pending is not None + assert not runtime.events + + +def test_idle_step_drains_pending_validation_and_release(runtime): + deferred_forward(runtime) + packet = NS( + _ascend_kv_guard_context={ + "pool_id": "pool", + "requests": {}, + "released_request_ids": ["r"], + "scheduled_requests": {}, + } + ) + runtime.wrap_execute(lambda output: None)(packet) + assert runtime._pending is None and not runtime._deferred_state + assert not runtime.state.requests + assert runtime.state.blocks[0, 1].valid_end == 20 + + +@pytest.mark.parametrize("failure", [False, True]) +def test_output_reader_and_next_forward_finish_only_once(runtime, failure): + started, release = Event(), Event() + calls = [] + + def complete(): + calls.append(True) + started.set() + assert release.wait(2) + if failure: + raise RuntimeError("audit failed") + + pending = runtime.subject._PendingAudit(complete) + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(pending.finish) for _ in range(2)] + assert started.wait(2) + release.set() + for future in futures: + if failure: + with pytest.raises(RuntimeError, match="audit failed"): + future.result() + else: + future.result() + assert calls == [True] and pending._complete is None + + +def test_snapshot_keeps_missing_layer_coverage(runtime): + bind(runtime) + runtime._defer_audit = True + runtime._scheduled = {"r": (0, 1)} + item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) + runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) + runtime._finish_pending() + assert runtime.state.blocks[0, 1].valid_end == 0 + assert any(event.get("reason") == "metadata_write_mapping_missing:layer1" for event in runtime.events) + + +@pytest.mark.parametrize("mode", ["FULL", "NONE"]) +@pytest.mark.parametrize("field,reason", [("table", "read_target_mismatch"), ("length", "read_length_mismatch")]) +def test_frozen_read_evidence_cannot_be_repaired_by_the_next_preparation(runtime, mode, field, reason): + bind(runtime, computed=19) + runtime._defer_audit = True + runtime._scheduled = {"r": (19, 1)} + slots = torch.tensor([147]) + table = torch.tensor([[2 if field == "table" else 1]]) + lengths = torch.tensor([21 if field == "length" else 20]) + item = metadata(slots, table, lengths) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + + def forward(count): + if mode != "FULL": + for layer in runtime.layer_groups: + runtime.write(layer, slots) + runtime.read(layer, table, lengths) + + runtime.wrap_forward(runner, forward, lambda: context)(1) + table[0, 0], lengths[0] = 1, 20 + runtime._finish_pending() + assert any(event.get("reason") == reason for event in runtime.events) diff --git a/tests/ut/debug/test_kv_guard_loads.py b/tests/ut/debug/test_kv_guard_loads.py new file mode 100644 index 000000000000..a30260104775 --- /dev/null +++ b/tests/ut/debug/test_kv_guard_loads.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Logical load snapshots and common connector boundaries; CPU contracts only.""" + +from types import SimpleNamespace as NS + +import pytest +import torch + +from . import complete_load, load_plan +from .test_kv_guard_runtime import bind, load, metadata +from .test_kv_guard_runtime import runtime_fixture as runtime_fixture +from .test_kv_guard_v2 import adapter_fixture as adapter_fixture +from .test_kv_guard_v2 import prepare + + +class Connector: + def __init__(self): + self.calls = [] + self.result = NS(finished_sending=set(), finished_recving=set(), failed_recving=set()) + self.invalid = set() + + def start_load_kv(self, *args, **kwargs): + self.calls.append(("start", args, kwargs)) + return "started" + + def wait_for_layer_load(self, layer_name, *args, **kwargs): + self.calls.append(("wait", layer_name, args, kwargs)) + return "ready" + + def get_transfer_results(self, *args, **kwargs): + self.calls.append(("results", args, kwargs)) + return self.result + + def get_block_ids_with_load_errors(self): + self.calls.append(("errors",)) + return self.invalid + + +@pytest.fixture +def connector(runtime): + connector = Connector() + runtime._connector = runtime.subject.ConnectorGuard(runtime, connector) + return connector + + +def packet(meta=None, scheduled=None, requests=None, released=(), replay=None): + scheduled = {} if scheduled is None else scheduled + context = { + "pool_id": "pool", + "requests": requests or {}, + "scheduled_requests": scheduled, + "released_request_ids": list(released), + "loads": [meta] if meta is not None else [], + "recompute_ranges": [replay] if replay is not None else [], + } + return NS(_ascend_kv_guard_context=context, total_num_scheduled_tokens=sum(n for _, n in scheduled.values())) + + +def poll(connector): + result = connector.get_transfer_results(set()) + invalid = connector.get_block_ids_with_load_errors() + return result, invalid + + +def reasons(runtime): + return {event.get("reason") for event in runtime.events} + + +@pytest.mark.parametrize( + "name, accepted", [("MooncakeConnector", True), ("MooncakeConnectorV1", True), ("AscendStoreConnector", False)] +) +def test_registered_ascend_mooncake_name_is_accepted(runtime, name, accepted): + impl = type("AscendAttentionBackendImpl", (), {})() + runtime.config = NS( + parallel_config=NS(), + model_config=NS(), + speculative_config=None, + kv_transfer_config=NS(kv_connector=name), + compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), + ) + runner = NS( + kv_cache_config=NS(kv_cache_groups=[NS(layer_names=["layer0"], kv_cache_spec=NS(block_size=128, head_size=64))]) + ) + if accepted: + runtime.configure(runner) + else: + with pytest.raises(ValueError, match="unadapted execution"): + runtime.configure(runner) + + +@pytest.mark.parametrize("local", [0, 19]) +def test_async_completion_uses_loaded_end_not_transport_local_offset(runtime, connector, local): + bind(runtime, computed=local) + meta = load_plan(runtime.state, end=64, start=local) + + def submit(output): + assert connector.start_load_kv("context", num_computed_tokens=local) == "started" + assert runtime.state.blocks[0, 1].valid_end == local + poll(connector) + + runtime.wrap_execute(submit)(packet(meta)) + connector.result.finished_recving = {"r"} + result, invalid = runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert result is connector.result and invalid is connector.invalid + assert runtime.state.blocks[0, 1].valid_end == 64 + assert not runtime.state.loads and not runtime.events + assert [call[0] for call in connector.calls] == ["start", "results", "errors", "results", "errors"] + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_sync_layer_waits_establish_group_readiness_only_after_every_layer(runtime, connector, mode): + bind(runtime) + meta = load_plan(runtime.state, end=64, is_async=False) + item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) + context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(count): + for index, layer in enumerate(runtime.layer_groups): + assert connector.wait_for_layer_load(layer) == "ready" + assert runtime.state.blocks[0, 1].valid_end == (0 if index == 0 else 64) + runtime.write(layer, item.slot_mapping) + runtime.read(layer, item.block_tables, [65]) + + forward = runtime.wrap_forward(runner, model, lambda: context) + + def execute(output): + connector.start_load_kv("context") + forward(1) + poll(connector) + + runtime.wrap_execute(execute)(packet(meta, {"r": (64, 1)})) + assert not runtime.events and not runtime.state.loads + assert runtime.state.blocks[0, 1].valid_end == 65 + + +@pytest.mark.parametrize("wait_late", [False, True]) +def test_missing_or_late_layer_wait_cannot_be_hidden_by_another_layer(runtime, connector, wait_late): + bind(runtime) + meta = load_plan(runtime.state, end=64, is_async=False) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="NONE") + + def model(count): + runtime.write("layer0", torch.tensor([192])) + runtime.read("layer0", torch.tensor([[1]]), [65]) + if wait_late: + connector.wait_for_layer_load("layer0") + connector.wait_for_layer_load("layer1") + runtime.write("layer1", torch.tensor([192])) + runtime.read("layer1", torch.tensor([[1]]), [65]) + + forward = runtime.wrap_forward(runner, model, lambda: context) + runtime.wrap_execute(lambda output: (connector.start_load_kv(None), forward(1)))(packet(meta, {"r": (64, 1)})) + assert "load_layer_not_ready" in reasons(runtime) + assert runtime.state.blocks[0, 1].valid_end == (64 if wait_late else 0) + assert not any(e["event"] == "guard.violation" for e in runtime.events) + + +def test_full_replay_without_load_completion_evidence_is_incomplete(runtime, connector): + bind(runtime) + meta = load_plan(runtime.state, end=64, is_async=False) + item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + forward = runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda n: None, lambda: context) + runtime.wrap_execute(lambda output: (connector.start_load_kv(None), forward(1)))(packet(meta, {"r": (64, 1)})) + assert "load_layer_not_ready" in reasons(runtime) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_receive_failure_preserves_local_prefix_and_allows_recompute(runtime, connector): + state = runtime.state + state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) + complete_load(state, end=128) + meta = load_plan(state, end=192, start=128) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.result.finished_recving = connector.result.failed_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert state.blocks[0, 1].valid_end == 128 + assert state.blocks[0, 2].valid_end == 0 + state.observe([{"group_id": 0, "slots": [256], "block_table": [[1, 2]], "seq_lens": [129]}], ["r"], "forward") + assert not runtime.events + + +def test_invalid_block_report_only_invalidates_actual_loaded_destinations(runtime, connector): + state = runtime.state + state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}}) + complete_load(state, end=128) + meta = load_plan(state, end=384, start=128) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.result.finished_recving = {"r"} + connector.invalid = {1, 2, 99} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert [state.blocks[0, n].valid_end for n in (1, 2, 3)] == [128, 0, 128] + assert not runtime.events + + +@pytest.mark.parametrize("invalid, expected", [({2}, [128, 0, 128]), ({1, 2, 3}, [128, 0, 0])]) +def test_early_invalid_blocks_survive_until_the_later_receive_completion(runtime, connector, invalid, expected): + state = runtime.state + state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}}) + complete_load(state, end=128) + meta = load_plan(state, end=384, start=128) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.invalid = invalid + runtime.wrap_execute(lambda output: poll(connector))(packet()) + connector.invalid = set() + connector.result.finished_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert [state.blocks[0, n].valid_end for n in (1, 2, 3)] == expected + assert not state.loads and not runtime.events + + +def test_partial_block_failure_invalidates_the_whole_destination_block(runtime, connector): + bind(runtime, computed=19) + meta = load_plan(runtime.state, end=64, start=19) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.result.finished_recving = connector.result.failed_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_start_exception_preserves_original_error_and_local_kv(runtime): + runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) + complete_load(runtime.state, end=128) + source = Connector() + error = RuntimeError("backend load failed") + + def fail(*args): + raise error + + source.start_load_kv = fail + runtime._connector = runtime.subject.ConnectorGuard(runtime, source) + meta = load_plan(runtime.state, end=192, start=128) + with pytest.raises(RuntimeError) as raised: + runtime.wrap_execute(lambda output: source.start_load_kv(None))(packet(meta)) + assert raised.value is error + assert runtime.state.blocks[0, 1].valid_end == 128 + assert runtime.state.blocks[0, 2].valid_end == 0 + assert not runtime._connector.inflight and not runtime.state.loads + + +@pytest.mark.parametrize("epoch", [1, 2]) +def test_old_completion_cannot_validate_reused_request_or_allocation(runtime, connector, epoch): + bind(runtime) + old = load_plan(runtime.state, end=64) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(old)) + runtime.wrap_execute(lambda output: None)(packet(requests={"r": [[(1, epoch)]]}, released=["r"])) + new = load_plan(runtime.state, end=32) + runtime.state.begin_load(new) + connector.result.finished_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert runtime.state.blocks[0, 1].valid_end == 0 + assert runtime.state.loads["r"]["load_id"] == new["load_id"] + assert not runtime.events # A known old submission may legally drain after cancellation. + + +def test_overlapping_request_id_is_reported_and_never_promotes_the_new_load(runtime, connector): + bind(runtime) + old = load_plan(runtime.state, end=64) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(old)) + new = load_plan(runtime.state, end=32) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))( + packet(new, released=["r"], requests={"r": [[(1, 2)]]}) + ) + connector.result.finished_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert "overlapping_load_request_id" in reasons(runtime) + assert runtime.state.blocks[0, 1].valid_end == 0 + + +def test_cancelled_request_can_drain_without_poisoning_later_requests(runtime, connector): + bind(runtime) + meta = load_plan(runtime.state, end=64) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + runtime.wrap_execute(lambda output: None)(packet(released=["r"])) + connector.result.finished_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert not runtime.events + assert not runtime._connector.inflight + runtime.state.bind({"requests": {"next": [[(1, 2)]]}}) + runtime.state.observe([{"group_id": 0, "slots": [128], "block_table": [[1]], "seq_lens": [1]}], ["next"], "forward") + assert runtime.state.blocks[0, 1].valid_end == 1 and not runtime.events + + +def test_completion_metadata_cannot_expand_the_registered_range(runtime): + bind(runtime) + meta = load_plan(runtime.state, end=64) + runtime.state.begin_load(meta) + runtime.state.load({**meta, "token_end": 128}) + assert runtime.state.blocks[0, 1].valid_end == 0 + assert "stale_load_completion" in reasons(runtime) + + +def test_completion_cannot_erase_a_detected_corruption_in_the_same_allocation(runtime, connector): + bind(runtime) + meta = load_plan(runtime.state, end=64) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + runtime.state.observe([{"group_id": 0, "slots": [142]}], [], "dummy") + connector.result.finished_recving = {"r"} + runtime.wrap_execute(lambda output: poll(connector))(packet()) + assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX + assert "load_destination_tainted" in reasons(runtime) + + +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_only_explicit_replay_allows_overwriting_a_loaded_tail(runtime, replay, mode): + bind(runtime, computed=128) + meta = load_plan(runtime.state, end=128, start=127) + meta.pop("load_id") + meta.pop("is_async") + item = metadata(torch.tensor([255]), torch.tensor([[1]]), [128]) + context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(count): + for layer in runtime.layer_groups: + runtime.write(layer, item.slot_mapping) + runtime.read(layer, item.block_tables, [128]) + + forward = runtime.wrap_forward(runner, model, lambda: context) + runtime.wrap_execute(lambda output: forward(1))(packet(scheduled={"r": (127, 1)}, replay=meta if replay else None)) + if replay: + assert runtime.state.blocks[0, 1].valid_end == 128 and not runtime.events + else: + assert "overwrite_valid_prefix" in reasons(runtime) + + +def test_replay_does_not_legalize_a_read_before_the_recomputed_write(runtime): + bind(runtime, computed=128) + replay = load_plan(runtime.state, end=128, start=127) + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(count): + runtime.read("layer0", torch.tensor([[1]]), [128]) + runtime.write("layer0", torch.tensor([255])) + + forward = runtime.wrap_forward(runner, model, lambda: NS(cudagraph_runtime_mode="NONE")) + runtime.wrap_execute(lambda output: forward(1))(packet(scheduled={"r": (127, 1)}, replay=replay)) + assert "read_unwritten_range" in reasons(runtime) + + +def test_post_forward_async_completion_does_not_legalize_an_earlier_read(runtime, connector): + bind(runtime) + meta = load_plan(runtime.state, end=64) + runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.result.finished_recving = {"r"} + runner = NS(input_batch=NS(req_ids=["r"])) + + def model(count): + runtime.read("layer0", torch.tensor([[1]]), [64]) + poll(connector) # V2 polls inside its forward scope. + assert runtime.state.blocks[0, 1].valid_end == 0 + + forward = runtime.wrap_forward(runner, model, lambda: NS(cudagraph_runtime_mode="NONE")) + runtime.wrap_execute(lambda output: forward(1))(packet()) + assert "read_unwritten_range" in reasons(runtime) + assert runtime.state.blocks[0, 1].valid_end == 64 + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("positional", [False, True]) +def test_scheduler_records_loading_bindings_even_without_scheduled_tokens(runtime, is_async, positional): + patch = load("guard_loading_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") + + class Pool: + blocks = [NS(block_id=n) for n in range(4)] + + def get_new_blocks(self, count): + return self.blocks[1 : count + 1] + + class Scheduler: + def __init__(self): + self.vllm_config = NS(additional_config={"kv_guard_config": {"enabled": True}}) + self.pool = Pool() + self.request = NS(request_id="r", num_computed_tokens=0) + self.kv_cache_manager = NS( + block_pool=self.pool, + allocate_slots=self.allocate, + get_block_ids=lambda request: [[1, 2]], + free=lambda request: None, + ) + + def allocate( + self, + request, + num_new_tokens, + num_new_computed_tokens=0, + new_computed_blocks=None, + num_lookahead_tokens=0, + num_external_computed_tokens=0, + delay_cache_blocks=False, + ): + self.pool.get_new_blocks(2) + return object() + + def schedule(self): + if positional: + self.kv_cache_manager.allocate_slots(self.request, 0, 19, None, 0, 173, is_async) + else: + self.kv_cache_manager.allocate_slots( + self.request, 0, 19, num_external_computed_tokens=173, delay_cache_blocks=is_async + ) + return NS( + scheduled_new_reqs=[], + scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), + num_scheduled_tokens={}, + finished_req_ids=set(), + ) + + patch.install_scheduler_guard(Scheduler) + scheduler = Scheduler() + context = scheduler.schedule()._ascend_kv_guard_context + assert context["requests"] == {"r": [[(1, 1), (2, 1)]]} + meta = context["loads"][0] + assert (meta["token_start"], meta["token_end"], meta["is_async"]) == (19, 192, is_async) + assert meta["block_table"] == context["requests"]["r"] + assert context["scheduled_requests"] == {} + runtime.state.bind(context) + assert runtime.state.blocks[0, 1].valid_end == 0 + assert runtime.state.loads["r"]["load_id"] == meta["load_id"] + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_v2_receive_only_completion_then_authorized_last_token_replay(adapter, connector, mode): + runtime = adapter.runtime + bind(runtime) + meta = load_plan(runtime.state, end=128) + adapter.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) + connector.result.finished_recving = {"r"} + adapter.wrap_execute(lambda output: poll(connector))(packet()) + assert runtime.state.blocks[0, 1].valid_end == 128 + replay = {**meta, "token_start": 127} + + def execute(output): + item = prepare(adapter, ["r"], [255], [[1]], [128], mode)["layer0"] + runtime.write("layer0", item.slot_mapping) + runtime.read("layer0", item.block_tables, [128]) + + adapter.wrap_execute(execute)(packet(scheduled={"r": (127, 1)}, replay=replay)) + assert runtime.state.blocks[0, 1].valid_end == 128 and not runtime.events + + +def test_v2_async_submission_can_follow_execute_in_sample_tokens(adapter, connector): + runtime = adapter.runtime + bind(runtime, "running", 1) + bind(runtime, "loading", 2) + meta = load_plan(runtime.state, "loading", end=64) + + def execute(output): + item = prepare(adapter, ["running"], [128], [[1]], [1])["layer0"] + runtime.write("layer0", item.slot_mapping) + runtime.read("layer0", item.block_tables, [1]) + return "awaiting sampling" + + assert adapter.wrap_execute(execute)(packet(meta, {"running": (0, 1)})) == "awaiting sampling" + assert runtime._connector.staged == [meta] and not runtime.events + # Actual MRV2 post_forward runs from sample_tokens, outside execute_model. + connector.start_load_kv(None) + connector.result.finished_recving = {"loading"} + poll(connector) + assert runtime.state.blocks[0, 2].valid_end == 64 + assert not runtime.state.loads and not runtime.events + + +def test_v2_unsubmitted_load_is_rejected_at_the_real_post_forward_boundary(adapter, connector): + runtime = adapter.runtime + bind(runtime, "running", 1) + bind(runtime, "loading", 2) + meta = load_plan(runtime.state, "loading", end=64) + + def execute(output): + prepare(adapter, ["running"], [128], [[1]], [1]) + + adapter.wrap_execute(execute)(packet(meta, {"running": (0, 1)})) + assert "load_not_submitted" not in reasons(runtime) + poll(connector) + assert "load_not_submitted" in reasons(runtime) + + +def test_receive_only_binding_is_not_required_in_the_worker_execution_batch(runtime): + waiting = packet(requests={"loading": [[(2, 1)]]}) + runner = NS(requests={}) + assert runtime.wrap_update(runner, lambda output: "unchanged")(waiting) == "unchanged" + assert not runtime.events + + +def test_scheduler_records_explicit_async_tail_replay_without_changing_request(runtime): + patch = load("guard_replay_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") + + class Scheduler: + def __init__(self): + self.vllm_config = NS(additional_config={"kv_guard_config": {"enabled": True}}) + self.request = NS(request_id="r", num_computed_tokens=128) + pool = NS(blocks=[NS(block_id=0), NS(block_id=1)], get_new_blocks=lambda count: [NS(block_id=1)]) + self.kv_cache_manager = NS(block_pool=pool, get_block_ids=lambda request: [[1]], free=lambda request: None) + + def _update_waiting_for_remote_kv(self, request): + request.num_computed_tokens -= 1 + return "unchanged" + + def schedule(self): + self.kv_cache_manager.block_pool.get_new_blocks(1) + assert self._update_waiting_for_remote_kv(self.request) == "unchanged" + return NS( + scheduled_new_reqs=[NS(req_id="r", num_computed_tokens=127)], + scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), + num_scheduled_tokens={"r": 1}, + finished_req_ids=set(), + ) + + patch.install_scheduler_guard(Scheduler) + scheduler = Scheduler() + context = scheduler.schedule()._ascend_kv_guard_context + assert scheduler.request.num_computed_tokens == 127 + assert context["scheduled_requests"] == {"r": (127, 1)} + assert context["recompute_ranges"] == [ + { + "pool_id": context["pool_id"], + "request_id": "r", + "block_table": [[(1, 1)]], + "token_start": 127, + "token_end": 128, + } + ] diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py index 1d0e94759072..375a6b7e0460 100644 --- a/tests/ut/debug/test_kv_guard_runtime.py +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -13,6 +13,8 @@ import pytest import torch +from . import complete_load + ROOT = Path(__file__).parents[3] @@ -39,6 +41,8 @@ def runtime_fixture(monkeypatch): monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_config", options) state = load("guard_runtime_state", "vllm_ascend/debug/kv_guard.py") monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard", state) + connector = load("guard_connector_subject", "vllm_ascend/debug/kv_guard_connector.py") + monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_connector", connector) subject = load("guard_runtime_subject", "vllm_ascend/debug/kv_guard_runtime.py") logger = ModuleType("vllm.logger") cast(Any, logger).logger = logging.Logger("guard-test") @@ -64,6 +68,8 @@ def bind(runtime, request="r", block=1, computed=0): runtime.state.bind( {"pool_id": "pool", "requests": {request: [[(block, 1)]]}, "scheduled_requests": {request: (computed, 0)}} ) + if computed: + complete_load(runtime.state, request, computed) def test_access_hooks_validate_case4_without_kv_copy_or_hash(runtime): @@ -133,7 +139,7 @@ def test_receive_only_step_binds_before_execute_and_cleans_finished(runtime): _ascend_kv_guard_context={"pool_id": "pool", "requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (0, 0)}} ) runtime.wrap_execute(lambda message: None)(output) - runtime.transfer("load", {"request_id": "r", "local_block_ids": [[1]], "num_computed_tokens": 19}) + complete_load(runtime.state) assert runtime.state.blocks[0, 1].valid_end == 19 runtime.wrap_execute(lambda message: None)( NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": ["r"], "scheduled_requests": {}}) @@ -462,6 +468,115 @@ def replay(n): assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) +@pytest.mark.parametrize( + "slots, table, lengths, use_fast", + [ + ([128, 129], [[1]], [2], True), + ([129, 128], [[1]], [2], False), + ([128, 128], [[1]], [2], False), + ([128, 129], [[2]], [2], False), + ([128, 129], [[1]], [1], False), + ([128, 129], [[1]], [3], False), + ], +) +def test_full_fast_and_complete_audits_agree(runtime, slots, table, lengths, use_fast): + outcomes = [] + for enabled in (False, True): + guard = runtime.subject.GuardRuntime(NS(additional_config={"kv_guard_config": {"enabled": True}})) + events = [] + guard.state.emit = lambda event, events=events, **fields: events.append({"event": event, **fields}) + guard.gap = lambda reason, events=events, **fields: events.append( + {"event": "guard.coverage_gap", "reason": reason, **fields} + ) + guard.state.configure({0: 128}) + guard.layer_groups = {"layer0": 0, "layer1": 0} + guard.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1)]]}}) + guard._scheduled = {"r": (0, 2)} + item = metadata(torch.tensor(slots), torch.tensor(table), torch.tensor(lengths)) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(guard.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=["r"])) + hits = [] + original = guard._try_fast_audit + + def attempt(runner, original=original, enabled=enabled, hits=hits): + used = original(runner) if enabled else False + hits.append(used) + return used + + guard._try_fast_audit = attempt + guard.wrap_forward(runner, lambda count: None, lambda context=context: context)(2) + assert hits == [enabled and use_fast] + outcomes.append((guard.state.blocks, events)) + assert outcomes[0] == outcomes[1] + + +def test_fast_candidate_waits_for_final_values_and_does_not_commit_during_preview(runtime, monkeypatch): + bind(runtime) + runtime._scheduled = {"r": (0, 1)} + item = metadata(torch.tensor([-1]), torch.tensor([[1]]), torch.tensor([0])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=["r"])) + preview = runtime._preview_append + seen = [] + + def staged(runner, plan): + assert item.slot_mapping.tolist() == [128] + result = preview(runner, plan) + assert runtime.state.blocks[0, 1].valid_end == 0 + seen.append(result) + return result + + monkeypatch.setattr(runtime, "_preview_append", staged) + monkeypatch.setattr(runtime._scratch, "observe", lambda *args: pytest.fail("normal plan unexpectedly fell back")) + + def forward(count): + assert not seen + item.slot_mapping[0], item.seq_lens[0] = 128, 1 + + runtime.wrap_forward(runner, forward, lambda: context)(1) + assert len(seen) == 1 and runtime.state.blocks[0, 1].valid_end == 1 + + +def test_fast_candidate_is_discarded_when_state_changes(runtime, monkeypatch): + bind(runtime) + runtime._scheduled = {"r": (0, 1)} + item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=["r"])) + preview = runtime._preview_append + + def changed(runner, plans): + candidate = preview(runner, plans) + runtime.state.blocks[0, 1].epoch = 2 + return candidate + + monkeypatch.setattr(runtime, "_preview_append", changed) + runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 0 + assert any(event.get("reason") == "allocation_changed_during_access_audit" for event in runtime.events) + + +def test_fast_preview_error_retains_complete_overwrite_diagnostic(runtime): + bind(runtime, computed=1) + runtime._scheduled = {"r": (0, 1)} + item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) + event = next(event for event in runtime.events if event.get("reason") == "overwrite_valid_prefix") + assert event["offsets"] == [0] and event["valid_end"] == 1 + assert event["evidence"] == "execution_plan" and event["layer"] == "layer0" + + +def test_fast_missing_layer_cannot_advance_validity(runtime): + bind(runtime) + runtime._scheduled = {"r": (0, 1)} + item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) + runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) + assert runtime.state.blocks[0, 1].valid_end == 0 + assert any(event["event"] == "guard.coverage_gap" for event in runtime.events) + + def test_current_full_graph_metadata_ignores_operator_callbacks(runtime): bind(runtime, computed=19) item = metadata(torch.tensor([147]), torch.tensor([[1]]), [20]) @@ -866,6 +981,7 @@ def test_shared_prefix_read_and_tail_append_remain_legal(runtime): "scheduled_requests": {"A": (128, 1), "B": (128, 1)}, } ) + complete_load(runtime.state, "A", 128) runtime._scheduled = {"A": (128, 1), "B": (128, 1)} runner = NS(input_batch=NS(req_ids=["B", "A"])) context = NS(cudagraph_runtime_mode="NONE") diff --git a/tests/ut/debug/test_kv_guard_scope.py b/tests/ut/debug/test_kv_guard_scope.py new file mode 100644 index 000000000000..da20203254bb --- /dev/null +++ b/tests/ut/debug/test_kv_guard_scope.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Outputs without a guaranteed sole consumer retain synchronous checking.""" + +from types import SimpleNamespace as NS + +import pytest + +from .test_kv_guard_runtime import runtime_fixture as runtime_fixture +from .test_kv_guard_v2 import adapter_fixture as adapter_fixture +from .test_kv_guard_v2 import message, prepare + + +@pytest.mark.parametrize("pooling,tp,pp", [(False, 1, 1), (True, 1, 1), (False, 2, 1), (False, 1, 2)]) +def test_only_single_worker_generation_can_defer_checks_to_sample_output(adapter, pooling, tp, pp): + runtime = adapter.runtime + runtime.config.scheduler_config = NS(async_scheduling=True) + adapter.runner.sample_tokens = lambda grammar: None + adapter.runner.is_pooling_model = pooling + runtime.config.parallel_config.tensor_parallel_size = tp + runtime.config.parallel_config.pipeline_parallel_size = pp + runtime.configure(adapter.runner) + assert runtime._defer_audit is (not pooling and tp == 1 and pp == 1) + + +@pytest.mark.parametrize("pooling,tp,pp", [(True, 1, 1), (False, 2, 1), (False, 1, 2)]) +def test_failure_is_detected_before_pool_or_nonreply_output_can_bypass_the_getter(adapter, pooling, tp, pp): + runtime = adapter.runtime + runtime.config.scheduler_config = NS(async_scheduling=True) + adapter.runner.sample_tokens = lambda grammar: None + adapter.runner.is_pooling_model = pooling + runtime.config.parallel_config.tensor_parallel_size = tp + runtime.config.parallel_config.pipeline_parallel_size = pp + runtime.action = "raise" + runtime.state.emit = runtime.report + runtime.configure(adapter.runner) + with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): + adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [129], [[1]], [1], "FULL"))(message()) + assert runtime._pending is None + + +@pytest.mark.parametrize("missing", ["pooling", "tp", "pp"]) +def test_unknown_output_layout_keeps_synchronous_validation(adapter, missing): + runtime = adapter.runtime + runtime.config.scheduler_config = NS(async_scheduling=True) + adapter.runner.sample_tokens = lambda grammar: None + adapter.runner.is_pooling_model = False + parallel = runtime.config.parallel_config + parallel.tensor_parallel_size = parallel.pipeline_parallel_size = 1 + if missing == "pooling": + del adapter.runner.is_pooling_model + else: + delattr(parallel, "tensor_parallel_size" if missing == "tp" else "pipeline_parallel_size") + runtime.configure(adapter.runner) + assert not runtime._defer_audit diff --git a/tests/ut/debug/test_kv_guard_v2.py b/tests/ut/debug/test_kv_guard_v2.py index de400b646133..9b14219acb25 100644 --- a/tests/ut/debug/test_kv_guard_v2.py +++ b/tests/ut/debug/test_kv_guard_v2.py @@ -8,12 +8,13 @@ import pytest import torch +from . import complete_load, load_plan from .test_kv_guard_runtime import bind, metadata from .test_kv_guard_runtime import runtime_fixture as runtime_fixture -@pytest.fixture -def adapter(runtime): +@pytest.fixture(name="adapter") +def adapter_fixture(runtime): subject = runtime.subject impl = type("AscendAttentionBackendImpl", (), {})() runtime.config = NS( @@ -190,6 +191,7 @@ def execute(output): def test_v2_clears_then_copies_after_physical_update(adapter): runtime = adapter.runtime runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}, "scheduled_requests": {"r": (4, 0)}}) + complete_load(runtime.state, end=4) packet = NS(new_block_ids_to_zero=[2], kv_cache_block_copies=[NS(src_block_id=1, dst_block_id=2)]) order = [] clear, copy = runtime.zero_blocks, runtime.copy_blocks @@ -312,11 +314,51 @@ def execute(output): adapter.wrap_execute(execute)(scheduler_output=message()) adapter.wrap_update(lambda output: "unchanged")(scheduler_output=NS()) + + +@pytest.mark.parametrize("async_output", [False, True]) +def test_v2_output_is_guarded_after_the_forward_scope_submits_its_audit(adapter, async_output): + runtime = adapter.runtime + runtime._defer_audit = True + consumed = [] + result = NS(get_output=lambda: consumed.append(True)) if async_output else "result" + + def execute(output): + prepare(adapter, ["r"], [128], [[1]], [1], "FULL") + return result + + output = adapter.wrap_execute(execute)(message()) + if async_output: + assert runtime.state.blocks[0, 1].valid_end == 0 + output.get_output() + assert consumed == [True] + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not runtime.events + + +@pytest.mark.parametrize("flag", ["dummy_run", "is_profile"]) +def test_v2_synthetic_execution_drains_the_real_pending_audit_first(adapter, flag): + runtime = adapter.runtime + runtime._defer_audit = True + + def real(output): + prepare(adapter, ["r"], [128], [[1]], [1], "FULL") + + adapter.wrap_execute(real)(message()) + assert runtime.state.blocks[0, 1].valid_end == 0 + + def synthetic(output, **kwargs): + assert runtime.state.blocks[0, 1].valid_end == 1 + prepare(adapter, ["dummy"], [-1], [[0]], [1], "FULL") + + adapter.wrap_execute(synthetic)(message(), **{flag: True}) + assert not runtime.events assert adapter.runtime.state.blocks[0, 1].valid_end == 1 assert not adapter.runtime.events def test_v2_pending_receive_also_blocks_graph_recapture(adapter): - adapter.runtime.state.load({"request_id": "incoming", "local_block_ids": [[1]], "num_computed_tokens": 19}) + adapter.runtime.state.bind({"pool_id": "pool", "requests": {"incoming": [[(1, 1)]]}}) + adapter.runtime.state.begin_load(load_plan(adapter.runtime.state, "incoming")) with pytest.raises(ValueError, match="recapture is unsupported"): adapter.wrap_capture(lambda: pytest.fail("capture must not run"))() diff --git a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py index 9a9f29a3cdc1..d49e02fbb027 100644 --- a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py +++ b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py @@ -5262,7 +5262,7 @@ def test_unused_remote_ports_are_notified_once_by_primary_receiver(primary, has_ @pytest.mark.parametrize("awaiting_kvs", [False, True]) @pytest.mark.parametrize("transfer_failed", [False, True]) -def test_guard_receive_callback_preserves_completion_contract(awaiting_kvs, transfer_failed): +def test_receive_completion_contract_is_preserved(awaiting_kvs, transfer_failed): receiver = KVCacheRecvingThread.__new__(KVCacheRecvingThread) receiver._is_failed_recv_request = MagicMock(side_effect=[False, transfer_failed]) receiver._transfer_kv_cache_all_groups = MagicMock() @@ -5277,7 +5277,6 @@ def test_guard_receive_callback_preserves_completion_contract(awaiting_kvs, tran receiver.request_queue = MagicMock() receiver._send_done_signal_to_free_remote_port = MagicMock() receiver._send_done_recv_signal = MagicMock() - receiver._kv_guard = MagicMock() meta = { "request_id": "r", "remote_request_id": "remote-r", @@ -5288,7 +5287,6 @@ def test_guard_receive_callback_preserves_completion_contract(awaiting_kvs, tran "awaiting_kvs": awaiting_kvs, } receiver._handle_request(meta) - receiver._kv_guard.transfer.assert_called_once_with("error" if transfer_failed else "load", meta) if transfer_failed: receiver._reformat_pending_kv_caches.assert_not_called() else: diff --git a/vllm_ascend/debug/kv_guard.py b/vllm_ascend/debug/kv_guard.py index fb805acbc5a9..54353de1d11a 100644 --- a/vllm_ascend/debug/kv_guard.py +++ b/vllm_ascend/debug/kv_guard.py @@ -33,7 +33,7 @@ def __init__(self, emit: Any): self.pool_id: str | None = None self.blocks: dict[tuple[int, int], Block] = {} self.requests: dict[str, BlockTable] = {} - self.pending_loads: dict[str, dict[str, Any]] = {} + self.loads: dict[str, dict[str, Any]] = {} self.group_sizes: dict[int, int] = {} self._lock = threading.RLock() @@ -60,8 +60,7 @@ def bind(self, context: dict[str, Any]) -> None: if pool != self.pool_id: self.blocks.clear() self.requests.clear() - if self.pool_id is not None: - self.pending_loads.clear() + self.loads.clear() self.pool_id = pool for request_id in context.get("released_request_ids", ()): for group, refs in enumerate(self.requests.pop(request_id, ())): @@ -69,7 +68,7 @@ def bind(self, context: dict[str, Any]) -> None: block = self.blocks.get((group, block_id)) if block is not None: block.owners.discard(request_id) - self.pending_loads.pop(request_id, None) + self.loads.pop(request_id, None) for request_id, groups in context.get("requests", {}).items(): keys = {(group, block) for group, refs in enumerate(groups) for block, _ in refs} for group, refs in enumerate(self.requests.get(request_id, ())): @@ -78,10 +77,8 @@ def bind(self, context: dict[str, Any]) -> None: if key not in keys and key in self.blocks: self.blocks[key].owners.discard(request_id) self.requests[request_id] = groups - computed = context.get("scheduled_requests", {}).get(request_id, (0, 0))[0] or 0 for group, refs in enumerate(groups): - size = self.group_sizes[group] - for logical, (block_id, epoch) in enumerate(refs): + for block_id, epoch in refs: if block_id == 0: continue if epoch <= 0: @@ -101,63 +98,102 @@ def bind(self, context: dict[str, Any]) -> None: block = Block(epoch) self.blocks[key] = block block.owners.add(request_id) - valid = min(size, max(0, computed - logical * size)) - # Reuse may establish validity, but never revive a tainted - # allocation or a destination awaiting a physical copy. - if ( - valid - and block.valid_end == 0 - and block_id not in context.get("uninitialized_block_ids", ()) - ): - block.valid_end = valid - pending = self.pending_loads.pop(request_id, None) - if pending is not None: - self._load(pending) + # A consumer's computed prefix is a read requirement, + # not evidence that this allocation has been written. + for meta in context.get("loads", ()): + self.begin_load(meta) - def load(self, meta: dict[str, Any]) -> None: + def begin_load(self, meta: dict[str, Any]) -> None: + """Register the scheduler's submission snapshot, before any receive.""" with self._lock: - if meta["request_id"] not in self.requests or not self.group_sizes: - self.pending_loads[meta["request_id"]] = { - key: meta[key] for key in ("request_id", "local_block_ids", "num_computed_tokens") if key in meta - } - else: - self._load(meta) + if meta["pool_id"] != self.pool_id or meta["request_id"] not in self.requests: + self.emit("guard.coverage_gap", reason="load_binding_missing", request_id=meta["request_id"]) + return + if not 0 <= meta["token_start"] < meta["token_end"]: + self.emit("guard.coverage_gap", reason="load_range_invalid", request_id=meta["request_id"]) + return + self.loads[meta["request_id"]] = {**meta, "block_table": [list(refs) for refs in meta["block_table"]]} - def _load(self, meta: dict[str, Any]) -> None: + def _range_targets(self, meta: dict[str, Any], group_ids=None, block_ids=None): request_id = meta["request_id"] - expected = { - (group, block): (logical, epoch) - for group, refs in enumerate(self.requests.get(request_id, ())) - for logical, (block, epoch) in enumerate(refs) - } - for group, ids in enumerate(meta.get("local_block_ids", ())): - for block_id in ids: - key = group, block_id - ref, block = expected.get(key), self.blocks.get(key) - if ref is None or block is None or request_id not in block.owners: - self._violation(key, block, "load", "unowned_destination", request_id=request_id) + if meta["pool_id"] != self.pool_id: + self.emit("guard.coverage_gap", reason="stale_pool_event", request_id=request_id) + return + current = self.requests.get(request_id, ()) + for group, refs in enumerate(meta["block_table"]): + if group_ids is not None and group not in group_ids: + continue + size = self.group_sizes[group] + first = meta["token_start"] // size + last = (meta["token_end"] + size - 1) // size + for logical, (block_id, epoch) in enumerate(refs[first:last], first): + start = max(0, meta["token_start"] - logical * size) + end = min(size, meta["token_end"] - logical * size) + if start >= end or (block_ids is not None and block_id not in block_ids): continue - logical, epoch = ref - if block.epoch != epoch: + key = group, block_id + block = self.blocks.get(key) + if block is not None and block.epoch != epoch: self._violation(key, block, "load", "stale_allocation_epoch", request_id=request_id) + elif ( + block is None + or group >= len(current) + or logical >= len(current[group]) + or current[group][logical] != (block_id, epoch) + or request_id not in block.owners + ): + self._violation(key, block, "load", "unowned_destination", request_id=request_id) + else: + yield key, block, start, end + + def _active_load(self, meta: dict[str, Any]) -> bool: + active = self.loads.get(meta["request_id"]) + if active == meta: + return True + self.emit( + "guard.coverage_gap", reason="stale_load_completion", request_id=meta["request_id"], load_id=meta["load_id"] + ) + return False + + def load(self, meta: dict[str, Any], *, group_ids=None, block_ids=None, finish: bool = True) -> None: + """Only a completed, matching submission can establish readable ranges.""" + with self._lock: + if not self._active_load(meta): + return + for key, block, start, end in self._range_targets(meta, group_ids, block_ids): + if block.valid_end == INVALID_PREFIX: + self.emit("guard.coverage_gap", reason="load_destination_tainted", group_id=key[0], block_id=key[1]) continue - computed = meta.get("num_computed_tokens") - if computed is None: - self.emit( - "guard.coverage_gap", reason="load_valid_length_missing", group_id=group, block_id=block_id - ) + if start > block.valid_end: + self.emit("guard.coverage_gap", reason="load_prefix_missing", group_id=key[0], block_id=key[1]) continue - size = self.group_sizes[group] - block.valid_end = min(size, max(0, computed - logical * size)) + block.valid_end = max(block.valid_end, end) + if finish: + self.finish_load(meta) + + def finish_load(self, meta: dict[str, Any]) -> None: + with self._lock: + if self.loads.get(meta["request_id"]) == meta: + self.loads.pop(meta["request_id"]) + + def fail_load(self, meta: dict[str, Any], *, block_ids=None, finish: bool = True) -> None: + with self._lock: + if not self._active_load(meta): + return + for _, block, _, _ in self._range_targets(meta, block_ids=block_ids): + # A block transfer can overwrite its local prefix before failing. + # Untouched local blocks stay valid; failed targets can be recomputed. + if block.valid_end != INVALID_PREFIX: + block.valid_end = 0 + if finish: + self.finish_load(meta) - def fail_load(self, request_id: str) -> None: + def recompute(self, meta: dict[str, Any]) -> None: + """Truncate only a framework-authorized replay, before auditing writes.""" with self._lock: - self.pending_loads.pop(request_id, None) - for group, refs in enumerate(self.requests.get(request_id, ())): - for block_id, epoch in refs: - block = self.blocks.get((group, block_id)) - if block is not None and block.epoch == epoch: - block.valid_end = INVALID_PREFIX + for _, block, start, _ in self._range_targets(meta): + if block.valid_end != INVALID_PREFIX: + block.valid_end = min(block.valid_end, start) def clear(self, group_id: int, block_id: int) -> None: with self._lock: diff --git a/vllm_ascend/debug/kv_guard_connector.py b/vllm_ascend/debug/kv_guard_connector.py new file mode 100644 index 000000000000..ca554b9ba52b --- /dev/null +++ b/vllm_ascend/debug/kv_guard_connector.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Translate the common connector lifecycle into logical KV readiness events. + +Submission snapshots come from the scheduler, never from backend-specific +transport offsets. No connector calls, waits or result polling are added. +""" + +from __future__ import annotations + +import functools +from typing import Any + + +class ConnectorGuard: + def __init__(self, runtime: Any, connector: Any): + self.runtime = runtime + self.state = runtime.state + self.staged: list[dict] = [] + self.inflight: dict[str, dict] = {} + self.sync: dict[str, dict] = {} + self.drained: list[dict] = [] + self.completed: list[tuple[dict, set[int]]] = [] + self.ready: dict[int, set[str]] = {} + self.ready_groups: dict[int, set[int]] = {} + self.failed: set[int] = set() + self.failed_blocks: dict[tuple[str, int], set[int]] = {} + self.post_forward_complete = False + self._wrap(connector, "start_load_kv", self.start) + self._wrap(connector, "wait_for_layer_load", self.wait) + self._wrap(connector, "get_transfer_results", self.results) + self._wrap(connector, "get_block_ids_with_load_errors", self.errors) + + @staticmethod + def _wrap(connector, name, callback): + original = getattr(connector, name) + + @functools.wraps(original) + def observed(*args, **kwargs): + return callback(original, *args, **kwargs) + + setattr(connector, name, observed) + + def stage(self, loads: list[dict]) -> None: + if self.staged or self.sync: + self.finish_step() + self.staged.extend(loads) + + def start(self, original, *args, **kwargs): + submitted, self.staged = self.staged, [] + accepted = [] + for meta in submitted: + request = meta["request_id"] + if request in self.inflight: + # Request-ID-only completions cannot distinguish overlapping incarnations. + self.runtime.gap("overlapping_load_request_id", request_id=request) + self.state.finish_load(meta) + continue + accepted.append(meta) + if meta["is_async"]: + self.inflight[request] = meta + else: + self.sync[request] = meta + self.runtime._check_failure() + try: + return original(*args, **kwargs) + except Exception: + for meta in accepted: + self.state.fail_load(meta) + self.inflight.pop(meta["request_id"], None) + raise + + def wait(self, original, layer_name, *args, **kwargs): + result = original(layer_name, *args, **kwargs) + group = self.runtime.layer_groups.get(layer_name) + if group is None: + return result + required = {name for name, index in self.runtime.layer_groups.items() if index == group} + for meta in self.sync.values(): + if meta["load_id"] in self.failed: + continue + ready = self.ready.setdefault(meta["load_id"], set()) + ready.add(layer_name) + groups = self.ready_groups.setdefault(meta["load_id"], set()) + if required <= ready and group not in groups: + self.state.load(meta, group_ids={group}, finish=False) + groups.add(group) + self.runtime._check_failure() + return result + + def check_ready(self, layer: str) -> None: + group = self.runtime.layer_groups.get(layer) + if group is None: + return + for meta in self.sync.values(): + if layer not in self.ready.get(meta["load_id"], ()): + self.runtime.gap("load_layer_not_ready", request_id=meta["request_id"], layer=layer) + self.runtime._missing_groups.add(group) + self.runtime._unready_load_groups.add(group) + + def before_audit(self) -> None: + for layer in self.runtime.layer_groups: + self.check_ready(layer) + + def results(self, original, *args, **kwargs): + result = original(*args, **kwargs) + failed = set(result.failed_recving or ()) + for request in set(result.finished_recving or ()) | failed: + meta = self.inflight.pop(request, None) + if meta is None: + meta = self.sync.get(request) + if meta is None: + continue + if request in failed: + self.failed.add(meta["load_id"]) + if not meta["is_async"]: + self.state.fail_load(meta, finish=False) + self.drained.append(meta) + return result + + def errors(self, original, *args, **kwargs): + invalid = original(*args, **kwargs) + if invalid: + # Core queries invalid blocks after transfer results; keep the snapshot until then. + for meta in (*self.inflight.values(), *self.sync.values()): + if self.state.loads.get(meta["request_id"]) != meta: + continue + affected = invalid & self._target_blocks(meta) + if affected: + self.failed_blocks.setdefault((meta["pool_id"], meta["load_id"]), set()).update(affected) + self.state.fail_load(meta, block_ids=affected, finish=False) + for meta in self.drained: + self.completed.append((meta, set(invalid or ()))) + self.drained.clear() + self.post_forward_complete = True + if self.runtime._phase is None: + self.finish_step() + self.runtime._check_failure() + return invalid + + def _target_blocks(self, meta: dict) -> set[int]: + targets = set() + for group, refs in enumerate(meta["block_table"]): + size = self.state.group_sizes[group] + first, last = meta["token_start"] // size, (meta["token_end"] + size - 1) // size + targets.update(block for block, _ in refs[first:last]) + return targets + + def commit_completions(self) -> None: + # MRV2 polls after forward but before its audit. These completions may + # authorize the next step, never an earlier read in the current step. + for meta, invalid in self.completed: + # Legacy connectors can report failed blocks in an earlier poll, + # then drain only a request-ID completion in a later step. + invalid |= self.failed_blocks.pop((meta["pool_id"], meta["load_id"]), set()) + if self.state.loads.get(meta["request_id"]) != meta: + continue # Cancellation may release ownership before the known old receive drains. + if meta["is_async"]: + if meta["load_id"] in self.failed: + self.state.fail_load(meta, finish=False) + else: + self.state.load(meta, block_ids=self._target_blocks(meta) - invalid, finish=False) + if invalid: + self.state.fail_load(meta, block_ids=invalid, finish=False) + if meta["is_async"] or meta["load_id"] in self.failed: + self.state.finish_load(meta) + self.completed.clear() + + def finish_step(self) -> None: + self.commit_completions() + for meta in self.staged: + self.runtime.gap("load_not_submitted", request_id=meta["request_id"]) + self.state.finish_load(meta) + for meta in self.sync.values(): + if meta["load_id"] not in self.failed: + required = set(self.runtime.layer_groups.values()) + if not required <= self.ready_groups.get(meta["load_id"], set()): + self.runtime.gap("load_completion_missing", request_id=meta["request_id"]) + self.state.finish_load(meta) + self.failed_blocks.pop((meta["pool_id"], meta["load_id"]), None) + self.staged.clear() + self.sync.clear() + self.ready.clear() + self.ready_groups.clear() + self.failed.clear() + self.post_forward_complete = False diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py index 79357d5f8c35..7571938039ca 100644 --- a/vllm_ascend/debug/kv_guard_runtime.py +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -14,12 +14,17 @@ import functools from collections import OrderedDict from contextlib import contextmanager +from copy import copy from dataclasses import dataclass, replace +from threading import Lock from types import SimpleNamespace from typing import Any, cast +import torch + from vllm_ascend.debug.kv_guard import INVALID_PREFIX, KVGuard from vllm_ascend.debug.kv_guard_config import KVGuardConfig +from vllm_ascend.debug.kv_guard_connector import ConnectorGuard _MAX_REPORT_KEYS = 1024 @@ -52,6 +57,29 @@ class KVGuardError(RuntimeError): """Stop returning model results after a violation or incomplete audit.""" +class _PendingAudit: + """One snapshot, consumed once by the output reader or next forward.""" + + def __init__(self, complete: Any): + self._complete = complete + self._lock = Lock() + self._done = False + self._error: BaseException | None = None + + def finish(self) -> None: + with self._lock: + if not self._done: + try: + self._complete() + except BaseException as exc: + self._error = exc + finally: + self._complete = None + self._done = True + if self._error is not None: + raise self._error + + class GuardRuntime: def __init__(self, vllm_config: Any): from vllm.logger import logger @@ -73,6 +101,14 @@ def __init__(self, vllm_config: Any): self._failure: str | None = None self.action = KVGuardConfig.from_vllm_config(vllm_config).action self._runner_adapter: _V2RunnerAdapter | None = None + self._host_buffers: list[Any] = [] + self._copy_event: Any = None + self._connector: ConnectorGuard | None = None + self._recomputes: list[dict] = [] + self._unready_load_groups: set[int] = set() + self._defer_audit = False + self._pending: _PendingAudit | None = None + self._deferred_state: list[tuple[Any, tuple]] = [] def _check_failure(self) -> None: if self.action == "raise" and self._failure is not None: @@ -85,31 +121,32 @@ def _emit_scratch(self, event: str, **fields: Any) -> None: def report(self, event: str, **fields: Any) -> None: failed = event in ("guard.violation", "guard.coverage_gap") if failed: - if self._failure is None: - self._failure = str(fields.get("reason", event)) - key = ( - event, - *( - fields.get(name) - for name in ( - "operation", - "reason", - "pool_id", - "group_id", - "block_id", - "alloc_epoch", - "layer", - "request_id", - ) - ), - tuple(fields.get("request_ids", ())), - ) - if key in self._reported: - self._reported.move_to_end(key) - return - self._reported[key] = None - if len(self._reported) > _MAX_REPORT_KEYS: - self._reported.popitem(last=False) + with self.state._lock: + if self._failure is None: + self._failure = str(fields.get("reason", event)) + key = ( + event, + *( + fields.get(name) + for name in ( + "operation", + "reason", + "pool_id", + "group_id", + "block_id", + "alloc_epoch", + "layer", + "request_id", + ) + ), + tuple(fields.get("request_ids", ())), + ) + if key in self._reported: + self._reported.move_to_end(key) + return + self._reported[key] = None + if len(self._reported) > _MAX_REPORT_KEYS: + self._reported.popitem(last=False) payload = " ".join(f"{key}={value}" for key, value in fields.items()) log = self.logger.error if failed else self.logger.info log("[KV_GUARD] %s %s", event, payload) @@ -128,7 +165,7 @@ def configure(self, runner: Any) -> None: getattr(parallel, "prefill_context_parallel_size", 1) != 1, getattr(parallel, "decode_context_parallel_size", 1) != 1, getattr(parallel, "enable_dbo", False), - transfer is not None and transfer.kv_connector != "MooncakeConnector", + transfer is not None and transfer.kv_connector not in ("MooncakeConnector", "MooncakeConnectorV1"), ) ) sizes = {} @@ -161,15 +198,28 @@ def configure(self, runner: Any) -> None: "(CP, DBO, speculation, encoder/decoder, connector)" ) self.state.configure(sizes) + self._defer_audit = ( + bool(getattr(getattr(config, "scheduler_config", None), "async_scheduling", False)) + and transfer is None + # Pooling bypasses sample_tokens; non-reply workers may discard outputs. + and getattr(runner, "is_pooling_model", None) is False + and getattr(parallel, "tensor_parallel_size", None) == 1 + and getattr(parallel, "pipeline_parallel_size", None) == 1 + and hasattr(runner, "sample_tokens") + ) if self._runner_adapter is not None: self._runner_adapter.configure() self.report("guard.enabled", failure_action=self.action) def write(self, layer: str, slots: Any, token_offset: int = 0) -> None: + if self._phase == "forward" and self._connector is not None and self._connector.sync: + self._connector.check_ready(layer) if self._phase is not None and not self._full_graph: self._accesses.append(WriteAccess(layer, slots, token_offset)) def read(self, layer: str, table: Any, lengths: Any, starts: Any = None, request_offset: int = 0) -> None: + if self._phase == "forward" and self._connector is not None and self._connector.sync: + self._connector.check_ready(layer) if self._phase is not None and not self._full_graph and table is not None: self._accesses.append(ReadAccess(layer, table, lengths, starts, request_offset)) @@ -265,7 +315,174 @@ def _expected_write_plan(self, runner: Any, group_id: int) -> ExpectedAccess | N query_ends.append(len(slots)) return ExpectedAccess(slots, query_ends, tables, lengths) + def _preview_append(self, runner: Any, expected: dict[int, ExpectedAccess]) -> Any: + """Prepare a silent, unpublished canonical WRITE-then-READ transaction.""" + baseline, proposed, bindings = {}, {}, {} + with self.state._lock: + pool = self.state.pool_id + for group_id, plan in expected.items(): + size = self.state.group_sizes[group_id] + intervals: dict[int, list[tuple[int, int]]] = {} + for row, request in enumerate(runner.input_batch.req_ids): + groups = self.state.requests[request] + bindings[request] = tuple(tuple(refs) for refs in groups) + refs = groups[group_id] + for block_id, epoch in refs[: len(plan.tables[row])]: + key = group_id, block_id + block = self.state.blocks.get(key) + if ( + block_id == 0 + or block is None + or epoch != block.epoch + or block.valid_end == INVALID_PREFIX + or request not in block.owners + ): + return None + baseline[key] = block.epoch, block.valid_end, frozenset(block.owners) + computed, count = self._scheduled[request] + for logical in range(computed // size, (computed + count - 1) // size + 1): + block_id, _ = refs[logical] + start = max(computed, logical * size) - logical * size + end = min(computed + count, (logical + 1) * size) - logical * size + intervals.setdefault(block_id, []).append((start, end)) + for block_id, spans in intervals.items(): + key = group_id, block_id + cursor = baseline[key][1] + for start, end in sorted(spans): + if start != cursor: + return None + cursor = end + proposed[key] = cursor + for row, length in enumerate(plan.lengths): + for logical, block_id in enumerate(plan.tables[row]): + key = group_id, block_id + if min(size, length - logical * size) > proposed.get(key, baseline[key][1]): + return None + return pool, baseline, proposed, bindings + + def _try_fast_audit(self, runner: Any) -> bool: + """Exact final metadata equality permits committing a precomputed state. + + shortcut: only canonical FULL plans; other plans retain the ordered audit. + """ + if ( + self._phase != "forward" + or not self._full_graph + or self._scheduled is None + or self._missing_groups + or self._unready_load_groups + or not self._metadata + or self._metadata.keys() != self.layer_groups.keys() + ): + return False + request_ids = list(runner.input_batch.req_ids) + if not request_ids or len(set(request_ids)) != len(request_ids): + return False + expected: dict[int, ExpectedAccess] = {} + inputs, checks, seen = {}, [], set() + for layer, item in self._metadata.items(): + group_id = self.layer_groups[layer] + if (group_id, id(item)) in seen: + continue + seen.add((group_id, id(item))) + if group_id not in expected: + plan = self._expected_write_plan(runner, group_id) + if plan is None: + return False + expected[group_id] = plan + plan = expected[group_id] + accesses = self._metadata_accesses(layer, item) + if len(accesses) != 2 or self._missing_groups: + return False + write, read = accesses + slots, table, lengths = write.slots, read.table, read.lengths + if not all( + isinstance(value, torch.Tensor) and value.dtype in (torch.int32, torch.int64) + for value in (slots, table, lengths) + ): + return False + rows, columns = len(request_ids), max(map(len, plan.tables), default=0) + if ( + slots.ndim != 1 + or slots.numel() != len(plan.slots) + or lengths.shape != (rows,) + or table.ndim != 2 + or table.shape[0] < rows + or table.shape[1] < columns + or any(value.device.type not in ("cpu", "npu") for value in (slots, table, lengths)) + ): + return False + keys = [] + for value in (slots, lengths, table[:rows, :columns]): + key = str(value.device), value.dtype, self._identity(value) + inputs.setdefault(key, value) + keys.append(key) + checks.append((plan, keys)) + + # Enqueue every copy before the CPU-only preview. No device values are + # inspected until the event, at the original post-forward audit point. + copied = {} + event = None + recorded = False + try: + for index, (key, value) in enumerate(inputs.items()): + if value.device.type == "cpu": + copied[key] = self._list(value) + continue + if self._copy_event is None: + self._copy_event = torch.npu.Event() + event = self._copy_event + if index >= len(self._host_buffers): + self._host_buffers.extend([None] * (index + 1 - len(self._host_buffers))) + host = self._host_buffers[index] + if host is None or host.shape != value.shape or host.dtype != value.dtype: + host = torch.empty(value.shape, dtype=value.dtype, device="cpu", pin_memory=True) + self._host_buffers[index] = host + host.copy_(value, non_blocking=True) + copied[key] = host + if event is not None: + event.record(torch.npu.current_stream()) + recorded = True + candidate = self._preview_append(runner, expected) + finally: + if event is not None: + # A partial enqueue failure must also drain outstanding DMA. + if not recorded: + event.record(torch.npu.current_stream()) + event.synchronize() + if candidate is None: + return False + values = {key: value.tolist() if isinstance(value, torch.Tensor) else value for key, value in copied.items()} + for plan, (slots, lengths, table) in checks: + if ( + values[slots] != plan.slots + or values[lengths] != plan.lengths + or any(row[: len(want)] != want for row, want in zip(values[table], plan.tables)) + ): + return False + pool, baseline, proposed, bindings = candidate + with self.state._lock: + unchanged = self.state.pool_id == pool and all( + tuple(tuple(refs) for refs in self.state.requests.get(request, ())) == groups + for request, groups in bindings.items() + ) + for key, before in baseline.items(): + block = self.state.blocks.get(key) + unchanged = ( + unchanged + and block is not None + and (block.epoch, block.valid_end, frozenset(block.owners)) == before + ) + if not unchanged: + self.gap("allocation_changed_during_access_audit") + return True + for key, valid in proposed.items(): + self.state.blocks[key].valid_end = valid + return True + def _audit(self, runner: Any) -> None: + if self._try_fast_audit(runner): + return phase = self._phase assert phase is not None # FULL requires this replay's final metadata. Never reuse a captured @@ -296,6 +513,8 @@ def _audit(self, runner: Any) -> None: if group_id is None: self.gap(f"unbound_access_layer:{layer}") continue + if group_id in self._unready_load_groups: + continue actual = observed.get(layer, []) if observed and not actual and not self._full_graph: self.gap(f"operator_layer_not_observed:{layer}") @@ -446,6 +665,137 @@ def observe_metadata(self, metadata: Any) -> None: self._missing_groups = set(self.layer_groups.values()) self.gap(f"metadata_collection_failed:{type(exc).__name__}") + def _finish_pending(self) -> None: + if self._pending is not None: + self._pending.finish() + self._pending = None + operations, self._deferred_state = self._deferred_state, [] + for operation, args in operations: + operation(*args) + self._check_failure() + + def _snapshot_audit(self, runner: Any) -> None: + """Freeze final evidence now; wait only at an output/forward boundary.""" + assert self._pending is None + job = copy(self) + job._scheduled = None if self._scheduled is None else dict(self._scheduled) + job._missing_groups = set(self._missing_groups) + job._unready_load_groups = set(self._unready_load_groups) + job.report, job.gap, job._check_failure = self.report, self.gap, self._check_failure + job._scratch = KVGuard(job._emit_scratch) + frozen_runner = SimpleNamespace(input_batch=SimpleNamespace(req_ids=list(runner.input_batch.req_ids))) + copies: dict[Any, Any] = {} + event = None + device_copies = 0 + + def freeze(value): + nonlocal event, device_copies + if isinstance(value, torch.Tensor): + key = value.device, value.dtype, self._identity(value) + if key not in copies: + if value.device.type == "cpu": + host = value.clone() + else: + if self._copy_event is None: + self._copy_event = torch.npu.Event() + event = self._copy_event + index = device_copies + device_copies += 1 + if index >= len(self._host_buffers): + self._host_buffers.append(None) + host = self._host_buffers[index] + if host is None or host.shape != value.shape or host.dtype != value.dtype: + host = torch.empty(value.shape, dtype=value.dtype, device="cpu", pin_memory=True) + self._host_buffers[index] = host + host.copy_(value, non_blocking=True) + copies[key] = host + return copies[key] + if isinstance(value, (list, tuple)): + return [freeze(item) for item in value] + return value + + try: + job._accesses = [ + WriteAccess(access.layer, freeze(access.slots), access.token_offset) + if isinstance(access, WriteAccess) + else ReadAccess( + access.layer, + freeze(access.table), + freeze(access.lengths), + freeze(access.starts), + access.request_offset, + ) + for access in self._accesses + ] + observed = {access.layer for access in self._accesses} + items, job._metadata = {}, {} + for layer, item in self._metadata.items(): + if not self._full_graph and layer in observed: + job._metadata[layer] = None # Only coverage is used when operator arguments were observed. + continue + key = id(item) + if key not in items: + slots = self._field(item, "slot_mapping") + count = getattr(item, "num_actual_tokens", None) + if slots is not None and count is not None: + slots = slots[:count] + items[key] = SimpleNamespace( + slot_mapping=freeze(slots), + block_tables=freeze(self._field(item, "block_tables", "block_table", "block_table_tensor")), + seq_lens=freeze(self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu")), + ) + job._metadata[layer] = items[key] + if event is not None: + event.record(torch.npu.current_stream()) + except Exception: + if event is not None: + # Even a partially failed snapshot must drain before buffer reuse. + event.record(torch.npu.current_stream()) + event.synchronize() + raise + + def complete(): + with torch.inference_mode(): + try: + if event is not None: + event.synchronize() + job._audit(frozen_runner) + except Exception as exc: + self.gap(f"access_audit_failed:{type(exc).__name__}") + self._check_failure() + + self._pending = _PendingAudit(complete) + + def _guard_output(self, output: Any) -> Any: + self._check_failure() + pending = self._pending + if pending is None: + return output + get_output = getattr(output, "get_output", None) + if get_output is None: + pending.finish() + self._check_failure() + return output + + @functools.wraps(get_output) + def checked(): + pending.finish() + self._check_failure() + return get_output() + + try: + output.get_output = checked + except (AttributeError, TypeError): + pending.finish() + return output + + def wrap_sample(self, sample: Any): + @functools.wraps(sample) + def observed(*args, **kwargs): + return self._guard_output(sample(*args, **kwargs)) + + return observed + @contextmanager def forward_scope( self, runner: Any, *, phase: str, metadata: Any = None, full_graph: bool = False, required: bool = True @@ -459,16 +809,30 @@ def forward_scope( yield if required or self._metadata or self._accesses: try: - self._audit(runner) + if phase == "forward" and self._connector is not None: + self._connector.before_audit() + if phase == "forward": + for meta in self._recomputes: + self.state.recompute(meta) + self._recomputes = [] + if self._defer_audit and phase == "forward" and self._connector is None: + self._snapshot_audit(runner) + else: + self._audit(runner) except Exception as exc: self.gap(f"access_audit_failed:{type(exc).__name__}") self._check_failure() finally: self._phase = None + if self._connector is not None: + self._connector.commit_completions() + if self._connector.post_forward_complete: + self._connector.finish_step() def wrap_forward(self, runner: Any, forward: Any, get_context: Any): @functools.wraps(forward) def observed(num_tokens, *args, **kwargs): + self._finish_pending() context = get_context() context._kv_guard = self with self.forward_scope( @@ -491,6 +855,8 @@ def observed(output, *args, **kwargs): self._check_failure() return result for request_id, groups in binding["requests"].items(): + if request_id not in binding.get("scheduled_requests", {}): + continue # Receive-only bindings are not runner batch entries. actual = runner.requests.get(request_id) expected = [[block for block, _ in refs] for refs in groups] if actual is None or [list(g) for g in actual.block_ids] != expected: @@ -506,10 +872,12 @@ def observed(output, *args, **kwargs): return observed @contextmanager - def execute_scope(self, output: Any, *, bind: bool = True): + def execute_scope(self, output: Any, *, bind: bool = True, finish_connector: bool = True): """Share scheduler binding and cleanup across both runner versions.""" self._check_failure() self._scheduled = None + self._unready_load_groups.clear() + self._recomputes = [] if bind: binding = getattr(output, "_ascend_kv_guard_context", None) if binding is not None: @@ -518,7 +886,22 @@ def execute_scope(self, output: Any, *, bind: bool = True): uninitialized.update( copy.dst_block_id for copy in (getattr(output, "kv_cache_block_copies", None) or ()) ) - self.state.bind({**binding, "uninitialized_block_ids": uninitialized}) + context = {**binding, "uninitialized_block_ids": uninitialized} + if self._pending is not None: + self._deferred_state.append((self.state.bind, (context,))) + else: + self.state.bind(context) + self._recomputes = binding.get("recompute_ranges", []) + if self._connector is None and getattr(self.config, "kv_transfer_config", None) is not None: + from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group + + if has_kv_transfer_group(): + self._connector = ConnectorGuard(self, get_kv_transfer_group()) + loads = binding.get("loads", []) + if self._connector is not None: + self._connector.stage(loads) + elif loads: + self.gap("load_connector_missing") self._scheduled = binding.get("scheduled_requests") if self._scheduled is None: self.gap("scheduler_write_plan_missing") @@ -527,18 +910,29 @@ def execute_scope(self, output: Any, *, bind: bool = True): try: self._check_failure() yield + self._check_failure() finally: + if self._deferred_state: + self._finish_pending() # Receive-only, idle and failed steps may never enter a forward. + if self._connector is not None and finish_connector: + self._connector.finish_step() self._scheduled = None + self._recomputes = [] + self._check_failure() def wrap_execute(self, execute: Any): @functools.wraps(execute) def observed(output, *args, **kwargs): with self.execute_scope(output): - return execute(output, *args, **kwargs) + result = execute(output, *args, **kwargs) + return self._guard_output(result) if result is not None else result return observed def zero_blocks(self, block_ids: list[int]) -> None: + if self._pending is not None: + self._deferred_state.append((self.zero_blocks, (list(block_ids),))) + return with self.state._lock: for group_id in self.state.group_sizes: for block_id in block_ids: @@ -555,6 +949,9 @@ def zero_blocks(self, block_ids: list[int]) -> None: self._check_failure() def copy_blocks(self, copies: Any) -> None: + if self._pending is not None: + self._deferred_state.append((self.copy_blocks, (list(copies),))) + return with self.state._lock: for copy in copies: if copy.dst_block_id == 0: @@ -571,15 +968,6 @@ def copy_blocks(self, copies: Any) -> None: self.state.copy_blocks(group_id, pairs) self._check_failure() - def transfer(self, event: str, meta: dict[str, Any]) -> None: - try: - if event == "load": - self.state.load(meta) - else: - self.state.fail_load(meta["request_id"]) - except Exception as exc: - self.gap(f"load_hook_failed:{type(exc).__name__}") - def get_guard(vllm_config: Any) -> GuardRuntime | None: if not KVGuardConfig.from_vllm_config(vllm_config).enabled: @@ -613,6 +1001,7 @@ def configure(self) -> None: def observed(input_batch, cudagraph_mode, *args, **kwargs): metadata = prepare(input_batch, cudagraph_mode, *args, **kwargs) if self.runtime._phase is not None: + self.runtime._finish_pending() self.input_batch = SimpleNamespace(req_ids=list(input_batch.req_ids[: input_batch.num_reqs])) self._prepared = True self.runtime._full_graph = "FULL" in str(cudagraph_mode) @@ -641,20 +1030,22 @@ def observed(scheduler_output, *args, **kwargs): required = bool(scheduler_output.total_num_scheduled_tokens) and not ( profile and flags.get("skip_attn_for_dummy_run") ) + if dummy or profile: + self.runtime._finish_pending() self.input_batch = SimpleNamespace(req_ids=[]) self._prepared = False # Synthetic dummy/profile outputs have no scheduler binding, but # their writes must still be checked against the shared live KV state. with ( - self.runtime.execute_scope(scheduler_output, bind=not (dummy or profile)), + self.runtime.execute_scope(scheduler_output, bind=not (dummy or profile), finish_connector=False), self.runtime.forward_scope(self, phase=phase, required=required), ): result = execute(scheduler_output, *args, **kwargs) if phase == "forward" and required and not self._prepared: self.runtime.gap("v2_attention_metadata_missing") self.runtime._missing_groups.update(self.runtime.layer_groups.values()) - return result + return self.runtime._guard_output(result) if result is not None else result return observed @@ -680,7 +1071,7 @@ def observed(*args, **kwargs): with self.runtime.state._lock: if ( self.runtime.state.requests - or self.runtime.state.pending_loads + or self.runtime.state.loads or any(block.valid_end != 0 for block in self.runtime.state.blocks.values()) ): raise ValueError( @@ -711,4 +1102,6 @@ def install_guard(runner: Any) -> GuardRuntime | None: cast(Any, runner)._model_forward = runtime.wrap_forward(runner, runner._model_forward, get_forward_context) cast(Any, runner).execute_model = runtime.wrap_execute(runner.execute_model) cast(Any, runner)._update_states = runtime.wrap_update(runner, runner._update_states) + if hasattr(runner, "sample_tokens"): + cast(Any, runner).sample_tokens = runtime.wrap_sample(runner.sample_tokens) return runtime diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py index 6e6c8063723c..de6166bda467 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py @@ -59,7 +59,6 @@ from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec, AscendSlidingWindowMLASpec -from vllm_ascend.debug.kv_guard_runtime import get_guard from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import global_te from vllm_ascend.distributed.kv_transfer.utils.utils import ( PD_QOS_DEFAULT, @@ -510,7 +509,6 @@ def __init__( self.request_task_counts_lock = threading.Lock() self.task_tracker = KVCacheTaskTracker() - self._kv_guard = get_guard(vllm_config) self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeAgentMetadata) @@ -777,9 +775,6 @@ def _handle_request(self, req_meta: dict[str, Any]): e, ) # Drain the failure with completion, after all local writes end. - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.transfer("error" if transfer_failed else "load", req_meta) if req_meta.get("awaiting_kvs", True): self.task_tracker.update_done_task_count(request_id) else: diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index f7f875a8f83f..e6fd376e96fb 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1686,6 +1686,7 @@ # Wrap the constructor once, then instrument allocation, release, and # schedule only for instances with additional_config.kv_guard_config.enabled. # Attach changed bindings and token ranges to SchedulerOutput. +# Include receive-only allocation snapshots and explicit prefix replay ranges. # Related PR (if no, explain why): # No upstream PR yet; this is the opt-in Ascend KV Guard integration. # Future Plan: diff --git a/vllm_ascend/patch/platform/patch_kv_guard.py b/vllm_ascend/patch/platform/patch_kv_guard.py index 1a64b60878d1..a21a41b7a9e3 100644 --- a/vllm_ascend/patch/platform/patch_kv_guard.py +++ b/vllm_ascend/patch/platform/patch_kv_guard.py @@ -7,6 +7,7 @@ """ import functools +import inspect import uuid from typing import Any, cast @@ -39,12 +40,85 @@ def allocated(*args, **kwargs): schedule = self.schedule mappings: dict[str, tuple] = {} released: set[str] = set() + loads: dict[str, dict] = {} + replays: dict[str, tuple[int, int]] = {} + load_sequence = 0 + + def block_refs(request_id): + return [ + [(block, counters[block]) for block in group] + for group in self.kv_cache_manager.get_block_ids(request_id) + ] + + allocate_slots = getattr(self.kv_cache_manager, "allocate_slots", None) + if allocate_slots is not None: + signature = inspect.signature(allocate_slots) + external_position = list(signature.parameters).index("num_external_computed_tokens") - 1 + + @functools.wraps(allocate_slots) + def slots_allocated(request, *args, **kwargs): + nonlocal load_sequence + external = kwargs.get( + "num_external_computed_tokens", args[external_position] if len(args) > external_position else 0 + ) + if not external: + return allocate_slots(request, *args, **kwargs) + arguments = signature.bind(request, *args, **kwargs).arguments + local_end = request.num_computed_tokens + arguments.get("num_new_computed_tokens", 0) + result = allocate_slots(request, *args, **kwargs) + if result is not None: + load_sequence += 1 + loads[request.request_id] = { + "pool_id": pool_id, + "load_id": load_sequence, + "request_id": request.request_id, + "block_table": block_refs(request.request_id), + "token_start": local_end, + "token_end": min( + local_end + external, getattr(self.kv_cache_manager, "max_model_len", local_end + external) + ), + "is_async": arguments.get("delay_cache_blocks", False), + } + return result + + self.kv_cache_manager.allocate_slots = slots_allocated + + def record_replay(request, start, end): + if start < end: + previous = replays.get(request.request_id, (start, end)) + replays[request.request_id] = min(start, previous[0]), max(end, previous[1]) + + mark_replay = getattr(self, "_mark_prefix_replay", None) + if mark_replay is not None: + + @functools.wraps(mark_replay) + def replay_marked(request, num_hit_tokens, *args, **kwargs): + count = mark_replay(request, num_hit_tokens, *args, **kwargs) + record_replay(request, num_hit_tokens - count, num_hit_tokens) + return count + + self._mark_prefix_replay = replay_marked + + update_waiting = getattr(self, "_update_waiting_for_remote_kv", None) + if update_waiting is not None: + + @functools.wraps(update_waiting) + def waiting_updated(request, *args, **kwargs): + before = request.num_computed_tokens + result = update_waiting(request, *args, **kwargs) + record_replay(request, request.num_computed_tokens, before) + return result + + self._update_waiting_for_remote_kv = waiting_updated + free = self.kv_cache_manager.free @functools.wraps(free) def freed(request, *args, **kwargs): result = free(request, *args, **kwargs) released.add(request.request_id) + loads.pop(request.request_id, None) + replays.pop(request.request_id, None) return result self.kv_cache_manager.free = freed @@ -58,9 +132,10 @@ def scheduled(*args, **kwargs): changed = {} for request_id in output.finished_req_ids | released: mappings.pop(request_id, None) - for request_id in output.num_scheduled_tokens: - blocks = [list(group) for group in self.kv_cache_manager.get_block_ids(request_id)] - refs = [[(block, counters[block]) for block in ids] for ids in blocks] + if request_id in output.finished_req_ids: + replays.pop(request_id, None) + for request_id in dict.fromkeys((*output.num_scheduled_tokens, *loads)): + refs = block_refs(request_id) signature = tuple(tuple(group) for group in refs) if mappings.get(request_id) != signature: changed[request_id] = refs @@ -74,6 +149,24 @@ def scheduled(*args, **kwargs): for request_id, count in output.num_scheduled_tokens.items() }, } + if loads: + output._ascend_kv_guard_context["loads"] = list(loads.values()) + loads.clear() + replay_ranges = [] + for request_id in output.num_scheduled_tokens: + replay = replays.pop(request_id, None) + if replay is not None: + replay_ranges.append( + { + "pool_id": pool_id, + "request_id": request_id, + "block_table": block_refs(request_id), + "token_start": replay[0], + "token_end": replay[1], + } + ) + if replay_ranges: + output._ascend_kv_guard_context["recompute_ranges"] = replay_ranges released.clear() return output From a0d8bef0b2efb7ed9da31ba683671e462e8e3080 Mon Sep 17 00:00:00 2001 From: TAJh Date: Fri, 6 Nov 2026 09:43:02 +0000 Subject: [PATCH 4/6] refactor: audit KV guard metadata once per cache group Signed-off-by: TAJh --- .../feature_guide/kv_cache_guard.md | 52 +- release_notes/fragments/17833.yaml | 4 +- tests/ut/debug/test_kv_guard_async.py | 30 +- tests/ut/debug/test_kv_guard_loads.py | 67 +- tests/ut/debug/test_kv_guard_runtime.py | 692 +++++------------- tests/ut/debug/test_kv_guard_v2.py | 7 - vllm_ascend/attention/attention_v1.py | 15 - vllm_ascend/attention/mla_v1.py | 26 - vllm_ascend/debug/kv_guard_connector.py | 2 +- vllm_ascend/debug/kv_guard_runtime.py | 531 +++----------- 10 files changed, 331 insertions(+), 1095 deletions(-) diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md index 3ed92851f7dc..4c4f8aba3504 100644 --- a/docs/source/user_guide/feature_guide/kv_cache_guard.md +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -2,7 +2,7 @@ KV Guard diagnoses address mappings that can cause cross-request writes, prefix overwrites, or reads from invalid KV ranges. It compares scheduled -destinations with execution metadata and observed operator arguments. +destinations with the runner's final execution metadata, once per KV cache group. The feature is disabled by default and supports adapted dense PA and MLA attention paths in Model Runner V1 and V2. @@ -62,7 +62,7 @@ slot of every scheduled token from those bindings and the actual batch order. | Logical block order | A request reads its own blocks in the wrong logical order. | | Initialized prefix | A read extends into unwritten data, or a write overwrites a valid prefix. | | Reserved storage | A write targets the null block instead of negative-slot padding. | -| Collection completeness | A scheduled write, binding, layer, or current graph plan is missing. | +| Collection completeness | A scheduled write, binding, cache group, or current graph plan is missing. | Each tracked block stores an allocation generation, a valid prefix length, and request owners. A zero prefix is unwritten; a positive prefix N makes `[0, N)` @@ -107,19 +107,19 @@ full-model incident. | Path | Audit evidence and boundary | | --- | --- | -| Eager | Ordered per-layer operator arguments, with adapted metadata plans where applicable. | -| FULL replay | The current execution's final attention metadata, inspected after forward and graph parameter update calls. Captured or previous-step plans are not reused as evidence. | -| Runner V1 | Bind before `execute_model`; collect through the existing forward context and attention hooks. | +| Eager and FULL replay | One final metadata record per KV cache group, inspected after forward and graph parameter updates. Captured or previous-step plans are not reused as evidence. | +| Runner V1 | Bind before `execute_model`; collect metadata through the existing forward context. | | Runner V2 | Bind before execute; obtain actual batch IDs and metadata after `model_state.prepare_attn`; audit at the end of execute. Persistent request-state indices are not batch positions. | | Dummy/profile execution | Check writes against existing KV state; ignore discarded scratch reads. | | Runner V2 startup capture | Require empty request/cache state; do not audit CPU metadata inside device graph capture. | -Canonical FULL plans can prepare a CPU validity transaction while the final -metadata copies complete. The transaction is committed only when every actual -plan exactly matches the independent scheduler-derived plan and the tracked -state is unchanged. Other plans and suspected violations use the complete audit. -This optimization preserves the audit and failure boundary, including its -device-completion wait. +Eager and FULL use the same checker. The runner selects one nonempty layer +metadata record for each cache group; it does not record or compare individual +layers' accesses. Only write-target block state is copied for the audit; +read-only block state is shared under the state lock. Validity is committed +only after checking the independent scheduler-derived plan and confirming the +write-target allocations are unchanged. Async snapshots still freeze every +step's selected metadata and wait for its copy before returning results. Live or cache-preserving Runner V2 recapture is rejected, including pending receives. Normal FULL replay remains audited. Allocator and kernel block @@ -134,10 +134,12 @@ this name and the legacy `MooncakeConnector` name pass the guard's connector che The shared load adapter is not an expansion of this connector support list; other backends still require end-to-end validation before enabling them. -The guard validates metadata and observed arguments. It does not inspect kernel -memory instructions, verify KV numerical values, or detect a kernel writing to -another address despite correct arguments. It cannot prevent or undo submitted -writes, prove transport byte integrity, or guarantee model output correctness. +The guard validates the declared metadata plan. It does not instrument attention +operators, detect per-layer argument differences or actual operator ordering, +inspect kernel memory instructions, or verify KV numerical values. Correct +metadata cannot establish that an operator used those addresses. The guard +cannot prevent or undo submitted writes, prove transport byte integrity, or +guarantee model output correctness. ## Read the logs @@ -154,20 +156,20 @@ Runner V2 may submit async loads during `sample_tokens` post-forward cleanup. Load submission and completion are checked at the connector cleanup boundary, after the separately audited `execute_model` forward. -Fields identify the layer, request, pool, cache group, block, allocation -generation, and relevant slot or range when available. Repeated failures are +Plan violations use `evidence=execution_plan` and identify the request, pool, +cache group, block, allocation generation, and relevant slot or range when +available. Connector readiness gaps may identify a missing layer wait; +attention accesses are not recorded per layer. Repeated failures are deduplicated with bounded state. Routine validity transitions are not logged. ## Validation and performance The focused CPU suite exercises allocation generations, ownership, shared -prefixes, exact destinations, read/write order, graph metadata freshness, +prefixes, exact destinations, declared read/write ranges, graph metadata freshness, failure latching, and the V1/V2 adapters: ```bash -pytest -q tests/ut/debug/test_kv_guard.py \ - tests/ut/debug/test_kv_guard_runtime.py tests/ut/debug/test_kv_guard_v2.py \ - tests/ut/debug/test_kv_guard_loads.py +pytest -q --confcutdir=tests/ut/debug tests/ut/debug ``` The single-card service regression starts separate disabled and enabled servers, @@ -179,7 +181,7 @@ pytest -q tests/e2e/pull_request/one_card/test_kv_cache_guard.py ``` CPU mocks establish checker behavior, not NPU execution correctness. Before -merge, record current-revision hardware runs for V1/V2 eager and FULL, prefix +merge, prioritize current-revision hardware runs for V2 eager and FULL, prefix reuse, observed preemption, Mooncake where used, and controlled fault injection. The service regression does not itself prove that a preemption or prefix-cache hit occurred; collect those metrics separately. Historical V1 results do not @@ -193,9 +195,9 @@ commands, and logs in the validation report. Do not infer device synchronization cost from CPU checker timings. The scheduler sends changed block/generation tables and compact request ranges. -Expected token slots are temporary and shared across layers in a cache group; +Expected token slots are temporary and derived once per cache group; the guard stores no persistent token history, KV snapshots, or content hashes. -Metadata transfers and access-pattern checks are deduplicated, but enabled mode +Metadata transfers are deduplicated, but enabled mode can still synchronize device metadata to CPU. Measure deployment overhead. ## Implementation guide and patch lifecycle @@ -206,7 +208,7 @@ can still synchronize device metadata to CPU. Measure deployment overhead. | `vllm_ascend/debug/kv_guard_config.py` | Validated per-engine options shared by the scheduler and workers. | | `vllm_ascend/debug/kv_guard_runtime.py` | Shared V1/V2 installation, scheduler binding, evidence collection, audit, and conservative state commit. Includes the V2 batch and capture adapter. | | `vllm_ascend/patch/platform/patch_kv_guard.py` | Conditional scheduler allocation and request-binding instrumentation. | -| Runner and attention implementations | Existing execution boundaries and ordered eager access arguments. | +| Runner implementations | Existing execution boundaries and one final metadata record per cache group. Attention implementations are unchanged. | | Mooncake receiver | Success/failure validity callbacks after receive completion and reformatting. | Both runners use one `GuardRuntime` and one address-checking state machine. diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml index e13c920457b4..5502170b23fd 100644 --- a/release_notes/fragments/17833.yaml +++ b/release_notes/fragments/17833.yaml @@ -7,6 +7,8 @@ entries: by default and configured per engine through additional_config.kv_guard_config (enabled and action). Both runners share one checker, which can log failures or stop returning worker results after a violation or incomplete audit. + The runner collects one final metadata plan per KV cache group; attention + operators are unchanged and per-layer accesses are not recorded. - category: bugfix text: >- KV guard no longer treats a consumer's computed-token count as proof of @@ -25,4 +27,4 @@ entries: text: >- Ordinary async generation with one model worker defers KV guard validation to output consumption or the next forward boundary, allowing address metadata copies to overlap - model execution while retaining every step's access checks. + model execution while retaining every step's metadata-plan checks. diff --git a/tests/ut/debug/test_kv_guard_async.py b/tests/ut/debug/test_kv_guard_async.py index 0c3a0c8af67f..c6727bad3e18 100644 --- a/tests/ut/debug/test_kv_guard_async.py +++ b/tests/ut/debug/test_kv_guard_async.py @@ -17,19 +17,16 @@ def deferred_forward(runtime, mode="FULL", slot=147): bind(runtime, computed=19) runtime._defer_audit = True runtime._scheduled = {"r": (19, 1)} - slots, table, lengths = torch.tensor([slot]), torch.tensor([[1]]), torch.tensor([20]) + slots, table, lengths = (torch.tensor([slot]), torch.tensor([[1]]), torch.tensor([20])) item = metadata(slots, table, lengths) runner = NS(input_batch=NS(req_ids=["r"])) context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) def forward(count): - if mode != "FULL": - for layer in runtime.layer_groups: - runtime.write(layer, slots) - runtime.read(layer, table, lengths) + pass runtime.wrap_forward(runner, forward, lambda: context)(1) - return slots, table, lengths + return (slots, table, lengths) @pytest.mark.parametrize("mode", ["FULL", "NONE"]) @@ -173,8 +170,10 @@ def complete(): assert calls == [True] and pending._complete is None -def test_snapshot_keeps_missing_layer_coverage(runtime): - bind(runtime) +def test_snapshot_keeps_missing_group_coverage(runtime): + runtime.state.configure({0: 128, 1: 128}) + runtime.layer_groups = {"layer0": 0, "layer1": 1} + runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1)], [(2, 1)]]}}) runtime._defer_audit = True runtime._scheduled = {"r": (0, 1)} item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) @@ -182,8 +181,12 @@ def test_snapshot_keeps_missing_layer_coverage(runtime): context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) runtime._finish_pending() - assert runtime.state.blocks[0, 1].valid_end == 0 - assert any(event.get("reason") == "metadata_write_mapping_missing:layer1" for event in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == 1 + assert runtime.state.blocks[1, 2].valid_end == 0 + assert any( + event.get("reason") == "metadata_write_mapping_missing" and event.get("group_id") == 1 + for event in runtime.events + ) @pytest.mark.parametrize("mode", ["FULL", "NONE"]) @@ -200,12 +203,9 @@ def test_frozen_read_evidence_cannot_be_repaired_by_the_next_preparation(runtime context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) def forward(count): - if mode != "FULL": - for layer in runtime.layer_groups: - runtime.write(layer, slots) - runtime.read(layer, table, lengths) + pass runtime.wrap_forward(runner, forward, lambda: context)(1) - table[0, 0], lengths[0] = 1, 20 + table[0, 0], lengths[0] = (1, 20) runtime._finish_pending() assert any(event.get("reason") == reason for event in runtime.events) diff --git a/tests/ut/debug/test_kv_guard_loads.py b/tests/ut/debug/test_kv_guard_loads.py index a30260104775..e0a3214b98b1 100644 --- a/tests/ut/debug/test_kv_guard_loads.py +++ b/tests/ut/debug/test_kv_guard_loads.py @@ -120,8 +120,6 @@ def model(count): for index, layer in enumerate(runtime.layer_groups): assert connector.wait_for_layer_load(layer) == "ready" assert runtime.state.blocks[0, 1].valid_end == (0 if index == 0 else 64) - runtime.write(layer, item.slot_mapping) - runtime.read(layer, item.block_tables, [65]) forward = runtime.wrap_forward(runner, model, lambda: context) @@ -131,30 +129,27 @@ def execute(output): poll(connector) runtime.wrap_execute(execute)(packet(meta, {"r": (64, 1)})) - assert not runtime.events and not runtime.state.loads + assert not runtime.events and (not runtime.state.loads) assert runtime.state.blocks[0, 1].valid_end == 65 -@pytest.mark.parametrize("wait_late", [False, True]) -def test_missing_or_late_layer_wait_cannot_be_hidden_by_another_layer(runtime, connector, wait_late): +@pytest.mark.parametrize("wait_both", [False, True]) +def test_declared_group_load_requires_all_layer_wait_results(runtime, connector, wait_both): bind(runtime) meta = load_plan(runtime.state, end=64, is_async=False) runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="NONE") + item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) def model(count): - runtime.write("layer0", torch.tensor([192])) - runtime.read("layer0", torch.tensor([[1]]), [65]) - if wait_late: + if wait_both: connector.wait_for_layer_load("layer0") connector.wait_for_layer_load("layer1") - runtime.write("layer1", torch.tensor([192])) - runtime.read("layer1", torch.tensor([[1]]), [65]) forward = runtime.wrap_forward(runner, model, lambda: context) runtime.wrap_execute(lambda output: (connector.start_load_kv(None), forward(1)))(packet(meta, {"r": (64, 1)})) - assert "load_layer_not_ready" in reasons(runtime) - assert runtime.state.blocks[0, 1].valid_end == (64 if wait_late else 0) + assert ("load_layer_not_ready" in reasons(runtime)) == (not wait_both) + assert runtime.state.blocks[0, 1].valid_end == (65 if wait_both else 0) assert not any(e["event"] == "guard.violation" for e in runtime.events) @@ -221,7 +216,7 @@ def test_partial_block_failure_invalidates_the_whole_destination_block(runtime, assert runtime.state.blocks[0, 1].valid_end == 0 -def test_start_exception_preserves_original_error_and_local_kv(runtime): +def test_start_exception_preserves_original_error_and_local_kv(runtime, monkeypatch): runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) complete_load(runtime.state, end=128) source = Connector() @@ -230,7 +225,7 @@ def test_start_exception_preserves_original_error_and_local_kv(runtime): def fail(*args): raise error - source.start_load_kv = fail + monkeypatch.setattr(source, "start_load_kv", fail) runtime._connector = runtime.subject.ConnectorGuard(runtime, source) meta = load_plan(runtime.state, end=192, start=128) with pytest.raises(RuntimeError) as raised: @@ -316,32 +311,16 @@ def test_only_explicit_replay_allows_overwriting_a_loaded_tail(runtime, replay, runner = NS(input_batch=NS(req_ids=["r"])) def model(count): - for layer in runtime.layer_groups: - runtime.write(layer, item.slot_mapping) - runtime.read(layer, item.block_tables, [128]) + pass forward = runtime.wrap_forward(runner, model, lambda: context) runtime.wrap_execute(lambda output: forward(1))(packet(scheduled={"r": (127, 1)}, replay=meta if replay else None)) if replay: - assert runtime.state.blocks[0, 1].valid_end == 128 and not runtime.events + assert runtime.state.blocks[0, 1].valid_end == 128 and (not runtime.events) else: assert "overwrite_valid_prefix" in reasons(runtime) -def test_replay_does_not_legalize_a_read_before_the_recomputed_write(runtime): - bind(runtime, computed=128) - replay = load_plan(runtime.state, end=128, start=127) - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(count): - runtime.read("layer0", torch.tensor([[1]]), [128]) - runtime.write("layer0", torch.tensor([255])) - - forward = runtime.wrap_forward(runner, model, lambda: NS(cudagraph_runtime_mode="NONE")) - runtime.wrap_execute(lambda output: forward(1))(packet(scheduled={"r": (127, 1)}, replay=replay)) - assert "read_unwritten_range" in reasons(runtime) - - def test_post_forward_async_completion_does_not_legalize_an_earlier_read(runtime, connector): bind(runtime) meta = load_plan(runtime.state, end=64) @@ -350,11 +329,12 @@ def test_post_forward_async_completion_does_not_legalize_an_earlier_read(runtime runner = NS(input_batch=NS(req_ids=["r"])) def model(count): - runtime.read("layer0", torch.tensor([[1]]), [64]) - poll(connector) # V2 polls inside its forward scope. + poll(connector) assert runtime.state.blocks[0, 1].valid_end == 0 - forward = runtime.wrap_forward(runner, model, lambda: NS(cudagraph_runtime_mode="NONE")) + item = metadata(torch.tensor([-1]), torch.tensor([[1]]), [64]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) + forward = runtime.wrap_forward(runner, model, lambda: context) runtime.wrap_execute(lambda output: forward(1))(packet()) assert "read_unwritten_range" in reasons(runtime) assert runtime.state.blocks[0, 1].valid_end == 64 @@ -435,12 +415,10 @@ def test_v2_receive_only_completion_then_authorized_last_token_replay(adapter, c replay = {**meta, "token_start": 127} def execute(output): - item = prepare(adapter, ["r"], [255], [[1]], [128], mode)["layer0"] - runtime.write("layer0", item.slot_mapping) - runtime.read("layer0", item.block_tables, [128]) + prepare(adapter, ["r"], [255], [[1]], [128], mode) adapter.wrap_execute(execute)(packet(scheduled={"r": (127, 1)}, replay=replay)) - assert runtime.state.blocks[0, 1].valid_end == 128 and not runtime.events + assert runtime.state.blocks[0, 1].valid_end == 128 and (not runtime.events) def test_v2_async_submission_can_follow_execute_in_sample_tokens(adapter, connector): @@ -450,19 +428,16 @@ def test_v2_async_submission_can_follow_execute_in_sample_tokens(adapter, connec meta = load_plan(runtime.state, "loading", end=64) def execute(output): - item = prepare(adapter, ["running"], [128], [[1]], [1])["layer0"] - runtime.write("layer0", item.slot_mapping) - runtime.read("layer0", item.block_tables, [1]) + prepare(adapter, ["running"], [128], [[1]], [1]) return "awaiting sampling" assert adapter.wrap_execute(execute)(packet(meta, {"running": (0, 1)})) == "awaiting sampling" - assert runtime._connector.staged == [meta] and not runtime.events - # Actual MRV2 post_forward runs from sample_tokens, outside execute_model. + assert runtime._connector.staged == [meta] and (not runtime.events) connector.start_load_kv(None) connector.result.finished_recving = {"loading"} poll(connector) assert runtime.state.blocks[0, 2].valid_end == 64 - assert not runtime.state.loads and not runtime.events + assert not runtime.state.loads and (not runtime.events) def test_v2_unsubmitted_load_is_rejected_at_the_real_post_forward_boundary(adapter, connector): diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py index 375a6b7e0460..615f22c6ee3a 100644 --- a/tests/ut/debug/test_kv_guard_runtime.py +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Runner plans, ordered operator arguments, graph replay and allocation transport.""" +"""Runner metadata per cache group, independent scheduler ranges and lifecycle.""" import importlib.util import logging @@ -72,66 +72,30 @@ def bind(runtime, request="r", block=1, computed=0): complete_load(runtime.state, request, computed) -def test_access_hooks_validate_case4_without_kv_copy_or_hash(runtime): - bind(runtime, computed=19) - slots, table = torch.tensor([142]), torch.tensor([[1]]) - runner, context = NS(input_batch=NS(req_ids=[])), NS(cudagraph_runtime_mode="NONE") - - def dummy(n): - runtime.write("layer0", slots) - runtime.read("layer0", table, [1]) - return "unchanged-model-result" - - model = runtime.wrap_forward(runner, dummy, lambda: context) - assert model(1, _kv_guard_phase="dummy") == "unchanged-model-result" - assert any(e.get("reason") == "dummy_write_to_live_block" and e["offsets"] == [14] for e in runtime.events) - runner.input_batch.req_ids = ["r"] - runtime.wrap_forward(runner, lambda n: runtime.read("layer0", table, [19]), lambda: context)(1) - assert any(e.get("reason") == "read_invalid_block" for e in runtime.events) - - -def test_full_graph_without_current_metadata_reports_a_gap(runtime): - bind(runtime, computed=19) - runner, context = NS(input_batch=NS(req_ids=[])), NS(cudagraph_runtime_mode="FULL") - - def replay(n): - runtime.write("layer0", torch.tensor([142])) - runtime.read("layer0", torch.tensor([[1]]), [19]) - - runtime.wrap_forward(runner, replay, lambda: context)(1, _kv_guard_phase="dummy") - assert any(e.get("reason") == "full_graph_metadata_missing" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 19 - - -def test_layer_hook_duplicates_do_not_look_like_repeated_writes(runtime): - bind(runtime) - slots, table = torch.arange(128, 147), torch.tensor([[1]]) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - for layer in ("layer0", "layer1"): - runtime.write(layer, slots[:19]) - runtime.read(layer, table, [19]) - - runtime.wrap_forward(runner, model, lambda: context)(19) - assert not any(e["event"] == "guard.violation" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 19 +def metadata(slots, table=None, lengths=None): + return NS(slot_mapping=slots, block_tables=table, seq_lens=lengths, num_actual_tokens=len(slots)) -def test_split_decode_prefill_write_ranges_and_read_request_offsets(runtime): - bind(runtime, "decode", 1, 19) - bind(runtime, "prefill", 2, 10) - slots = torch.tensor([147, 266, 267]) - runner, context = NS(input_batch=NS(req_ids=["decode", "prefill"])), NS(cudagraph_runtime_mode="NONE") +def test_group_audit_shares_read_only_blocks_and_rejects_a_changed_write_allocation(runtime, monkeypatch): + runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) + complete_load(runtime.state, end=128) + runtime._scheduled = {"r": (128, 1)} + item = metadata(torch.tensor([256]), torch.tensor([[1, 2]]), torch.tensor([129])) + runner = NS(input_batch=NS(req_ids=["r"])) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) + guard_type = type(runtime.state) + observe = guard_type.observe - def model(n): - runtime.write("layer0", slots[:1]) - runtime.write("layer0", slots[1:]) - runtime.read("layer0", torch.tensor([[1]]), [20]) - runtime.read("layer0", torch.tensor([[2]]), [12], request_offset=1) + def observe_with_allocation_change(shadow, *args): + assert shadow.blocks[0, 1] is runtime.state.blocks[0, 1] + assert shadow.blocks[0, 2] is not runtime.state.blocks[0, 2] + observe(shadow, *args) + runtime.state.blocks[0, 2].epoch = 2 - runtime.wrap_forward(runner, model, lambda: context)(3) - assert not any(e["event"] == "guard.violation" for e in runtime.events) + monkeypatch.setattr(guard_type, "observe", observe_with_allocation_change) + runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) + assert runtime.state.blocks[0, 2].valid_end == 0 + assert any(event.get("reason") == "allocation_changed_during_access_audit" for event in runtime.events) def test_receive_only_step_binds_before_execute_and_cleans_finished(runtime): @@ -181,44 +145,6 @@ def test_guard_rejects_invalid_configuration(runtime, options): assert not hasattr(config, "_ascend_kv_guard") -def test_v1_installs_the_shared_runtime_at_existing_boundaries(runtime, monkeypatch): - context = NS(cudagraph_runtime_mode="NONE") - forward_context = ModuleType("vllm.forward_context") - cast(Any, forward_context).get_forward_context = lambda: context - monkeypatch.setitem(sys.modules, "vllm.forward_context", forward_context) - calls = [] - result = object() - runner = NS(vllm_config=runtime.config, input_batch=NS(req_ids=["r"]), requests={"r": NS(block_ids=[[1]])}) - - def forward(num_tokens, *, marker): - assert marker is result - assert runtime._phase == "forward" - runtime.write("layer0", torch.tensor([128])) - return marker - - def update(output): - calls.append("update") - - def execute(output, *, marker): - assert runtime._scheduled == {"r": (0, 1)} - runner._update_states(output) - return runner._model_forward(1, marker=marker) - - runner._model_forward = forward - runner._update_states = update - runner.execute_model = execute - assert runtime.subject.install_guard(runner) is runtime - packet = NS( - _ascend_kv_guard_context={"pool_id": "pool", "requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (0, 1)}} - ) - assert runner.execute_model(packet, marker=result) is result - assert context._kv_guard is runtime - assert calls == ["update"] - assert runtime.state.blocks[0, 1].valid_end == 1 - assert runtime._phase is None and runtime._scheduled is None - assert not runtime.events - - @pytest.mark.parametrize("kernel_size, accepted", [(128, True), (16, False)]) def test_unadapted_kernel_block_coordinates_are_rejected(runtime, kernel_size, accepted): impl = type("AscendAttentionBackendImpl", (), {})() @@ -237,7 +163,7 @@ def test_unadapted_kernel_block_coordinates_are_rejected(runtime, kernel_size, a ) if accepted: runtime.configure(runner) - assert impl._kv_guard is runtime + assert not hasattr(impl, "_kv_guard") else: with pytest.raises(ValueError, match="unadapted"): runtime.configure(runner) @@ -348,7 +274,6 @@ def operation(**kwargs): def replay(count): bound.apply("stream") - assert not runtime._accesses return "result" assert runtime.wrap_forward(runner, replay, lambda: context)(2) == "result" @@ -387,49 +312,6 @@ def test_repeated_violations_are_deduplicated_without_hiding_other_blocks(runtim assert [level for level, args in messages].count("error") == 2 -def test_eager_raw_key_read_does_not_require_paged_cache_read_metadata(runtime): - bind(runtime) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.write("layer0", torch.tensor([128])) - runtime.read("layer0", None, [1]) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not any(e["event"] == "guard.coverage_gap" for e in runtime.events) - - -def test_same_layer_double_write_is_not_deduplicated(runtime): - bind(runtime) - slots = torch.tensor([128]) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.write("layer0", slots) - runtime.write("layer0", slots) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert any(e.get("reason") == "overwrite_valid_prefix" for e in runtime.events) - - -def test_equivalent_cloned_slots_across_layers_remain_one_logical_write(runtime): - bind(runtime) - slots = torch.tensor([128]) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.write("layer0", slots) - runtime.write("layer1", slots.clone()) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert not any(e["event"] == "guard.violation" for e in runtime.events) - - -def metadata(slots, table=None, lengths=None): - return NS(slot_mapping=slots, block_tables=table, seq_lens=lengths, num_actual_tokens=len(slots)) - - def test_runner_plan_detects_case4_without_any_operator_hook(runtime): bind(runtime, computed=19) item = metadata(torch.tensor([142]), torch.tensor([[1]]), torch.tensor([19])) @@ -468,95 +350,7 @@ def replay(n): assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) -@pytest.mark.parametrize( - "slots, table, lengths, use_fast", - [ - ([128, 129], [[1]], [2], True), - ([129, 128], [[1]], [2], False), - ([128, 128], [[1]], [2], False), - ([128, 129], [[2]], [2], False), - ([128, 129], [[1]], [1], False), - ([128, 129], [[1]], [3], False), - ], -) -def test_full_fast_and_complete_audits_agree(runtime, slots, table, lengths, use_fast): - outcomes = [] - for enabled in (False, True): - guard = runtime.subject.GuardRuntime(NS(additional_config={"kv_guard_config": {"enabled": True}})) - events = [] - guard.state.emit = lambda event, events=events, **fields: events.append({"event": event, **fields}) - guard.gap = lambda reason, events=events, **fields: events.append( - {"event": "guard.coverage_gap", "reason": reason, **fields} - ) - guard.state.configure({0: 128}) - guard.layer_groups = {"layer0": 0, "layer1": 0} - guard.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1)]]}}) - guard._scheduled = {"r": (0, 2)} - item = metadata(torch.tensor(slots), torch.tensor(table), torch.tensor(lengths)) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(guard.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=["r"])) - hits = [] - original = guard._try_fast_audit - - def attempt(runner, original=original, enabled=enabled, hits=hits): - used = original(runner) if enabled else False - hits.append(used) - return used - - guard._try_fast_audit = attempt - guard.wrap_forward(runner, lambda count: None, lambda context=context: context)(2) - assert hits == [enabled and use_fast] - outcomes.append((guard.state.blocks, events)) - assert outcomes[0] == outcomes[1] - - -def test_fast_candidate_waits_for_final_values_and_does_not_commit_during_preview(runtime, monkeypatch): - bind(runtime) - runtime._scheduled = {"r": (0, 1)} - item = metadata(torch.tensor([-1]), torch.tensor([[1]]), torch.tensor([0])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=["r"])) - preview = runtime._preview_append - seen = [] - - def staged(runner, plan): - assert item.slot_mapping.tolist() == [128] - result = preview(runner, plan) - assert runtime.state.blocks[0, 1].valid_end == 0 - seen.append(result) - return result - - monkeypatch.setattr(runtime, "_preview_append", staged) - monkeypatch.setattr(runtime._scratch, "observe", lambda *args: pytest.fail("normal plan unexpectedly fell back")) - - def forward(count): - assert not seen - item.slot_mapping[0], item.seq_lens[0] = 128, 1 - - runtime.wrap_forward(runner, forward, lambda: context)(1) - assert len(seen) == 1 and runtime.state.blocks[0, 1].valid_end == 1 - - -def test_fast_candidate_is_discarded_when_state_changes(runtime, monkeypatch): - bind(runtime) - runtime._scheduled = {"r": (0, 1)} - item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=["r"])) - preview = runtime._preview_append - - def changed(runner, plans): - candidate = preview(runner, plans) - runtime.state.blocks[0, 1].epoch = 2 - return candidate - - monkeypatch.setattr(runtime, "_preview_append", changed) - runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 0 - assert any(event.get("reason") == "allocation_changed_during_access_audit" for event in runtime.events) - - -def test_fast_preview_error_retains_complete_overwrite_diagnostic(runtime): +def test_metadata_retains_complete_overwrite_diagnostic(runtime): bind(runtime, computed=1) runtime._scheduled = {"r": (0, 1)} item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) @@ -564,27 +358,7 @@ def test_fast_preview_error_retains_complete_overwrite_diagnostic(runtime): runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) event = next(event for event in runtime.events if event.get("reason") == "overwrite_valid_prefix") assert event["offsets"] == [0] and event["valid_end"] == 1 - assert event["evidence"] == "execution_plan" and event["layer"] == "layer0" - - -def test_fast_missing_layer_cannot_advance_validity(runtime): - bind(runtime) - runtime._scheduled = {"r": (0, 1)} - item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) - runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 0 - assert any(event["event"] == "guard.coverage_gap" for event in runtime.events) - - -def test_current_full_graph_metadata_ignores_operator_callbacks(runtime): - bind(runtime, computed=19) - item = metadata(torch.tensor([147]), torch.tensor([[1]]), [20]) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor([142])), lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 20 - assert not any(e["event"] == "guard.violation" for e in runtime.events) + assert event["evidence"] == "execution_plan" and event["group_id"] == 0 def test_unadapted_ring_cache_is_rejected_instead_of_silently_reducing_checks(runtime): @@ -603,71 +377,6 @@ def test_unadapted_ring_cache_is_rejected_instead_of_silently_reducing_checks(ru runtime.configure(runner) -def test_missing_layer_metadata_does_not_advance_group_validity(runtime): - bind(runtime) - item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) - assert any(e.get("reason") == "metadata_write_mapping_missing:layer1" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_operator_read_before_write_cannot_be_legalized_by_later_write(runtime): - bind(runtime) - table = torch.tensor([[1]]) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.read("layer0", table, [1]) - runtime.write("layer0", torch.tensor([128])) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert any(e.get("reason") == "read_unwritten_range" for e in runtime.events) - - -def test_later_illegal_write_does_not_retroactively_invalidate_earlier_read(runtime): - bind(runtime, computed=19) - runner, context = NS(input_batch=NS(req_ids=["r"])), NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.read("layer0", torch.tensor([[1]]), [19]) - runtime.write("layer0", torch.tensor([142])) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert any(e.get("reason") == "overwrite_valid_prefix" for e in runtime.events) - assert not any(e.get("reason") == "read_invalid_block" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX - - -def test_one_layers_write_cannot_hide_another_layers_missing_write(runtime): - bind(runtime) - slots, table = torch.tensor([128]), torch.tensor([[1]]) - item = metadata(slots, table, [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(n): - runtime.write("layer0", slots) - runtime.read("layer0", table, [1]) - runtime.read("layer1", table, [1]) - - runtime.wrap_forward(runner, model, lambda: context)(1) - violation = next(e for e in runtime.events if e.get("reason") == "read_unwritten_range") - assert violation["layer"] == "layer1" and violation["evidence"] == "operator_arguments" - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_missing_known_operator_layer_is_reported_and_not_committed(runtime): - bind(runtime) - item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - runtime.wrap_forward(runner, lambda n: runtime.write("layer0", item.slot_mapping), lambda: context)(1) - assert any(e.get("reason") == "operator_layer_not_observed:layer1" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 0 - - def test_failed_forward_preserves_exception_and_does_not_commit_planned_writes(runtime): bind(runtime) item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) @@ -726,87 +435,6 @@ class OtherAttentionImpl: assert not hasattr(impl, "_kv_guard") -def test_eager_uses_operator_arguments_without_building_unused_plans(runtime, monkeypatch): - bind(runtime) - slots, table = torch.tensor([128]), torch.tensor([[1]]) - item = metadata(slots, table, [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - - def unused_plan(*args): - raise AssertionError("eager must not materialize an unused metadata plan") - - monkeypatch.setattr(runtime, "_metadata_accesses", unused_plan) - - def model(n): - for layer in ("layer0", "layer1"): - runtime.write(layer, slots) - runtime.read(layer, table, [1]) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - -def test_full_metadata_skips_redundant_hooks_and_builds_one_shared_plan(runtime, monkeypatch): - bind(runtime) - slots, table = torch.tensor([128]), torch.tensor([[1]]) - item = metadata(slots, table, [1]) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - parsed = [] - original = runtime._metadata_accesses - - def parse(layer, value): - parsed.append(layer) - return original(layer, value) - - monkeypatch.setattr(runtime, "_metadata_accesses", parse) - - def replay(n): - for layer in ("layer0", "layer1"): - runtime.write(layer, slots) - runtime.read(layer, table, [1]) - assert not runtime._accesses - - runtime.wrap_forward(runner, replay, lambda: context)(1) - assert len(parsed) == 1 - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - -def test_unique_layer_patterns_share_owner_snapshot_and_reuse_checker(runtime, monkeypatch): - bind(runtime) - slots, table = torch.tensor([128]), torch.tensor([[1]]) - context = NS(cudagraph_runtime_mode="NONE") - runner = NS(input_batch=NS(req_ids=["r"])) - owner_copies = [] - original = runtime.subject.replace - scratch = runtime._scratch - - def copied(block, **fields): - if "owners" in fields: - owner_copies.append(block) - return original(block, **fields) - - monkeypatch.setattr(runtime.subject, "replace", copied) - - def unexpected_checker(*args): - raise AssertionError("the scratch checker must be reused") - - monkeypatch.setattr(runtime.subject, "KVGuard", unexpected_checker) - - def model(n): - for layer in ("layer0", "layer1"): - runtime.write(layer, slots.clone()) - runtime.read(layer, table, [1]) - - runtime.wrap_forward(runner, model, lambda: context)(1) - assert runtime._scratch is scratch and len(owner_copies) == 1 - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - def test_execute_and_update_bind_once_without_skipping_receive_only_or_next_step(runtime, monkeypatch): output = NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": []}) runner = NS(requests={}) @@ -835,96 +463,9 @@ def test_shared_unsupported_metadata_reports_missing_read_coverage_for_each_grou context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) runner = NS(input_batch=NS(req_ids=[])) runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) - assert runtime._missing_groups == {0, 1} - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -@pytest.mark.parametrize("swapped", [False, True]) -def test_scheduler_plan_detects_cross_request_slot_swaps(runtime, mode, swapped): - bind(runtime, "A", 1) - bind(runtime, "B", 2) - slots = torch.tensor([256, 128] if swapped else [128, 256]) - table = torch.tensor([[1], [2]]) - item = metadata(slots, table, [1, 1]) - context = NS(cudagraph_runtime_mode=mode, attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["A", "B"])) - output = NS(_ascend_kv_guard_context={"requests": {}, "scheduled_requests": {"A": (0, 1), "B": (0, 1)}}) - - def model(n): - for layer in runtime.layer_groups: - runtime.write(layer, slots) - runtime.read(layer, table, [1, 1]) - - forward = runtime.wrap_forward(runner, model, lambda: context) - runtime.wrap_execute(lambda message: forward(2))(output) - violations = [e for e in runtime.events if e["event"] == "guard.violation"] - if swapped: - mismatch = next(e for e in violations if e["reason"] == "write_target_mismatch") - assert mismatch["request_id"] == "A" and mismatch["slot"] == 256 and mismatch["expected_slot"] == 128 - assert all(b.valid_end == runtime.subject.INVALID_PREFIX for b in runtime.state.blocks.values()) - else: - assert not violations - assert runtime._scheduled is None - - -def test_split_mixed_batch_uses_explicit_token_offset_and_runner_request_order(runtime): - bind(runtime, "prefill", 2, 10) - bind(runtime, "decode", 1, 19) - runtime._scheduled = {"prefill": (10, 2), "decode": (19, 1)} - runner = NS(input_batch=NS(req_ids=["decode", "prefill"])) - context = NS(cudagraph_runtime_mode="NONE") - - def model(n): - runtime.write("layer0", torch.tensor([266, 267]), token_offset=1) - runtime.write("layer0", torch.tensor([147])) - runtime.read("layer0", torch.tensor([[1], [2]]), [20, 12]) - - runtime.wrap_forward(runner, model, lambda: context)(3) - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - -@pytest.mark.parametrize( - "slots, mismatch", [([128, 129, -1], False), ([129, 128, -1], True), ([128, -1, -1], True), ([128, 129, 256], True)] -) -def test_exact_token_positions_distinguish_permutation_missing_write_and_padding(runtime, slots, mismatch): - bind(runtime) - runtime._scheduled = {"r": (0, 2)} - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="NONE") - runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor(slots)), lambda: context)(3) - assert any(e.get("reason") == "write_target_mismatch" for e in runtime.events) == mismatch - - -def test_missing_scheduler_range_does_not_advance_validity(runtime): - bind(runtime) - runtime._scheduled = {"r": (None, 1)} - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="NONE") - runtime.wrap_forward(runner, lambda n: runtime.write("layer0", torch.tensor([128])), lambda: context)(1) - assert any(e.get("reason") == "scheduler_write_range_missing" and e["request_id"] == "r" for e in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_raise_mode_stops_output_and_remains_latched_after_caller_catches(runtime): - bind(runtime, computed=19) - runtime.action = "raise" - runtime.state.emit = runtime.report - runner = NS(input_batch=NS(req_ids=[])) - context = NS(cudagraph_runtime_mode="NONE") - completed = [] - - def model(n): - runtime.write("layer0", torch.tensor([142])) - completed.append("kernel submitted") - return "must not escape" - - forward = runtime.wrap_forward(runner, model, lambda: context) - with pytest.raises(runtime.subject.KVGuardError, match="dummy_write_to_live_block"): - forward(1, _kv_guard_phase="dummy") - assert completed == ["kernel submitted"] # Auditing cannot undo submitted operations. - with pytest.raises(runtime.subject.KVGuardError): - forward(1, _kv_guard_phase="dummy") - assert completed == ["kernel submitted"] and runtime._phase is None + assert { + event.get("group_id") for event in runtime.events if event.get("reason") == "metadata_read_range_missing" + } == {0, 1} def test_raise_mode_stops_on_coverage_gap_before_execution(runtime): @@ -962,65 +503,158 @@ def execute(message): assert any(e.get("reason") == "clear_live_block" for e in runtime.events) -def test_read_block_order_within_one_request_is_checked_against_scheduler(runtime): - runtime.state.bind( - {"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}, "scheduled_requests": {"r": (256, 1)}} +def test_copy_to_null_or_unbound_block_cannot_silently_pass(runtime): + bind(runtime, computed=19) + runtime.copy_blocks([NS(src_block_id=1, dst_block_id=0), NS(src_block_id=1, dst_block_id=99)]) + assert any(e.get("reason") == "write_null_block" and e["operation"] == "copy" for e in runtime.events) + assert any(e.get("reason") == "copy_destination_unbound" and e["block_id"] == 99 for e in runtime.events) + + +def test_scheduler_optional_zero_and_copy_fields_accept_none(runtime): + output = NS( + _ascend_kv_guard_context={"requests": {}, "scheduled_requests": {}}, + new_block_ids_to_zero=None, + kv_cache_block_copies=None, ) - runtime._scheduled = {"r": (256, 1)} - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="NONE") - runtime.wrap_forward(runner, lambda n: runtime.read("layer0", torch.tensor([[2, 1]]), [256]), lambda: context)(1) - assert any(e.get("reason") == "read_target_mismatch" for e in runtime.events) + assert runtime.wrap_execute(lambda message: "executed")(output) == "executed" + assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) -def test_shared_prefix_read_and_tail_append_remain_legal(runtime): - runtime.state.bind( - { - "pool_id": "pool", - "requests": {"A": [[(1, 1), (2, 1)]], "B": [[(1, 1), (3, 1)]]}, - "scheduled_requests": {"A": (128, 1), "B": (128, 1)}, - } +def run_plan(runtime, slots, table, lengths, *, ids=("r",), mode="FULL", phase="forward", forward=None): + item = metadata(torch.tensor(slots), torch.tensor(table), torch.tensor(lengths)) + context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) + runner = NS(input_batch=NS(req_ids=list(ids))) + return runtime.wrap_forward(runner, forward or (lambda count: None), lambda: context)( + len(slots), _kv_guard_phase=phase ) - complete_load(runtime.state, "A", 128) - runtime._scheduled = {"A": (128, 1), "B": (128, 1)} - runner = NS(input_batch=NS(req_ids=["B", "A"])) - context = NS(cudagraph_runtime_mode="NONE") - def model(n): - runtime.write("layer0", torch.tensor([384, 256])) - runtime.read("layer0", torch.tensor([[1, 3], [1, 2]]), [129, 129]) - runtime.wrap_forward(runner, model, lambda: context)(2) - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +def test_group_metadata_is_checked_once_without_attention_hooks(runtime, monkeypatch, mode): + bind(runtime) + runtime._scheduled = {"r": (0, 1)} + calls = [] + original = runtime.subject.KVGuard.observe + + def checked(state, groups, requests, phase): + calls.append((len(groups), list(requests))) + return original(state, groups, requests, phase) + + monkeypatch.setattr(runtime.subject.KVGuard, "observe", checked) + run_plan(runtime, [128], [[1]], [1], mode=mode) + assert calls == [(1, ["r"])] + assert runtime.state.blocks[0, 1].valid_end == 1 + assert not runtime.events + assert set(runtime._metadata) == {0} + assert not hasattr(runtime, "write") and not hasattr(runtime, "read") + assert not hasattr(runtime, "_accesses") + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +@pytest.mark.parametrize( + "slots,table,lengths,reason", + [ + ([128, 129], [[1]], [2], None), + ([129, 128], [[1]], [2], "write_target_mismatch"), + ([128, 128], [[1]], [2], "write_target_mismatch"), + ([128, 129], [[2]], [2], "read_target_mismatch"), + ([128, 129], [[1]], [1], None), + ([128, 129], [[1]], [3], "read_length_mismatch"), + ], +) +def test_metadata_constraints_share_one_checker(runtime, mode, slots, table, lengths, reason): + bind(runtime) + runtime._scheduled = {"r": (0, 2)} + run_plan(runtime, slots, table, lengths, mode=mode) + violations = [event for event in runtime.events if event["event"] == "guard.violation"] + assert any(event.get("reason") == reason for event in violations) if reason else not violations + assert all(event["evidence"] == "execution_plan" and "layer" not in event for event in violations) + + +@pytest.mark.parametrize("mode", ["NONE", "FULL"]) +@pytest.mark.parametrize("swapped", [False, True]) +def test_group_plan_checks_actual_request_order_and_cross_request_slots(runtime, mode, swapped): + bind(runtime, "A", 1) + bind(runtime, "B", 2) + runtime._scheduled = {"A": (0, 1), "B": (0, 1)} + run_plan(runtime, [128, 256] if swapped else [256, 128], [[2], [1]], [1, 1], ids=("B", "A"), mode=mode) + assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == swapped -def test_truncated_write_observation_is_incomplete_even_if_read_only_uses_old_prefix(runtime): +@pytest.mark.parametrize("slots,mismatch", [([147, 148], False), ([148, 147], True), ([147, -1], True)]) +def test_group_metadata_distinguishes_exact_token_positions_and_padding(runtime, slots, mismatch): bind(runtime, computed=19) runtime._scheduled = {"r": (19, 2)} - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="NONE") + run_plan(runtime, slots, [[1]], [21]) + assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == mismatch - def model(n): - runtime.write("layer0", torch.tensor([147])) - runtime.read("layer0", torch.tensor([[1]]), [19]) - runtime.wrap_forward(runner, model, lambda: context)(2) - assert any(e.get("reason") == "write_tokens_not_observed" for e in runtime.events) +def test_truncated_group_write_cannot_advance_validity(runtime): + bind(runtime, computed=19) + runtime._scheduled = {"r": (19, 2)} + run_plan(runtime, [147], [[1]], [19]) + assert any(event.get("reason") == "write_tokens_not_observed" for event in runtime.events) assert runtime.state.blocks[0, 1].valid_end == 19 -def test_copy_to_null_or_unbound_block_cannot_silently_pass(runtime): +def test_computed_range_does_not_initialize_unwritten_blocks(runtime): + bind(runtime) + runtime._scheduled = {"r": (19, 1)} + run_plan(runtime, [147], [[1]], [20]) + assert any(event.get("reason") == "noncontiguous_write" for event in runtime.events) + assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX + + +def test_read_block_order_is_checked_within_one_request(runtime): + runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) + complete_load(runtime.state, "r", 254) + runtime._scheduled = {"r": (254, 1)} + run_plan(runtime, [382], [[2, 1]], [255]) + assert any(event.get("reason") == "read_target_mismatch" for event in runtime.events) + + +def test_shared_prefix_and_distinct_tail_blocks_remain_legal(runtime): + runtime.state.bind({"pool_id": "pool", "requests": {"A": [[(1, 1), (2, 1)]], "B": [[(1, 1), (3, 1)]]}}) + complete_load(runtime.state, "A", 128) + runtime._scheduled = {"A": (128, 1), "B": (128, 1)} + run_plan(runtime, [384, 256], [[1, 3], [1, 2]], [129, 129], ids=("B", "A")) + assert not runtime.events + + +def test_dummy_metadata_write_still_taints_a_live_block(runtime): bind(runtime, computed=19) - runtime.copy_blocks([NS(src_block_id=1, dst_block_id=0), NS(src_block_id=1, dst_block_id=99)]) - assert any(e.get("reason") == "write_null_block" and e["operation"] == "copy" for e in runtime.events) - assert any(e.get("reason") == "copy_destination_unbound" and e["block_id"] == 99 for e in runtime.events) + run_plan(runtime, [142], [[1]], [19], ids=(), mode="NONE", phase="dummy") + assert any(event.get("reason") == "dummy_write_to_live_block" for event in runtime.events) + run_plan(runtime, [-1], [[1]], [19], mode="NONE") + assert any(event.get("reason") == "read_invalid_block" for event in runtime.events) -def test_scheduler_optional_zero_and_copy_fields_accept_none(runtime): - output = NS( - _ascend_kv_guard_context={"requests": {}, "scheduled_requests": {}}, - new_block_ids_to_zero=None, - kv_cache_block_copies=None, +def test_only_group_coverage_is_required_not_every_layer(runtime): + bind(runtime) + item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) + context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer1": item}) + runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) + assert not runtime.events and runtime.state.blocks[0, 1].valid_end == 1 + + +def test_missing_cache_group_is_an_explicit_coverage_gap(runtime): + runtime.state.configure({0: 128, 1: 128}) + runtime.layer_groups = {"layer0": 0, "layer1": 1} + item = metadata(torch.tensor([-1]), torch.tensor([[0]]), [0]) + context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) + runtime.wrap_forward(NS(input_batch=NS(req_ids=[])), lambda count: None, lambda: context)(1) + assert any( + event.get("reason") == "metadata_write_mapping_missing" and event.get("group_id") == 1 + for event in runtime.events ) - assert runtime.wrap_execute(lambda message: "executed")(output) == "executed" - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) + + +def test_raise_mode_stops_results_and_latches_metadata_failure(runtime): + runtime.action = "raise" + runtime.state.emit = runtime.report + bind(runtime, computed=19) + runtime._scheduled = {"r": (19, 1)} + with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): + run_plan(runtime, [148], [[1]], [20]) + with pytest.raises(runtime.subject.KVGuardError): + runtime.wrap_execute(lambda output: None)(NS()) diff --git a/tests/ut/debug/test_kv_guard_v2.py b/tests/ut/debug/test_kv_guard_v2.py index 9b14219acb25..510def7d25ea 100644 --- a/tests/ut/debug/test_kv_guard_v2.py +++ b/tests/ut/debug/test_kv_guard_v2.py @@ -74,9 +74,6 @@ def test_sorted_v2_batch_uses_batch_ids_instead_of_request_state_indices(adapter def execute(output): prepare(adapter, ["B", "A"], [256, 257, 128], [[2], [1]], [2, 1], mode) - if mode == "NONE": - runtime.write("layer0", torch.tensor([256, 257, 128])) - runtime.read("layer0", torch.tensor([[2], [1]]), [2, 1]) return "unchanged" assert adapter.wrap_execute(execute)(packet) == "unchanged" @@ -99,11 +96,8 @@ def test_v2_full_replay_observes_current_metadata_after_parameter_update(adapter def execute(output): item = prepare(adapter, ["r"], [128, -1], [[1], [0]], [1, 0], "FULL")["layer0"] - # Replay/parameter updates mutate the same buffers after collection. item.slot_mapping[0] = 256 item.block_tables[0, 0] = 2 - runtime.write("layer0", torch.tensor([128])) - assert runtime._accesses == [] return "replayed" assert adapter.wrap_execute(execute)(packet) == "replayed" @@ -179,7 +173,6 @@ def execute(output, *args): def test_v2_failed_forward_does_not_commit_writes_or_leak_scope(adapter): def execute(output): prepare(adapter, ["r"], [128], [[1]], [1]) - adapter.runtime.write("layer0", torch.tensor([128])) raise RuntimeError("original failure") with pytest.raises(RuntimeError, match="original failure"): diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py index e79d35a48c68..d6341b4ee1b7 100644 --- a/vllm_ascend/attention/attention_v1.py +++ b/vllm_ascend/attention/attention_v1.py @@ -954,9 +954,6 @@ def forward_fused_infer_attention( key, value, block_size, block_table, actual_seq_lengths_kv = self._get_fia_params( key, value, attn_metadata, kv_cache ) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.read(self._layer_name, block_table, actual_seq_lengths_kv) num_tokens = attn_metadata.actual_seq_lengths_q[-1] query = query[:num_tokens] if ( @@ -1119,9 +1116,6 @@ def _forward_fia_chunked_prefill_split( seq_lens_list = attn_metadata.seq_lens_list num_tokens = int(actual_seq_qlen[-1]) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.read(self._layer_name, block_table, seq_lens_list) # decode part if num_decode_tokens > 0: decode_out, _ = DeviceOperator.npu_fused_infer_attention_score( @@ -1187,9 +1181,6 @@ def forward_paged_attention( attn_metadata: AscendMetadata, output: torch.Tensor | None = None, ) -> torch.Tensor: - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.read(self._layer_name, attn_metadata.block_tables, attn_metadata.seq_lens) if _EXTRA_CTX.capturing: return self.full_graph_pa(query, attn_metadata, output) torch_npu._npu_paged_attention( @@ -1243,9 +1234,6 @@ def do_kv_cache_update( if self.key_cache is None: self.key_cache, self.value_cache = kv_cache[0], kv_cache[1] - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.write(layer.layer_name, slot_mapping) DeviceOperator.reshape_and_cache( key=key, value=value, @@ -1280,9 +1268,6 @@ def reshape_and_cache( # quick fix to make sure slots is int32 for cross attention case. # see: https://github.com/vllm-project/vllm/blob/ce88756b967c2c5006746a424c15dd59a284ed8c/vllm/model_executor/layers/attention/cross_attention.py#L117 slots_to_cache = slots[: attn_metadata.num_actual_tokens] if not encoder_decoder else slots.to(torch.int32) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.write(self._layer_name, slots_to_cache) DeviceOperator.reshape_and_cache( key=key_to_cache, value=value_to_cache, diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py index 01bf85edec2c..ac024d1f1af9 100644 --- a/vllm_ascend/attention/mla_v1.py +++ b/vllm_ascend/attention/mla_v1.py @@ -1609,15 +1609,6 @@ def _compute_prefill_context( pe_cache = cache_k_pe k_pe = torch.empty(toks, num_heads, rope_dim, dtype=q_pe.dtype, device=q_pe.device) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.read( - self._kv_guard_layer_name, - prefill_metadata.block_table, - context_seq_len_npu, - prefill_metadata.chunked_context.starts[i], - request_offset=attn_metadata.num_decodes, - ) DeviceOperator.kv_cache_load( cache_kv_c, pe_cache, @@ -1839,9 +1830,6 @@ def exec_kv_decode( slots: torch.Tensor, return_current_kv: bool = False, ): - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.write(self._kv_guard_layer_name, slots) if self._uses_strided_kv_adapter() and any(not cache.is_contiguous() for cache in kv_cache): current_k_pe, current_k_nope = self._exec_strided_kv_writer(kv_no_split, cos, sin, kv_cache, slots) if return_current_kv: @@ -2671,13 +2659,6 @@ def exec_kv_prefill( *, attn_metadata: AscendMLAMetadata | None = None, ): - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.write( - self._kv_guard_layer_name, - slots, - token_offset=attn_metadata.num_decode_tokens if attn_metadata else 0, - ) pcp_local_range = None if self.use_pcp: assert attn_metadata is not None @@ -2787,9 +2768,6 @@ def _forward_decode( dequant_scale_q_nope = decode_preprocess_res.dequant_scale_q_nope decode_meta = attn_metadata.decode assert decode_meta is not None - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.read(self._kv_guard_layer_name, decode_meta.block_table, decode_meta.seq_lens_list) # TODO: The CANN package is expected to support num_heads that are not # powers of 2 in 2026 Q2. Once supported, all padding operations under # `if self.head_padding > 0` in this function can be removed. @@ -3066,9 +3044,6 @@ def mla_preprocess_only_decode(self, hidden_states, kv_cache, attn_metadata): from cann_ops_transformer import mla_prolog # type: ignore[import-not-found,import-untyped] # noqa: PLC0415 bsz = attn_metadata.num_decode_tokens - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.write(self._kv_guard_layer_name, attn_metadata.slot_mapping[:bsz]) history_slots = attn_metadata.slot_mapping[:bsz] return_current_kv = self._decode_requires_current_kv(attn_metadata) if return_current_kv: @@ -3439,7 +3414,6 @@ def forward( # Profiling run. return output.fill_(0) - self._kv_guard_layer_name = layer_name num_actual_tokens = self.get_num_actual_tokens(attn_metadata) assert ( attn_metadata.num_decodes is not None diff --git a/vllm_ascend/debug/kv_guard_connector.py b/vllm_ascend/debug/kv_guard_connector.py index ca554b9ba52b..f5b54ebd83be 100644 --- a/vllm_ascend/debug/kv_guard_connector.py +++ b/vllm_ascend/debug/kv_guard_connector.py @@ -139,7 +139,7 @@ def errors(self, original, *args, **kwargs): return invalid def _target_blocks(self, meta: dict) -> set[int]: - targets = set() + targets: set[int] = set() for group, refs in enumerate(meta["block_table"]): size = self.state.group_sizes[group] first, last = meta["token_start"] // size, (meta["token_end"] + size - 1) // size diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py index 7571938039ca..901fe3693bc8 100644 --- a/vllm_ascend/debug/kv_guard_runtime.py +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -3,8 +3,8 @@ """Worker L1 guard with one runner execution-plan collection point. No KV snapshots, hashes, general-purpose trace or offline reader. -Adapted per-layer metadata supplies declared access plans. Operator hooks -preserve argument order; scheduler ranges independently identify each token. +Runner metadata supplies one declared access plan per KV cache group. +Scheduler ranges independently identify each token; attention is not instrumented. Only small metadata is copied, after the forward and graph parameter updates; the checker neither inspects KV values nor proves kernel memory accesses. """ @@ -29,22 +29,6 @@ _MAX_REPORT_KEYS = 1024 -@dataclass(slots=True) -class WriteAccess: - layer: str - slots: Any - token_offset: int = 0 - - -@dataclass(slots=True) -class ReadAccess: - layer: str - table: Any - lengths: Any - starts: Any = None - request_offset: int = 0 - - @dataclass(slots=True) class ExpectedAccess: slots: list[int] @@ -88,11 +72,8 @@ def __init__(self, vllm_config: Any): self.config = vllm_config self.state = KVGuard(self.report) self._scratch = KVGuard(self._emit_scratch) - self._audit_layer = "" - self._audit_evidence = "" self.layer_groups: dict[str, int] = {} - self._accesses: list[WriteAccess | ReadAccess] = [] - self._metadata: dict[str, Any] = {} + self._metadata: dict[int, Any] = {} self._missing_groups: set[int] = set() self._phase: str | None = None self._full_graph = False @@ -116,7 +97,7 @@ def _check_failure(self) -> None: def _emit_scratch(self, event: str, **fields: Any) -> None: if event == "guard.violation": - self.state.emit(event, layer=self._audit_layer, evidence=self._audit_evidence, **fields) + self.state.emit(event, evidence="execution_plan", **fields) def report(self, event: str, **fields: Any) -> None: failed = event in ("guard.violation", "guard.coverage_gap") @@ -188,9 +169,7 @@ def configure(self, runner: Any) -> None: raise ValueError(f"KV L1 guard: unadapted cache validity rules: {type(layer_spec).__name__}") layer = config.compilation_config.static_forward_context.get(name) impl = getattr(layer, "impl", None) - if impl is not None and type(impl).__name__ in ("AscendAttentionBackendImpl", "AscendMLAImpl"): - impl._kv_guard = self - else: + if impl is None or type(impl).__name__ not in ("AscendAttentionBackendImpl", "AscendMLAImpl"): raise ValueError(f"KV L1 guard: unadapted attention implementation: {type(impl).__name__}") if unsupported: raise ValueError( @@ -211,18 +190,6 @@ def configure(self, runner: Any) -> None: self._runner_adapter.configure() self.report("guard.enabled", failure_action=self.action) - def write(self, layer: str, slots: Any, token_offset: int = 0) -> None: - if self._phase == "forward" and self._connector is not None and self._connector.sync: - self._connector.check_ready(layer) - if self._phase is not None and not self._full_graph: - self._accesses.append(WriteAccess(layer, slots, token_offset)) - - def read(self, layer: str, table: Any, lengths: Any, starts: Any = None, request_offset: int = 0) -> None: - if self._phase == "forward" and self._connector is not None and self._connector.sync: - self._connector.check_ready(layer) - if self._phase is not None and not self._full_graph and table is not None: - self._accesses.append(ReadAccess(layer, table, lengths, starts, request_offset)) - @staticmethod def _list(value: Any) -> list: if isinstance(value, (list, tuple)): @@ -234,40 +201,21 @@ def _field(metadata: Any, *names: str) -> Any: return next((value for name in names if (value := getattr(metadata, name, None)) is not None), None) def collect_metadata(self, metadata: Any) -> None: - """Keep references/coverage only; materialize the selected source later.""" - self._metadata, self._missing_groups = {}, set() + self._metadata = {} + self._missing_groups = set() if metadata is None: - return # Profiling and eager operator-only tests may have no attention plan. + return if not isinstance(metadata, dict): self.gap("microbatch_metadata_unadapted") - self._missing_groups.update(self.layer_groups.values()) + self._missing_groups.update(self.state.group_sizes) return - for layer, group in self.layer_groups.items(): - item = metadata.get(layer) - if item is None: - self.gap(f"metadata_write_mapping_missing:{layer}") - self._missing_groups.add(group) - continue - self._metadata[layer] = item - - def _metadata_accesses(self, layer: str, item: Any) -> list[WriteAccess | ReadAccess]: - slots = self._field(item, "slot_mapping") - if slots is None: - self.gap(f"metadata_write_mapping_missing:{layer}") - self._missing_groups.add(self.layer_groups[layer]) - return [] - count = getattr(item, "num_actual_tokens", None) - if count is not None: - slots = slots[:count] - accesses: list[WriteAccess | ReadAccess] = [WriteAccess(layer, slots)] - table = self._field(item, "block_tables", "block_table", "block_table_tensor") - lengths = self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") - if table is None or lengths is None: - self.gap(f"metadata_read_range_missing:{layer}") - self._missing_groups.add(self.layer_groups[layer]) - else: - accesses.append(ReadAccess(layer, table, lengths)) - return accesses + for name, item in metadata.items(): + group = self.layer_groups.get(name) + if group is not None and item is not None: + self._metadata.setdefault(group, item) + for group in self.state.group_sizes.keys() - self._metadata.keys(): + self.gap("metadata_write_mapping_missing", group_id=group) + self._missing_groups.add(group) @staticmethod def _identity(value: Any) -> Any: @@ -315,347 +263,88 @@ def _expected_write_plan(self, runner: Any, group_id: int) -> ExpectedAccess | N query_ends.append(len(slots)) return ExpectedAccess(slots, query_ends, tables, lengths) - def _preview_append(self, runner: Any, expected: dict[int, ExpectedAccess]) -> Any: - """Prepare a silent, unpublished canonical WRITE-then-READ transaction.""" - baseline, proposed, bindings = {}, {}, {} - with self.state._lock: - pool = self.state.pool_id - for group_id, plan in expected.items(): - size = self.state.group_sizes[group_id] - intervals: dict[int, list[tuple[int, int]]] = {} - for row, request in enumerate(runner.input_batch.req_ids): - groups = self.state.requests[request] - bindings[request] = tuple(tuple(refs) for refs in groups) - refs = groups[group_id] - for block_id, epoch in refs[: len(plan.tables[row])]: - key = group_id, block_id - block = self.state.blocks.get(key) - if ( - block_id == 0 - or block is None - or epoch != block.epoch - or block.valid_end == INVALID_PREFIX - or request not in block.owners - ): - return None - baseline[key] = block.epoch, block.valid_end, frozenset(block.owners) - computed, count = self._scheduled[request] - for logical in range(computed // size, (computed + count - 1) // size + 1): - block_id, _ = refs[logical] - start = max(computed, logical * size) - logical * size - end = min(computed + count, (logical + 1) * size) - logical * size - intervals.setdefault(block_id, []).append((start, end)) - for block_id, spans in intervals.items(): - key = group_id, block_id - cursor = baseline[key][1] - for start, end in sorted(spans): - if start != cursor: - return None - cursor = end - proposed[key] = cursor - for row, length in enumerate(plan.lengths): - for logical, block_id in enumerate(plan.tables[row]): - key = group_id, block_id - if min(size, length - logical * size) > proposed.get(key, baseline[key][1]): - return None - return pool, baseline, proposed, bindings - - def _try_fast_audit(self, runner: Any) -> bool: - """Exact final metadata equality permits committing a precomputed state. - - shortcut: only canonical FULL plans; other plans retain the ordered audit. - """ - if ( - self._phase != "forward" - or not self._full_graph - or self._scheduled is None - or self._missing_groups - or self._unready_load_groups - or not self._metadata - or self._metadata.keys() != self.layer_groups.keys() - ): - return False - request_ids = list(runner.input_batch.req_ids) - if not request_ids or len(set(request_ids)) != len(request_ids): - return False - expected: dict[int, ExpectedAccess] = {} - inputs, checks, seen = {}, [], set() - for layer, item in self._metadata.items(): - group_id = self.layer_groups[layer] - if (group_id, id(item)) in seen: - continue - seen.add((group_id, id(item))) - if group_id not in expected: - plan = self._expected_write_plan(runner, group_id) - if plan is None: - return False - expected[group_id] = plan - plan = expected[group_id] - accesses = self._metadata_accesses(layer, item) - if len(accesses) != 2 or self._missing_groups: - return False - write, read = accesses - slots, table, lengths = write.slots, read.table, read.lengths - if not all( - isinstance(value, torch.Tensor) and value.dtype in (torch.int32, torch.int64) - for value in (slots, table, lengths) - ): - return False - rows, columns = len(request_ids), max(map(len, plan.tables), default=0) - if ( - slots.ndim != 1 - or slots.numel() != len(plan.slots) - or lengths.shape != (rows,) - or table.ndim != 2 - or table.shape[0] < rows - or table.shape[1] < columns - or any(value.device.type not in ("cpu", "npu") for value in (slots, table, lengths)) - ): - return False - keys = [] - for value in (slots, lengths, table[:rows, :columns]): - key = str(value.device), value.dtype, self._identity(value) - inputs.setdefault(key, value) - keys.append(key) - checks.append((plan, keys)) - - # Enqueue every copy before the CPU-only preview. No device values are - # inspected until the event, at the original post-forward audit point. - copied = {} - event = None - recorded = False - try: - for index, (key, value) in enumerate(inputs.items()): - if value.device.type == "cpu": - copied[key] = self._list(value) - continue - if self._copy_event is None: - self._copy_event = torch.npu.Event() - event = self._copy_event - if index >= len(self._host_buffers): - self._host_buffers.extend([None] * (index + 1 - len(self._host_buffers))) - host = self._host_buffers[index] - if host is None or host.shape != value.shape or host.dtype != value.dtype: - host = torch.empty(value.shape, dtype=value.dtype, device="cpu", pin_memory=True) - self._host_buffers[index] = host - host.copy_(value, non_blocking=True) - copied[key] = host - if event is not None: - event.record(torch.npu.current_stream()) - recorded = True - candidate = self._preview_append(runner, expected) - finally: - if event is not None: - # A partial enqueue failure must also drain outstanding DMA. - if not recorded: - event.record(torch.npu.current_stream()) - event.synchronize() - if candidate is None: - return False - values = {key: value.tolist() if isinstance(value, torch.Tensor) else value for key, value in copied.items()} - for plan, (slots, lengths, table) in checks: - if ( - values[slots] != plan.slots - or values[lengths] != plan.lengths - or any(row[: len(want)] != want for row, want in zip(values[table], plan.tables)) - ): - return False - pool, baseline, proposed, bindings = candidate - with self.state._lock: - unchanged = self.state.pool_id == pool and all( - tuple(tuple(refs) for refs in self.state.requests.get(request, ())) == groups - for request, groups in bindings.items() - ) - for key, before in baseline.items(): - block = self.state.blocks.get(key) - unchanged = ( - unchanged - and block is not None - and (block.epoch, block.valid_end, frozenset(block.owners)) == before - ) - if not unchanged: - self.gap("allocation_changed_during_access_audit") - return True - for key, valid in proposed.items(): - self.state.blocks[key].valid_end = valid - return True - def _audit(self, runner: Any) -> None: - if self._try_fast_audit(runner): - return phase = self._phase assert phase is not None - # FULL requires this replay's final metadata. Never reuse a captured - # address plan when the current metadata is missing. - if self._full_graph and not self._metadata: - if phase != "warmup": - self.gap("full_graph_metadata_missing") - return - observed: dict[str, list[WriteAccess | ReadAccess]] = {} - for access in self._accesses: - observed.setdefault(access.layer, []).append(access) - if not self._metadata and not observed: + if not self._metadata: if phase != "warmup": - self.gap("execution_access_plan_missing") + self.gap("full_graph_metadata_missing" if self._full_graph else "execution_access_plan_missing") return - - copies: dict[Any, list] = {} - results: dict[int, list[dict]] = {} - audited: set[Any] = set() - templates: dict[tuple[int, int], list[WriteAccess | ReadAccess]] = {} - baseline: dict = {} + request_ids = list(runner.input_batch.req_ids) if phase == "forward" else [] shadow = self._scratch shadow.configure(self.state.group_sizes) - pool_id = shadow.pool_id = self.state.pool_id - expected: dict[int, ExpectedAccess | None] = {} - for layer in dict.fromkeys((*self._metadata, *observed)): - group_id = self.layer_groups.get(layer) - if group_id is None: - self.gap(f"unbound_access_layer:{layer}") + shadow.emit = self._emit_scratch + for group, item in self._metadata.items(): + if group in self._unready_load_groups: continue - if group_id in self._unready_load_groups: - continue - actual = observed.get(layer, []) - if observed and not actual and not self._full_graph: - self.gap(f"operator_layer_not_observed:{layer}") - self._missing_groups.add(group_id) - if actual and not self._full_graph: - accesses = actual - evidence = "operator_arguments" - elif layer in self._metadata: - item = self._metadata[layer] - template_key = id(item), group_id - if template_key not in templates: - templates[template_key] = self._metadata_accesses(layer, item) - accesses = templates[template_key] - evidence = "execution_plan" + slots = self._field(item, "slot_mapping") + table = self._field(item, "block_tables", "block_table", "block_table_tensor") + lengths = self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") + if slots is None: + self.gap("metadata_write_mapping_missing", group_id=group) + self._missing_groups.add(group) + slots = [] else: - accesses = actual - evidence = "operator_arguments" - signature = ( - group_id, - evidence, - tuple( - ("write", self._identity(a.slots), a.token_offset) - if isinstance(a, WriteAccess) - else ( - "read", - self._identity(a.table), - self._identity(a.lengths), - self._identity(a.starts), - a.request_offset, + count = getattr(item, "num_actual_tokens", None) + slots = self._list(slots if count is None else slots[:count]) + slots = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] + size = self.state.group_sizes[group] + if phase != "forward": + lengths, table = [], [] # Discarded scratch reads do not establish cache validity. + elif table is None or lengths is None: + self.gap("metadata_read_range_missing", group_id=group) + self._missing_groups.add(group) + lengths, table = [], [] + else: + lengths = self._list(lengths) + columns = max(((length + size - 1) // size for length in lengths), default=0) + table = ( + [row[:columns] for row in table[: len(lengths)]] + if isinstance(table, (list, tuple)) + else self._list(table[: len(lengths), :columns]) + ) + entry = {"group_id": group, "slots": slots, "block_table": table, "seq_lens": lengths} + writes = {(group, slot // size) for slot in slots if slot >= 0} + keys = writes | {(group, block) for row in table for block in row if block >= 0} + # Metadata is already on the host. Keep state inspection and commit atomic, + # sharing read-only blocks and cloning only blocks this step may write. + with self.state._lock: + plan = self._expected_write_plan(runner, group) + if plan is not None: + entry.update( + expected_slots=plan.slots, + query_ends=plan.query_ends, + expected_block_table=plan.tables, + expected_seq_lens=plan.lengths, ) - for a in accesses - ), - ) - if signature in audited: - continue - - def values(value): - key = self._identity(value) - if key not in copies: - copies[key] = self._list(value) - return copies[key] - - groups: list[tuple[dict[str, Any], int]] = [] - write_ranges: list[tuple[int, int]] = [] - keys: set[tuple[int, int]] = set() - write_keys: set[tuple[int, int]] = set() - size = self.state.group_sizes[group_id] - if group_id not in expected: - expected[group_id] = self._expected_write_plan(runner, group_id) - plan = expected[group_id] - for access in accesses: - entry: dict[str, Any] = { - "group_id": group_id, - "slots": [], - "block_table": [], - "seq_lens": [], + if len(slots) < len(plan.slots): + self.gap("write_tokens_not_observed", group_id=group, first_missing_token=len(slots)) + self._missing_groups.add(group) + pool = shadow.pool_id = self.state.pool_id + before = { + key: (block, block.epoch, block.valid_end, frozenset(block.owners)) + for key in writes + if (block := self.state.blocks.get(key)) is not None } - if isinstance(access, WriteAccess): - slots = values(access.slots) - entry["slots"] = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] - if plan is not None: - entry["expected_slots"], entry["query_ends"] = plan.slots, plan.query_ends - entry["token_offset"] = access.token_offset - write_ranges.append((access.token_offset, access.token_offset + len(entry["slots"]))) - write_keys.update((group_id, slot // size) for slot in entry["slots"] if slot >= 0) - keys.update(write_keys) - offset = 0 - else: - if self._phase in ("dummy", "warmup"): - # Dummy outputs are discarded. Their writes can corrupt - # live KV, but reading reserved scratch is expected. - continue - lengths = values(access.lengths) - if access.starts is not None: - lengths = [a + b for a, b in zip(lengths, values(access.starts))] - entry["seq_lens"] = lengths - needed = max(((length + size - 1) // size for length in lengths), default=0) - if isinstance(access.table, (list, tuple)): - entry["block_table"] = [row[:needed] for row in access.table[: len(lengths)]] - else: - entry["block_table"] = values(access.table[: len(lengths), :needed]) - keys.update((group_id, block) for row in entry["block_table"] for block in row if block >= 0) - offset = access.request_offset - if plan is not None: - entry["expected_block_table"] = plan.tables[offset:] - entry["expected_seq_lens"] = plan.lengths[offset:] - groups.append((entry, offset)) - - if plan is not None: - covered = 0 - for start, end in sorted(write_ranges): - if start > covered: - break - covered = max(covered, end) - if covered < len(plan.slots): - self.gap("write_tokens_not_observed", layer=layer, first_missing_token=covered) - self._missing_groups.add(group_id) - - with self.state._lock: - # Owners are immutable within the audit. Copy them once per - # touched block; read-only state can also be shared by layers. - for key in keys: - if key not in baseline and key in self.state.blocks: - initial = self.state.blocks[key] - baseline[key] = replace(initial, owners=set(initial.owners)) shadow.blocks = { - key: replace(baseline[key]) if key in write_keys else baseline[key] + key: replace(block) if key in writes else block for key in keys - if key in baseline + if (block := self.state.blocks.get(key)) is not None } - self._audit_layer, self._audit_evidence = layer, evidence - request_ids = list(runner.input_batch.req_ids) if self._phase == "forward" else [] - for entry, offset in groups: - shadow.observe([entry], request_ids[offset:], phase) - audited.add(signature) - results.setdefault(group_id, []).append(shadow.blocks) - - # Each layer starts from the same pre-forward validity. Another layer's - # write cannot legalize an earlier read or hide a missing write. Merge - # conservatively; no validity is advanced on failed/partial collection. - with self.state._lock: - for group_id, versions_by_layer in results.items(): - touched = {key for versions in versions_by_layer for key in versions} - for key in touched: - block = self.state.blocks.get(key) - if block is None: - continue - versions = [versions.get(key, baseline.get(key, block)) for versions in versions_by_layer] - if self.state.pool_id != pool_id or any(b.epoch != block.epoch for b in versions): - self.gap("allocation_changed_during_access_audit") - continue - if any(b.valid_end == INVALID_PREFIX for b in versions): + shadow.observe([entry], request_ids, phase) + if self.state.pool_id != pool or any( + self.state.blocks.get(key) is not block + or (block.epoch, block.valid_end, frozenset(block.owners)) != (epoch, valid, owners) + for key, (block, epoch, valid, owners) in before.items() + ): + self.gap("allocation_changed_during_access_audit") + continue + for key in before: + block, updated = self.state.blocks[key], shadow.blocks[key] + if updated.valid_end == INVALID_PREFIX: block.valid_end = INVALID_PREFIX - elif ( - block.valid_end != INVALID_PREFIX - and group_id not in self._missing_groups - and self._phase == "forward" - ): - valid = min(b.valid_end for b in versions) - if valid > block.valid_end: - block.valid_end = valid + elif phase == "forward" and group not in self._missing_groups and block.valid_end != INVALID_PREFIX: + block.valid_end = max(block.valid_end, updated.valid_end) def observe_metadata(self, metadata: Any) -> None: try: @@ -667,8 +356,10 @@ def observe_metadata(self, metadata: Any) -> None: def _finish_pending(self) -> None: if self._pending is not None: - self._pending.finish() - self._pending = None + try: + self._pending.finish() + finally: + self._pending = None operations, self._deferred_state = self._deferred_state, [] for operation, args in operations: operation(*args) @@ -681,7 +372,8 @@ def _snapshot_audit(self, runner: Any) -> None: job._scheduled = None if self._scheduled is None else dict(self._scheduled) job._missing_groups = set(self._missing_groups) job._unready_load_groups = set(self._unready_load_groups) - job.report, job.gap, job._check_failure = self.report, self.gap, self._check_failure + reporting = cast(Any, job) + reporting.report, reporting.gap, reporting._check_failure = self.report, self.gap, self._check_failure job._scratch = KVGuard(job._emit_scratch) frozen_runner = SimpleNamespace(input_batch=SimpleNamespace(req_ids=list(runner.input_batch.req_ids))) copies: dict[Any, Any] = {} @@ -715,36 +407,15 @@ def freeze(value): return value try: - job._accesses = [ - WriteAccess(access.layer, freeze(access.slots), access.token_offset) - if isinstance(access, WriteAccess) - else ReadAccess( - access.layer, - freeze(access.table), - freeze(access.lengths), - freeze(access.starts), - access.request_offset, + job._metadata = {} + for group, item in self._metadata.items(): + slots = self._field(item, "slot_mapping") + count = getattr(item, "num_actual_tokens", None) + job._metadata[group] = SimpleNamespace( + slot_mapping=freeze(slots if slots is None or count is None else slots[:count]), + block_tables=freeze(self._field(item, "block_tables", "block_table", "block_table_tensor")), + seq_lens=freeze(self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu")), ) - for access in self._accesses - ] - observed = {access.layer for access in self._accesses} - items, job._metadata = {}, {} - for layer, item in self._metadata.items(): - if not self._full_graph and layer in observed: - job._metadata[layer] = None # Only coverage is used when operator arguments were observed. - continue - key = id(item) - if key not in items: - slots = self._field(item, "slot_mapping") - count = getattr(item, "num_actual_tokens", None) - if slots is not None and count is not None: - slots = slots[:count] - items[key] = SimpleNamespace( - slot_mapping=freeze(slots), - block_tables=freeze(self._field(item, "block_tables", "block_table", "block_table_tensor")), - seq_lens=freeze(self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu")), - ) - job._metadata[layer] = items[key] if event is not None: event.record(torch.npu.current_stream()) except Exception: @@ -803,11 +474,10 @@ def forward_scope( """Audit one completed forward; never advance validity on exceptions.""" self._check_failure() self._phase, self._full_graph = phase, full_graph - self._accesses = [] self.observe_metadata(metadata) try: yield - if required or self._metadata or self._accesses: + if required or self._metadata: try: if phase == "forward" and self._connector is not None: self._connector.before_audit() @@ -815,8 +485,10 @@ def forward_scope( for meta in self._recomputes: self.state.recompute(meta) self._recomputes = [] - if self._defer_audit and phase == "forward" and self._connector is None: + if phase == "forward": self._snapshot_audit(runner) + if not self._defer_audit: + self._finish_pending() else: self._audit(runner) except Exception as exc: @@ -834,7 +506,6 @@ def wrap_forward(self, runner: Any, forward: Any, get_context: Any): def observed(num_tokens, *args, **kwargs): self._finish_pending() context = get_context() - context._kv_guard = self with self.forward_scope( runner, phase=kwargs.pop("_kv_guard_phase", "forward"), From 1199135a92f40eab4bb040c3a33004cf71d63d02 Mon Sep 17 00:00:00 2001 From: TAJh Date: Sat, 7 Nov 2026 10:38:40 +0000 Subject: [PATCH 5/6] refactor(debug): consume standard scheduler output for KV checks Keep the cache-group metadata audit while removing scheduler, allocator, connector and result-method patches. Bind expected addresses from standard SchedulerOutput and call the guard from native Ascend runner boundaries. Limit scope to metadata-plan consistency and complete audits synchronously. Signed-off-by: TAJh --- .../feature_guide/kv_cache_guard.md | 290 ++---- release_notes/fragments/17833.yaml | 35 +- tests/ut/debug/__init__.py | 22 - tests/ut/debug/conftest.py | 67 ++ tests/ut/debug/helpers.py | 40 + tests/ut/debug/test_kv_guard.py | 331 ++----- tests/ut/debug/test_kv_guard_async.py | 211 ----- tests/ut/debug/test_kv_guard_loads.py | 502 ----------- tests/ut/debug/test_kv_guard_runtime.py | 796 ++++------------ tests/ut/debug/test_kv_guard_scope.py | 152 +++- tests/ut/debug/test_kv_guard_v2.py | 357 -------- .../kv_p2p/test_mooncake_connector.py | 39 - vllm_ascend/debug/kv_guard.py | 425 ++------- vllm_ascend/debug/kv_guard_config.py | 2 +- vllm_ascend/debug/kv_guard_connector.py | 186 ---- vllm_ascend/debug/kv_guard_runtime.py | 852 ++++-------------- .../kv_transfer/kv_p2p/mooncake_connector.py | 3 +- vllm_ascend/patch/__init__.py | 17 - vllm_ascend/patch/platform/__init__.py | 3 - vllm_ascend/patch/platform/patch_kv_guard.py | 182 ---- vllm_ascend/worker/v1/model_runner.py | 56 +- vllm_ascend/worker/v2/model_runner.py | 20 +- vllm_ascend/worker/v2/model_states/default.py | 9 + 23 files changed, 831 insertions(+), 3766 deletions(-) create mode 100644 tests/ut/debug/conftest.py create mode 100644 tests/ut/debug/helpers.py delete mode 100644 tests/ut/debug/test_kv_guard_async.py delete mode 100644 tests/ut/debug/test_kv_guard_loads.py delete mode 100644 tests/ut/debug/test_kv_guard_v2.py delete mode 100644 vllm_ascend/debug/kv_guard_connector.py delete mode 100644 vllm_ascend/patch/platform/patch_kv_guard.py diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md index 4c4f8aba3504..fc96e882a087 100644 --- a/docs/source/user_guide/feature_guide/kv_cache_guard.md +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -1,227 +1,117 @@ # KV cache address guard -KV Guard diagnoses address mappings that can cause cross-request writes, -prefix overwrites, or reads from invalid KV ranges. It compares scheduled -destinations with the runner's final execution metadata, once per KV cache group. -The feature is disabled by default and supports adapted dense PA and MLA -attention paths in Model Runner V1 and V2. +KV Guard compares the scheduler's declared token destinations with the runner's +final attention metadata, once per KV cache group. Eager and FULL graph modes +use the same metadata-plan checks. Attention implementations are unchanged; +operator arguments and per-layer execution order are not recorded. ## Enable the guard -Use the same per-engine configuration for Model Runner V1 and V2: +The guard is disabled by default. Enable it for an engine with: ```bash -vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true}}' +vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' ``` -| `kv_guard_config` option | Type | Default | Meaning | +| Option | Type | Default | Meaning | | --- | --- | --- | --- | -| `enabled` | bool | `false` | Enable KV address diagnostics for this engine. | +| `enabled` | bool | `false` | Enable worker metadata-plan checks. | | `action` | `"log"` or `"raise"` | `"log"` | Report failures or stop returning worker results. | -`action="log"` reports failures while preserving diagnostic execution. To stop -returning worker results after a violation or incomplete audit, select `raise`: +The Python API accepts the same `additional_config` dictionary. Unknown fields, +non-boolean `enabled` values and unknown actions are rejected. Each runner owns +its guard; the configuration object does not hold mutable checking state. -```bash -vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' -``` +`raise` latches the first violation or incomplete audit. Catching an exception +does not reset the guard; restart the worker before reuse. There is no collective +failure coordination across workers. -For the Python API, pass the same dictionary through `additional_config`: +## Independent execution plan -```python -from vllm import LLM +The worker consumes the unmodified `SchedulerOutput` before its runner updates: -llm = LLM(model="MODEL", additional_config={"kv_guard_config": {"enabled": True, "action": "raise"}}) -``` +- New requests supply complete block tables and computed-token counts. +- Cached requests supply block-table increments, or replacements for resumed + requests, and current computed-token counts. +- Scheduled token counts define the token positions authorized for this step. +- Finished and preempted request IDs retire the corresponding bindings. + +The guard keeps its own request-to-block mirror. It derives expected slots, +read block tables and sequence lengths from these scheduler messages and the +actual batch order. It does not use the runner's address metadata to construct +its expected plan. + +For each KV cache group, the guard selects the first non-null metadata item +whose layer belongs to that group, matching the current group-level audit. +Other layers in that group are not independently audited. A missing group, +incomplete plan or incorrect batch binding reports a coverage gap. + +Checks compare exact token slots, logical read block order and declared read +lengths. For example, swapping A's and B's slots is an error even if both target +blocks exist. Negative slot padding and zero-length padded request rows are +allowed; padding must not authorize extra writes or reads. + +FULL replay uses this execution's final metadata after graph parameter updates. +Captured or previous-step plans cannot substitute for the current plan. V2 +uses the actual batch request order, not persistent request-state indices. + +Runner and model-state implementations call the guard explicitly. No Scheduler +or allocator patch is installed, and runner, connector and asynchronous output +methods are not replaced at runtime. + +## Coverage boundaries + +Dense PA and MLA append-only cache paths in Runner V1/V2 are adapted. The guard +checks consistency between scheduler messages and selected metadata. It does +not prove that a particular layer executed, that operator arguments matched +the metadata, or that a kernel used those addresses. + +Allocator generations, initialized KV contents, historical prefix writes and +transport integrity are outside this revision's scope. A computed prefix is +a framework input, not proof that its bytes were initialized correctly. Prefix +reuse retains address-plan checks without validating the reused contents. + +KV connectors, CP, DBO, speculative decoding, encoder/decoder execution, +sliding/circular caches, compressed/state caches and unadapted attention or +model-state implementations are rejected. There is no remote-load submission, +completion, cancellation or failed-load recovery guarantee. + +Startup graph capture does not read device metadata on the CPU. Recapture with +live request state is rejected. Dummy writes can be checked against request +blocks currently known to the mirror; this is not an allocator-liveness proof. + +The guard cannot prevent or undo writes already submitted to the device, verify +KV numerical values, or establish model-output correctness. + +## Failure reporting and performance + +Logs use `[KV_GUARD]`. Plan mismatches use `guard.violation` with +`evidence=execution_plan` and a cache group when available. +`guard.coverage_gap` means required evidence is missing, not a successful check. +Reports are deduplicated with bounded state. + +Audits complete synchronously before guarded execution returns. Copying address +metadata to CPU can synchronize the NPU. No deferred audit is attached to an +asynchronous output. Performance results for earlier deferred implementations +do not apply; measure this revision on the intended workload. + +## Validation -The scheduler and workers read this configuration independently for each engine. -Omitting `kv_guard_config` or setting `enabled=false` disables instrumentation. -Invalid option names, non-boolean `enabled` values, and unknown actions are rejected. -KV guard environment variables are no longer used. - -The first failure is latched. Catching the exception does not reset the worker; -restart it before reuse. This is a worker-level stop, not cancellation of one -request, and it adds no collective failure coordination across ranks. - -## What is checked - -An allocated block is not necessarily the correct destination for a token. -For example, if requests A and B own blocks 1 and 2, swapping their write slots -still points to allocated memory. Ordinary allocation logs cannot establish -that each token used its intended destination. - -The scheduler supplies changed request-to-block bindings, allocation generations, -released requests, and each request's computed and scheduled token counts through -the existing scheduler output. The worker independently derives the expected -slot of every scheduled token from those bindings and the actual batch order. - -| Constraint | Failure example | -| --- | --- | -| Exact token destination | Slots are swapped between requests or between tokens in one request. | -| Allocation and ownership | A real request accesses another request's block or a recycled allocation. | -| Logical block order | A request reads its own blocks in the wrong logical order. | -| Initialized prefix | A read extends into unwritten data, or a write overwrites a valid prefix. | -| Reserved storage | A write targets the null block instead of negative-slot padding. | -| Collection completeness | A scheduled write, binding, cache group, or current graph plan is missing. | - -Each tracked block stores an allocation generation, a valid prefix length, and -request owners. A zero prefix is unwritten; a positive prefix N makes `[0, N)` -readable; an invalid prefix prevents subsequent reads from being accepted. -Completed loads and contiguous appends establish readable prefixes. Completion -and preemption release ownership while preserving reusable cached prefixes. -The scheduler's computed-token count describes a consumer's read requirement; -it does not establish validity. A same-step prefix producer must first supply -legal write evidence, while an already valid prefix still rejects overwrites. -Physical clearing, block copies, and completed receives update validity at their -existing lifecycle boundaries. The scheduler records external-load destinations, -allocation generations, and logical token ranges when slots are allocated, -including requests that are still waiting and have no scheduled tokens. -Connector transport offsets are not interpreted as initialized-prefix lengths. - -The common connector methods provide completion evidence. Async receive results -authorize later forwards. Sync loads become readable for a cache group only -after its required layer waits; a missing wait reports incomplete coverage. -FULL replay without observable sync-load completion also reports incomplete -coverage. No extra connector polling or device waits are introduced. -Failed loads discard only their destination blocks, preserving untouched local -prefix blocks. A failed transfer can have overwritten the beginning of a partial -destination block, so that entire block must be recomputed. -Block failures remain attached to their submission until its completion is -drained, even if the connector reports failures and completion in different steps. -Completions must match the submission snapshot and current allocation. Reusing -a request ID while its old receive is outstanding reports incomplete coverage, -because the connector's request-ID-only result cannot distinguish both loads. -Known old submissions may drain after cancellation without changing a new -allocation or treating cancellation itself as a guard failure. -Framework-authorized prefix replay truncates the specified validity range before -write auditing; an ordinary prefix overwrite still fails. Transport drain and -completion results are returned unchanged. - -For example, with block size 128 and a loaded 19-token prefix in block 1, the -next token should write slot 147. A dummy write to slot 142 overwrites offset 14; -the guard diagnoses the write and marks the block invalid for subsequent reads. -This example illustrates the metadata checks, not reproduction of a particular -full-model incident. - -## Execution paths and supported configurations - -| Path | Audit evidence and boundary | -| --- | --- | -| Eager and FULL replay | One final metadata record per KV cache group, inspected after forward and graph parameter updates. Captured or previous-step plans are not reused as evidence. | -| Runner V1 | Bind before `execute_model`; collect metadata through the existing forward context. | -| Runner V2 | Bind before execute; obtain actual batch IDs and metadata after `model_state.prepare_attn`; audit at the end of execute. Persistent request-state indices are not batch positions. | -| Dummy/profile execution | Check writes against existing KV state; ignore discarded scratch reads. | -| Runner V2 startup capture | Require empty request/cache state; do not audit CPU metadata inside device graph capture. | - -Eager and FULL use the same checker. The runner selects one nonempty layer -metadata record for each cache group; it does not record or compare individual -layers' accesses. Only write-target block state is copied for the audit; -read-only block state is shared under the state lock. Validity is committed -only after checking the independent scheduler-derived plan and confirming the -write-target allocations are unchanged. Async snapshots still freeze every -step's selected metadata and wait for its copy before returning results. - -Live or cache-preserving Runner V2 recapture is rejected, including pending -receives. Normal FULL replay remains audited. Allocator and kernel block -coordinates must match. - -The guard explicitly rejects CP, DBO, speculation, encoder/decoder execution, -sliding or circular caches, other compressed or state-cache semantics, unknown -attention implementations, and connectors other than Mooncake. These paths -require their own validity rules before they can be supported. -The Ascend block-transfer implementation is registered as `MooncakeConnectorV1`; -this name and the legacy `MooncakeConnector` name pass the guard's connector check. -The shared load adapter is not an expansion of this connector support list; -other backends still require end-to-end validation before enabling them. - -The guard validates the declared metadata plan. It does not instrument attention -operators, detect per-layer argument differences or actual operator ordering, -inspect kernel memory instructions, or verify KV numerical values. Correct -metadata cannot establish that an operator used those addresses. The guard -cannot prevent or undo submitted writes, prove transport byte integrity, or -guarantee model output correctness. - -## Read the logs - -Logs use the `[KV_GUARD]` prefix: - -- `guard.enabled`: the configured failure action. -- `guard.violation`: an observed constraint violation, with an operation and - reason such as `write_target_mismatch`, `overwrite_valid_prefix`, or - `read_unwritten_range`. -- `guard.coverage_gap`: required evidence was missing, so the audit was - incomplete. This is not a successful check or an exemption from validation. - -Runner V2 may submit async loads during `sample_tokens` post-forward cleanup. -Load submission and completion are checked at the connector cleanup boundary, -after the separately audited `execute_model` forward. - -Plan violations use `evidence=execution_plan` and identify the request, pool, -cache group, block, allocation generation, and relevant slot or range when -available. Connector readiness gaps may identify a missing layer wait; -attention accesses are not recorded per layer. Repeated failures are -deduplicated with bounded state. Routine validity transitions are not logged. - -## Validation and performance - -The focused CPU suite exercises allocation generations, ownership, shared -prefixes, exact destinations, declared read/write ranges, graph metadata freshness, -failure latching, and the V1/V2 adapters: +Run focused tests with: ```bash pytest -q --confcutdir=tests/ut/debug tests/ut/debug ``` -The single-card service regression starts separate disabled and enabled servers, -uses `raise`, repeats a shared prompt, submits concurrent requests, and compares -greedy outputs for both runners in eager and FULL modes: +The service regression compares disabled/enabled greedy outputs for V1/V2, +eager/FULL, repeated prompts and concurrent requests: ```bash pytest -q tests/e2e/pull_request/one_card/test_kv_cache_guard.py ``` -CPU mocks establish checker behavior, not NPU execution correctness. Before -merge, prioritize current-revision hardware runs for V2 eager and FULL, prefix -reuse, observed preemption, Mooncake where used, and controlled fault injection. -The service regression does not itself prove that a preemption or prefix-cache -hit occurred; collect those metrics separately. Historical V1 results do not -validate a rebased revision or Runner V2. - -Benchmark disabled and enabled configurations on the same hardware, model, -request trace, sequence lengths, concurrency, seed, and graph mode. Record -successful requests, throughput, TTFT, ITL, CPU usage, and run-to-run variation. -Fix the source commits, actual module load paths, software versions, exact -commands, and logs in the validation report. Do not infer device synchronization -cost from CPU checker timings. - -The scheduler sends changed block/generation tables and compact request ranges. -Expected token slots are temporary and derived once per cache group; -the guard stores no persistent token history, KV snapshots, or content hashes. -Metadata transfers are deduplicated, but enabled mode -can still synchronize device metadata to CPU. Measure deployment overhead. - -## Implementation guide and patch lifecycle - -| Component | Responsibility | -| --- | --- | -| `vllm_ascend/debug/kv_guard.py` | Allocation, ownership, prefix validity, and read/write rules. | -| `vllm_ascend/debug/kv_guard_config.py` | Validated per-engine options shared by the scheduler and workers. | -| `vllm_ascend/debug/kv_guard_runtime.py` | Shared V1/V2 installation, scheduler binding, evidence collection, audit, and conservative state commit. Includes the V2 batch and capture adapter. | -| `vllm_ascend/patch/platform/patch_kv_guard.py` | Conditional scheduler allocation and request-binding instrumentation. | -| Runner implementations | Existing execution boundaries and one final metadata record per cache group. Attention implementations are unchanged. | -| Mooncake receiver | Success/failure validity callbacks after receive completion and reformatting. | - -Both runners use one `GuardRuntime` and one address-checking state machine. -The internal V2 adapter only handles sorted batch IDs, metadata timing, and -capture boundaries; it does not implement another checker. - -The scheduler constructor is wrapped once; allocation, release, and schedule -methods are instrumented only for engines with `kv_guard_config.enabled=true`. -Current public hooks do not provide both allocation generations and the -independent scheduler-to-worker request contract. Its removal path is an upstream allocation-lifecycle and scheduler -output extension that supplies the same evidence without wrapping instance -methods. Changes to those upstream contracts require adapter and lifecycle -regressions, alongside NPU tests for affected execution paths. - -This feature's current scope is address-legality diagnostics. Content tracing, -KV dumps, automatic recovery, and additional cache semantics are separate work. +CPU tests establish checker behavior, not NPU service correctness. Record +exact-head NPU service runs, current replay metadata and controlled plan faults +before merge. Standalone NPU metadata or synthetic graph checks do not replace +a full model-service run. Allocator-generation or transport guarantees require +an upstream vLLM lifecycle contract rather than Scheduler patches in Ascend. diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml index 5502170b23fd..ff08aa0f3b79 100644 --- a/release_notes/fragments/17833.yaml +++ b/release_notes/fragments/17833.yaml @@ -2,29 +2,12 @@ pr: 17833 entries: - category: feature text: >- - Added opt-in KV cache address-legality diagnostics for adapted dense PA - and MLA attention paths in Model Runner V1 and V2. The guard is disabled - by default and configured per engine through additional_config.kv_guard_config - (enabled and action). Both runners share one checker, which can log failures - or stop returning worker results after a violation or incomplete audit. - The runner collects one final metadata plan per KV cache group; attention - operators are unchanged and per-layer accesses are not recorded. - - category: bugfix - text: >- - KV guard no longer treats a consumer's computed-token count as proof of - initialized cache bytes. Same-step shared-prefix writes are checked against - recorded validity, which continues to reject overwrites of valid prefixes. - External-load readiness now uses submission-time logical ranges and - allocation generations rather than transport-local computed-token offsets. - Failed loads preserve unaffected local prefixes, and explicit framework - tail replay is distinguished from an illegal prefix overwrite. - Failed-block records are retained across connector polls until receive - completion, and the registered Ascend MooncakeConnectorV1 name is accepted. - Runner V2 async loads are finalized at connector cleanup, including when - submission is deferred from execute_model to sample_tokens. - Cancelled submissions can drain without modifying a later allocation. - - category: performance - text: >- - Ordinary async generation with one model worker defers KV guard validation to output consumption - or the next forward boundary, allowing address metadata copies to overlap - model execution while retaining every step's metadata-plan checks. + Added opt-in KV cache metadata-plan diagnostics for dense PA and MLA paths + in Runner V1/V2. Each cache group is checked once using final attention + metadata in both eager and FULL modes; attention implementations are + unchanged and operator/per-layer accesses are not recorded. Standard + scheduler messages provide the independent expected address plan, without + Scheduler/allocator patches or runtime method replacement. Per-engine + additional_config.kv_guard_config supports log or sticky raise actions. + Audits complete synchronously. Allocator generations, initialized-prefix + validity and KV connectors are outside this feature's scope. diff --git a/tests/ut/debug/__init__.py b/tests/ut/debug/__init__.py index 44d0f6a6de8e..208f01a7cb5e 100644 --- a/tests/ut/debug/__init__.py +++ b/tests/ut/debug/__init__.py @@ -1,24 +1,2 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -def load_plan(guard, request="r", end=19, start=0, *, is_async=True): - guard._test_load_id = getattr(guard, "_test_load_id", 0) + 1 - return { - "pool_id": guard.pool_id, - "load_id": guard._test_load_id, - "request_id": request, - "block_table": [list(refs) for refs in guard.requests[request]], - "token_start": start, - "token_end": end, - "is_async": is_async, - } - - -def complete_load(guard, request="r", end=19, start=0): - meta = load_plan(guard, request, end, start) - guard.begin_load(meta) - guard.load(meta) - return meta diff --git a/tests/ut/debug/conftest.py b/tests/ut/debug/conftest.py new file mode 100644 index 000000000000..20708c921499 --- /dev/null +++ b/tests/ut/debug/conftest.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load the device-independent guard without importing the NPU plugin.""" + +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType +from types import SimpleNamespace as NS +from typing import Any, cast + +import pytest + +ROOT = Path(__file__).parents[3] + + +@pytest.fixture +def subject(monkeypatch): + for name in ("vllm_ascend", "vllm_ascend.debug", "vllm"): + module = ModuleType(name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, name, module) + logger = ModuleType("vllm.logger") + cast(Any, logger).logger = logging.getLogger("kv-guard-test") + monkeypatch.setitem(sys.modules, "vllm.logger", logger) + for name, path in ( + ("vllm_ascend.config_utils", "config_utils.py"), + ("vllm_ascend.debug.kv_guard_config", "debug/kv_guard_config.py"), + ("vllm_ascend.debug.kv_guard", "debug/kv_guard.py"), + ("vllm_ascend.debug.kv_guard_runtime", "debug/kv_guard_runtime.py"), + ): + spec = importlib.util.spec_from_file_location(name, ROOT / "vllm_ascend" / path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def config(): + return NS( + additional_config={"kv_guard_config": {"enabled": True}}, + model_config=NS(is_encoder_decoder=False), + parallel_config=NS(prefill_context_parallel_size=1, decode_context_parallel_size=1, enable_dbo=False), + compilation_config=NS(static_forward_context={}), + speculative_config=None, + kv_transfer_config=None, + ) + + +@pytest.fixture +def runtime(subject, config): + guard = subject.create_guard(config) + guard.events = [] + report = guard.report + + def recorded(event, **fields): + guard.events.append({"event": event, **fields}) + report(event, **fields) + + guard.report = recorded + guard.state.emit = recorded + guard.state.configure({0: 4}) + guard.layer_groups = {"layer0": 0} + return guard diff --git a/tests/ut/debug/helpers.py b/tests/ut/debug/helpers.py new file mode 100644 index 000000000000..e25851ee73f4 --- /dev/null +++ b/tests/ut/debug/helpers.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace as NS + +import torch + + +def message(*, new=(), cached=(), resumed=(), counts=None, finished=(), preempted=()): + return NS( + scheduled_new_reqs=[ + NS(req_id=req, block_ids=groups, num_computed_tokens=computed) for req, groups, computed in new + ], + scheduled_cached_reqs=NS( + req_ids=[req for req, _, _ in cached], + new_block_ids=[groups for _, groups, _ in cached], + num_computed_tokens=[computed for _, _, computed in cached], + resumed_req_ids=set(resumed), + ), + num_scheduled_tokens=counts or {}, + finished_req_ids=set(finished), + preempted_req_ids=set(preempted), + ) + + +def metadata(slots, tables, lengths): + return NS( + slot_mapping=torch.tensor(slots), + block_tables=torch.tensor(tables), + seq_lens=torch.tensor(lengths), + num_actual_tokens=len(slots), + ) + + +def observe(guard, request_ids, slots, tables, lengths, *, full_graph=False): + guard.observe_metadata({"layer0": metadata(slots, tables, lengths)}, request_ids, full_graph) + + +def reasons(guard): + return [event.get("reason") for event in guard.events] diff --git a/tests/ut/debug/test_kv_guard.py b/tests/ut/debug/test_kv_guard.py index d3d09eafbe1f..9cd366386e1e 100644 --- a/tests/ut/debug/test_kv_guard.py +++ b/tests/ut/debug/test_kv_guard.py @@ -1,283 +1,58 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""L1 transitions, invalid accesses and exact Case4 address regression.""" - -import importlib.util -import sys -from pathlib import Path +"""Shadow bindings use standard scheduler messages, never runner tables.""" import pytest -from . import complete_load, load_plan - -spec = importlib.util.spec_from_file_location( - "kv_guard_subject", Path(__file__).parents[3] / "vllm_ascend/debug/kv_guard.py" -) -assert spec is not None and spec.loader is not None -subject = importlib.util.module_from_spec(spec) -sys.modules[spec.name] = subject -spec.loader.exec_module(subject) - - -@pytest.fixture -def guard(): - events = [] - guard = subject.KVGuard(lambda event, **fields: events.append({"event": event, **fields})) - guard.configure({0: 128}) - guard.events = events - return guard - - -def bind(guard, request="r", computed=0, epoch=1, pool="pool", ids=(1,)): - guard.bind( - { - "pool_id": pool, - "requests": {request: [[(block, epoch) for block in ids]]}, - "scheduled_requests": {request: (computed, 0)}, - } - ) - if computed: - complete_load(guard, request, computed) - - -def access(guard, slots=(), length=0, ids=(1,), request_ids=("r",), phase="forward"): - guard.observe( - [{"group_id": 0, "slots": list(slots), "block_table": [list(ids)], "seq_lens": [length]}], - list(request_ids), - phase, - ) - - -def violations(guard): - return [e for e in guard.events if e["event"] == "guard.violation"] - - -def test_partial_fill_seal_and_read(guard): - bind(guard) - access(guard, range(128, 147), 19) - assert guard.blocks[0, 1].valid_end == 19 - access(guard, range(147, 256), 128) - assert guard.blocks[0, 1].valid_end == 128 - assert not violations(guard) - - -@pytest.mark.parametrize("reverse", [False, True]) -def test_same_step_prefix_consumer_cannot_prevalidate_producers_write(guard, reverse): - items = [("producer", [[(1, 1), (2, 1)]]), ("consumer", [[(1, 1), (3, 1)]])] - guard.bind( - { - "pool_id": "pool", - "requests": dict(reversed(items) if reverse else items), - "scheduled_requests": {"producer": (0, 256), "consumer": (128, 128)}, - } - ) - assert guard.blocks[0, 1].valid_end == 0 - guard.observe( - [ - { - "group_id": 0, - "slots": list(range(128, 384)) + list(range(384, 512)), - "block_table": [[1, 2], [1, 3]], - "seq_lens": [256, 256], - } - ], - ["producer", "consumer"], - "forward", - ) - assert not violations(guard) - assert all(guard.blocks[0, block].valid_end == 128 for block in (1, 2, 3)) - - -def test_computed_without_write_evidence_does_not_make_unknown_prefix_readable(guard): - guard.bind({"requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (128, 0)}}) - access(guard, (), 128) - assert violations(guard)[-1]["reason"] == "read_unwritten_range" - - -def test_recorded_prefix_survives_release_and_same_epoch_reuse(guard): - bind(guard) - access(guard, range(128, 256), 128) - guard.bind({"released_request_ids": ["r"]}) - guard.bind({"requests": {"next": [[(1, 1)]]}, "scheduled_requests": {"next": (128, 0)}}) - access(guard, (), 128, request_ids=("next",)) - assert not violations(guard) - access(guard, [128], 128, request_ids=("next",)) - assert violations(guard)[0]["reason"] == "overwrite_valid_prefix" - - -def test_unknown_read_and_skipped_write(guard): - bind(guard) - access(guard, (), 1) - assert violations(guard)[-1]["reason"] == "read_unwritten_range" - access(guard, [130]) - assert violations(guard)[-1]["reason"] == "noncontiguous_write" - assert guard.blocks[0, 1].valid_end == subject.INVALID_PREFIX - - -def test_case4_loaded_prefix_dummy_write_then_dirty_read_without_tensors(guard): - bind(guard) - complete_load(guard) - assert guard.blocks[0, 1].valid_end == 19 - access(guard, [142], 1, request_ids=(), phase="dummy") - dirty_write = violations(guard)[0] - assert dirty_write["reason"] == "dummy_write_to_live_block" - assert dirty_write["block_id"] == 1 and dirty_write["offsets"] == [14] - assert dirty_write["valid_end"] == 19 and dirty_write["owner_request_ids"] == ["r"] - access(guard, (), 19) - assert violations(guard)[-1]["reason"] == "read_invalid_block" - - -def test_loaded_tail_allows_append_but_not_overwrite(guard): - bind(guard, computed=19) - access(guard, [147], 20) - assert not violations(guard) - access(guard, [142], 20) - assert violations(guard)[0]["reason"] == "overwrite_valid_prefix" - - -def test_reallocation_epoch_clears_old_validity_and_taint(guard): - bind(guard, computed=19) - access(guard, [142], phase="dummy", request_ids=()) - bind(guard, request="new", epoch=2) - assert guard.blocks[0, 1].valid_end == 0 - assert guard.blocks[0, 1].owners == {"new"} - guard.events.clear() - access(guard, [128], 1, request_ids=("new",)) - assert not violations(guard) - - -def test_shared_prefix_and_finished_request(guard): - bind(guard, computed=128) - bind(guard, request="shared", computed=128) - access(guard, (), 128, request_ids=("shared",)) - assert not violations(guard) - guard.bind({"released_request_ids": ["r"], "requests": {}}) - assert guard.blocks[0, 1].owners == {"shared"} - access(guard, (), 128) - assert violations(guard)[-1]["reason"] == "unowned_read" - - -def test_cache_lookup_does_not_revive_tainted_same_epoch(guard): - bind(guard, computed=19) - access(guard, [142], phase="dummy", request_ids=()) - guard.bind({"requests": {"prefix-hit": [[(1, 1)]]}, "scheduled_requests": {"prefix-hit": (19, 0)}}) - assert guard.blocks[0, 1].valid_end == subject.INVALID_PREFIX - - -def test_unregistered_completion_cannot_establish_validity(guard): - bind(guard) - guard.load(load_plan(guard)) - assert guard.blocks[0, 1].valid_end == 0 - access(guard, (), 19) - assert violations(guard)[-1]["reason"] == "read_unwritten_range" - assert any(e.get("reason") == "stale_load_completion" for e in guard.events) - - -def test_old_completion_never_applied_to_different_request(guard): - bind(guard, request="old") - meta = load_plan(guard, "old") - guard.begin_load(meta) - guard.bind({"released_request_ids": ["old"]}) - bind(guard, request="new", epoch=2) - guard.load(meta) - assert guard.blocks[0, 1].valid_end == 0 - assert not guard.loads - - -def test_no_per_token_state_log_and_negative_padding(guard): - bind(guard) - for offset in range(64): - access(guard, [128 + offset, -1], offset + 1) - assert not violations(guard) - assert len([e for e in guard.events if e["event"] == "guard.state"]) == 0 - - -def test_null_block_write_is_not_padding_and_taints_later_read(guard): - access(guard, [-1, 0, 14], ids=(0,), request_ids=(), phase="dummy") - assert violations(guard)[0]["reason"] == "write_null_block" - assert violations(guard)[0]["offsets"] == [0, 14] - access(guard, (), 1, ids=(0,)) - assert violations(guard)[-1]["reason"] == "read_null_block" - guard.clear(0, 0) - access(guard, (), 1, ids=(0,)) - assert violations(guard)[-1]["reason"] == "read_null_block" - - -def test_group_pool_and_allocation_isolation(guard): - guard.configure({0: 128, 1: 128}) - bind(guard, computed=19) - guard.bind({"pool_id": "pool", "requests": {"r2": [[], [(1, 1)]]}, "scheduled_requests": {"r2": (7, 0)}}) - complete_load(guard, "r2", 7) - assert guard.blocks[0, 1].valid_end == 19 - assert guard.blocks[1, 1].valid_end == 7 - bind(guard, pool="new-pool") - assert guard.blocks[0, 1].valid_end == 0 - assert (1, 1) not in guard.blocks - - -def test_copy_clear_and_read_invalid_destination(guard): - bind(guard, computed=19) - bind(guard, request="dest", ids=(2,)) - guard.copy_blocks(0, [(1, 2)]) - assert guard.blocks[0, 2].valid_end == 19 - access(guard, (), 19, ids=(2,), request_ids=("dest",)) - assert not violations(guard) - guard.clear(0, 1) - guard.copy_blocks(0, [(1, 2)]) - assert violations(guard)[-1]["reason"] == "invalid_copy_source_or_destination" - access(guard, (), 19, ids=(2,), request_ids=("dest",)) - assert violations(guard)[-1]["reason"] == "read_invalid_block" - - -def test_recycling_a_live_block_is_a_lifecycle_violation(guard): - bind(guard, computed=19) - bind(guard, request="new", epoch=2) - assert violations(guard)[0]["reason"] == "recycle_live_block" - - -def test_preempted_request_release_allows_reallocation_without_false_violation(guard): - bind(guard, computed=19) - guard.bind({"released_request_ids": ["r"]}) - access(guard, (), 19) - assert violations(guard)[-1]["reason"] == "unowned_read" - guard.events.clear() - bind(guard, request="new", epoch=2) - access(guard, [128], 1, request_ids=("new",)) - assert not violations(guard) - - -def test_block_copy_uses_original_sources_for_overlapping_pairs(guard): - bind(guard, computed=19) - bind(guard, request="b", computed=7, ids=(2,)) - bind(guard, request="c", ids=(3,)) - guard.copy_blocks(0, [(1, 2), (2, 3)]) - assert guard.blocks[0, 2].valid_end == 19 - assert guard.blocks[0, 3].valid_end == 7 - assert not violations(guard) - - -def test_missing_allocation_generation_does_not_make_a_block_writable(guard): - bind(guard, epoch=0) - assert (0, 1) not in guard.blocks - assert any(e.get("reason") == "allocation_generation_missing" for e in guard.events) - access(guard, [128]) - assert violations(guard)[-1]["reason"] == "unallocated_block" - - -def test_copying_an_empty_prefix_does_not_make_bytes_readable(guard): - bind(guard) - bind(guard, request="dest", ids=(2,)) - guard.copy_blocks(0, [(1, 2)]) - assert not violations(guard) - access(guard, (), 1, ids=(2,), request_ids=("dest",)) - assert violations(guard)[-1]["reason"] == "read_unwritten_range" - - -def test_replaced_pool_drops_old_loads(guard): - bind(guard) - meta = load_plan(guard) - guard.begin_load(meta) - bind(guard, request="next", pool="replacement") - guard.load(meta) - assert not guard.loads - assert guard.blocks[0, 1].valid_end == 0 +from .helpers import message, reasons + + +def test_new_plan_copies_scheduler_blocks_and_accepts_framework_prefix(runtime): + blocks = ([2, 3],) + runtime.state.bind(message(new=[("r", blocks, 3)], counts={"r": 3})) + blocks[0][0] = 99 + plan = runtime.state.expected(0, ["r"]) + assert plan.slots == [11, 12, 13] + assert plan.tables == [[2, 3]] + assert plan.lengths == [6] + assert not runtime.events + + +def test_append_none_resume_preempt_finish_and_request_id_reuse(runtime): + state = runtime.state + state.bind(message(new=[("r", ([2],), 0)], counts={"r": 3})) + state.bind(message(cached=[("r", ([3],), 3)], counts={"r": 3})) + assert state.expected(0, ["r"]).slots == [11, 12, 13] + state.bind(message(cached=[("r", None, 6)], counts={"r": 1})) + assert state.expected(0, ["r"]).slots == [14] + state.bind(message(preempted=["r"])) + assert state.requests == {} + state.bind(message(cached=[("r", ([7],), 0)], resumed=["r"], counts={"r": 2})) + assert state.expected(0, ["r"]).slots == [28, 29] + state.bind(message(finished=["r"])) + assert state.requests == {} + state.bind(message(new=[("r", ([11],), 0)], counts={"r": 1})) + assert state.expected(0, ["r"]).slots == [44] + + +def test_multiple_groups_and_actual_batch_reordering(runtime): + state = runtime.state + state.configure({0: 4, 1: 8}) + state.bind(message(new=[("a", ([2], [5]), 1), ("b", ([3], [7]), 0)], counts={"a": 2, "b": 1})) + assert state.expected(0, ["b", "a"]).slots == [12, 9, 10] + assert state.expected(1, ["b", "a"]).slots == [56, 41, 42] + + +@pytest.mark.parametrize("blocks,computed,count", [([], 0, 1), ([0], 0, 1), ([2], -1, 1), ([2], 0, 0)]) +def test_incomplete_or_invalid_scheduler_plan_is_a_gap(runtime, blocks, computed, count): + runtime.state.bind(message(new=[("r", (blocks,), computed)], counts={"r": count})) + assert runtime.state.expected(0, ["r"]) is None + assert any(event["event"] == "guard.coverage_gap" for event in runtime.events) + + +def test_missing_cached_or_resumed_binding_is_not_reconstructed_from_runner(runtime): + runtime.state.bind(message(cached=[("r", ([2],), 0)], counts={"r": 1})) + assert runtime.state.expected(0, ["r"]) is None + assert "cached_request_binding_missing" in reasons(runtime) + runtime.state.bind(message(cached=[("r", None, 0)], resumed=["r"], counts={"r": 1})) + assert "resume_block_mapping_missing" in reasons(runtime) diff --git a/tests/ut/debug/test_kv_guard_async.py b/tests/ut/debug/test_kv_guard_async.py deleted file mode 100644 index c6727bad3e18..000000000000 --- a/tests/ut/debug/test_kv_guard_async.py +++ /dev/null @@ -1,211 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Deferred evidence must remain exact and precede output/state consumption.""" - -from concurrent.futures import ThreadPoolExecutor -from threading import Event -from types import SimpleNamespace as NS - -import pytest -import torch - -from .test_kv_guard_runtime import bind, metadata -from .test_kv_guard_runtime import runtime_fixture as runtime_fixture - - -def deferred_forward(runtime, mode="FULL", slot=147): - bind(runtime, computed=19) - runtime._defer_audit = True - runtime._scheduled = {"r": (19, 1)} - slots, table, lengths = (torch.tensor([slot]), torch.tensor([[1]]), torch.tensor([20])) - item = metadata(slots, table, lengths) - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - - def forward(count): - pass - - runtime.wrap_forward(runner, forward, lambda: context)(1) - return (slots, table, lengths) - - -@pytest.mark.parametrize("mode", ["FULL", "NONE"]) -@pytest.mark.parametrize("slot", [147, 148]) -def test_deferred_audit_freezes_final_values_without_publishing_validity(runtime, mode, slot): - slots, table, lengths = deferred_forward(runtime, mode, slot) - assert runtime.state.blocks[0, 1].valid_end == 19 - slots[0], table[0, 0], lengths[0] = 0, 2, 99 - runtime._finish_pending() - mismatch = any(event.get("reason") == "write_target_mismatch" for event in runtime.events) - assert mismatch == (slot != 147) - if slot == 147: - assert runtime.state.blocks[0, 1].valid_end == 20 - assert not runtime.events - - -@pytest.mark.parametrize("mode", ["FULL", "NONE"]) -def test_invalid_snapshot_cannot_escape_through_async_output(runtime, mode): - runtime.action = "raise" - runtime.state.emit = runtime.report - deferred_forward(runtime, mode, 148) - consumed = [] - output = NS(get_output=lambda: consumed.append(True)) - assert runtime._guard_output(output) is output - with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): - output.get_output() - assert not consumed - with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): - output.get_output() - assert not consumed - - -def test_direct_output_finishes_audit_before_return(runtime): - deferred_forward(runtime) - marker = object() - assert runtime._guard_output(marker) is marker - assert runtime.state.blocks[0, 1].valid_end == 20 - - -def test_completed_output_still_honors_a_later_sticky_failure(runtime): - runtime.action = "raise" - deferred_forward(runtime) - consumed = [] - output = runtime._guard_output(NS(get_output=lambda: consumed.append(True))) - runtime._finish_pending() - runtime._failure = "later_step_failed" - with pytest.raises(runtime.subject.KVGuardError, match="later_step_failed"): - output.get_output() - assert not consumed - - -def test_readonly_output_fallback_cannot_bypass_a_sticky_failure(runtime): - class Output: - __slots__ = () - - def get_output(self): - return "result" - - runtime.action = "raise" - deferred_forward(runtime) - assert runtime._guard_output(Output()).get_output() == "result" - runtime._failure = "later_step_failed" - with pytest.raises(runtime.subject.KVGuardError, match="later_step_failed"): - runtime._guard_output(Output()) - - -def test_next_binding_and_clear_follow_the_predecessor_audit(runtime): - deferred_forward(runtime) - pending = runtime._pending - consumed = [] - output = runtime._guard_output(NS(get_output=lambda: consumed.append(True))) - original_bind = runtime.state.bind - calls = [] - - def checked_bind(context): - assert runtime.state.blocks[0, 1].valid_end == 20 - calls.append("bind") - original_bind(context) - - runtime.state.bind = checked_bind - packet = NS( - _ascend_kv_guard_context={ - "pool_id": "pool", - "requests": {"next": [[(2, 1)]]}, - "released_request_ids": ["r"], - "scheduled_requests": {"next": (0, 1)}, - } - ) - with runtime.execute_scope(packet): - runtime.zero_blocks([2]) - assert not calls and "r" in runtime.state.requests - runtime._finish_pending() - assert calls == ["bind"] and "r" not in runtime.state.requests - assert runtime.state.blocks[0, 2].valid_end == 0 - # A late reader reuses the completed decision, not the next step's buffers/state. - runtime.state.blocks[0, 1].epoch = 9 - output.get_output() - assert consumed == [True] and pending is not None - assert not runtime.events - - -def test_idle_step_drains_pending_validation_and_release(runtime): - deferred_forward(runtime) - packet = NS( - _ascend_kv_guard_context={ - "pool_id": "pool", - "requests": {}, - "released_request_ids": ["r"], - "scheduled_requests": {}, - } - ) - runtime.wrap_execute(lambda output: None)(packet) - assert runtime._pending is None and not runtime._deferred_state - assert not runtime.state.requests - assert runtime.state.blocks[0, 1].valid_end == 20 - - -@pytest.mark.parametrize("failure", [False, True]) -def test_output_reader_and_next_forward_finish_only_once(runtime, failure): - started, release = Event(), Event() - calls = [] - - def complete(): - calls.append(True) - started.set() - assert release.wait(2) - if failure: - raise RuntimeError("audit failed") - - pending = runtime.subject._PendingAudit(complete) - with ThreadPoolExecutor(max_workers=2) as executor: - futures = [executor.submit(pending.finish) for _ in range(2)] - assert started.wait(2) - release.set() - for future in futures: - if failure: - with pytest.raises(RuntimeError, match="audit failed"): - future.result() - else: - future.result() - assert calls == [True] and pending._complete is None - - -def test_snapshot_keeps_missing_group_coverage(runtime): - runtime.state.configure({0: 128, 1: 128}) - runtime.layer_groups = {"layer0": 0, "layer1": 1} - runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1)], [(2, 1)]]}}) - runtime._defer_audit = True - runtime._scheduled = {"r": (0, 1)} - item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) - runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) - runtime._finish_pending() - assert runtime.state.blocks[0, 1].valid_end == 1 - assert runtime.state.blocks[1, 2].valid_end == 0 - assert any( - event.get("reason") == "metadata_write_mapping_missing" and event.get("group_id") == 1 - for event in runtime.events - ) - - -@pytest.mark.parametrize("mode", ["FULL", "NONE"]) -@pytest.mark.parametrize("field,reason", [("table", "read_target_mismatch"), ("length", "read_length_mismatch")]) -def test_frozen_read_evidence_cannot_be_repaired_by_the_next_preparation(runtime, mode, field, reason): - bind(runtime, computed=19) - runtime._defer_audit = True - runtime._scheduled = {"r": (19, 1)} - slots = torch.tensor([147]) - table = torch.tensor([[2 if field == "table" else 1]]) - lengths = torch.tensor([21 if field == "length" else 20]) - item = metadata(slots, table, lengths) - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - - def forward(count): - pass - - runtime.wrap_forward(runner, forward, lambda: context)(1) - table[0, 0], lengths[0] = (1, 20) - runtime._finish_pending() - assert any(event.get("reason") == reason for event in runtime.events) diff --git a/tests/ut/debug/test_kv_guard_loads.py b/tests/ut/debug/test_kv_guard_loads.py deleted file mode 100644 index e0a3214b98b1..000000000000 --- a/tests/ut/debug/test_kv_guard_loads.py +++ /dev/null @@ -1,502 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Logical load snapshots and common connector boundaries; CPU contracts only.""" - -from types import SimpleNamespace as NS - -import pytest -import torch - -from . import complete_load, load_plan -from .test_kv_guard_runtime import bind, load, metadata -from .test_kv_guard_runtime import runtime_fixture as runtime_fixture -from .test_kv_guard_v2 import adapter_fixture as adapter_fixture -from .test_kv_guard_v2 import prepare - - -class Connector: - def __init__(self): - self.calls = [] - self.result = NS(finished_sending=set(), finished_recving=set(), failed_recving=set()) - self.invalid = set() - - def start_load_kv(self, *args, **kwargs): - self.calls.append(("start", args, kwargs)) - return "started" - - def wait_for_layer_load(self, layer_name, *args, **kwargs): - self.calls.append(("wait", layer_name, args, kwargs)) - return "ready" - - def get_transfer_results(self, *args, **kwargs): - self.calls.append(("results", args, kwargs)) - return self.result - - def get_block_ids_with_load_errors(self): - self.calls.append(("errors",)) - return self.invalid - - -@pytest.fixture -def connector(runtime): - connector = Connector() - runtime._connector = runtime.subject.ConnectorGuard(runtime, connector) - return connector - - -def packet(meta=None, scheduled=None, requests=None, released=(), replay=None): - scheduled = {} if scheduled is None else scheduled - context = { - "pool_id": "pool", - "requests": requests or {}, - "scheduled_requests": scheduled, - "released_request_ids": list(released), - "loads": [meta] if meta is not None else [], - "recompute_ranges": [replay] if replay is not None else [], - } - return NS(_ascend_kv_guard_context=context, total_num_scheduled_tokens=sum(n for _, n in scheduled.values())) - - -def poll(connector): - result = connector.get_transfer_results(set()) - invalid = connector.get_block_ids_with_load_errors() - return result, invalid - - -def reasons(runtime): - return {event.get("reason") for event in runtime.events} - - -@pytest.mark.parametrize( - "name, accepted", [("MooncakeConnector", True), ("MooncakeConnectorV1", True), ("AscendStoreConnector", False)] -) -def test_registered_ascend_mooncake_name_is_accepted(runtime, name, accepted): - impl = type("AscendAttentionBackendImpl", (), {})() - runtime.config = NS( - parallel_config=NS(), - model_config=NS(), - speculative_config=None, - kv_transfer_config=NS(kv_connector=name), - compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), - ) - runner = NS( - kv_cache_config=NS(kv_cache_groups=[NS(layer_names=["layer0"], kv_cache_spec=NS(block_size=128, head_size=64))]) - ) - if accepted: - runtime.configure(runner) - else: - with pytest.raises(ValueError, match="unadapted execution"): - runtime.configure(runner) - - -@pytest.mark.parametrize("local", [0, 19]) -def test_async_completion_uses_loaded_end_not_transport_local_offset(runtime, connector, local): - bind(runtime, computed=local) - meta = load_plan(runtime.state, end=64, start=local) - - def submit(output): - assert connector.start_load_kv("context", num_computed_tokens=local) == "started" - assert runtime.state.blocks[0, 1].valid_end == local - poll(connector) - - runtime.wrap_execute(submit)(packet(meta)) - connector.result.finished_recving = {"r"} - result, invalid = runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert result is connector.result and invalid is connector.invalid - assert runtime.state.blocks[0, 1].valid_end == 64 - assert not runtime.state.loads and not runtime.events - assert [call[0] for call in connector.calls] == ["start", "results", "errors", "results", "errors"] - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_sync_layer_waits_establish_group_readiness_only_after_every_layer(runtime, connector, mode): - bind(runtime) - meta = load_plan(runtime.state, end=64, is_async=False) - item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) - context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(count): - for index, layer in enumerate(runtime.layer_groups): - assert connector.wait_for_layer_load(layer) == "ready" - assert runtime.state.blocks[0, 1].valid_end == (0 if index == 0 else 64) - - forward = runtime.wrap_forward(runner, model, lambda: context) - - def execute(output): - connector.start_load_kv("context") - forward(1) - poll(connector) - - runtime.wrap_execute(execute)(packet(meta, {"r": (64, 1)})) - assert not runtime.events and (not runtime.state.loads) - assert runtime.state.blocks[0, 1].valid_end == 65 - - -@pytest.mark.parametrize("wait_both", [False, True]) -def test_declared_group_load_requires_all_layer_wait_results(runtime, connector, wait_both): - bind(runtime) - meta = load_plan(runtime.state, end=64, is_async=False) - runner = NS(input_batch=NS(req_ids=["r"])) - item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) - - def model(count): - if wait_both: - connector.wait_for_layer_load("layer0") - connector.wait_for_layer_load("layer1") - - forward = runtime.wrap_forward(runner, model, lambda: context) - runtime.wrap_execute(lambda output: (connector.start_load_kv(None), forward(1)))(packet(meta, {"r": (64, 1)})) - assert ("load_layer_not_ready" in reasons(runtime)) == (not wait_both) - assert runtime.state.blocks[0, 1].valid_end == (65 if wait_both else 0) - assert not any(e["event"] == "guard.violation" for e in runtime.events) - - -def test_full_replay_without_load_completion_evidence_is_incomplete(runtime, connector): - bind(runtime) - meta = load_plan(runtime.state, end=64, is_async=False) - item = metadata(torch.tensor([192]), torch.tensor([[1]]), [65]) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - forward = runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda n: None, lambda: context) - runtime.wrap_execute(lambda output: (connector.start_load_kv(None), forward(1)))(packet(meta, {"r": (64, 1)})) - assert "load_layer_not_ready" in reasons(runtime) - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_receive_failure_preserves_local_prefix_and_allows_recompute(runtime, connector): - state = runtime.state - state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) - complete_load(state, end=128) - meta = load_plan(state, end=192, start=128) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.result.finished_recving = connector.result.failed_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert state.blocks[0, 1].valid_end == 128 - assert state.blocks[0, 2].valid_end == 0 - state.observe([{"group_id": 0, "slots": [256], "block_table": [[1, 2]], "seq_lens": [129]}], ["r"], "forward") - assert not runtime.events - - -def test_invalid_block_report_only_invalidates_actual_loaded_destinations(runtime, connector): - state = runtime.state - state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}}) - complete_load(state, end=128) - meta = load_plan(state, end=384, start=128) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.result.finished_recving = {"r"} - connector.invalid = {1, 2, 99} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert [state.blocks[0, n].valid_end for n in (1, 2, 3)] == [128, 0, 128] - assert not runtime.events - - -@pytest.mark.parametrize("invalid, expected", [({2}, [128, 0, 128]), ({1, 2, 3}, [128, 0, 0])]) -def test_early_invalid_blocks_survive_until_the_later_receive_completion(runtime, connector, invalid, expected): - state = runtime.state - state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1), (3, 1)]]}}) - complete_load(state, end=128) - meta = load_plan(state, end=384, start=128) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.invalid = invalid - runtime.wrap_execute(lambda output: poll(connector))(packet()) - connector.invalid = set() - connector.result.finished_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert [state.blocks[0, n].valid_end for n in (1, 2, 3)] == expected - assert not state.loads and not runtime.events - - -def test_partial_block_failure_invalidates_the_whole_destination_block(runtime, connector): - bind(runtime, computed=19) - meta = load_plan(runtime.state, end=64, start=19) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.result.finished_recving = connector.result.failed_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_start_exception_preserves_original_error_and_local_kv(runtime, monkeypatch): - runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) - complete_load(runtime.state, end=128) - source = Connector() - error = RuntimeError("backend load failed") - - def fail(*args): - raise error - - monkeypatch.setattr(source, "start_load_kv", fail) - runtime._connector = runtime.subject.ConnectorGuard(runtime, source) - meta = load_plan(runtime.state, end=192, start=128) - with pytest.raises(RuntimeError) as raised: - runtime.wrap_execute(lambda output: source.start_load_kv(None))(packet(meta)) - assert raised.value is error - assert runtime.state.blocks[0, 1].valid_end == 128 - assert runtime.state.blocks[0, 2].valid_end == 0 - assert not runtime._connector.inflight and not runtime.state.loads - - -@pytest.mark.parametrize("epoch", [1, 2]) -def test_old_completion_cannot_validate_reused_request_or_allocation(runtime, connector, epoch): - bind(runtime) - old = load_plan(runtime.state, end=64) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(old)) - runtime.wrap_execute(lambda output: None)(packet(requests={"r": [[(1, epoch)]]}, released=["r"])) - new = load_plan(runtime.state, end=32) - runtime.state.begin_load(new) - connector.result.finished_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert runtime.state.blocks[0, 1].valid_end == 0 - assert runtime.state.loads["r"]["load_id"] == new["load_id"] - assert not runtime.events # A known old submission may legally drain after cancellation. - - -def test_overlapping_request_id_is_reported_and_never_promotes_the_new_load(runtime, connector): - bind(runtime) - old = load_plan(runtime.state, end=64) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(old)) - new = load_plan(runtime.state, end=32) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))( - packet(new, released=["r"], requests={"r": [[(1, 2)]]}) - ) - connector.result.finished_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert "overlapping_load_request_id" in reasons(runtime) - assert runtime.state.blocks[0, 1].valid_end == 0 - - -def test_cancelled_request_can_drain_without_poisoning_later_requests(runtime, connector): - bind(runtime) - meta = load_plan(runtime.state, end=64) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - runtime.wrap_execute(lambda output: None)(packet(released=["r"])) - connector.result.finished_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert not runtime.events - assert not runtime._connector.inflight - runtime.state.bind({"requests": {"next": [[(1, 2)]]}}) - runtime.state.observe([{"group_id": 0, "slots": [128], "block_table": [[1]], "seq_lens": [1]}], ["next"], "forward") - assert runtime.state.blocks[0, 1].valid_end == 1 and not runtime.events - - -def test_completion_metadata_cannot_expand_the_registered_range(runtime): - bind(runtime) - meta = load_plan(runtime.state, end=64) - runtime.state.begin_load(meta) - runtime.state.load({**meta, "token_end": 128}) - assert runtime.state.blocks[0, 1].valid_end == 0 - assert "stale_load_completion" in reasons(runtime) - - -def test_completion_cannot_erase_a_detected_corruption_in_the_same_allocation(runtime, connector): - bind(runtime) - meta = load_plan(runtime.state, end=64) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - runtime.state.observe([{"group_id": 0, "slots": [142]}], [], "dummy") - connector.result.finished_recving = {"r"} - runtime.wrap_execute(lambda output: poll(connector))(packet()) - assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX - assert "load_destination_tainted" in reasons(runtime) - - -@pytest.mark.parametrize("replay", [False, True]) -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_only_explicit_replay_allows_overwriting_a_loaded_tail(runtime, replay, mode): - bind(runtime, computed=128) - meta = load_plan(runtime.state, end=128, start=127) - meta.pop("load_id") - meta.pop("is_async") - item = metadata(torch.tensor([255]), torch.tensor([[1]]), [128]) - context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(count): - pass - - forward = runtime.wrap_forward(runner, model, lambda: context) - runtime.wrap_execute(lambda output: forward(1))(packet(scheduled={"r": (127, 1)}, replay=meta if replay else None)) - if replay: - assert runtime.state.blocks[0, 1].valid_end == 128 and (not runtime.events) - else: - assert "overwrite_valid_prefix" in reasons(runtime) - - -def test_post_forward_async_completion_does_not_legalize_an_earlier_read(runtime, connector): - bind(runtime) - meta = load_plan(runtime.state, end=64) - runtime.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.result.finished_recving = {"r"} - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(count): - poll(connector) - assert runtime.state.blocks[0, 1].valid_end == 0 - - item = metadata(torch.tensor([-1]), torch.tensor([[1]]), [64]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item}) - forward = runtime.wrap_forward(runner, model, lambda: context) - runtime.wrap_execute(lambda output: forward(1))(packet()) - assert "read_unwritten_range" in reasons(runtime) - assert runtime.state.blocks[0, 1].valid_end == 64 - - -@pytest.mark.parametrize("is_async", [False, True]) -@pytest.mark.parametrize("positional", [False, True]) -def test_scheduler_records_loading_bindings_even_without_scheduled_tokens(runtime, is_async, positional): - patch = load("guard_loading_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") - - class Pool: - blocks = [NS(block_id=n) for n in range(4)] - - def get_new_blocks(self, count): - return self.blocks[1 : count + 1] - - class Scheduler: - def __init__(self): - self.vllm_config = NS(additional_config={"kv_guard_config": {"enabled": True}}) - self.pool = Pool() - self.request = NS(request_id="r", num_computed_tokens=0) - self.kv_cache_manager = NS( - block_pool=self.pool, - allocate_slots=self.allocate, - get_block_ids=lambda request: [[1, 2]], - free=lambda request: None, - ) - - def allocate( - self, - request, - num_new_tokens, - num_new_computed_tokens=0, - new_computed_blocks=None, - num_lookahead_tokens=0, - num_external_computed_tokens=0, - delay_cache_blocks=False, - ): - self.pool.get_new_blocks(2) - return object() - - def schedule(self): - if positional: - self.kv_cache_manager.allocate_slots(self.request, 0, 19, None, 0, 173, is_async) - else: - self.kv_cache_manager.allocate_slots( - self.request, 0, 19, num_external_computed_tokens=173, delay_cache_blocks=is_async - ) - return NS( - scheduled_new_reqs=[], - scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), - num_scheduled_tokens={}, - finished_req_ids=set(), - ) - - patch.install_scheduler_guard(Scheduler) - scheduler = Scheduler() - context = scheduler.schedule()._ascend_kv_guard_context - assert context["requests"] == {"r": [[(1, 1), (2, 1)]]} - meta = context["loads"][0] - assert (meta["token_start"], meta["token_end"], meta["is_async"]) == (19, 192, is_async) - assert meta["block_table"] == context["requests"]["r"] - assert context["scheduled_requests"] == {} - runtime.state.bind(context) - assert runtime.state.blocks[0, 1].valid_end == 0 - assert runtime.state.loads["r"]["load_id"] == meta["load_id"] - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_v2_receive_only_completion_then_authorized_last_token_replay(adapter, connector, mode): - runtime = adapter.runtime - bind(runtime) - meta = load_plan(runtime.state, end=128) - adapter.wrap_execute(lambda output: connector.start_load_kv(None))(packet(meta)) - connector.result.finished_recving = {"r"} - adapter.wrap_execute(lambda output: poll(connector))(packet()) - assert runtime.state.blocks[0, 1].valid_end == 128 - replay = {**meta, "token_start": 127} - - def execute(output): - prepare(adapter, ["r"], [255], [[1]], [128], mode) - - adapter.wrap_execute(execute)(packet(scheduled={"r": (127, 1)}, replay=replay)) - assert runtime.state.blocks[0, 1].valid_end == 128 and (not runtime.events) - - -def test_v2_async_submission_can_follow_execute_in_sample_tokens(adapter, connector): - runtime = adapter.runtime - bind(runtime, "running", 1) - bind(runtime, "loading", 2) - meta = load_plan(runtime.state, "loading", end=64) - - def execute(output): - prepare(adapter, ["running"], [128], [[1]], [1]) - return "awaiting sampling" - - assert adapter.wrap_execute(execute)(packet(meta, {"running": (0, 1)})) == "awaiting sampling" - assert runtime._connector.staged == [meta] and (not runtime.events) - connector.start_load_kv(None) - connector.result.finished_recving = {"loading"} - poll(connector) - assert runtime.state.blocks[0, 2].valid_end == 64 - assert not runtime.state.loads and (not runtime.events) - - -def test_v2_unsubmitted_load_is_rejected_at_the_real_post_forward_boundary(adapter, connector): - runtime = adapter.runtime - bind(runtime, "running", 1) - bind(runtime, "loading", 2) - meta = load_plan(runtime.state, "loading", end=64) - - def execute(output): - prepare(adapter, ["running"], [128], [[1]], [1]) - - adapter.wrap_execute(execute)(packet(meta, {"running": (0, 1)})) - assert "load_not_submitted" not in reasons(runtime) - poll(connector) - assert "load_not_submitted" in reasons(runtime) - - -def test_receive_only_binding_is_not_required_in_the_worker_execution_batch(runtime): - waiting = packet(requests={"loading": [[(2, 1)]]}) - runner = NS(requests={}) - assert runtime.wrap_update(runner, lambda output: "unchanged")(waiting) == "unchanged" - assert not runtime.events - - -def test_scheduler_records_explicit_async_tail_replay_without_changing_request(runtime): - patch = load("guard_replay_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") - - class Scheduler: - def __init__(self): - self.vllm_config = NS(additional_config={"kv_guard_config": {"enabled": True}}) - self.request = NS(request_id="r", num_computed_tokens=128) - pool = NS(blocks=[NS(block_id=0), NS(block_id=1)], get_new_blocks=lambda count: [NS(block_id=1)]) - self.kv_cache_manager = NS(block_pool=pool, get_block_ids=lambda request: [[1]], free=lambda request: None) - - def _update_waiting_for_remote_kv(self, request): - request.num_computed_tokens -= 1 - return "unchanged" - - def schedule(self): - self.kv_cache_manager.block_pool.get_new_blocks(1) - assert self._update_waiting_for_remote_kv(self.request) == "unchanged" - return NS( - scheduled_new_reqs=[NS(req_id="r", num_computed_tokens=127)], - scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), - num_scheduled_tokens={"r": 1}, - finished_req_ids=set(), - ) - - patch.install_scheduler_guard(Scheduler) - scheduler = Scheduler() - context = scheduler.schedule()._ascend_kv_guard_context - assert scheduler.request.num_computed_tokens == 127 - assert context["scheduled_requests"] == {"r": (127, 1)} - assert context["recompute_ranges"] == [ - { - "pool_id": context["pool_id"], - "request_id": "r", - "block_table": [[(1, 1)]], - "token_start": 127, - "token_end": 128, - } - ] diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py index 615f22c6ee3a..b0c843c6dfe6 100644 --- a/tests/ut/debug/test_kv_guard_runtime.py +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -1,660 +1,204 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Runner metadata per cache group, independent scheduler ranges and lifecycle.""" +"""One current execution plan per group is compared with an independent mirror.""" -import importlib.util -import logging -import sys -from pathlib import Path -from types import ModuleType from types import SimpleNamespace as NS -from typing import Any, cast import pytest import torch -from . import complete_load - -ROOT = Path(__file__).parents[3] - - -def load(name, path): - spec = importlib.util.spec_from_file_location(name, ROOT / path) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture(name="runtime") -def runtime_fixture(monkeypatch): - package = ModuleType("vllm_ascend") - package.__path__ = [] - debug = ModuleType("vllm_ascend.debug") - debug.__path__ = [] - monkeypatch.setitem(sys.modules, "vllm_ascend", package) - monkeypatch.setitem(sys.modules, "vllm_ascend.debug", debug) - utils = load("guard_config_utils", "vllm_ascend/config_utils.py") - monkeypatch.setitem(sys.modules, "vllm_ascend.config_utils", utils) - options = load("guard_config_subject", "vllm_ascend/debug/kv_guard_config.py") - monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_config", options) - state = load("guard_runtime_state", "vllm_ascend/debug/kv_guard.py") - monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard", state) - connector = load("guard_connector_subject", "vllm_ascend/debug/kv_guard_connector.py") - monkeypatch.setitem(sys.modules, "vllm_ascend.debug.kv_guard_connector", connector) - subject = load("guard_runtime_subject", "vllm_ascend/debug/kv_guard_runtime.py") - logger = ModuleType("vllm.logger") - cast(Any, logger).logger = logging.Logger("guard-test") - vllm = ModuleType("vllm") - vllm.__path__ = [] - monkeypatch.setitem(sys.modules, "vllm", vllm) - monkeypatch.setitem(sys.modules, "vllm.logger", logger) - config = NS(additional_config={"kv_guard_config": {"enabled": True}}) - runtime = subject.get_guard(config) - assert subject.get_guard(config) is runtime - runtime.subject = subject - runtime.state.configure({0: 128}) - runtime.layer_groups = {"layer0": 0, "layer1": 0} - runtime.events = [] - runtime.state.emit = lambda event, **fields: runtime.events.append({"event": event, **fields}) - runtime.gap = lambda reason, **fields: runtime.events.append( - {"event": "guard.coverage_gap", "reason": reason, **fields} - ) - return runtime - - -def bind(runtime, request="r", block=1, computed=0): - runtime.state.bind( - {"pool_id": "pool", "requests": {request: [[(block, 1)]]}, "scheduled_requests": {request: (computed, 0)}} - ) - if computed: - complete_load(runtime.state, request, computed) - - -def metadata(slots, table=None, lengths=None): - return NS(slot_mapping=slots, block_tables=table, seq_lens=lengths, num_actual_tokens=len(slots)) - - -def test_group_audit_shares_read_only_blocks_and_rejects_a_changed_write_allocation(runtime, monkeypatch): - runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) - complete_load(runtime.state, end=128) - runtime._scheduled = {"r": (128, 1)} - item = metadata(torch.tensor([256]), torch.tensor([[1, 2]]), torch.tensor([129])) - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) - guard_type = type(runtime.state) - observe = guard_type.observe - - def observe_with_allocation_change(shadow, *args): - assert shadow.blocks[0, 1] is runtime.state.blocks[0, 1] - assert shadow.blocks[0, 2] is not runtime.state.blocks[0, 2] - observe(shadow, *args) - runtime.state.blocks[0, 2].epoch = 2 - - monkeypatch.setattr(guard_type, "observe", observe_with_allocation_change) - runtime.wrap_forward(runner, lambda count: None, lambda: context)(1) - assert runtime.state.blocks[0, 2].valid_end == 0 - assert any(event.get("reason") == "allocation_changed_during_access_audit" for event in runtime.events) - - -def test_receive_only_step_binds_before_execute_and_cleans_finished(runtime): - output = NS( - _ascend_kv_guard_context={"pool_id": "pool", "requests": {"r": [[(1, 1)]]}, "scheduled_requests": {"r": (0, 0)}} - ) - runtime.wrap_execute(lambda message: None)(output) - complete_load(runtime.state) - assert runtime.state.blocks[0, 1].valid_end == 19 - runtime.wrap_execute(lambda message: None)( - NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": ["r"], "scheduled_requests": {}}) - ) - assert runtime.state.blocks[0, 1].owners == set() - - -def test_disabled_guard_does_not_create_runtime(runtime, monkeypatch): - monkeypatch.setenv("VLLM_ASCEND_KV_GUARD", "1") - monkeypatch.setenv("VLLM_ASCEND_KV_GUARD_ACTION", "raise") - config = NS() - assert runtime.subject.get_guard(config) is None - assert not hasattr(config, "_ascend_kv_guard") - - -@pytest.mark.parametrize("use_v2", [False, True]) -def test_guard_configuration_is_per_engine(runtime, use_v2, monkeypatch): - monkeypatch.setenv("VLLM_ASCEND_KV_GUARD", "0") - configurations = [ - NS(use_v2_model_runner=use_v2, additional_config={"kv_guard_config": options}) - for options in ({"enabled": True}, {"enabled": False}, {"enabled": True, "action": "raise"}) - ] - log_guard, disabled_guard, raise_guard = [runtime.subject.get_guard(config) for config in configurations] - assert log_guard is not None and log_guard.action == "log" - assert disabled_guard is None - assert raise_guard is not None and raise_guard.action == "raise" - assert log_guard is not raise_guard - assert runtime.subject.get_guard(configurations[0]) is log_guard - assert runtime.subject.get_guard(configurations[2]) is raise_guard +from .helpers import message, metadata, observe, reasons -@pytest.mark.parametrize( - "options", [None, [], True, {"enabled": "false"}, {"enabled": 1}, {"action": "stop"}, {"enable": True}] -) -def test_guard_rejects_invalid_configuration(runtime, options): - config = NS(additional_config={"kv_guard_config": options}) - with pytest.raises(ValueError): - runtime.subject.get_guard(config) - assert not hasattr(config, "_ascend_kv_guard") - - -@pytest.mark.parametrize("kernel_size, accepted", [(128, True), (16, False)]) -def test_unadapted_kernel_block_coordinates_are_rejected(runtime, kernel_size, accepted): - impl = type("AscendAttentionBackendImpl", (), {})() - runtime.config = NS( - parallel_config=NS(), - kv_transfer_config=None, - model_config=NS(), - speculative_config=None, - compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), - ) - runner = NS( - input_batch=NS(block_table=NS(block_tables=[NS(block_size=kernel_size)])), - kv_cache_config=NS( - kv_cache_groups=[NS(kv_cache_spec=NS(block_size=128, head_size=64), layer_names=["layer0"])] - ), - ) - if accepted: - runtime.configure(runner) - assert not hasattr(impl, "_kv_guard") - else: - with pytest.raises(ValueError, match="unadapted"): - runtime.configure(runner) - - -@pytest.mark.parametrize("use_v2", [False, True]) -def test_scheduler_preserves_output_and_sends_only_changed_generations(runtime, use_v2): - patch = load("guard_scheduler_subject", "vllm_ascend/patch/platform/patch_kv_guard.py") - - class Pool: - blocks = [NS(block_id=i) for i in range(3)] - null_block = blocks[0] - - def get_new_blocks(self, count): - return [self.blocks[1]] - - class Scheduler: - def __init__(self, vllm_config): - self.vllm_config = vllm_config - self.pool = Pool() - self.mapping = {"r": [[1]]} - self.kv_cache_manager = NS( - block_pool=self.pool, get_block_ids=lambda request: self.mapping[request], free=lambda request: None - ) - self.output = NS( - scheduled_new_reqs=[NS(req_id="r", num_computed_tokens=0)], - scheduled_cached_reqs=NS(req_ids=[], num_computed_tokens=[]), - num_scheduled_tokens={"r": 1}, - finished_req_ids=set(), - ) - - def schedule(self): - return self.output - - patch.install_scheduler_guard(Scheduler) - patch.install_scheduler_guard(Scheduler) - scheduler = Scheduler(NS(use_v2_model_runner=use_v2, additional_config={"kv_guard_config": {"enabled": True}})) - for config in (NS(additional_config=None), NS(additional_config={"kv_guard_config": {"enabled": False}})): - disabled = Scheduler(vllm_config=config) - assert "schedule" not in vars(disabled) - assert "get_new_blocks" not in vars(disabled.pool) - assert disabled.schedule() is disabled.output - assert not hasattr(disabled.output, "_ascend_kv_guard_context") - scheduler.pool.get_new_blocks(1) - assert scheduler.schedule() is scheduler.output - binding = scheduler.output._ascend_kv_guard_context - assert binding["requests"]["r"] == [[(1, 1)]] - assert binding["scheduled_requests"] == {"r": (0, 1)} - scheduler.schedule() - assert scheduler.output._ascend_kv_guard_context["requests"] == {} - scheduler.pool.get_new_blocks(1) - scheduler.schedule() - assert scheduler.output._ascend_kv_guard_context["requests"]["r"] == [[(1, 2)]] - scheduler.kv_cache_manager.free(NS(request_id="r")) - scheduler.schedule() - assert scheduler.output._ascend_kv_guard_context["released_request_ids"] == ["r"] - assert "r" in scheduler.output._ascend_kv_guard_context["requests"] - scheduler.output.finished_req_ids = {"r"} - scheduler.output.num_scheduled_tokens = {} - scheduler.schedule() - assert scheduler.output._ascend_kv_guard_context == { - "pool_id": binding["pool_id"], - "requests": {}, - "released_request_ids": ["r"], - "scheduled_requests": {}, - } - - -@pytest.mark.parametrize("slots, mismatch", [([128, 129], False), ([129, 128], True)]) -def test_full_graph_update_audits_final_metadata_without_guard_callbacks(runtime, monkeypatch, slots, mismatch): - order: list[Any] = [] - monkeypatch.setattr( - torch, - "npu", - NS( - NPUGraph=object, - graph_task_update_begin=lambda *args: order.append("begin"), - graph_task_update_end=lambda *args: order.append("end"), - ), - raising=False, - ) - monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) - utils = ModuleType("vllm_ascend.utils") - cast(Any, utils).weak_ref_tensors = lambda value: value - monkeypatch.setitem(sys.modules, "vllm_ascend.utils", utils) - graph = load("guard_graph_update_subject", "vllm_ascend/compilation/updatable_graph.py") - bind(runtime) - runtime.layer_groups = {"layer0": 0} - runtime._scheduled = {"r": (0, 2)} - item = metadata(torch.tensor([-1, -1]), torch.tensor([[1]]), torch.tensor([0])) - runner = NS(input_batch=NS(req_ids=["r"])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) - - def operation(**kwargs): - order.append(("operation", kwargs["context_lens"].tolist())) - item.slot_mapping[:] = torch.tensor(slots) - item.seq_lens[:] = kwargs["context_lens"] - - task = graph.GraphUpdateTask( - operation=operation, - kwargs={"block_table": item.block_tables, "context_lens": torch.tensor([0])}, - provider=NS(layer_name="layer0"), - provider_index=0, - handle="handle", - event=NS(record=lambda stream: order.append("event")), - ) - bound = task.bind({"context_lens": torch.tensor([2])}) - - def replay(count): - bound.apply("stream") - return "result" - - assert runtime.wrap_forward(runner, replay, lambda: context)(2) == "result" - assert order == ["begin", ("operation", [2]), "end", "event"] - assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == mismatch - assert not any(event["event"] == "guard.coverage_gap" for event in runtime.events) - if not mismatch: - assert runtime.state.blocks[0, 1].valid_end == 2 - - -def test_v1_dummy_wrapper_preserves_decoder_replay_argument(runtime): - calls = [] - runner = NS(input_batch=NS(req_ids=[])) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata=None) - - def forward(count, *, decoder_replay_dummy): - calls.append((count, decoder_replay_dummy)) - return "unchanged" - - wrapped = runtime.wrap_forward(runner, forward, lambda: context) - assert wrapped(8, decoder_replay_dummy=True, _kv_guard_phase="warmup") == "unchanged" - assert calls == [(8, True)] +@pytest.mark.parametrize("full_graph", [False, True]) +def test_prefix_append_is_a_metadata_plan_for_both_runner_modes(runtime, full_graph): + with runtime.execution(NS(), message(new=[("r", ([2, 3],), 3)], counts={"r": 3})): + observe(runtime, ["r"], [11, 12, 13], [[2, 3]], [6], full_graph=full_graph) assert not runtime.events + assert not hasattr(runtime, "write") + assert not hasattr(runtime, "read") -def test_repeated_violations_are_deduplicated_without_hiding_other_blocks(runtime): - messages = [] - runtime.logger = NS( - info=lambda *args: messages.append(("info", args)), - warning=lambda *args: messages.append(("warning", args)), - error=lambda *args: messages.append(("error", args)), - ) - runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=1) - runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=1) - runtime.report("guard.violation", operation="read", reason="read_invalid_block", block_id=2) - assert [level for level, args in messages].count("error") == 2 - - -def test_runner_plan_detects_case4_without_any_operator_hook(runtime): - bind(runtime, computed=19) - item = metadata(torch.tensor([142]), torch.tensor([[1]]), torch.tensor([19])) - runner = NS(input_batch=NS(req_ids=[])) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runtime.wrap_forward(runner, lambda n: "result", lambda: context)(1, _kv_guard_phase="dummy") - violation = next(e for e in runtime.events if e.get("reason") == "dummy_write_to_live_block") - assert violation["evidence"] == "execution_plan" and violation["offsets"] == [14] - assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX - - -def test_final_metadata_is_read_after_forward_parameter_updates_and_copied_once(runtime, monkeypatch): - bind(runtime) - slots, table, lengths = torch.tensor([-1]), torch.tensor([[1]]), torch.tensor([0]) - item = metadata(slots, table, lengths) - runner = NS(input_batch=NS(req_ids=["r"]), kv_caches=object()) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) - copies = [] - original = runtime._list +@pytest.mark.parametrize("lengths,valid", [([1, 0, 0], True), ([1, 1, 0], False), ([], False)]) +def test_full_graph_padding_requires_zero_lengths_after_real_rows(runtime, lengths, valid): + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata({"layer0": metadata([8], [[2], [0], [0]], lengths)}, ["r"], True) + assert ("read_length_mismatch" not in reasons(runtime)) is valid - def record_copy(tensor): - copies.append(tensor.data_ptr()) - return original(tensor) - - monkeypatch.setattr(runtime, "_list", record_copy) - - def replay(n): - assert not copies # No D2H before asynchronous graph parameter updates. - assert runtime.state.blocks[0, 1].valid_end == 0 - slots[0], lengths[0] = 128, 1 - return "result" - - assert runtime.wrap_forward(runner, replay, lambda: context)(1) == "result" - assert runtime.state.blocks[0, 1].valid_end == 1 - assert len(copies) == 3 and len(set(copies)) == 3 - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - -def test_metadata_retains_complete_overwrite_diagnostic(runtime): - bind(runtime, computed=1) - runtime._scheduled = {"r": (0, 1)} - item = metadata(torch.tensor([128]), torch.tensor([[1]]), torch.tensor([1])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) - event = next(event for event in runtime.events if event.get("reason") == "overwrite_valid_prefix") - assert event["offsets"] == [0] and event["valid_end"] == 1 - assert event["evidence"] == "execution_plan" and event["group_id"] == 0 - - -def test_unadapted_ring_cache_is_rejected_instead_of_silently_reducing_checks(runtime): - spec = type("CircularBufferSpec", (), {"block_size": 128, "head_size": 64})() - runtime.config = NS( - parallel_config=NS(), - kv_transfer_config=None, - model_config=NS(is_hybrid=True), - speculative_config=None, - compilation_config=NS(static_forward_context={"ring": NS(kv_cache=object())}), - ) - runner = NS( - input_batch=NS(req_ids=[]), kv_cache_config=NS(kv_cache_groups=[NS(kv_cache_spec=spec, layer_names=["ring"])]) - ) - with pytest.raises(ValueError, match="unadapted cache validity"): - runtime.configure(runner) - - -def test_failed_forward_preserves_exception_and_does_not_commit_planned_writes(runtime): - bind(runtime) - item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - - def model(n): - raise RuntimeError("model failure") - - with pytest.raises(RuntimeError, match="model failure"): - runtime.wrap_forward(runner, model, lambda: context)(1) - assert runtime.state.blocks[0, 1].valid_end == 0 and runtime._phase is None - - -def test_only_needed_block_table_rows_and_columns_are_copied(runtime, monkeypatch): - bind(runtime) - table = torch.full((32, 1024), 99) - table[0, 0] = 1 - item = metadata(torch.tensor([128]), table, [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=["r"])) - shapes = [] - original = runtime._list - def record_shape(value): - shapes.append(tuple(getattr(value, "shape", ()))) - return original(value) +def test_corrupted_runner_table_cannot_change_shadow_plan(runtime): + runner = NS(input_batch=NS(req_ids=["r"], block_table=[[99]])) + with runtime.execution(runner, message(new=[("r", ([2],), 0)], counts={"r": 1})): + observe(runtime, ["r"], [396], [[99]], [1]) + assert "write_target_mismatch" in reasons(runtime) + assert "read_block_mapping_mismatch" in reasons(runtime) + assert runtime.state.requests["r"] == ([2],) - monkeypatch.setattr(runtime, "_list", record_shape) - runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) - assert (1, 1) in shapes and (32, 1024) not in shapes - assert not any(e["event"] == "guard.violation" for e in runtime.events) - - -def test_unknown_backend_is_rejected_without_mutating_it(runtime): - class OtherAttentionImpl: - __slots__ = () - - impl = OtherAttentionImpl() - runtime.layer_groups.clear() - runtime.config = NS( - parallel_config=NS(), - kv_transfer_config=None, - model_config=NS(), - speculative_config=None, - compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), - ) - runner = NS( - input_batch=NS(req_ids=["r"]), - kv_cache_config=NS( - kv_cache_groups=[NS(kv_cache_spec=NS(block_size=128, head_size=64), layer_names=["layer0"])] - ), - ) - with pytest.raises(ValueError, match="unadapted attention implementation"): - runtime.configure(runner) - assert not hasattr(impl, "_kv_guard") - - -def test_execute_and_update_bind_once_without_skipping_receive_only_or_next_step(runtime, monkeypatch): - output = NS(_ascend_kv_guard_context={"requests": {}, "released_request_ids": []}) - runner = NS(requests={}) - bound = [] - original = runtime.state.bind - - def record(binding): - bound.append(binding) - return original(binding) - - monkeypatch.setattr(runtime.state, "bind", record) - update = runtime.wrap_update(runner, lambda value: None) - execute = runtime.wrap_execute(update) - execute(output) - assert len(bound) == 1 and runtime._scheduled is None - execute(output) - assert len(bound) == 2 and runtime._scheduled is None - runtime.wrap_execute(lambda value: None)(output) - assert len(bound) == 3 and runtime._scheduled is None - - -def test_shared_unsupported_metadata_reports_missing_read_coverage_for_each_group(runtime): - runtime.layer_groups = {"layer0": 0, "layer1": 1} - runtime.state.configure({0: 128, 1: 128}) - item = metadata(torch.tensor([-1])) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item, "layer1": item}) - runner = NS(input_batch=NS(req_ids=[])) - runtime.wrap_forward(runner, lambda n: None, lambda: context)(1) - assert { - event.get("group_id") for event in runtime.events if event.get("reason") == "metadata_read_range_missing" - } == {0, 1} +@pytest.mark.parametrize("full_graph", [False, True]) +def test_swapped_requests_and_wrong_offsets_are_detected(runtime, full_graph): + output = message(new=[("a", ([2],), 0), ("b", ([3],), 0)], counts={"a": 1, "b": 1}) + with runtime.execution(NS(), output): + observe(runtime, ["b", "a"], [8, 12], [[3], [2]], [1, 1], full_graph=full_graph) + assert "write_target_mismatch" in reasons(runtime) -def test_raise_mode_stops_on_coverage_gap_before_execution(runtime): - runtime.action = "raise" - runtime.gap = runtime.subject.GuardRuntime.gap.__get__(runtime) - called = [] - with pytest.raises(runtime.subject.KVGuardError, match="scheduler_write_plan_missing"): - runtime.wrap_execute(lambda message: called.append(True))(NS(_ascend_kv_guard_context={"requests": {}})) - assert not called - - -def test_physical_zero_and_copy_preserve_only_proven_validity(runtime): - bind(runtime, "source", 1, 19) - output = NS( - _ascend_kv_guard_context={ - "pool_id": "pool", - "requests": {"dest": [[(2, 1)]]}, - "scheduled_requests": {"dest": (19, 0)}, - }, - new_block_ids_to_zero=[2], - kv_cache_block_copies=[NS(src_block_id=1, dst_block_id=2)], - ) - - def execute(message): - assert runtime.state.blocks[0, 2].valid_end == 0 - runtime.zero_blocks(message.new_block_ids_to_zero) - assert runtime.state.blocks[0, 2].valid_end == 0 - runtime.copy_blocks(message.kv_cache_block_copies) - assert runtime.state.blocks[0, 2].valid_end == 19 - - runtime.wrap_execute(execute)(output) - assert not any(e["event"] == "guard.violation" for e in runtime.events) - runtime.zero_blocks([2]) - assert runtime.state.blocks[0, 2].valid_end == runtime.subject.INVALID_PREFIX - assert any(e.get("reason") == "clear_live_block" for e in runtime.events) - - -def test_copy_to_null_or_unbound_block_cannot_silently_pass(runtime): - bind(runtime, computed=19) - runtime.copy_blocks([NS(src_block_id=1, dst_block_id=0), NS(src_block_id=1, dst_block_id=99)]) - assert any(e.get("reason") == "write_null_block" and e["operation"] == "copy" for e in runtime.events) - assert any(e.get("reason") == "copy_destination_unbound" and e["block_id"] == 99 for e in runtime.events) - - -def test_scheduler_optional_zero_and_copy_fields_accept_none(runtime): - output = NS( - _ascend_kv_guard_context={"requests": {}, "scheduled_requests": {}}, - new_block_ids_to_zero=None, - kv_cache_block_copies=None, - ) - assert runtime.wrap_execute(lambda message: "executed")(output) == "executed" - assert not any(e["event"] in ("guard.violation", "guard.coverage_gap") for e in runtime.events) - - -def run_plan(runtime, slots, table, lengths, *, ids=("r",), mode="FULL", phase="forward", forward=None): - item = metadata(torch.tensor(slots), torch.tensor(table), torch.tensor(lengths)) - context = NS(cudagraph_runtime_mode=mode, attn_metadata=dict.fromkeys(runtime.layer_groups, item)) - runner = NS(input_batch=NS(req_ids=list(ids))) - return runtime.wrap_forward(runner, forward or (lambda count: None), lambda: context)( - len(slots), _kv_guard_phase=phase - ) - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_group_metadata_is_checked_once_without_attention_hooks(runtime, monkeypatch, mode): - bind(runtime) - runtime._scheduled = {"r": (0, 1)} - calls = [] - original = runtime.subject.KVGuard.observe - - def checked(state, groups, requests, phase): - calls.append((len(groups), list(requests))) - return original(state, groups, requests, phase) - - monkeypatch.setattr(runtime.subject.KVGuard, "observe", checked) - run_plan(runtime, [128], [[1]], [1], mode=mode) - assert calls == [(1, ["r"])] - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not runtime.events - assert set(runtime._metadata) == {0} - assert not hasattr(runtime, "write") and not hasattr(runtime, "read") - assert not hasattr(runtime, "_accesses") - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) @pytest.mark.parametrize( - "slots,table,lengths,reason", + "slots,tables,lengths,reason", [ - ([128, 129], [[1]], [2], None), - ([129, 128], [[1]], [2], "write_target_mismatch"), - ([128, 128], [[1]], [2], "write_target_mismatch"), - ([128, 129], [[2]], [2], "read_target_mismatch"), - ([128, 129], [[1]], [1], None), - ([128, 129], [[1]], [3], "read_length_mismatch"), + ([10], [[2]], [2], "write_target_mismatch"), + ([9], [[3]], [2], "read_block_mapping_mismatch"), + ([9], [[2]], [1], "read_length_mismatch"), + ([9], [[2]], [3], "read_length_mismatch"), + ([], [[2]], [2], "write_tokens_not_observed"), ], ) -def test_metadata_constraints_share_one_checker(runtime, mode, slots, table, lengths, reason): - bind(runtime) - runtime._scheduled = {"r": (0, 2)} - run_plan(runtime, slots, table, lengths, mode=mode) - violations = [event for event in runtime.events if event["event"] == "guard.violation"] - assert any(event.get("reason") == reason for event in violations) if reason else not violations - assert all(event["evidence"] == "execution_plan" and "layer" not in event for event in violations) - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -@pytest.mark.parametrize("swapped", [False, True]) -def test_group_plan_checks_actual_request_order_and_cross_request_slots(runtime, mode, swapped): - bind(runtime, "A", 1) - bind(runtime, "B", 2) - runtime._scheduled = {"A": (0, 1), "B": (0, 1)} - run_plan(runtime, [128, 256] if swapped else [256, 128], [[2], [1]], [1, 1], ids=("B", "A"), mode=mode) - assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == swapped - - -@pytest.mark.parametrize("slots,mismatch", [([147, 148], False), ([148, 147], True), ([147, -1], True)]) -def test_group_metadata_distinguishes_exact_token_positions_and_padding(runtime, slots, mismatch): - bind(runtime, computed=19) - runtime._scheduled = {"r": (19, 2)} - run_plan(runtime, slots, [[1]], [21]) - assert any(event.get("reason") == "write_target_mismatch" for event in runtime.events) == mismatch - - -def test_truncated_group_write_cannot_advance_validity(runtime): - bind(runtime, computed=19) - runtime._scheduled = {"r": (19, 2)} - run_plan(runtime, [147], [[1]], [19]) - assert any(event.get("reason") == "write_tokens_not_observed" for event in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == 19 - - -def test_computed_range_does_not_initialize_unwritten_blocks(runtime): - bind(runtime) - runtime._scheduled = {"r": (19, 1)} - run_plan(runtime, [147], [[1]], [20]) - assert any(event.get("reason") == "noncontiguous_write" for event in runtime.events) - assert runtime.state.blocks[0, 1].valid_end == runtime.subject.INVALID_PREFIX - - -def test_read_block_order_is_checked_within_one_request(runtime): - runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}}) - complete_load(runtime.state, "r", 254) - runtime._scheduled = {"r": (254, 1)} - run_plan(runtime, [382], [[2, 1]], [255]) - assert any(event.get("reason") == "read_target_mismatch" for event in runtime.events) - - -def test_shared_prefix_and_distinct_tail_blocks_remain_legal(runtime): - runtime.state.bind({"pool_id": "pool", "requests": {"A": [[(1, 1), (2, 1)]], "B": [[(1, 1), (3, 1)]]}}) - complete_load(runtime.state, "A", 128) - runtime._scheduled = {"A": (128, 1), "B": (128, 1)} - run_plan(runtime, [384, 256], [[1, 3], [1, 2]], [129, 129], ids=("B", "A")) +def test_exact_slot_read_table_and_read_length(runtime, slots, tables, lengths, reason): + with runtime.execution(NS(), message(new=[("r", ([2],), 1)], counts={"r": 1})): + observe(runtime, ["r"], slots, tables, lengths) + assert reason in reasons(runtime) + + +def test_same_group_uses_only_first_known_non_none_metadata(runtime): + runtime.layer_groups = {"layer0": 0, "layer1": 0, "layer2": 0} + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata( + {"unknown": object(), "layer0": None, "layer1": metadata([8], [[2]], [1]), "layer2": object()}, + ["r"], + ) assert not runtime.events -def test_dummy_metadata_write_still_taints_a_live_block(runtime): - bind(runtime, computed=19) - run_plan(runtime, [142], [[1]], [19], ids=(), mode="NONE", phase="dummy") - assert any(event.get("reason") == "dummy_write_to_live_block" for event in runtime.events) - run_plan(runtime, [-1], [[1]], [19], mode="NONE") - assert any(event.get("reason") == "read_invalid_block" for event in runtime.events) +def test_first_selected_invalid_metadata_does_not_fall_back_to_later_good_layer(runtime): + runtime.layer_groups = {"layer0": 0, "layer1": 0} + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata({"layer0": NS(), "layer1": metadata([8], [[2]], [1])}, ["r"]) + assert "metadata_write_mapping_missing" in reasons(runtime) + assert "metadata_read_range_missing" in reasons(runtime) + + +def test_other_layers_of_selected_group_do_not_require_metadata(runtime): + runtime.layer_groups = {"layer0": 0, "layer1": 0} + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + observe(runtime, ["r"], [8], [[2]], [1]) + assert not runtime.events -def test_only_group_coverage_is_required_not_every_layer(runtime): - bind(runtime) - item = metadata(torch.tensor([128]), torch.tensor([[1]]), [1]) - context = NS(cudagraph_runtime_mode="NONE", attn_metadata={"layer1": item}) - runtime.wrap_forward(NS(input_batch=NS(req_ids=["r"])), lambda count: None, lambda: context)(1) - assert not runtime.events and runtime.state.blocks[0, 1].valid_end == 1 +def test_each_group_gets_its_own_expected_block_coordinates(runtime): + runtime.state.configure({0: 4, 1: 8}) + runtime.layer_groups = {"layer0": 0, "layer1": 1} + output = message(new=[("r", ([2], [5]), 1)], counts={"r": 1}) + with runtime.execution(NS(), output): + runtime.observe_metadata({"layer0": metadata([9], [[2]], [2]), "layer1": metadata([41], [[5]], [2])}, ["r"]) + assert not runtime.events -def test_missing_cache_group_is_an_explicit_coverage_gap(runtime): - runtime.state.configure({0: 128, 1: 128}) +def test_missing_group_reports_only_that_group(runtime): + runtime.state.configure({0: 4, 1: 8}) runtime.layer_groups = {"layer0": 0, "layer1": 1} - item = metadata(torch.tensor([-1]), torch.tensor([[0]]), [0]) - context = NS(cudagraph_runtime_mode="FULL", attn_metadata={"layer0": item}) - runtime.wrap_forward(NS(input_batch=NS(req_ids=[])), lambda count: None, lambda: context)(1) - assert any( - event.get("reason") == "metadata_write_mapping_missing" and event.get("group_id") == 1 - for event in runtime.events - ) + with runtime.execution(NS(), message(new=[("r", ([2], [5]), 0)], counts={"r": 1})): + observe(runtime, ["r"], [8], [[2]], [1]) + assert runtime.events == [{"event": "guard.coverage_gap", "reason": "metadata_group_missing", "group_id": 1}] + + +def test_metadata_is_read_once_per_group_only_after_final_parameters_change(runtime, monkeypatch): + runtime.layer_groups = {"layer0": 0, "layer1": 0} + current = metadata([99], [[2]], [1]) + reads = [] + original = runtime._list + + def recorded(value): + reads.append(value) + return original(value) + + monkeypatch.setattr(runtime, "_list", recorded) + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata({"layer0": current, "layer1": current}, ["r"], True) + assert reads == [] + current.slot_mapping.copy_(torch.tensor([8])) + assert sum(value.data_ptr() == current.slot_mapping.data_ptr() for value in reads) == 1 + assert not runtime.events + + +def test_only_actual_slot_prefix_is_materialized_from_padded_buffer(runtime, monkeypatch): + current = metadata([8, 999, 999, 999], [[2]], [1]) + current.num_actual_tokens = 1 + materialized = [] + original = runtime._list + + def recorded(value): + if value.data_ptr() == current.slot_mapping.data_ptr(): + materialized.append(value.numel()) + return original(value) + + monkeypatch.setattr(runtime, "_list", recorded) + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata({"layer0": current}, ["r"], True) + assert materialized == [1] + assert not runtime.events + + +def test_full_graph_requires_fresh_metadata_every_step(runtime): + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + observe(runtime, ["r"], [8], [[2]], [1], full_graph=True) + with runtime.execution(NS(input_batch=NS(req_ids=["r"])), message(cached=[("r", None, 1)], counts={"r": 1})): + pass + assert "metadata_group_missing" in reasons(runtime) + + +@pytest.mark.parametrize("batch", [[], ["r", "r"], ["other"]]) +def test_batch_identity_mismatch_is_a_coverage_gap(runtime, batch): + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + observe(runtime, batch, [8], [[2]], [1]) + assert "batch_request_mismatch" in reasons(runtime) + + +def test_no_token_step_can_release_without_a_forward(runtime): + runtime.state.requests = {"r": ([2],)} + with runtime.execution(NS(), message(finished=["r"])): + pass + assert runtime.state.requests == {} + assert not runtime.events + + +def test_no_token_step_does_not_accept_an_active_metadata_plan(runtime): + with runtime.execution(NS(), message()): + observe(runtime, [], [8], [[2]], [1]) + assert "unscheduled_attention_plan" in reasons(runtime) + + +def test_native_exception_preserved_without_audit_and_metadata_references_cleared(runtime): + with ( + pytest.raises(RuntimeError, match="native failure"), + runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})), + ): + observe(runtime, ["r"], [99], [[2]], [1]) + raise RuntimeError("native failure") + assert not runtime.events + assert runtime._phase is None + assert not runtime._metadata -def test_raise_mode_stops_results_and_latches_metadata_failure(runtime): +def test_raise_is_synchronous_and_sticky_without_modifying_output_getter(runtime, subject): runtime.action = "raise" - runtime.state.emit = runtime.report - bind(runtime, computed=19) - runtime._scheduled = {"r": (19, 1)} - with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): - run_plan(runtime, [148], [[1]], [20]) - with pytest.raises(runtime.subject.KVGuardError): - runtime.wrap_execute(lambda output: None)(NS()) + output = NS(get_output=lambda: "result") + original = output.get_output + with ( + pytest.raises(subject.KVGuardError, match="write_target_mismatch"), + runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})), + ): + observe(runtime, ["r"], [9], [[2]], [1], full_graph=True) + assert output.get_output is original + with pytest.raises(subject.KVGuardError, match="restart the worker"): + runtime.check_failure() + + +def test_log_reports_group_execution_plan_evidence(runtime, caplog): + with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): + observe(runtime, ["r"], [9], [[2]], [1]) + result = "native result" + runtime.check_failure() + assert result == "native result" + assert "group_id=0" in caplog.text + assert "evidence=execution_plan" in caplog.text + assert "layer=" not in caplog.text diff --git a/tests/ut/debug/test_kv_guard_scope.py b/tests/ut/debug/test_kv_guard_scope.py index da20203254bb..1f5e8fc0558e 100644 --- a/tests/ut/debug/test_kv_guard_scope.py +++ b/tests/ut/debug/test_kv_guard_scope.py @@ -1,55 +1,113 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Outputs without a guaranteed sole consumer retain synchronous checking.""" +"""Per-runner ownership, explicit unsupported modes, and capture isolation.""" from types import SimpleNamespace as NS import pytest -from .test_kv_guard_runtime import runtime_fixture as runtime_fixture -from .test_kv_guard_v2 import adapter_fixture as adapter_fixture -from .test_kv_guard_v2 import message, prepare - - -@pytest.mark.parametrize("pooling,tp,pp", [(False, 1, 1), (True, 1, 1), (False, 2, 1), (False, 1, 2)]) -def test_only_single_worker_generation_can_defer_checks_to_sample_output(adapter, pooling, tp, pp): - runtime = adapter.runtime - runtime.config.scheduler_config = NS(async_scheduling=True) - adapter.runner.sample_tokens = lambda grammar: None - adapter.runner.is_pooling_model = pooling - runtime.config.parallel_config.tensor_parallel_size = tp - runtime.config.parallel_config.pipeline_parallel_size = pp - runtime.configure(adapter.runner) - assert runtime._defer_audit is (not pooling and tp == 1 and pp == 1) - - -@pytest.mark.parametrize("pooling,tp,pp", [(True, 1, 1), (False, 2, 1), (False, 1, 2)]) -def test_failure_is_detected_before_pool_or_nonreply_output_can_bypass_the_getter(adapter, pooling, tp, pp): - runtime = adapter.runtime - runtime.config.scheduler_config = NS(async_scheduling=True) - adapter.runner.sample_tokens = lambda grammar: None - adapter.runner.is_pooling_model = pooling - runtime.config.parallel_config.tensor_parallel_size = tp - runtime.config.parallel_config.pipeline_parallel_size = pp - runtime.action = "raise" - runtime.state.emit = runtime.report - runtime.configure(adapter.runner) - with pytest.raises(runtime.subject.KVGuardError, match="write_target_mismatch"): - adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [129], [[1]], [1], "FULL"))(message()) - assert runtime._pending is None - - -@pytest.mark.parametrize("missing", ["pooling", "tp", "pp"]) -def test_unknown_output_layout_keeps_synchronous_validation(adapter, missing): - runtime = adapter.runtime - runtime.config.scheduler_config = NS(async_scheduling=True) - adapter.runner.sample_tokens = lambda grammar: None - adapter.runner.is_pooling_model = False - parallel = runtime.config.parallel_config - parallel.tensor_parallel_size = parallel.pipeline_parallel_size = 1 - if missing == "pooling": - del adapter.runner.is_pooling_model +from .helpers import metadata, reasons + + +def runner_for(config, impl_name="AscendAttentionBackendImpl"): + impl = type(impl_name, (), {})() + config.compilation_config.static_forward_context = {"layer0": NS(impl=impl)} + spec = NS(block_size=4, head_size=16) + return NS(kv_cache_config=NS(kv_cache_groups=[NS(kv_cache_spec=spec, layer_names=["layer0"])])), impl + + +def test_disabled_has_no_runtime_or_config_side_effect(subject, config): + config.additional_config = {} + assert subject.create_guard(config) is None + assert not hasattr(config, "_ascend_kv_guard") + + +def test_same_config_creates_independent_per_runner_state(subject, config): + first, second = subject.create_guard(config), subject.create_guard(config) + first.state.requests["r"] = ([2],) + assert first is not second + assert second.state.requests == {} + assert not hasattr(config, "_ascend_kv_guard") + + +@pytest.mark.parametrize( + "options", [{"enabled": "yes"}, {"enabled": 1}, {"action": "ignore"}, {"enabled": True, "typo": 1}] +) +def test_config_rejects_invalid_options(subject, config, options): + config.additional_config = {"kv_guard_config": options} + with pytest.raises(ValueError): + subject.create_guard(config) + + +@pytest.mark.parametrize("impl_name", ["AscendAttentionBackendImpl", "AscendMLAImpl"]) +def test_attention_only_validated_and_declared_model_state_bound(runtime, config, impl_name): + runner, impl = runner_for(config, impl_name) + runner.model_state = NS(kv_guard=None) + runner.block_tables = NS(block_sizes=[4], kernel_block_sizes=[4]) + runtime.configure(runner) + assert not hasattr(impl, "kv_guard") + assert runner.model_state.kv_guard is runtime + assert runtime.layer_groups == {"layer0": 0} + + +def test_empty_cache_groups_cannot_claim_guard_enabled(runtime): + with pytest.raises(ValueError, match="no supported attention layers"): + runtime.configure(NS(kv_cache_config=NS(kv_cache_groups=[]))) + + +@pytest.mark.parametrize( + "mode", ["connector", "cp", "dbo", "speculation", "encoder", "sliding", "compressed", "coordinates", "model_state"] +) +def test_unsupported_modes_fail_before_execution(runtime, config, mode): + runner, _ = runner_for(config) + spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec + if mode == "connector": + config.kv_transfer_config = NS(kv_connector="MooncakeConnector") + elif mode == "cp": + config.parallel_config.decode_context_parallel_size = 2 + elif mode == "dbo": + config.parallel_config.enable_dbo = True + elif mode == "speculation": + config.speculative_config = NS() + elif mode == "encoder": + config.model_config.is_encoder_decoder = True + elif mode == "sliding": + spec.sliding_window = 16 + elif mode == "compressed": + spec.compress_ratio = 4 + elif mode == "coordinates": + runner.block_tables = NS(block_sizes=[4], kernel_block_sizes=[2]) else: - delattr(parallel, "tensor_parallel_size" if missing == "tp" else "pipeline_parallel_size") - runtime.configure(adapter.runner) - assert not runtime._defer_audit + runner.model_state = NS() + with pytest.raises(ValueError, match="unsupported"): + runtime.configure(runner) + + +@pytest.mark.parametrize("full_graph", [False, True]) +def test_dummy_live_write_uses_current_group_plan(runtime, full_graph): + runtime.state.requests = {"r": ([2],)} + with runtime.dummy_scope(NS()): + runtime.observe_metadata({"layer0": metadata([9], [[2]], [1])}, [], full_graph) + assert "dummy_write_to_live_block" in reasons(runtime) + assert runtime.state.requests == {"r": ([2],)} + + +def test_dummy_scratch_is_allowed(runtime): + runtime.state.requests = {"r": ([2],)} + with runtime.dummy_scope(NS()): + runtime.observe_metadata({"layer0": metadata([-1, 0, 12], [[0]], [0])}, []) + assert not runtime.events + + +@pytest.mark.parametrize("phase", ["warmup", "capture"]) +def test_startup_capture_never_materializes_tensor_metadata(runtime, monkeypatch, phase): + monkeypatch.setattr(runtime, "_list", lambda _: pytest.fail("CPU metadata read during capture")) + with runtime.dummy_scope(NS(), phase=phase): + runtime.observe_metadata({"layer0": metadata([8], [[2]], [1])}, [], True) + assert not runtime.events + + +def test_live_recapture_is_explicitly_rejected(runtime): + runtime.state.requests = {"r": ([2],)} + with pytest.raises(ValueError, match="no live requests"): + runtime.check_capture() diff --git a/tests/ut/debug/test_kv_guard_v2.py b/tests/ut/debug/test_kv_guard_v2.py deleted file mode 100644 index 510def7d25ea..000000000000 --- a/tests/ut/debug/test_kv_guard_v2.py +++ /dev/null @@ -1,357 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""MRV2 batch ordering, graph updates, dummy writes and lifecycle integration.""" - -from types import SimpleNamespace as NS -from typing import Any - -import pytest -import torch - -from . import complete_load, load_plan -from .test_kv_guard_runtime import bind, metadata -from .test_kv_guard_runtime import runtime_fixture as runtime_fixture - - -@pytest.fixture(name="adapter") -def adapter_fixture(runtime): - subject = runtime.subject - impl = type("AscendAttentionBackendImpl", (), {})() - runtime.config = NS( - use_v2_model_runner=True, - additional_config={"kv_guard_config": {"enabled": True}}, - parallel_config=NS(), - model_config=NS(), - speculative_config=None, - kv_transfer_config=None, - compilation_config=NS(static_forward_context={"layer0": NS(impl=impl)}), - _ascend_kv_guard=runtime, - ) - runner: Any = NS( - vllm_config=runtime.config, - block_tables=NS(block_sizes=[128], kernel_block_sizes=[128]), - kv_cache_config=NS( - kv_cache_groups=[NS(layer_names=["layer0"], kv_cache_spec=NS(block_size=128, head_size=64))] - ), - model_state=NS(prepare_attn=lambda batch, mode: runner.next_metadata), - next_metadata={}, - ) - runtime.layer_groups.clear() - result = subject._V2RunnerAdapter(runner, runtime) - runtime._runner_adapter = result - runtime.configure(runner) - result.subject = subject - return result - - -def message(requests=None, scheduled=None, released=()): - scheduled = {"r": (0, 1)} if scheduled is None else scheduled - return NS( - total_num_scheduled_tokens=sum(count for _, count in scheduled.values()), - _ascend_kv_guard_context={ - "pool_id": "pool", - "requests": {"r": [[(1, 1)]]} if requests is None else requests, - "scheduled_requests": scheduled, - "released_request_ids": list(released), - }, - ) - - -def prepare(adapter, ids, slots, tables, lengths, mode="NONE"): - adapter.runner.next_metadata = {"layer0": metadata(torch.tensor(slots), torch.tensor(tables), lengths)} - return adapter.runner.model_state.prepare_attn(NS(req_ids=ids, num_reqs=len(ids)), mode) - - -def reasons(runtime): - return {event.get("reason") for event in runtime.events} - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_sorted_v2_batch_uses_batch_ids_instead_of_request_state_indices(adapter, mode): - runtime = adapter.runtime - packet = message(requests={"A": [[(1, 1)]], "B": [[(2, 1)]]}, scheduled={"A": (0, 1), "B": (0, 2)}) - adapter.runner.req_states = NS(req_id_to_index={"A": 9, "B": 3}) - - def execute(output): - prepare(adapter, ["B", "A"], [256, 257, 128], [[2], [1]], [2, 1], mode) - return "unchanged" - - assert adapter.wrap_execute(execute)(packet) == "unchanged" - assert not runtime.events - assert runtime.state.blocks[0, 2].valid_end == 2 - assert runtime.state.blocks[0, 1].valid_end == 1 - assert runtime._phase is None and runtime._scheduled is None - - -@pytest.mark.parametrize("mode", ["NONE", "FULL"]) -def test_v2_cross_request_swaps_are_not_legalized_by_allocated_blocks(adapter, mode): - packet = message(requests={"A": [[(1, 1)]], "B": [[(2, 1)]]}, scheduled={"A": (0, 1), "B": (0, 1)}) - adapter.wrap_execute(lambda output: prepare(adapter, ["A", "B"], [256, 128], [[1], [2]], [1, 1], mode))(packet) - assert "write_target_mismatch" in reasons(adapter.runtime) - - -def test_v2_full_replay_observes_current_metadata_after_parameter_update(adapter): - runtime = adapter.runtime - packet = message(requests={"r": [[(2, 1)]]}) - - def execute(output): - item = prepare(adapter, ["r"], [128, -1], [[1], [0]], [1, 0], "FULL")["layer0"] - item.slot_mapping[0] = 256 - item.block_tables[0, 0] = 2 - return "replayed" - - assert adapter.wrap_execute(execute)(packet) == "replayed" - assert not runtime.events - assert runtime.state.blocks[0, 2].valid_end == 1 - - -def test_v2_full_cannot_reuse_last_steps_metadata(adapter): - forward = adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [128], [[1]], [1], "FULL")) - forward(message()) - adapter.runner.model_state.prepare_attn = lambda batch, mode: None - adapter.wrap_execute(lambda output: None)(message(requests={}, scheduled={"r": (1, 1)})) - assert "v2_attention_metadata_missing" in reasons(adapter.runtime) - assert adapter.runtime.state.blocks[0, 1].valid_end == 1 - - -@pytest.mark.parametrize("ids", [["r", "r"], [], ["other"]]) -def test_v2_missing_duplicate_or_unscheduled_batch_ids_are_coverage_gaps(adapter, ids): - adapter.wrap_execute(lambda output: prepare(adapter, ids, [128], [[1]], [1]))(message()) - assert "v2_batch_request_mismatch" in reasons(adapter.runtime) - assert adapter.runtime.state.blocks[0, 1].valid_end <= 0 - - -def test_v2_receive_only_and_finished_steps_do_not_require_attention_metadata(adapter): - sentinel = object() - execute = adapter.wrap_execute(lambda output: sentinel) - assert execute(message(scheduled={})) is sentinel - assert "r" in adapter.runtime.state.requests - assert execute(message(requests={}, scheduled={}, released=["r"])) is sentinel - assert adapter.runtime.state.requests == {} - assert not adapter.runtime.events - - -@pytest.mark.parametrize("profile", [False, True]) -def test_v2_dummy_or_profile_writes_still_detect_case4_without_a_scheduler_binding(adapter, profile): - runtime = adapter.runtime - bind(runtime, computed=19) - packet = NS(total_num_scheduled_tokens=1) - adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [142], [[1]], [19]))( - packet, dummy_run=True, is_profile=profile - ) - assert "dummy_write_to_live_block" in reasons(runtime) - assert "scheduler_binding_missing" not in reasons(runtime) - assert runtime.state.blocks[0, 1].valid_end < 0 - adapter.wrap_execute(lambda output: prepare(adapter, ["r"], [147], [[1]], [20]))( - message(requests={}, scheduled={"r": (19, 1)}) - ) - assert "read_invalid_block" in reasons(runtime) - - -def test_v2_dummy_scratch_reads_are_ignored_but_null_writes_are_reported(adapter): - packet = NS(total_num_scheduled_tokens=1) - adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [-1], [[0]], [1]))(packet, dummy_run=True) - assert not adapter.runtime.events - adapter.wrap_execute(lambda output, **kwargs: prepare(adapter, ["dummy"], [0], [[0]], [1]))(packet, dummy_run=True) - assert "write_null_block" in reasons(adapter.runtime) - - -def test_v2_profile_without_attention_and_positional_flags_preserve_arguments(adapter): - token = object() - calls = [] - - def execute(output, *args): - calls.append((output, args)) - return token - - packet = NS(total_num_scheduled_tokens=16) - assert adapter.wrap_execute(execute)(packet, token, True, True, True) is token - assert calls == [(packet, (token, True, True, True))] - assert not adapter.runtime.events - - -def test_v2_failed_forward_does_not_commit_writes_or_leak_scope(adapter): - def execute(output): - prepare(adapter, ["r"], [128], [[1]], [1]) - raise RuntimeError("original failure") - - with pytest.raises(RuntimeError, match="original failure"): - adapter.wrap_execute(execute)(message()) - assert adapter.runtime.state.blocks[0, 1].valid_end == 0 - assert adapter.runtime._phase is None and adapter.runtime._scheduled is None - - -def test_v2_clears_then_copies_after_physical_update(adapter): - runtime = adapter.runtime - runtime.state.bind({"pool_id": "pool", "requests": {"r": [[(1, 1), (2, 1)]]}, "scheduled_requests": {"r": (4, 0)}}) - complete_load(runtime.state, end=4) - packet = NS(new_block_ids_to_zero=[2], kv_cache_block_copies=[NS(src_block_id=1, dst_block_id=2)]) - order = [] - clear, copy = runtime.zero_blocks, runtime.copy_blocks - - def cleared(blocks): - order.append("clear") - clear(blocks) - - def copied(blocks): - order.append("copy") - copy(blocks) - - runtime.zero_blocks = cleared - runtime.copy_blocks = copied - adapter.wrap_update(lambda output: order.append("physical"))(packet) - assert order == ["physical", "clear", "copy"] - assert runtime.state.blocks[0, 2].valid_end == 4 - adapter.wrap_update(lambda output: None)(NS(new_block_ids_to_zero=None, kv_cache_block_copies=None)) - - -def test_v2_startup_capture_does_not_copy_metadata_inside_capture(adapter): - calls = [] - adapter.runtime.observe_metadata = lambda metadata: calls.append(metadata) - capture = adapter.wrap_capture(lambda: prepare(adapter, ["dummy"], [-1], [[0]], [1])) - capture() - assert calls == [] and not adapter.runtime.events - - -@pytest.mark.parametrize("released", [False, True]) -def test_v2_live_or_prefix_preserving_recapture_is_rejected(adapter, released): - bind(adapter.runtime, computed=19) - if released: - adapter.runtime.state.bind({"released_request_ids": ["r"]}) - with pytest.raises(ValueError, match="recapture is unsupported"): - adapter.wrap_capture(lambda: pytest.fail("capture must not run"))() - - -def test_v2_kernel_allocator_coordinate_mismatch_is_rejected(adapter): - adapter.runner.block_tables.kernel_block_sizes = [64] - with pytest.raises(ValueError, match="allocator/kernel"): - adapter.configure() - - -def test_v2_install_disabled_does_not_wrap_methods(adapter): - assert adapter.subject.install_guard(NS(vllm_config=NS(use_v2_model_runner=True))) is None - - -def enable_raise(runtime): - runtime.action = "raise" - append = runtime.state.emit - - def report(event, **fields): - append(event, **fields) - runtime.report(event, **fields) - - runtime.state.emit = report - runtime.gap = lambda reason, **fields: report("guard.coverage_gap", reason=reason, **fields) - - -def test_v2_raise_stops_output_and_latches_before_the_next_model_call(adapter): - enable_raise(adapter.runtime) - calls = [] - - def execute(output): - prepare(adapter, ["r"], [129], [[1]], [1]) - calls.append(True) - return "must not escape" - - forward = adapter.wrap_execute(execute) - for _ in range(2): - with pytest.raises(adapter.runtime.subject.KVGuardError): - forward(message()) - assert calls == [True] - - -def test_v2_missing_binding_and_batch_gap_stop_before_kernel_submission(adapter): - enable_raise(adapter.runtime) - calls = [] - with pytest.raises(adapter.runtime.subject.KVGuardError, match="scheduler_binding_missing"): - adapter.wrap_execute(lambda output: calls.append(True))(NS(total_num_scheduled_tokens=1)) - assert calls == [] - - -def test_v2_wrong_batch_ids_raise_before_the_model_runs(adapter): - enable_raise(adapter.runtime) - calls = [] - - def execute(output): - prepare(adapter, ["unknown"], [128], [[1]], [1]) - calls.append(True) - - with pytest.raises(adapter.runtime.subject.KVGuardError, match="v2_batch_request_mismatch"): - adapter.wrap_execute(execute)(message()) - assert calls == [] - - -def test_v2_install_wraps_real_runner_boundaries_and_reinitializes_model_state(adapter): - runner = adapter.runner - runner.execute_model = lambda output: prepare(adapter, ["r"], [128], [[1]], [1]) - runner.update_requests = lambda output: None - runner.capture_model = lambda: "capture" - installed = adapter.subject.install_guard(runner) - assert installed is not None - # Profiling can replace the model-state object before real cache setup. - runner.model_state = NS(prepare_attn=lambda batch, mode: runner.next_metadata) - installed.configure(runner) - wrapped = runner.model_state.prepare_attn - installed.configure(runner) - assert runner.model_state.prepare_attn is wrapped - assert runner.capture_model() == "capture" - runner.execute_model(message()) - assert adapter.runtime.state.blocks[0, 1].valid_end == 1 - assert not adapter.runtime.events - - -def test_v2_keyword_calls_preserve_upstream_parameter_names(adapter): - def execute(output): - adapter.runner.next_metadata = {"layer0": metadata(torch.tensor([128]), torch.tensor([[1]]), [1])} - return adapter.runner.model_state.prepare_attn(input_batch=NS(req_ids=["r"], num_reqs=1), cudagraph_mode="NONE") - - adapter.wrap_execute(execute)(scheduler_output=message()) - adapter.wrap_update(lambda output: "unchanged")(scheduler_output=NS()) - - -@pytest.mark.parametrize("async_output", [False, True]) -def test_v2_output_is_guarded_after_the_forward_scope_submits_its_audit(adapter, async_output): - runtime = adapter.runtime - runtime._defer_audit = True - consumed = [] - result = NS(get_output=lambda: consumed.append(True)) if async_output else "result" - - def execute(output): - prepare(adapter, ["r"], [128], [[1]], [1], "FULL") - return result - - output = adapter.wrap_execute(execute)(message()) - if async_output: - assert runtime.state.blocks[0, 1].valid_end == 0 - output.get_output() - assert consumed == [True] - assert runtime.state.blocks[0, 1].valid_end == 1 - assert not runtime.events - - -@pytest.mark.parametrize("flag", ["dummy_run", "is_profile"]) -def test_v2_synthetic_execution_drains_the_real_pending_audit_first(adapter, flag): - runtime = adapter.runtime - runtime._defer_audit = True - - def real(output): - prepare(adapter, ["r"], [128], [[1]], [1], "FULL") - - adapter.wrap_execute(real)(message()) - assert runtime.state.blocks[0, 1].valid_end == 0 - - def synthetic(output, **kwargs): - assert runtime.state.blocks[0, 1].valid_end == 1 - prepare(adapter, ["dummy"], [-1], [[0]], [1], "FULL") - - adapter.wrap_execute(synthetic)(message(), **{flag: True}) - assert not runtime.events - assert adapter.runtime.state.blocks[0, 1].valid_end == 1 - assert not adapter.runtime.events - - -def test_v2_pending_receive_also_blocks_graph_recapture(adapter): - adapter.runtime.state.bind({"pool_id": "pool", "requests": {"incoming": [[(1, 1)]]}}) - adapter.runtime.state.begin_load(load_plan(adapter.runtime.state, "incoming")) - with pytest.raises(ValueError, match="recapture is unsupported"): - adapter.wrap_capture(lambda: pytest.fail("capture must not run"))() diff --git a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py index af5c9c8fa4ff..dcc6de552a7d 100644 --- a/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py +++ b/tests/ut/distributed/kv_transfer/kv_p2p/test_mooncake_connector.py @@ -5436,45 +5436,6 @@ def test_unused_remote_ports_are_notified_once_by_primary_receiver(primary, has_ assert sent == ([("r", "idle-peer", 7000, ports)] if primary and has_ports else []) -@pytest.mark.parametrize("awaiting_kvs", [False, True]) -@pytest.mark.parametrize("transfer_failed", [False, True]) -def test_receive_completion_contract_is_preserved(awaiting_kvs, transfer_failed): - receiver = KVCacheRecvingThread.__new__(KVCacheRecvingThread) - receiver._is_failed_recv_request = MagicMock(side_effect=[False, transfer_failed]) - receiver._transfer_kv_cache_all_groups = MagicMock() - receiver._mark_request_task_done = MagicMock(return_value=True) - receiver.pending_reformat_lock = threading.Lock() - receiver.pending_reformat = {"r": object()} - receiver.proc_not_transfer_request_lock = threading.Lock() - receiver.proc_not_transfer_request = {} - receiver._reformat_pending_kv_caches = MagicMock() - receiver._clear_failed_recv_request = MagicMock() - receiver.task_tracker = MagicMock() - receiver.request_queue = MagicMock() - receiver._send_done_signal_to_free_remote_port = MagicMock() - receiver._send_done_recv_signal = MagicMock() - meta = { - "request_id": "r", - "remote_request_id": "remote-r", - "remote_host": "host", - "remote_handshake_port": 1, - "remote_port_send_num": {}, - "all_task_done": True, - "awaiting_kvs": awaiting_kvs, - } - receiver._handle_request(meta) - if transfer_failed: - receiver._reformat_pending_kv_caches.assert_not_called() - else: - receiver._reformat_pending_kv_caches.assert_called_once_with("r") - if awaiting_kvs: - receiver.task_tracker.update_done_task_count.assert_called_once_with("r") - else: - receiver.task_tracker.update_done_task_count.assert_called_once_with("r", report_completion=False) - receiver.request_queue.task_done.assert_called_once() - receiver._send_done_recv_signal.assert_called_once() - - @pytest.mark.parametrize("pulls", [1, 2, 4]) def test_head_reformat_preserves_token_head_order_at_scatter_boundary(monkeypatch, pulls): """Validate PD's CPU reshape arithmetic, not the external NPU scatter op.""" diff --git a/vllm_ascend/debug/kv_guard.py b/vllm_ascend/debug/kv_guard.py index 54353de1d11a..6cac42e51c93 100644 --- a/vllm_ascend/debug/kv_guard.py +++ b/vllm_ascend/debug/kv_guard.py @@ -1,373 +1,94 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""KV address legality: allocation generation, valid prefix and request owners. +"""Independent request/block mirror built only from standard SchedulerOutput. -A prefix of zero is unwritten; INVALID_PREFIX is tainted. Loading and appending -produce the same readable prefix, so their provenance is not persistent state. +Computed prefixes are framework inputs, not proof of KV contents or allocator +generations. This mirror checks execution plans, not the scheduler's allocator. """ from __future__ import annotations -import threading -from bisect import bisect_right -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any -INVALID_PREFIX = -1 -# Group and logical block index are the positions in this nested table. -BlockTable = list[list[tuple[int, int]]] - @dataclass(slots=True) -class Block: - epoch: int - valid_end: int = 0 - owners: set[str] = field(default_factory=set) +class AccessPlan: + slots: list[int] + tables: list[list[int]] + lengths: list[int] class KVGuard: - """One worker's pool, with request bindings and a reverse ownership index.""" - def __init__(self, emit: Any): self.emit = emit - self.pool_id: str | None = None - self.blocks: dict[tuple[int, int], Block] = {} - self.requests: dict[str, BlockTable] = {} - self.loads: dict[str, dict[str, Any]] = {} + self.requests: dict[str, tuple[list[int], ...]] = {} + self.scheduled: dict[str, tuple[int, int]] = {} self.group_sizes: dict[int, int] = {} - self._lock = threading.RLock() def configure(self, group_sizes: dict[int, int]) -> None: - self.group_sizes = group_sizes - - def _violation(self, key: tuple[int, int], block: Block | None, operation: str, reason: str, **fields: Any) -> None: - self.emit( - "guard.violation", - operation=operation, - reason=reason, - group_id=key[0], - block_id=key[1], - pool_id=self.pool_id, - alloc_epoch=block.epoch if block else None, - valid_end=block.valid_end if block else 0, - owner_request_ids=sorted(block.owners) if block else [], - **fields, - ) - - def bind(self, context: dict[str, Any]) -> None: - with self._lock: - pool = context.get("pool_id", self.pool_id) - if pool != self.pool_id: - self.blocks.clear() - self.requests.clear() - self.loads.clear() - self.pool_id = pool - for request_id in context.get("released_request_ids", ()): - for group, refs in enumerate(self.requests.pop(request_id, ())): - for block_id, _ in refs: - block = self.blocks.get((group, block_id)) - if block is not None: - block.owners.discard(request_id) - self.loads.pop(request_id, None) - for request_id, groups in context.get("requests", {}).items(): - keys = {(group, block) for group, refs in enumerate(groups) for block, _ in refs} - for group, refs in enumerate(self.requests.get(request_id, ())): - for block_id, _ in refs: - key = group, block_id - if key not in keys and key in self.blocks: - self.blocks[key].owners.discard(request_id) - self.requests[request_id] = groups - for group, refs in enumerate(groups): - for block_id, epoch in refs: - if block_id == 0: - continue - if epoch <= 0: - self.emit( - "guard.coverage_gap", - reason="allocation_generation_missing", - group_id=group, - block_id=block_id, - request_id=request_id, - ) - continue - key = group, block_id - block = self.blocks.get(key) - if block is None or block.epoch != epoch: - if block is not None and block.owners: - self._violation(key, block, "recycle", "recycle_live_block", request_id=request_id) - block = Block(epoch) - self.blocks[key] = block - block.owners.add(request_id) - # A consumer's computed prefix is a read requirement, - # not evidence that this allocation has been written. - for meta in context.get("loads", ()): - self.begin_load(meta) - - def begin_load(self, meta: dict[str, Any]) -> None: - """Register the scheduler's submission snapshot, before any receive.""" - with self._lock: - if meta["pool_id"] != self.pool_id or meta["request_id"] not in self.requests: - self.emit("guard.coverage_gap", reason="load_binding_missing", request_id=meta["request_id"]) - return - if not 0 <= meta["token_start"] < meta["token_end"]: - self.emit("guard.coverage_gap", reason="load_range_invalid", request_id=meta["request_id"]) - return - self.loads[meta["request_id"]] = {**meta, "block_table": [list(refs) for refs in meta["block_table"]]} - - def _range_targets(self, meta: dict[str, Any], group_ids=None, block_ids=None): - request_id = meta["request_id"] - if meta["pool_id"] != self.pool_id: - self.emit("guard.coverage_gap", reason="stale_pool_event", request_id=request_id) - return - current = self.requests.get(request_id, ()) - for group, refs in enumerate(meta["block_table"]): - if group_ids is not None and group not in group_ids: - continue - size = self.group_sizes[group] - first = meta["token_start"] // size - last = (meta["token_end"] + size - 1) // size - for logical, (block_id, epoch) in enumerate(refs[first:last], first): - start = max(0, meta["token_start"] - logical * size) - end = min(size, meta["token_end"] - logical * size) - if start >= end or (block_ids is not None and block_id not in block_ids): - continue - key = group, block_id - block = self.blocks.get(key) - if block is not None and block.epoch != epoch: - self._violation(key, block, "load", "stale_allocation_epoch", request_id=request_id) - elif ( - block is None - or group >= len(current) - or logical >= len(current[group]) - or current[group][logical] != (block_id, epoch) - or request_id not in block.owners - ): - self._violation(key, block, "load", "unowned_destination", request_id=request_id) + self.group_sizes = dict(group_sizes) + + @staticmethod + def _copy_groups(groups: Any) -> tuple[list[int], ...]: + return tuple(list(group) for group in groups) + + def bind(self, output: Any) -> None: + """Apply the scheduler's append/replace/release contract before execution.""" + for request_id in output.finished_req_ids | (output.preempted_req_ids or set()): + self.requests.pop(request_id, None) + computed = {} + for request in output.scheduled_new_reqs: + self.requests[request.req_id] = self._copy_groups(request.block_ids) + computed[request.req_id] = request.num_computed_tokens + cached = output.scheduled_cached_reqs + for index, request_id in enumerate(cached.req_ids): + groups = cached.new_block_ids[index] + if request_id in cached.resumed_req_ids: + if groups is None: + self.emit("guard.coverage_gap", reason="resume_block_mapping_missing", request_id=request_id) + self.requests.pop(request_id, None) else: - yield key, block, start, end - - def _active_load(self, meta: dict[str, Any]) -> bool: - active = self.loads.get(meta["request_id"]) - if active == meta: - return True - self.emit( - "guard.coverage_gap", reason="stale_load_completion", request_id=meta["request_id"], load_id=meta["load_id"] - ) - return False - - def load(self, meta: dict[str, Any], *, group_ids=None, block_ids=None, finish: bool = True) -> None: - """Only a completed, matching submission can establish readable ranges.""" - with self._lock: - if not self._active_load(meta): - return - for key, block, start, end in self._range_targets(meta, group_ids, block_ids): - if block.valid_end == INVALID_PREFIX: - self.emit("guard.coverage_gap", reason="load_destination_tainted", group_id=key[0], block_id=key[1]) - continue - if start > block.valid_end: - self.emit("guard.coverage_gap", reason="load_prefix_missing", group_id=key[0], block_id=key[1]) - continue - block.valid_end = max(block.valid_end, end) - if finish: - self.finish_load(meta) - - def finish_load(self, meta: dict[str, Any]) -> None: - with self._lock: - if self.loads.get(meta["request_id"]) == meta: - self.loads.pop(meta["request_id"]) - - def fail_load(self, meta: dict[str, Any], *, block_ids=None, finish: bool = True) -> None: - with self._lock: - if not self._active_load(meta): - return - for _, block, _, _ in self._range_targets(meta, block_ids=block_ids): - # A block transfer can overwrite its local prefix before failing. - # Untouched local blocks stay valid; failed targets can be recomputed. - if block.valid_end != INVALID_PREFIX: - block.valid_end = 0 - if finish: - self.finish_load(meta) - - def recompute(self, meta: dict[str, Any]) -> None: - """Truncate only a framework-authorized replay, before auditing writes.""" - with self._lock: - for _, block, start, _ in self._range_targets(meta): - if block.valid_end != INVALID_PREFIX: - block.valid_end = min(block.valid_end, start) - - def clear(self, group_id: int, block_id: int) -> None: - with self._lock: - block = self.blocks.get((group_id, block_id)) - if block is not None: - block.valid_end = INVALID_PREFIX - - def copy_blocks(self, group_id: int, copies: list[tuple[int, int]]) -> None: - """Match the operator's snapshot-all-sources-before-writing semantics.""" - with self._lock: - source_ends = { - source: block.valid_end - for source, _ in copies - if (block := self.blocks.get((group_id, source))) is not None - } - for source_id, destination_id in copies: - valid = source_ends.get(source_id, INVALID_PREFIX) - key = group_id, destination_id - target = self.blocks.get(key) - if valid == INVALID_PREFIX or target is None or not target.owners: - self._violation( - key, target, "copy", "invalid_copy_source_or_destination", source_block_id=source_id - ) - if target is not None: - target.valid_end = INVALID_PREFIX + self.requests[request_id] = self._copy_groups(groups) + elif request_id not in self.requests: + self.emit("guard.coverage_gap", reason="cached_request_binding_missing", request_id=request_id) + elif groups is not None: + current = self.requests[request_id] + if len(current) != len(groups): + self.emit("guard.coverage_gap", reason="cache_group_count_mismatch", request_id=request_id) + self.requests.pop(request_id, None) else: - target.valid_end = valid - - def observe(self, groups: list[dict[str, Any]], request_ids: list[str], phase: str) -> None: - """Audit WRITE_SLOT then READ_RANGE using attention's consumed metadata. - - An illegal write taints its destination so subsequent reads diagnose - an invalid prefix, even when the block address itself remains legal. - """ - with self._lock: - for group in groups: - group_id = group["group_id"] - size = self.group_sizes[group_id] - ranges: dict[int, list[int]] = {} - expected = group.get("expected_slots") - query_ends = group.get("query_ends", ()) - token_offset = group.get("token_offset", 0) - mismatched: set[int] = set() - slots = group["slots"] - # List equality handles the common valid path in C. Only an - # actual mismatch needs per-token Python diagnostics. - count = min(len(slots), max(0, len(expected) - token_offset)) if expected is not None else 0 - if expected is not None and ( - slots[:count] != expected[token_offset : token_offset + count] - or any(slot >= 0 for slot in slots[count:]) - ): - for index, slot in enumerate(slots): - token_index = token_offset + index - target = expected[token_index] if token_index < len(expected) else None - if (target is not None and slot == target) or (target is None and slot < 0): - continue - destination = slot if slot >= 0 else target - assert destination is not None - block_id = destination // size - key = group_id, block_id - block = self.blocks.get(key) - if block_id in mismatched: - continue - self._violation( - key, - block, - "write", - "write_target_mismatch", - token_index=token_index, - slot=slot, - expected_slot=target, - request_id=request_ids[bisect_right(query_ends, token_index)] - if query_ends and token_index < query_ends[-1] - else None, - phase=phase, - ) - mismatched.add(block_id) - if block is not None: - block.valid_end = INVALID_PREFIX - for slot in slots: - if slot < 0: - continue - ranges.setdefault(slot // size, []).append(slot % size) - for block_id, offsets in ranges.items(): - key, offsets = (group_id, block_id), sorted(offsets) - block = self.blocks.get(key) - if block_id in mismatched: - continue - if block_id == 0: - self._violation(key, None, "write", "write_null_block", offsets=offsets, phase=phase) - continue - if phase == "warmup" and block is None: - continue # Capture may use explicitly allocated scratch. - reason = None - if block is None: - reason = "unallocated_block" - elif phase in ("dummy", "warmup"): - reason = "dummy_write_to_live_block" - elif not block.owners.intersection(request_ids): - reason = "unowned_write" - elif block.valid_end == INVALID_PREFIX: - reason = "write_invalid_block" - elif offsets[0] < block.valid_end: - reason = "overwrite_valid_prefix" - elif offsets != list(range(block.valid_end, block.valid_end + len(offsets))): - reason = "noncontiguous_write" - if reason is not None: - self._violation( - key, block, "write", reason, offsets=offsets, phase=phase, request_ids=request_ids - ) - if block is not None: - block.valid_end = INVALID_PREFIX - elif block is not None: - block.valid_end = offsets[-1] + 1 - for row, length in enumerate(group.get("seq_lens", ())): - if length <= 0: - continue - request_id = request_ids[row] if row < len(request_ids) else None - table = group["block_table"][row] - expected_tables = group.get("expected_block_table") - expected_lengths = group.get("expected_seq_lens") - if expected_lengths is not None and ( - row >= len(expected_lengths) or length > expected_lengths[row] - ): - self.emit( - "guard.violation", - operation="read", - reason="read_length_mismatch", - group_id=group_id, - request_id=request_id, - length=length, - phase=phase, - ) - continue - needed = (length + size - 1) // size - if needed > len(table): - self.emit( - "guard.violation", - operation="read", - reason="short_block_table", - group_id=group_id, - phase=phase, - ) - continue - for logical, block_id in enumerate(table[:needed]): - if expected_tables is not None and block_id != expected_tables[row][logical]: - self._violation( - (group_id, block_id), - self.blocks.get((group_id, block_id)), - "read", - "read_target_mismatch", - request_id=request_id, - logical_block_index=logical, - expected_block_id=expected_tables[row][logical], - phase=phase, - ) - continue - if block_id == 0: - if request_id is not None: - reason = "read_null_block" - self._violation( - (group_id, block_id), None, "read", reason, phase=phase, request_id=request_id - ) - continue - key, end = (group_id, block_id), min(size, length - logical * size) - block = self.blocks.get(key) - if block is None: - reason = "read_unallocated_block" - elif block.valid_end == INVALID_PREFIX: - reason = "read_invalid_block" - elif request_id not in block.owners: - reason = "unowned_read" - elif end > block.valid_end: - reason = "read_unwritten_range" - else: - continue - self._violation(key, block, "read", reason, end_offset=end, phase=phase, request_id=request_id) + self.requests[request_id] = tuple(old + list(new) for old, new in zip(current, groups)) + computed[request_id] = cached.num_computed_tokens[index] + self.scheduled = { + request_id: (computed.get(request_id, -1), count) + for request_id, count in output.num_scheduled_tokens.items() + } + + def expected(self, group_id: int, request_ids: list[str]) -> AccessPlan | None: + slots: list[int] = [] + tables: list[list[int]] = [] + lengths: list[int] = [] + size = self.group_sizes[group_id] + for request_id in request_ids: + computed, count = self.scheduled.get(request_id, (-1, 0)) + groups = self.requests.get(request_id, ()) + needed = (computed + count + size - 1) // size + if computed < 0 or count <= 0 or group_id >= len(groups) or len(groups[group_id]) < needed: + self.emit( + "guard.coverage_gap", reason="scheduler_plan_incomplete", request_id=request_id, group_id=group_id + ) + return None + blocks = groups[group_id][:needed] + if any(block <= 0 for block in blocks): + self.emit( + "guard.coverage_gap", reason="scheduler_block_id_invalid", request_id=request_id, group_id=group_id + ) + return None + tables.append(list(blocks)) + lengths.append(computed + count) + slots.extend(blocks[token // size] * size + token % size for token in range(computed, computed + count)) + return AccessPlan(slots, tables, lengths) + + def live_blocks(self, group_id: int) -> set[int]: + return {block for groups in self.requests.values() if group_id < len(groups) for block in groups[group_id]} diff --git a/vllm_ascend/debug/kv_guard_config.py b/vllm_ascend/debug/kv_guard_config.py index 93fecd8a7a29..42cb5b42de33 100644 --- a/vllm_ascend/debug/kv_guard_config.py +++ b/vllm_ascend/debug/kv_guard_config.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Per-engine KV guard options, usable before AscendConfig initialization.""" +"""Per-engine options for worker KV metadata-plan diagnostics.""" from typing import Any, Literal diff --git a/vllm_ascend/debug/kv_guard_connector.py b/vllm_ascend/debug/kv_guard_connector.py deleted file mode 100644 index f5b54ebd83be..000000000000 --- a/vllm_ascend/debug/kv_guard_connector.py +++ /dev/null @@ -1,186 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Translate the common connector lifecycle into logical KV readiness events. - -Submission snapshots come from the scheduler, never from backend-specific -transport offsets. No connector calls, waits or result polling are added. -""" - -from __future__ import annotations - -import functools -from typing import Any - - -class ConnectorGuard: - def __init__(self, runtime: Any, connector: Any): - self.runtime = runtime - self.state = runtime.state - self.staged: list[dict] = [] - self.inflight: dict[str, dict] = {} - self.sync: dict[str, dict] = {} - self.drained: list[dict] = [] - self.completed: list[tuple[dict, set[int]]] = [] - self.ready: dict[int, set[str]] = {} - self.ready_groups: dict[int, set[int]] = {} - self.failed: set[int] = set() - self.failed_blocks: dict[tuple[str, int], set[int]] = {} - self.post_forward_complete = False - self._wrap(connector, "start_load_kv", self.start) - self._wrap(connector, "wait_for_layer_load", self.wait) - self._wrap(connector, "get_transfer_results", self.results) - self._wrap(connector, "get_block_ids_with_load_errors", self.errors) - - @staticmethod - def _wrap(connector, name, callback): - original = getattr(connector, name) - - @functools.wraps(original) - def observed(*args, **kwargs): - return callback(original, *args, **kwargs) - - setattr(connector, name, observed) - - def stage(self, loads: list[dict]) -> None: - if self.staged or self.sync: - self.finish_step() - self.staged.extend(loads) - - def start(self, original, *args, **kwargs): - submitted, self.staged = self.staged, [] - accepted = [] - for meta in submitted: - request = meta["request_id"] - if request in self.inflight: - # Request-ID-only completions cannot distinguish overlapping incarnations. - self.runtime.gap("overlapping_load_request_id", request_id=request) - self.state.finish_load(meta) - continue - accepted.append(meta) - if meta["is_async"]: - self.inflight[request] = meta - else: - self.sync[request] = meta - self.runtime._check_failure() - try: - return original(*args, **kwargs) - except Exception: - for meta in accepted: - self.state.fail_load(meta) - self.inflight.pop(meta["request_id"], None) - raise - - def wait(self, original, layer_name, *args, **kwargs): - result = original(layer_name, *args, **kwargs) - group = self.runtime.layer_groups.get(layer_name) - if group is None: - return result - required = {name for name, index in self.runtime.layer_groups.items() if index == group} - for meta in self.sync.values(): - if meta["load_id"] in self.failed: - continue - ready = self.ready.setdefault(meta["load_id"], set()) - ready.add(layer_name) - groups = self.ready_groups.setdefault(meta["load_id"], set()) - if required <= ready and group not in groups: - self.state.load(meta, group_ids={group}, finish=False) - groups.add(group) - self.runtime._check_failure() - return result - - def check_ready(self, layer: str) -> None: - group = self.runtime.layer_groups.get(layer) - if group is None: - return - for meta in self.sync.values(): - if layer not in self.ready.get(meta["load_id"], ()): - self.runtime.gap("load_layer_not_ready", request_id=meta["request_id"], layer=layer) - self.runtime._missing_groups.add(group) - self.runtime._unready_load_groups.add(group) - - def before_audit(self) -> None: - for layer in self.runtime.layer_groups: - self.check_ready(layer) - - def results(self, original, *args, **kwargs): - result = original(*args, **kwargs) - failed = set(result.failed_recving or ()) - for request in set(result.finished_recving or ()) | failed: - meta = self.inflight.pop(request, None) - if meta is None: - meta = self.sync.get(request) - if meta is None: - continue - if request in failed: - self.failed.add(meta["load_id"]) - if not meta["is_async"]: - self.state.fail_load(meta, finish=False) - self.drained.append(meta) - return result - - def errors(self, original, *args, **kwargs): - invalid = original(*args, **kwargs) - if invalid: - # Core queries invalid blocks after transfer results; keep the snapshot until then. - for meta in (*self.inflight.values(), *self.sync.values()): - if self.state.loads.get(meta["request_id"]) != meta: - continue - affected = invalid & self._target_blocks(meta) - if affected: - self.failed_blocks.setdefault((meta["pool_id"], meta["load_id"]), set()).update(affected) - self.state.fail_load(meta, block_ids=affected, finish=False) - for meta in self.drained: - self.completed.append((meta, set(invalid or ()))) - self.drained.clear() - self.post_forward_complete = True - if self.runtime._phase is None: - self.finish_step() - self.runtime._check_failure() - return invalid - - def _target_blocks(self, meta: dict) -> set[int]: - targets: set[int] = set() - for group, refs in enumerate(meta["block_table"]): - size = self.state.group_sizes[group] - first, last = meta["token_start"] // size, (meta["token_end"] + size - 1) // size - targets.update(block for block, _ in refs[first:last]) - return targets - - def commit_completions(self) -> None: - # MRV2 polls after forward but before its audit. These completions may - # authorize the next step, never an earlier read in the current step. - for meta, invalid in self.completed: - # Legacy connectors can report failed blocks in an earlier poll, - # then drain only a request-ID completion in a later step. - invalid |= self.failed_blocks.pop((meta["pool_id"], meta["load_id"]), set()) - if self.state.loads.get(meta["request_id"]) != meta: - continue # Cancellation may release ownership before the known old receive drains. - if meta["is_async"]: - if meta["load_id"] in self.failed: - self.state.fail_load(meta, finish=False) - else: - self.state.load(meta, block_ids=self._target_blocks(meta) - invalid, finish=False) - if invalid: - self.state.fail_load(meta, block_ids=invalid, finish=False) - if meta["is_async"] or meta["load_id"] in self.failed: - self.state.finish_load(meta) - self.completed.clear() - - def finish_step(self) -> None: - self.commit_completions() - for meta in self.staged: - self.runtime.gap("load_not_submitted", request_id=meta["request_id"]) - self.state.finish_load(meta) - for meta in self.sync.values(): - if meta["load_id"] not in self.failed: - required = set(self.runtime.layer_groups.values()) - if not required <= self.ready_groups.get(meta["load_id"], set()): - self.runtime.gap("load_completion_missing", request_id=meta["request_id"]) - self.state.finish_load(meta) - self.failed_blocks.pop((meta["pool_id"], meta["load_id"]), None) - self.staged.clear() - self.sync.clear() - self.ready.clear() - self.ready_groups.clear() - self.failed.clear() - self.post_forward_complete = False diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py index 901fe3693bc8..6927a4cda2e8 100644 --- a/vllm_ascend/debug/kv_guard_runtime.py +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -1,67 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Worker L1 guard with one runner execution-plan collection point. +"""Compare one attention execution plan per cache group with SchedulerOutput. -No KV snapshots, hashes, general-purpose trace or offline reader. -Runner metadata supplies one declared access plan per KV cache group. -Scheduler ranges independently identify each token; attention is not instrumented. -Only small metadata is copied, after the forward and graph parameter updates; -the checker neither inspects KV values nor proves kernel memory accesses. +Eager and FULL replay use this step's final metadata. The checker does not +observe operator arguments, kernel memory accesses, KV contents or generations. """ from __future__ import annotations -import functools from collections import OrderedDict from contextlib import contextmanager -from copy import copy -from dataclasses import dataclass, replace -from threading import Lock -from types import SimpleNamespace -from typing import Any, cast +from typing import Any -import torch - -from vllm_ascend.debug.kv_guard import INVALID_PREFIX, KVGuard +from vllm_ascend.debug.kv_guard import KVGuard from vllm_ascend.debug.kv_guard_config import KVGuardConfig -from vllm_ascend.debug.kv_guard_connector import ConnectorGuard _MAX_REPORT_KEYS = 1024 -@dataclass(slots=True) -class ExpectedAccess: - slots: list[int] - query_ends: list[int] - tables: list[list[int]] - lengths: list[int] - - class KVGuardError(RuntimeError): - """Stop returning model results after a violation or incomplete audit.""" - - -class _PendingAudit: - """One snapshot, consumed once by the output reader or next forward.""" - - def __init__(self, complete: Any): - self._complete = complete - self._lock = Lock() - self._done = False - self._error: BaseException | None = None - - def finish(self) -> None: - with self._lock: - if not self._done: - try: - self._complete() - except BaseException as exc: - self._error = exc - finally: - self._complete = None - self._done = True - if self._error is not None: - raise self._error + """A violation or missing coverage prevents further model output.""" class GuardRuntime: @@ -70,64 +28,31 @@ def __init__(self, vllm_config: Any): self.logger = logger self.config = vllm_config + self.action = KVGuardConfig.from_vllm_config(vllm_config).action self.state = KVGuard(self.report) - self._scratch = KVGuard(self._emit_scratch) self.layer_groups: dict[str, int] = {} - self._metadata: dict[int, Any] = {} - self._missing_groups: set[int] = set() - self._phase: str | None = None - self._full_graph = False - self._reported: OrderedDict[tuple, None] = OrderedDict() - self._scheduled: dict[str, tuple[int | None, int]] | None = None self._failure: str | None = None - self.action = KVGuardConfig.from_vllm_config(vllm_config).action - self._runner_adapter: _V2RunnerAdapter | None = None - self._host_buffers: list[Any] = [] - self._copy_event: Any = None - self._connector: ConnectorGuard | None = None - self._recomputes: list[dict] = [] - self._unready_load_groups: set[int] = set() - self._defer_audit = False - self._pending: _PendingAudit | None = None - self._deferred_state: list[tuple[Any, tuple]] = [] - - def _check_failure(self) -> None: - if self.action == "raise" and self._failure is not None: - raise KVGuardError(f"KV guard stopped model output: {self._failure}; restart the worker before reuse") + self._reported: OrderedDict[tuple, None] = OrderedDict() + self._phase: str | None = None + self._reset_step() - def _emit_scratch(self, event: str, **fields: Any) -> None: - if event == "guard.violation": - self.state.emit(event, evidence="execution_plan", **fields) + def _reset_step(self) -> None: + self._metadata: dict[int, Any] = {} + self._request_ids: list[str] | None = None def report(self, event: str, **fields: Any) -> None: failed = event in ("guard.violation", "guard.coverage_gap") if failed: - with self.state._lock: - if self._failure is None: - self._failure = str(fields.get("reason", event)) - key = ( - event, - *( - fields.get(name) - for name in ( - "operation", - "reason", - "pool_id", - "group_id", - "block_id", - "alloc_epoch", - "layer", - "request_id", - ) - ), - tuple(fields.get("request_ids", ())), - ) - if key in self._reported: - self._reported.move_to_end(key) - return - self._reported[key] = None - if len(self._reported) > _MAX_REPORT_KEYS: - self._reported.popitem(last=False) + fields.setdefault("evidence", "execution_plan") + if self._failure is None: + self._failure = str(fields.get("reason", event)) + key = event, tuple((name, repr(value)) for name, value in sorted(fields.items())) + if key in self._reported: + self._reported.move_to_end(key) + return + self._reported[key] = None + if len(self._reported) > _MAX_REPORT_KEYS: + self._reported.popitem(last=False) payload = " ".join(f"{key}={value}" for key, value in fields.items()) log = self.logger.error if failed else self.logger.info log("[KV_GUARD] %s %s", event, payload) @@ -135,61 +60,71 @@ def report(self, event: str, **fields: Any) -> None: def gap(self, reason: str, **fields: Any) -> None: self.report("guard.coverage_gap", reason=reason, **fields) + def check_failure(self) -> None: + if self.action == "raise" and self._failure is not None: + raise KVGuardError(f"KV guard stopped model output: {self._failure}; restart the worker before reuse") + def configure(self, runner: Any) -> None: - config = self.config - parallel = config.parallel_config - transfer = config.kv_transfer_config - unsupported = any( + config, parallel = self.config, self.config.parallel_config + if any( ( + config.kv_transfer_config is not None, getattr(config.model_config, "is_encoder_decoder", False), config.speculative_config is not None, getattr(parallel, "prefill_context_parallel_size", 1) != 1, getattr(parallel, "decode_context_parallel_size", 1) != 1, getattr(parallel, "enable_dbo", False), - transfer is not None and transfer.kv_connector not in ("MooncakeConnector", "MooncakeConnectorV1"), ) - ) + ): + raise ValueError("KV guard: unsupported connector, CP, DBO, speculation or encoder/decoder mode") sizes = {} - tables = getattr(getattr(getattr(runner, "input_batch", None), "block_table", None), "block_tables", ()) + self.layer_groups = {} for group_id, group in enumerate(runner.kv_cache_config.kv_cache_groups): - spec = group.kv_cache_spec - sizes[group_id] = spec.block_size - if tables and tables[group_id].block_size != spec.block_size: - unsupported = True + sizes[group_id] = group.kv_cache_spec.block_size for name in group.layer_names: - self.layer_groups[name] = group_id - layer_spec = getattr(spec, "kv_cache_specs", {}).get(name, spec) - append = ( - hasattr(layer_spec, "head_size") - and getattr(layer_spec, "sliding_window", None) is None - and getattr(layer_spec, "compress_ratio", 1) == 1 - and type(layer_spec).__name__ != "CircularBufferSpec" - ) - if not append: - raise ValueError(f"KV L1 guard: unadapted cache validity rules: {type(layer_spec).__name__}") + spec = getattr(group.kv_cache_spec, "kv_cache_specs", {}).get(name, group.kv_cache_spec) + if ( + not hasattr(spec, "head_size") + or getattr(spec, "sliding_window", None) is not None + or getattr(spec, "attention_chunk_size", None) is not None + or getattr(spec, "compress_ratio", 1) != 1 + or type(spec).__name__ == "CircularBufferSpec" + ): + raise ValueError(f"KV guard: unsupported cache specification: {type(spec).__name__}") layer = config.compilation_config.static_forward_context.get(name) impl = getattr(layer, "impl", None) - if impl is None or type(impl).__name__ not in ("AscendAttentionBackendImpl", "AscendMLAImpl"): - raise ValueError(f"KV L1 guard: unadapted attention implementation: {type(impl).__name__}") - if unsupported: - raise ValueError( - "KV L1 guard: unadapted execution mode or block coordinates " - "(CP, DBO, speculation, encoder/decoder, connector)" - ) + if type(impl).__name__ not in ("AscendAttentionBackendImpl", "AscendMLAImpl"): + raise ValueError(f"KV guard: unsupported attention implementation: {type(impl).__name__}") + self.layer_groups[name] = group_id + if not self.layer_groups: + raise ValueError("KV guard: no supported attention layers") + tables = getattr(runner, "block_tables", None) + if tables is not None and ( + list(tables.block_sizes) != list(sizes.values()) or list(tables.kernel_block_sizes) != list(sizes.values()) + ): + raise ValueError("KV guard: unsupported allocator/kernel block coordinates") + batch_tables = getattr(getattr(getattr(runner, "input_batch", None), "block_table", None), "block_tables", ()) + if batch_tables and any(table.block_size != sizes[index] for index, table in enumerate(batch_tables)): + raise ValueError("KV guard: unsupported allocator/kernel block coordinates") + model_state = getattr(runner, "model_state", None) + if model_state is not None and not hasattr(model_state, "kv_guard"): + raise ValueError("KV guard: unsupported model state") self.state.configure(sizes) - self._defer_audit = ( - bool(getattr(getattr(config, "scheduler_config", None), "async_scheduling", False)) - and transfer is None - # Pooling bypasses sample_tokens; non-reply workers may discard outputs. - and getattr(runner, "is_pooling_model", None) is False - and getattr(parallel, "tensor_parallel_size", None) == 1 - and getattr(parallel, "pipeline_parallel_size", None) == 1 - and hasattr(runner, "sample_tokens") - ) - if self._runner_adapter is not None: - self._runner_adapter.configure() + if model_state is not None: + model_state.kv_guard = self self.report("guard.enabled", failure_action=self.action) + def observe_metadata(self, metadata: Any, request_ids: list[str] | None = None, full_graph: bool = False) -> None: + if self._phase not in ("forward", "dummy"): + return + self._metadata = {} + self._request_ids = list(request_ids) if request_ids is not None else None + if isinstance(metadata, dict): + for name, item in metadata.items(): + group = self.layer_groups.get(name) + if group is not None and item is not None: + self._metadata.setdefault(group, item) + @staticmethod def _list(value: Any) -> list: if isinstance(value, (list, tuple)): @@ -200,579 +135,116 @@ def _list(value: Any) -> list: def _field(metadata: Any, *names: str) -> Any: return next((value for name in names if (value := getattr(metadata, name, None)) is not None), None) - def collect_metadata(self, metadata: Any) -> None: - self._metadata = {} - self._missing_groups = set() - if metadata is None: - return - if not isinstance(metadata, dict): - self.gap("microbatch_metadata_unadapted") - self._missing_groups.update(self.state.group_sizes) - return - for name, item in metadata.items(): - group = self.layer_groups.get(name) - if group is not None and item is not None: - self._metadata.setdefault(group, item) - for group in self.state.group_sizes.keys() - self._metadata.keys(): - self.gap("metadata_write_mapping_missing", group_id=group) - self._missing_groups.add(group) - - @staticmethod - def _identity(value: Any) -> Any: + def _slots(self, metadata: Any, group: int) -> list[int] | None: + value = self._field(metadata, "slot_mapping") if value is None: + self.gap("metadata_write_mapping_missing", group_id=group) return None - if isinstance(value, (list, tuple)): - return tuple(GuardRuntime._identity(item) for item in value) - if isinstance(value, (int, float, str)): - return value - return value.data_ptr(), tuple(value.shape), tuple(value.stride()) - - def _expected_write_plan(self, runner: Any, group_id: int) -> ExpectedAccess | None: - """Derive destinations from the scheduler, never from the audited mapping. - - Per-token lists exist only for this forward and are shared by all layers - in a cache group. No persistent token-ID history or KV tensor is copied. - """ - if self._phase != "forward" or self._scheduled is None: - return None - size = self.state.group_sizes[group_id] - slots: list[int] = [] - query_ends: list[int] = [] - tables: list[list[int]] = [] - lengths: list[int] = [] - for request_id in runner.input_batch.req_ids: - computed, count = self._scheduled.get(request_id, (None, 0)) - if computed is None or computed < 0 or count <= 0: - self.gap("scheduler_write_range_missing", request_id=request_id) - self._missing_groups.add(group_id) - return None - groups = self.state.requests.get(request_id, ()) - refs = groups[group_id] if group_id < len(groups) else () - needed = (computed + count + size - 1) // size - if len(refs) < needed: - self.gap("scheduler_read_blocks_missing", request_id=request_id, group_id=group_id) - self._missing_groups.add(group_id) - return None - tables.append([block for block, _ in refs[:needed]]) - lengths.append(computed + count) - for logical in range(computed // size, (computed + count - 1) // size + 1): - block_id, _ = refs[logical] - start = max(computed, logical * size) - logical * size - end = min(computed + count, (logical + 1) * size) - logical * size - slots.extend(range(block_id * size + start, block_id * size + end)) - query_ends.append(len(slots)) - return ExpectedAccess(slots, query_ends, tables, lengths) + count = getattr(metadata, "num_actual_tokens", None) + slots = self._list(value if count is None else value[:count]) + slots = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] + return slots def _audit(self, runner: Any) -> None: - phase = self._phase - assert phase is not None - if not self._metadata: - if phase != "warmup": - self.gap("full_graph_metadata_missing" if self._full_graph else "execution_access_plan_missing") + if not self.state.scheduled: + for group, metadata in self._metadata.items(): + slots = self._slots(metadata, group) + lengths = self._field(metadata, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") + if (slots is not None and any(slot >= 0 for slot in slots)) or ( + lengths is not None and any(self._list(lengths)) + ): + self.report("guard.violation", reason="unscheduled_attention_plan", group_id=group) return - request_ids = list(runner.input_batch.req_ids) if phase == "forward" else [] - shadow = self._scratch - shadow.configure(self.state.group_sizes) - shadow.emit = self._emit_scratch - for group, item in self._metadata.items(): - if group in self._unready_load_groups: + request_ids = self._request_ids + if request_ids is None: + request_ids = list(getattr(getattr(runner, "input_batch", None), "req_ids", ())) + for group in self.state.group_sizes: + metadata = self._metadata.get(group) + if metadata is None: + self.gap("metadata_group_missing", group_id=group) continue - slots = self._field(item, "slot_mapping") - table = self._field(item, "block_tables", "block_table", "block_table_tensor") - lengths = self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") - if slots is None: - self.gap("metadata_write_mapping_missing", group_id=group) - self._missing_groups.add(group) - slots = [] - else: - count = getattr(item, "num_actual_tokens", None) - slots = self._list(slots if count is None else slots[:count]) - slots = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] - size = self.state.group_sizes[group] - if phase != "forward": - lengths, table = [], [] # Discarded scratch reads do not establish cache validity. - elif table is None or lengths is None: + if len(request_ids) != len(set(request_ids)) or set(request_ids) != set(self.state.scheduled): + self.gap("batch_request_mismatch", group_id=group) + continue + plan = self.state.expected(group, request_ids) + if plan is None: + continue + slots = self._slots(metadata, group) + if slots is not None: + if slots[: len(plan.slots)] != plan.slots or any(slot >= 0 for slot in slots[len(plan.slots) :]): + self.report("guard.violation", reason="write_target_mismatch", group_id=group) + if len(slots) < len(plan.slots): + self.gap("write_tokens_not_observed", group_id=group) + table = self._field(metadata, "block_tables", "block_table", "block_table_tensor") + lengths = self._field(metadata, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu") + if table is None or lengths is None: self.gap("metadata_read_range_missing", group_id=group) - self._missing_groups.add(group) - lengths, table = [], [] - else: - lengths = self._list(lengths) - columns = max(((length + size - 1) // size for length in lengths), default=0) - table = ( - [row[:columns] for row in table[: len(lengths)]] - if isinstance(table, (list, tuple)) - else self._list(table[: len(lengths), :columns]) - ) - entry = {"group_id": group, "slots": slots, "block_table": table, "seq_lens": lengths} - writes = {(group, slot // size) for slot in slots if slot >= 0} - keys = writes | {(group, block) for row in table for block in row if block >= 0} - # Metadata is already on the host. Keep state inspection and commit atomic, - # sharing read-only blocks and cloning only blocks this step may write. - with self.state._lock: - plan = self._expected_write_plan(runner, group) - if plan is not None: - entry.update( - expected_slots=plan.slots, - query_ends=plan.query_ends, - expected_block_table=plan.tables, - expected_seq_lens=plan.lengths, - ) - if len(slots) < len(plan.slots): - self.gap("write_tokens_not_observed", group_id=group, first_missing_token=len(slots)) - self._missing_groups.add(group) - pool = shadow.pool_id = self.state.pool_id - before = { - key: (block, block.epoch, block.valid_end, frozenset(block.owners)) - for key in writes - if (block := self.state.blocks.get(key)) is not None - } - shadow.blocks = { - key: replace(block) if key in writes else block - for key in keys - if (block := self.state.blocks.get(key)) is not None - } - shadow.observe([entry], request_ids, phase) - if self.state.pool_id != pool or any( - self.state.blocks.get(key) is not block - or (block.epoch, block.valid_end, frozenset(block.owners)) != (epoch, valid, owners) - for key, (block, epoch, valid, owners) in before.items() - ): - self.gap("allocation_changed_during_access_audit") - continue - for key in before: - block, updated = self.state.blocks[key], shadow.blocks[key] - if updated.valid_end == INVALID_PREFIX: - block.valid_end = INVALID_PREFIX - elif phase == "forward" and group not in self._missing_groups and block.valid_end != INVALID_PREFIX: - block.valid_end = max(block.valid_end, updated.valid_end) - - def observe_metadata(self, metadata: Any) -> None: - try: - self.collect_metadata(metadata) - except Exception as exc: - self._metadata = {} - self._missing_groups = set(self.layer_groups.values()) - self.gap(f"metadata_collection_failed:{type(exc).__name__}") - - def _finish_pending(self) -> None: - if self._pending is not None: - try: - self._pending.finish() - finally: - self._pending = None - operations, self._deferred_state = self._deferred_state, [] - for operation, args in operations: - operation(*args) - self._check_failure() - - def _snapshot_audit(self, runner: Any) -> None: - """Freeze final evidence now; wait only at an output/forward boundary.""" - assert self._pending is None - job = copy(self) - job._scheduled = None if self._scheduled is None else dict(self._scheduled) - job._missing_groups = set(self._missing_groups) - job._unready_load_groups = set(self._unready_load_groups) - reporting = cast(Any, job) - reporting.report, reporting.gap, reporting._check_failure = self.report, self.gap, self._check_failure - job._scratch = KVGuard(job._emit_scratch) - frozen_runner = SimpleNamespace(input_batch=SimpleNamespace(req_ids=list(runner.input_batch.req_ids))) - copies: dict[Any, Any] = {} - event = None - device_copies = 0 - - def freeze(value): - nonlocal event, device_copies - if isinstance(value, torch.Tensor): - key = value.device, value.dtype, self._identity(value) - if key not in copies: - if value.device.type == "cpu": - host = value.clone() - else: - if self._copy_event is None: - self._copy_event = torch.npu.Event() - event = self._copy_event - index = device_copies - device_copies += 1 - if index >= len(self._host_buffers): - self._host_buffers.append(None) - host = self._host_buffers[index] - if host is None or host.shape != value.shape or host.dtype != value.dtype: - host = torch.empty(value.shape, dtype=value.dtype, device="cpu", pin_memory=True) - self._host_buffers[index] = host - host.copy_(value, non_blocking=True) - copies[key] = host - return copies[key] - if isinstance(value, (list, tuple)): - return [freeze(item) for item in value] - return value - - try: - job._metadata = {} - for group, item in self._metadata.items(): - slots = self._field(item, "slot_mapping") - count = getattr(item, "num_actual_tokens", None) - job._metadata[group] = SimpleNamespace( - slot_mapping=freeze(slots if slots is None or count is None else slots[:count]), - block_tables=freeze(self._field(item, "block_tables", "block_table", "block_table_tensor")), - seq_lens=freeze(self._field(item, "seq_lens_gpu", "seq_lens", "seq_lens_list", "seq_lens_cpu")), - ) - if event is not None: - event.record(torch.npu.current_stream()) - except Exception: - if event is not None: - # Even a partially failed snapshot must drain before buffer reuse. - event.record(torch.npu.current_stream()) - event.synchronize() - raise - - def complete(): - with torch.inference_mode(): - try: - if event is not None: - event.synchronize() - job._audit(frozen_runner) - except Exception as exc: - self.gap(f"access_audit_failed:{type(exc).__name__}") - self._check_failure() - - self._pending = _PendingAudit(complete) - - def _guard_output(self, output: Any) -> Any: - self._check_failure() - pending = self._pending - if pending is None: - return output - get_output = getattr(output, "get_output", None) - if get_output is None: - pending.finish() - self._check_failure() - return output - - @functools.wraps(get_output) - def checked(): - pending.finish() - self._check_failure() - return get_output() - - try: - output.get_output = checked - except (AttributeError, TypeError): - pending.finish() - return output - - def wrap_sample(self, sample: Any): - @functools.wraps(sample) - def observed(*args, **kwargs): - return self._guard_output(sample(*args, **kwargs)) - - return observed + continue + lengths = self._list(lengths) + rows = len(plan.lengths) + if lengths[:rows] != plan.lengths or any(lengths[rows:]): + self.report("guard.violation", reason="read_length_mismatch", group_id=group) + columns = max((len(row) for row in plan.tables), default=0) + table = ( + [list(row[:columns]) for row in table[:rows]] + if isinstance(table, (list, tuple)) + else self._list(table[:rows, :columns]) + ) + if len(table) < rows or any( + actual[: len(expected)] != expected for actual, expected in zip(table, plan.tables) + ): + self.report("guard.violation", reason="read_block_mapping_mismatch", group_id=group) @contextmanager - def forward_scope( - self, runner: Any, *, phase: str, metadata: Any = None, full_graph: bool = False, required: bool = True - ): - """Audit one completed forward; never advance validity on exceptions.""" - self._check_failure() - self._phase, self._full_graph = phase, full_graph - self.observe_metadata(metadata) + def execution(self, runner: Any, scheduler_output: Any): + self.check_failure() + self.state.bind(scheduler_output) + self._reset_step() + self._phase = "forward" try: + self.check_failure() yield - if required or self._metadata: - try: - if phase == "forward" and self._connector is not None: - self._connector.before_audit() - if phase == "forward": - for meta in self._recomputes: - self.state.recompute(meta) - self._recomputes = [] - if phase == "forward": - self._snapshot_audit(runner) - if not self._defer_audit: - self._finish_pending() - else: - self._audit(runner) - except Exception as exc: - self.gap(f"access_audit_failed:{type(exc).__name__}") - self._check_failure() + try: + self._audit(runner) + except Exception as exc: + self.gap(f"access_audit_failed:{type(exc).__name__}") + self.check_failure() finally: self._phase = None - if self._connector is not None: - self._connector.commit_completions() - if self._connector.post_forward_complete: - self._connector.finish_step() - - def wrap_forward(self, runner: Any, forward: Any, get_context: Any): - @functools.wraps(forward) - def observed(num_tokens, *args, **kwargs): - self._finish_pending() - context = get_context() - with self.forward_scope( - runner, - phase=kwargs.pop("_kv_guard_phase", "forward"), - metadata=getattr(context, "attn_metadata", None), - full_graph="FULL" in str(context.cudagraph_runtime_mode), - ): - return forward(num_tokens, *args, **kwargs) + self._reset_step() - return observed - - def wrap_update(self, runner: Any, update: Any): - @functools.wraps(update) - def observed(output, *args, **kwargs): - result = update(output, *args, **kwargs) - binding = getattr(output, "_ascend_kv_guard_context", None) - if binding is None: - self.gap("scheduler_binding_missing") - self._check_failure() - return result - for request_id, groups in binding["requests"].items(): - if request_id not in binding.get("scheduled_requests", {}): - continue # Receive-only bindings are not runner batch entries. - actual = runner.requests.get(request_id) - expected = [[block for block, _ in refs] for refs in groups] - if actual is None or [list(g) for g in actual.block_ids] != expected: - self.report( - "guard.violation", - operation="bind", - reason="worker_block_mapping_mismatch", - request_id=request_id, - ) - self._check_failure() - return result - - return observed + def check_capture(self) -> None: + self.check_failure() + if self.state.requests: + raise ValueError("KV guard: capture requires no live requests") @contextmanager - def execute_scope(self, output: Any, *, bind: bool = True, finish_connector: bool = True): - """Share scheduler binding and cleanup across both runner versions.""" - self._check_failure() - self._scheduled = None - self._unready_load_groups.clear() - self._recomputes = [] - if bind: - binding = getattr(output, "_ascend_kv_guard_context", None) - if binding is not None: - # Receive-only and idle DP steps can return before _update_states. - uninitialized = set(getattr(output, "new_block_ids_to_zero", None) or ()) - uninitialized.update( - copy.dst_block_id for copy in (getattr(output, "kv_cache_block_copies", None) or ()) - ) - context = {**binding, "uninitialized_block_ids": uninitialized} - if self._pending is not None: - self._deferred_state.append((self.state.bind, (context,))) - else: - self.state.bind(context) - self._recomputes = binding.get("recompute_ranges", []) - if self._connector is None and getattr(self.config, "kv_transfer_config", None) is not None: - from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group - - if has_kv_transfer_group(): - self._connector = ConnectorGuard(self, get_kv_transfer_group()) - loads = binding.get("loads", []) - if self._connector is not None: - self._connector.stage(loads) - elif loads: - self.gap("load_connector_missing") - self._scheduled = binding.get("scheduled_requests") - if self._scheduled is None: - self.gap("scheduler_write_plan_missing") - else: - self.gap("scheduler_binding_missing") + def dummy_scope(self, runner: Any, phase: str = "dummy"): + self.check_failure() + if phase == "capture": + self.check_capture() + previous = self._phase, self._metadata, self._request_ids + self._reset_step() + self._phase = phase try: - self._check_failure() yield - self._check_failure() - finally: - if self._deferred_state: - self._finish_pending() # Receive-only, idle and failed steps may never enter a forward. - if self._connector is not None and finish_connector: - self._connector.finish_step() - self._scheduled = None - self._recomputes = [] - self._check_failure() - - def wrap_execute(self, execute: Any): - @functools.wraps(execute) - def observed(output, *args, **kwargs): - with self.execute_scope(output): - result = execute(output, *args, **kwargs) - return self._guard_output(result) if result is not None else result - - return observed - - def zero_blocks(self, block_ids: list[int]) -> None: - if self._pending is not None: - self._deferred_state.append((self.zero_blocks, (list(block_ids),))) - return - with self.state._lock: - for group_id in self.state.group_sizes: - for block_id in block_ids: - if block_id == 0: - self.state.clear(group_id, block_id) + if phase == "dummy": + for group, size in self.state.group_sizes.items(): + metadata = self._metadata.get(group) + if metadata is None: + if self.state.requests: + self.gap("dummy_write_plan_missing", group_id=group) continue - block = self.state.blocks.get((group_id, block_id)) - # A fresh allocation is already UNKNOWN. Clearing it must - # not turn its first legitimate write into an invalid write. - if block is not None and block.valid_end != 0: - if block.valid_end > 0 and block.owners: - self.state._violation((group_id, block_id), block, "clear", "clear_live_block") - self.state.clear(group_id, block_id) - self._check_failure() - - def copy_blocks(self, copies: Any) -> None: - if self._pending is not None: - self._deferred_state.append((self.copy_blocks, (list(copies),))) - return - with self.state._lock: - for copy in copies: - if copy.dst_block_id == 0: - for group_id in self.state.group_sizes: - self.state._violation((group_id, 0), None, "copy", "write_null_block") - elif not any((group_id, copy.dst_block_id) in self.state.blocks for group_id in self.state.group_sizes): - self.gap("copy_destination_unbound", block_id=copy.dst_block_id) - for group_id in self.state.group_sizes: - pairs = [ - (copy.src_block_id, copy.dst_block_id) - for copy in copies - if (group_id, copy.dst_block_id) in self.state.blocks - ] - self.state.copy_blocks(group_id, pairs) - self._check_failure() + slots = self._slots(metadata, group) + live = self.state.live_blocks(group) + if slots is not None and any(slot >= 0 and slot // size in live for slot in slots): + self.report("guard.violation", reason="dummy_write_to_live_block", group_id=group) + self.check_failure() + finally: + self._phase, self._metadata, self._request_ids = previous -def get_guard(vllm_config: Any) -> GuardRuntime | None: +def create_guard(vllm_config: Any) -> GuardRuntime | None: if not KVGuardConfig.from_vllm_config(vllm_config).enabled: return None - guard = getattr(vllm_config, "_ascend_kv_guard", None) - if guard is None: - guard = GuardRuntime(vllm_config) - vllm_config._ascend_kv_guard = guard - return guard - - -class _V2RunnerAdapter: - """Only V2 batch ordering, metadata timing and capture need adaptation.""" - - def __init__(self, runner: Any, runtime: GuardRuntime): - self.runner = runner - self.runtime = runtime - self.input_batch = SimpleNamespace(req_ids=[]) - self._prepared = False - - def configure(self) -> None: - tables = self.runner.block_tables - sizes = [group.kv_cache_spec.block_size for group in self.runner.kv_cache_config.kv_cache_groups] - if list(tables.block_sizes) != sizes or list(tables.kernel_block_sizes) != sizes: - raise ValueError("KV guard: unadapted MRV2 allocator/kernel block coordinates") - prepare = self.runner.model_state.prepare_attn - if getattr(prepare, "_ascend_kv_guard", False): - return - - @functools.wraps(prepare) - def observed(input_batch, cudagraph_mode, *args, **kwargs): - metadata = prepare(input_batch, cudagraph_mode, *args, **kwargs) - if self.runtime._phase is not None: - self.runtime._finish_pending() - self.input_batch = SimpleNamespace(req_ids=list(input_batch.req_ids[: input_batch.num_reqs])) - self._prepared = True - self.runtime._full_graph = "FULL" in str(cudagraph_mode) - self.runtime.observe_metadata(metadata) - scheduled = self.runtime._scheduled - if self.runtime._phase == "forward" and scheduled is not None: - ids = self.input_batch.req_ids - if len(ids) != len(set(ids)) or set(ids) != set(scheduled): - self.runtime.gap("v2_batch_request_mismatch") - self.runtime._missing_groups.update(self.runtime.layer_groups.values()) - self.runtime._check_failure() - return metadata - - cast(Any, observed)._ascend_kv_guard = True - self.runner.model_state.prepare_attn = observed - - def wrap_execute(self, execute: Any): - @functools.wraps(execute) - def observed(scheduler_output, *args, **kwargs): - # Forward positional and keyword arguments unchanged across vLLM lanes. - flags = dict(zip(("intermediate_tensors", "dummy_run", "skip_attn_for_dummy_run", "is_profile"), args)) - flags.update(kwargs) - dummy = bool(flags.get("dummy_run", False)) - profile = bool(flags.get("is_profile", False)) - phase = "warmup" if profile else "dummy" if dummy else "forward" - required = bool(scheduler_output.total_num_scheduled_tokens) and not ( - profile and flags.get("skip_attn_for_dummy_run") - ) - if dummy or profile: - self.runtime._finish_pending() - self.input_batch = SimpleNamespace(req_ids=[]) - self._prepared = False - - # Synthetic dummy/profile outputs have no scheduler binding, but - # their writes must still be checked against the shared live KV state. - with ( - self.runtime.execute_scope(scheduler_output, bind=not (dummy or profile), finish_connector=False), - self.runtime.forward_scope(self, phase=phase, required=required), - ): - result = execute(scheduler_output, *args, **kwargs) - if phase == "forward" and required and not self._prepared: - self.runtime.gap("v2_attention_metadata_missing") - self.runtime._missing_groups.update(self.runtime.layer_groups.values()) - return self.runtime._guard_output(result) if result is not None else result - - return observed - - def wrap_update(self, update: Any): - @functools.wraps(update) - def observed(scheduler_output, *args, **kwargs): - result = update(scheduler_output, *args, **kwargs) - # The parent has completed the physical operations at this point. - zeroed = getattr(scheduler_output, "new_block_ids_to_zero", None) - if zeroed: - self.runtime.zero_blocks(zeroed) - copies = getattr(scheduler_output, "kv_cache_block_copies", None) - if copies: - self.runtime.copy_blocks(copies) - return result - - return observed - - def wrap_capture(self, capture: Any): - @functools.wraps(capture) - def observed(*args, **kwargs): - self.runtime._check_failure() - with self.runtime.state._lock: - if ( - self.runtime.state.requests - or self.runtime.state.loads - or any(block.valid_end != 0 for block in self.runtime.state.blocks.values()) - ): - raise ValueError( - "KV guard: capture requires empty KV state; live/cache-preserving recapture is unsupported" - ) - # Startup capture has no request KV to protect. Do not synchronize - # metadata to the CPU inside a device graph capture. Real replays - # are audited by wrap_execute using their current metadata. - return capture(*args, **kwargs) - - return observed - - -def install_guard(runner: Any) -> GuardRuntime | None: - """Install one per-engine guard with only the runner-specific hooks.""" - runtime = get_guard(runner.vllm_config) - if runtime is None: - return None - if getattr(runner.vllm_config, "use_v2_model_runner", False): - adapter = _V2RunnerAdapter(runner, runtime) - runtime._runner_adapter = adapter - cast(Any, runner).execute_model = adapter.wrap_execute(runner.execute_model) - cast(Any, runner).update_requests = adapter.wrap_update(runner.update_requests) - cast(Any, runner).capture_model = adapter.wrap_capture(runner.capture_model) - else: - from vllm.forward_context import get_forward_context - - cast(Any, runner)._model_forward = runtime.wrap_forward(runner, runner._model_forward, get_forward_context) - cast(Any, runner).execute_model = runtime.wrap_execute(runner.execute_model) - cast(Any, runner)._update_states = runtime.wrap_update(runner, runner._update_states) - if hasattr(runner, "sample_tokens"): - cast(Any, runner).sample_tokens = runtime.wrap_sample(runner.sample_tokens) - return runtime + return GuardRuntime(vllm_config) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py index 65792897fd6e..bbe1321c9a51 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_connector.py @@ -756,8 +756,7 @@ def _handle_request(self, req_meta: dict[str, Any]): finally: all_tasks_done = self._mark_request_task_done(request_id, all_task_done) if all_tasks_done: - transfer_failed = transfer_failed or self._is_failed_recv_request(request_id) - if transfer_failed: + if transfer_failed or self._is_failed_recv_request(request_id): with self.pending_reformat_lock: self.pending_reformat.pop(request_id, None) else: diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 816a60b17423..57cbd3b3f0c8 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -1870,20 +1870,3 @@ # Remove this patch when upstream provides a backend capability hook for # enabling MoE sequence parallelism with DP=1. # -# ** File: platform/patch_kv_guard.py ** -# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# 1. `vllm.v1.core.sched.scheduler.Scheduler` -# Why: -# Optional KV address diagnostics need allocation generations and request -# bindings independently of worker attention metadata. -# How: -# Wrap the constructor once, then instrument allocation, release, and -# schedule only for instances with additional_config.kv_guard_config.enabled. -# Attach changed bindings and token ranges to SchedulerOutput. -# Include receive-only allocation snapshots and explicit prefix replay ranges. -# Related PR (if no, explain why): -# No upstream PR yet; this is the opt-in Ascend KV Guard integration. -# Future Plan: -# Replace the wrappers with upstream allocation lifecycle and scheduler -# output extension points that provide the same independent evidence. -# diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index 30236e978fbe..7ce09f058feb 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -25,9 +25,6 @@ import vllm_ascend.patch.platform.patch_pp_mtp # noqa import vllm_ascend.patch.platform.patch_use_v2_model_runner # noqa from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile -from vllm_ascend.patch.platform.patch_kv_guard import apply_patch as _apply_kv_guard - -_apply_kv_guard() if get_current_hardware_profile().supports(HardwareCapability.STANDARD_MAMBA_PATCH): import vllm_ascend.patch.platform.patch_mamba_config # noqa diff --git a/vllm_ascend/patch/platform/patch_kv_guard.py b/vllm_ascend/patch/platform/patch_kv_guard.py deleted file mode 100644 index a21a41b7a9e3..000000000000 --- a/vllm_ascend/patch/platform/patch_kv_guard.py +++ /dev/null @@ -1,182 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Transport allocation generations to the Worker L1 state machine. - -No lifecycle trace, extra publisher, token hashes or per-token log messages. -Changed bindings, releases and compact per-request write ranges use the existing channel. -""" - -import functools -import inspect -import uuid -from typing import Any, cast - -from vllm_ascend.debug.kv_guard_config import KVGuardConfig - - -def install_scheduler_guard(scheduler_cls): - if getattr(scheduler_cls.__init__, "_ascend_kv_guard", False): - return - original_init = scheduler_cls.__init__ - - @functools.wraps(original_init) - def initialize(self, *args, **kwargs): - original_init(self, *args, **kwargs) - if not KVGuardConfig.from_vllm_config(self.vllm_config).enabled: - return - pool = self.kv_cache_manager.block_pool - counters = [0] * len(pool.blocks) - pool_id = uuid.uuid4().hex - allocate = pool.get_new_blocks - - @functools.wraps(allocate) - def allocated(*args, **kwargs): - blocks = allocate(*args, **kwargs) - for block in blocks: - counters[block.block_id] += 1 - return blocks - - pool.get_new_blocks = allocated - schedule = self.schedule - mappings: dict[str, tuple] = {} - released: set[str] = set() - loads: dict[str, dict] = {} - replays: dict[str, tuple[int, int]] = {} - load_sequence = 0 - - def block_refs(request_id): - return [ - [(block, counters[block]) for block in group] - for group in self.kv_cache_manager.get_block_ids(request_id) - ] - - allocate_slots = getattr(self.kv_cache_manager, "allocate_slots", None) - if allocate_slots is not None: - signature = inspect.signature(allocate_slots) - external_position = list(signature.parameters).index("num_external_computed_tokens") - 1 - - @functools.wraps(allocate_slots) - def slots_allocated(request, *args, **kwargs): - nonlocal load_sequence - external = kwargs.get( - "num_external_computed_tokens", args[external_position] if len(args) > external_position else 0 - ) - if not external: - return allocate_slots(request, *args, **kwargs) - arguments = signature.bind(request, *args, **kwargs).arguments - local_end = request.num_computed_tokens + arguments.get("num_new_computed_tokens", 0) - result = allocate_slots(request, *args, **kwargs) - if result is not None: - load_sequence += 1 - loads[request.request_id] = { - "pool_id": pool_id, - "load_id": load_sequence, - "request_id": request.request_id, - "block_table": block_refs(request.request_id), - "token_start": local_end, - "token_end": min( - local_end + external, getattr(self.kv_cache_manager, "max_model_len", local_end + external) - ), - "is_async": arguments.get("delay_cache_blocks", False), - } - return result - - self.kv_cache_manager.allocate_slots = slots_allocated - - def record_replay(request, start, end): - if start < end: - previous = replays.get(request.request_id, (start, end)) - replays[request.request_id] = min(start, previous[0]), max(end, previous[1]) - - mark_replay = getattr(self, "_mark_prefix_replay", None) - if mark_replay is not None: - - @functools.wraps(mark_replay) - def replay_marked(request, num_hit_tokens, *args, **kwargs): - count = mark_replay(request, num_hit_tokens, *args, **kwargs) - record_replay(request, num_hit_tokens - count, num_hit_tokens) - return count - - self._mark_prefix_replay = replay_marked - - update_waiting = getattr(self, "_update_waiting_for_remote_kv", None) - if update_waiting is not None: - - @functools.wraps(update_waiting) - def waiting_updated(request, *args, **kwargs): - before = request.num_computed_tokens - result = update_waiting(request, *args, **kwargs) - record_replay(request, request.num_computed_tokens, before) - return result - - self._update_waiting_for_remote_kv = waiting_updated - - free = self.kv_cache_manager.free - - @functools.wraps(free) - def freed(request, *args, **kwargs): - result = free(request, *args, **kwargs) - released.add(request.request_id) - loads.pop(request.request_id, None) - replays.pop(request.request_id, None) - return result - - self.kv_cache_manager.free = freed - - @functools.wraps(schedule) - def scheduled(*args, **kwargs): - output = schedule(*args, **kwargs) - computed = {r.req_id: r.num_computed_tokens for r in output.scheduled_new_reqs} - cached = output.scheduled_cached_reqs - computed.update(zip(cached.req_ids, cached.num_computed_tokens)) - changed = {} - for request_id in output.finished_req_ids | released: - mappings.pop(request_id, None) - if request_id in output.finished_req_ids: - replays.pop(request_id, None) - for request_id in dict.fromkeys((*output.num_scheduled_tokens, *loads)): - refs = block_refs(request_id) - signature = tuple(tuple(group) for group in refs) - if mappings.get(request_id) != signature: - changed[request_id] = refs - mappings[request_id] = signature - output._ascend_kv_guard_context = { - "pool_id": pool_id, - "requests": changed, - "released_request_ids": sorted(output.finished_req_ids | released), - "scheduled_requests": { - request_id: (computed.get(request_id), count) - for request_id, count in output.num_scheduled_tokens.items() - }, - } - if loads: - output._ascend_kv_guard_context["loads"] = list(loads.values()) - loads.clear() - replay_ranges = [] - for request_id in output.num_scheduled_tokens: - replay = replays.pop(request_id, None) - if replay is not None: - replay_ranges.append( - { - "pool_id": pool_id, - "request_id": request_id, - "block_table": block_refs(request_id), - "token_start": replay[0], - "token_end": replay[1], - } - ) - if replay_ranges: - output._ascend_kv_guard_context["recompute_ranges"] = replay_ranges - released.clear() - return output - - self.schedule = scheduled - - cast(Any, initialize)._ascend_kv_guard = True - scheduler_cls.__init__ = initialize - - -def apply_patch(): - from vllm.v1.core.sched.scheduler import Scheduler - - install_scheduler_guard(Scheduler) diff --git a/vllm_ascend/worker/v1/model_runner.py b/vllm_ascend/worker/v1/model_runner.py index d08c40bb53d0..99b1335787e5 100644 --- a/vllm_ascend/worker/v1/model_runner.py +++ b/vllm_ascend/worker/v1/model_runner.py @@ -152,7 +152,7 @@ ) from vllm_ascend.compilation.breakable_aclgraph import BreakableACLGraphWrapper from vllm_ascend.core.kv_cache_interface import is_circular_kv_cache_spec -from vllm_ascend.debug.kv_guard_runtime import install_guard +from vllm_ascend.debug.kv_guard_runtime import GuardRuntime, create_guard from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h import get_prebound_copy_sfa_slots from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise.layerwise_cache_layout import ( @@ -379,6 +379,8 @@ class ExecuteModelState(NamedTuple): class NPUModelRunner(GPUModelRunner): + _kv_guard: GuardRuntime | None = None + # vLLM #51718 describes compatible KV cache groups as views over one # standardized backing allocation. The default runner preserves that # contract for its layer/block-compact Attention+Mamba path. @@ -398,7 +400,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - self._kv_guard = install_guard(self) + self._kv_guard = create_guard(vllm_config) self.lookback_token_ids: CpuGpuBuffer | None = None self.device_metadata_executor: DeviceMetadataExecutor | None = None @@ -1031,9 +1033,6 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.copy_blocks(scheduler_output.kv_cache_block_copies) scheduler_output = replace( scheduler_output, new_block_ids_to_zero=[], @@ -2354,6 +2353,15 @@ def execute_model( self, scheduler_output: "SchedulerOutput", intermediate_tensors: IntermediateTensors | None = None, + ) -> ModelRunnerOutput | IntermediateTensors | None: + guard = self._kv_guard + with guard.execution(self, scheduler_output) if guard is not None else nullcontext(): + return self._execute_model(scheduler_output, intermediate_tensors) + + def _execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: IntermediateTensors | None = None, ) -> ModelRunnerOutput | IntermediateTensors | None: self._cpp_execution_time_ms = None profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config @@ -2888,6 +2896,8 @@ def _skip_drafting( def sample_tokens( self, grammar_output: "GrammarOutput | None" ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + if self._kv_guard is not None: + self._kv_guard.check_failure() kv_connector_output = self.kv_connector_output self.kv_connector_output = None pp = get_pp_group() @@ -3579,6 +3589,12 @@ def _model_forward( if executor is not None and executor.submission_in_flight: executor.release() + if self._kv_guard is not None: + self._kv_guard.observe_metadata( + forward_context.attn_metadata, + request_ids=list(self.input_batch.req_ids), + full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL, + ) return hidden_states def _prepare_device_metadata_for_forward( @@ -4581,18 +4597,17 @@ def dummy_drafter_compute_logits(hidden_states): if not is_graph_capturing and self.ascend_config.enable_force_eplb \ and self.vllm_config.model_config.is_moe: build_force_eplb_topk(self.device, self.max_num_tokens) - guard_kwargs: dict[str, Any] = {} - if getattr(self, "_kv_guard", None) is not None: - guard_kwargs["_kv_guard_phase"] = "warmup" if is_profile or is_graph_capturing else "dummy" - outputs = self._model_forward( - num_tokens_padded, - input_ids, - positions, - intermediate_tensors, - inputs_embeds, - decoder_replay_dummy=True, - **guard_kwargs, - ) + guard = self._kv_guard + phase = "warmup" if is_profile or is_graph_capturing else "dummy" + with guard.dummy_scope(self, phase=phase) if guard is not None else nullcontext(): + outputs = self._model_forward( + num_tokens_padded, + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + decoder_replay_dummy=True, + ) if active_device_metadata_executor is not None and active_device_metadata_executor.submission_in_flight: active_device_metadata_executor.release() self.kvpp.complete_forward() @@ -5084,9 +5099,6 @@ def _fill_c8_mxfp_v_scale_caches(self, kv_caches: dict[str, Any]) -> None: def _zero_block_ids(self, block_ids: list[int]) -> None: """Prepare freshly assigned attention pages, including static C8 scales.""" super()._zero_block_ids(block_ids) - guard = getattr(self, "_kv_guard", None) - if guard is not None: - guard.zero_blocks(block_ids) # MRV1 can receive new attention block IDs without a KV zeroer. Scale # restoration must still run: Mamba owns the same physical pages while # those IDs are assigned to its cache groups. @@ -5119,7 +5131,7 @@ def initialize_kv_cache_tensors( kv_caches = self._reshape_kv_cache_tensors( kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes ) - guard = getattr(self, "_kv_guard", None) + guard = self._kv_guard if guard is not None: guard.configure(self) @@ -6957,6 +6969,8 @@ def _check_and_update_cudagraph_mode( def capture_model(self) -> int: """Capture NPU graphs and return actual graph pool memory bytes consumed.""" + if self._kv_guard is not None: + self._kv_guard.check_capture() parent_module_name = _get_gpu_model_runner_module_name(self) self._engram_capture_active = True try: diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index cfe20cc29940..225b95ad68f7 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -76,7 +76,7 @@ _finish_profiling_chunk_timing, _start_profiling_chunk_timing, ) -from vllm_ascend.debug.kv_guard_runtime import install_guard +from vllm_ascend.debug.kv_guard_runtime import GuardRuntime, create_guard from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h.dspark_context import ( bind_dspark_context_receiver, configure_dspark_kv_transfer, @@ -141,6 +141,8 @@ class NPUModelRunner(GPUModelRunner): """Model runner for Ascend NPUs.""" + _kv_guard: GuardRuntime | None = None + # vLLM #51718 overlays hybrid Attention/Mamba groups in one standardized # backing allocation. Ascend MRV2 preserves that allocation and exposes # backend-specific views; MLA may use a page-strided fused/component view. @@ -289,7 +291,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.debugger = AclGraphDumper(dump_cfg) self.runtime_guard = RuntimeGuardProcessor.bind(self) - self._kv_guard = install_guard(self) + self._kv_guard = create_guard(vllm_config) @property def pcp_manager_cls(self) -> type[AscendPCPManager]: @@ -323,6 +325,8 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: @runtime_guard_sample_tokens def sample_tokens(self, grammar_output): + if self._kv_guard is not None: + self._kv_guard.check_failure() if ( lmhead_tp_enable() and self.prompt_logprobs_worker is not None @@ -520,7 +524,7 @@ def initialize_kv_cache( self.model_state.kvpp_runtime = self.kvpp self._fill_c8_mxfp_v_scale_caches() - guard = getattr(self, "_kv_guard", None) + guard = self._kv_guard if guard is not None: guard.configure(self) @@ -583,6 +587,8 @@ def _register_sparse_kv_caches(self, kv_caches: dict[str, Any]) -> None: @torch.inference_mode() def capture_model(self, *, profile_only: bool = False) -> int: """Capture decoder and multimodal encoder graphs on Ascend.""" + if self._kv_guard is not None: + self._kv_guard.check_capture() if hasattr(self, "model_state") and self.model_state.supports_mm_inputs: encoder_runner = getattr(self.model_state, "encoder_runner", None) encoder_manager = getattr(encoder_runner, "cudagraph_manager", None) @@ -641,7 +647,13 @@ def execute_model( ) profile_kwargs = {"randomize_inputs": randomize_inputs} forward_failed = True - with dp_coordination_context: + guard_scope = nullcontext() + if self._kv_guard is not None: + if dummy_run or is_profile: + guard_scope = self._kv_guard.dummy_scope(self, phase="warmup" if is_profile else "dummy") + else: + guard_scope = self._kv_guard.execution(self, scheduler_output) + with guard_scope, dp_coordination_context: try: output = super().execute_model( scheduler_output, diff --git a/vllm_ascend/worker/v2/model_states/default.py b/vllm_ascend/worker/v2/model_states/default.py index 73a290369f97..972535950fd4 100644 --- a/vllm_ascend/worker/v2/model_states/default.py +++ b/vllm_ascend/worker/v2/model_states/default.py @@ -62,6 +62,7 @@ def get_extra_common_attn_kwargs(self, cache_group_idx: int, num_reqs: int) -> d if TYPE_CHECKING: from vllm_ascend.attention.dsa_v41 import AscendDSAV41Metadata, AscendDSAV41MetadataBuilder from vllm_ascend.attention.utils import AscendCommonAttentionMetadata + from vllm_ascend.debug.kv_guard_runtime import GuardRuntime from vllm_ascend.models.deepseek_v41.decoder_replay_layers import AscendDecoderReplayLayers from vllm_ascend.worker.device_metadata import TargetDeviceMetadata from vllm_ascend.worker.v2.kvpp import KVPPRuntime @@ -71,6 +72,7 @@ def get_extra_common_attn_kwargs(self, cache_group_idx: int, num_reqs: int) -> d class AscendModelState(DefaultModelState): """Model state for Ascend NPUs.""" + kv_guard: "GuardRuntime | None" = None pcp_manager: "AscendPCPManager | None" = None pcp_context: "AscendPCPAttentionContext | None" = None kvpp_runtime: "KVPPRuntime | None" = None @@ -84,6 +86,7 @@ def finish_execution(self, *, failed: bool) -> None: def __init__(self, vllm_config, model, encoder_cache, device): super().__init__(vllm_config, model, encoder_cache, device) + self.kv_guard = None self._offload_req_ids: CpuGpuBuffer | None = None self._offload_token_to_req: CpuGpuBuffer | None = None self._offload_pool_slots: CpuGpuBuffer | None = None @@ -573,4 +576,10 @@ def prepare_attn( if getattr(metadata, "copy_sfa_copy_src_offsets", None) is not None: manager.restore_copy_sfa_tails(metadata) break + if self.kv_guard is not None and not for_capture: + self.kv_guard.observe_metadata( + self.attn_metadata, + request_ids=list(input_batch.req_ids[: input_batch.num_reqs]), + full_graph=cudagraph_mode == CUDAGraphMode.FULL, + ) return self.attn_metadata From 02ddd24301005b5b0a6fa23109a7884e7b692a4b Mon Sep 17 00:00:00 2001 From: TAJh Date: Sat, 7 Nov 2026 10:59:32 +0000 Subject: [PATCH 6/6] refactor(debug): limit KV metadata checks to Model Runner V2 Remove the V1 runner integration and compatibility fallbacks. Require explicit V2 batch identities and model state, reject enabled guard configuration with V1, and restrict the service regression and documentation to V2. Signed-off-by: TAJh --- .../configuration/additional_config.md | 2 +- .../feature_guide/kv_cache_guard.md | 11 +++-- release_notes/fragments/17833.yaml | 2 +- .../one_card/test_kv_cache_guard.py | 5 +- tests/ut/debug/test_kv_guard_runtime.py | 13 ++++-- tests/ut/debug/test_kv_guard_scope.py | 24 ++++++++-- tests/ut/test_ascend_config.py | 30 ++++++++++++ vllm_ascend/ascend_config.py | 2 + vllm_ascend/debug/kv_guard_runtime.py | 19 ++++---- vllm_ascend/worker/v1/model_runner.py | 46 ++++--------------- 10 files changed, 89 insertions(+), 65 deletions(-) diff --git a/docs/source/user_guide/configuration/additional_config.md b/docs/source/user_guide/configuration/additional_config.md index 258ae463297d..279b2d467b00 100644 --- a/docs/source/user_guide/configuration/additional_config.md +++ b/docs/source/user_guide/configuration/additional_config.md @@ -38,7 +38,7 @@ The following table lists additional configuration options available in vLLM Asc |-------------------------------------|------|---------|-----------------------------------------------------------------------------------------------------------| | `xlite_graph_config` | dict | `{}` | Configuration options for Xlite graph mode | | `finegrained_tp_config` | dict | `{}` | Configuration options for fine-grained tensor parallelism. See [Fine-Grained Tensor Parallelism](../feature_guide/Fine_grained_TP.md). | -| `kv_guard_config` | dict | `{}` | Per-engine KV cache address diagnostics for Model Runner V1 and V2. `enabled` defaults to `false`; `action` accepts `"log"` (default) or `"raise"`. See [KV cache address guard](../feature_guide/kv_cache_guard.md). | +| `kv_guard_config` | dict | `{}` | Per-engine KV cache address diagnostics for Model Runner V2. `enabled` defaults to `false`; `action` accepts `"log"` (default) or `"raise"`. See [KV cache address guard](../feature_guide/kv_cache_guard.md). | | `ascend_compilation_config` | dict | `{}` | Configuration options for ascend compilation | | `ascend_warmup_config` | dict | `{}` | Configuration options for startup warmup that overlaps weight loading | | `eplb_config` | dict | `{}` | Runner-specific EPLB extensions. See [Expert Parallelism Load Balancer](../feature_guide/expert_parallelism_load_balancer.md). | diff --git a/docs/source/user_guide/feature_guide/kv_cache_guard.md b/docs/source/user_guide/feature_guide/kv_cache_guard.md index fc96e882a087..30a47d700eec 100644 --- a/docs/source/user_guide/feature_guide/kv_cache_guard.md +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -1,16 +1,16 @@ # KV cache address guard -KV Guard compares the scheduler's declared token destinations with the runner's +KV Guard in Model Runner V2 compares the scheduler's declared token destinations with the runner's final attention metadata, once per KV cache group. Eager and FULL graph modes use the same metadata-plan checks. Attention implementations are unchanged; operator arguments and per-layer execution order are not recorded. ## Enable the guard -The guard is disabled by default. Enable it for an engine with: +The guard is disabled by default and requires Model Runner V2. Enable it for an engine with: ```bash -vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' +VLLM_USE_V2_MODEL_RUNNER=1 vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' ``` | Option | Type | Default | Meaning | @@ -21,6 +21,7 @@ vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action The Python API accepts the same `additional_config` dictionary. Unknown fields, non-boolean `enabled` values and unknown actions are rejected. Each runner owns its guard; the configuration object does not hold mutable checking state. +Enabling the guard with Model Runner V1 raises a configuration error. `raise` latches the first violation or incomplete audit. Catching an exception does not reset the guard; restart the worker before reuse. There is no collective @@ -61,7 +62,7 @@ methods are not replaced at runtime. ## Coverage boundaries -Dense PA and MLA append-only cache paths in Runner V1/V2 are adapted. The guard +Dense PA and MLA append-only cache paths in Runner V2 are adapted. The guard checks consistency between scheduler messages and selected metadata. It does not prove that a particular layer executed, that operator arguments matched the metadata, or that a kernel used those addresses. @@ -103,7 +104,7 @@ Run focused tests with: pytest -q --confcutdir=tests/ut/debug tests/ut/debug ``` -The service regression compares disabled/enabled greedy outputs for V1/V2, +The service regression compares disabled/enabled greedy outputs for Runner V2, eager/FULL, repeated prompts and concurrent requests: ```bash diff --git a/release_notes/fragments/17833.yaml b/release_notes/fragments/17833.yaml index ff08aa0f3b79..3f10abeb3d0e 100644 --- a/release_notes/fragments/17833.yaml +++ b/release_notes/fragments/17833.yaml @@ -3,7 +3,7 @@ entries: - category: feature text: >- Added opt-in KV cache metadata-plan diagnostics for dense PA and MLA paths - in Runner V1/V2. Each cache group is checked once using final attention + in Runner V2. Each cache group is checked once using final attention metadata in both eager and FULL modes; attention implementations are unchanged and operator/per-layer accesses are not recorded. Standard scheduler messages provide the independent expected address plan, without diff --git a/tests/e2e/pull_request/one_card/test_kv_cache_guard.py b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py index 0de7114cc7d6..c68f37af22fe 100644 --- a/tests/e2e/pull_request/one_card/test_kv_cache_guard.py +++ b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py @@ -8,7 +8,6 @@ @pytest.mark.e2e_model("Qwen/Qwen3-0.6B") -@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) @pytest.mark.parametrize( "graph_mode", [ @@ -39,7 +38,7 @@ ], ) @wait_until_npu_memory_free() -def test_guard_preserves_sequential_and_concurrent_generation(use_v2, graph_mode): +def test_guard_preserves_sequential_and_concurrent_generation(graph_mode): model = "Qwen/Qwen3-0.6B" prompt = "Explain how a key value cache works. " * 40 arguments = [ @@ -65,7 +64,7 @@ def test_guard_preserves_sequential_and_concurrent_generation(use_v2, graph_mode results = [] for enabled in (False, True): port = get_open_port() - environment = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2))} + environment = {"VLLM_USE_V2_MODEL_RUNNER": "1"} with ( RemoteOpenAIServer( model, diff --git a/tests/ut/debug/test_kv_guard_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py index b0c843c6dfe6..8983daff3c29 100644 --- a/tests/ut/debug/test_kv_guard_runtime.py +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize("full_graph", [False, True]) -def test_prefix_append_is_a_metadata_plan_for_both_runner_modes(runtime, full_graph): +def test_v2_prefix_append_is_a_metadata_plan_for_eager_and_full(runtime, full_graph): with runtime.execution(NS(), message(new=[("r", ([2, 3],), 3)], counts={"r": 3})): observe(runtime, ["r"], [11, 12, 13], [[2, 3]], [6], full_graph=full_graph) assert not runtime.events @@ -27,7 +27,7 @@ def test_full_graph_padding_requires_zero_lengths_after_real_rows(runtime, lengt def test_corrupted_runner_table_cannot_change_shadow_plan(runtime): - runner = NS(input_batch=NS(req_ids=["r"], block_table=[[99]])) + runner = NS(block_tables=NS(block_table=[[99]])) with runtime.execution(runner, message(new=[("r", ([2],), 0)], counts={"r": 1})): observe(runtime, ["r"], [396], [[99]], [1]) assert "write_target_mismatch" in reasons(runtime) @@ -141,7 +141,7 @@ def recorded(value): def test_full_graph_requires_fresh_metadata_every_step(runtime): with runtime.execution(NS(), message(new=[("r", ([2],), 0)], counts={"r": 1})): observe(runtime, ["r"], [8], [[2]], [1], full_graph=True) - with runtime.execution(NS(input_batch=NS(req_ids=["r"])), message(cached=[("r", None, 1)], counts={"r": 1})): + with runtime.execution(NS(), message(cached=[("r", None, 1)], counts={"r": 1})): pass assert "metadata_group_missing" in reasons(runtime) @@ -153,6 +153,13 @@ def test_batch_identity_mismatch_is_a_coverage_gap(runtime, batch): assert "batch_request_mismatch" in reasons(runtime) +def test_actual_batch_ids_must_be_explicitly_observed(runtime): + runner = NS(input_batch=NS(req_ids=["r"])) + with runtime.execution(runner, message(new=[("r", ([2],), 0)], counts={"r": 1})): + runtime.observe_metadata({"layer0": metadata([8], [[2]], [1])}) + assert "batch_request_ids_missing" in reasons(runtime) + + def test_no_token_step_can_release_without_a_forward(runtime): runtime.state.requests = {"r": ([2],)} with runtime.execution(NS(), message(finished=["r"])): diff --git a/tests/ut/debug/test_kv_guard_scope.py b/tests/ut/debug/test_kv_guard_scope.py index 1f5e8fc0558e..004491bf611c 100644 --- a/tests/ut/debug/test_kv_guard_scope.py +++ b/tests/ut/debug/test_kv_guard_scope.py @@ -13,7 +13,11 @@ def runner_for(config, impl_name="AscendAttentionBackendImpl"): impl = type(impl_name, (), {})() config.compilation_config.static_forward_context = {"layer0": NS(impl=impl)} spec = NS(block_size=4, head_size=16) - return NS(kv_cache_config=NS(kv_cache_groups=[NS(kv_cache_spec=spec, layer_names=["layer0"])])), impl + return NS( + kv_cache_config=NS(kv_cache_groups=[NS(kv_cache_spec=spec, layer_names=["layer0"])]), + model_state=NS(kv_guard=None), + block_tables=NS(block_sizes=[4], kernel_block_sizes=[4]), + ), impl def test_disabled_has_no_runtime_or_config_side_effect(subject, config): @@ -42,14 +46,28 @@ def test_config_rejects_invalid_options(subject, config, options): @pytest.mark.parametrize("impl_name", ["AscendAttentionBackendImpl", "AscendMLAImpl"]) def test_attention_only_validated_and_declared_model_state_bound(runtime, config, impl_name): runner, impl = runner_for(config, impl_name) - runner.model_state = NS(kv_guard=None) - runner.block_tables = NS(block_sizes=[4], kernel_block_sizes=[4]) runtime.configure(runner) assert not hasattr(impl, "kv_guard") assert runner.model_state.kv_guard is runtime assert runtime.layer_groups == {"layer0": 0} +@pytest.mark.parametrize("model_state", [None, NS()]) +def test_v2_model_state_must_declare_guard_member(runtime, config, model_state): + runner, impl = runner_for(config) + runner.model_state = model_state + with pytest.raises(ValueError, match="unsupported V2 model state"): + runtime.configure(runner) + assert not hasattr(impl, "kv_guard") + + +def test_runner_without_v2_model_state_is_rejected(runtime, config): + runner, _ = runner_for(config) + del runner.model_state + with pytest.raises(ValueError, match="unsupported V2 model state"): + runtime.configure(runner) + + def test_empty_cache_groups_cannot_claim_guard_enabled(runtime): with pytest.raises(ValueError, match="no supported attention layers"): runtime.configure(NS(kv_cache_config=NS(kv_cache_groups=[]))) diff --git a/tests/ut/test_ascend_config.py b/tests/ut/test_ascend_config.py index 3087e413b8e2..6789fc36d18b 100644 --- a/tests/ut/test_ascend_config.py +++ b/tests/ut/test_ascend_config.py @@ -52,6 +52,7 @@ init_ascend_config, is_mega_moe_supported, ) +from vllm_ascend.debug.kv_guard_config import KVGuardConfig from vllm_ascend.device.hardware import AscendDeviceType from vllm_ascend.device.hardware_profile import get_hardware_profile from vllm_ascend.utils import clear_enable_sp, enable_sp, shared_expert_dp_enabled @@ -2233,6 +2234,35 @@ def test_config_factory_keeps_kvpp_enabled(self, _check_config): self.assertEqual(actual.kvpp_config.size, 4) +class TestKVGuardRunnerSelection(TestBase): + def test_enabled_rejects_model_runner_v1(self): + config = VllmConfig(device_config=DeviceConfig(device="cpu")) + ascend_config = AscendConfig( + kv_guard_config=KVGuardConfig(enabled=True), sparse_kv_offload_config=SparseKVOffloadConfig() + ) + with ( + patch.object(_VllmConfig, "use_v2_model_runner", new_callable=PropertyMock, return_value=False), + self.assertRaisesRegex(ValueError, "requires Model Runner V2"), + ): + ascend_config.derive_and_validate(config) + + def test_enabled_accepts_model_runner_v2(self): + config = VllmConfig(device_config=DeviceConfig(device="cpu")) + ascend_config = AscendConfig( + kv_guard_config=KVGuardConfig(enabled=True), sparse_kv_offload_config=SparseKVOffloadConfig() + ) + with patch.object(_VllmConfig, "use_v2_model_runner", new_callable=PropertyMock, return_value=True): + ascend_config.derive_and_validate(config) + self.assertTrue(ascend_config.kv_guard_config.enabled) + + def test_disabled_preserves_model_runner_v1(self): + config = VllmConfig(device_config=DeviceConfig(device="cpu")) + ascend_config = AscendConfig(sparse_kv_offload_config=SparseKVOffloadConfig()) + with patch.object(_VllmConfig, "use_v2_model_runner", new_callable=PropertyMock, return_value=False): + ascend_config.derive_and_validate(config) + self.assertFalse(ascend_config.kv_guard_config.enabled) + + class TestEngramSharedMemoryDefaults(TestBase): def test_shared_tables_are_derived_for_mrv2_dp_or_pcp(self): topologies = ( diff --git a/vllm_ascend/ascend_config.py b/vllm_ascend/ascend_config.py index a546f9c4d2a5..f8e69b112c44 100644 --- a/vllm_ascend/ascend_config.py +++ b/vllm_ascend/ascend_config.py @@ -692,6 +692,8 @@ def _validate_user_input_ranges(self): # the max_num_batched_tokens that sequence-parallel writeback corrected). def derive_and_validate(self, vllm_config: VllmConfig) -> AscendConfig: vc = vllm_config + if self.kv_guard_config.enabled and not vc.use_v2_model_runner: + raise ValueError("KV guard requires Model Runner V2") engram_config = getattr(vc, "engram_config", None) if ( engram_config is not None diff --git a/vllm_ascend/debug/kv_guard_runtime.py b/vllm_ascend/debug/kv_guard_runtime.py index 6927a4cda2e8..ae19b4a6a828 100644 --- a/vllm_ascend/debug/kv_guard_runtime.py +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -103,15 +103,11 @@ def configure(self, runner: Any) -> None: list(tables.block_sizes) != list(sizes.values()) or list(tables.kernel_block_sizes) != list(sizes.values()) ): raise ValueError("KV guard: unsupported allocator/kernel block coordinates") - batch_tables = getattr(getattr(getattr(runner, "input_batch", None), "block_table", None), "block_tables", ()) - if batch_tables and any(table.block_size != sizes[index] for index, table in enumerate(batch_tables)): - raise ValueError("KV guard: unsupported allocator/kernel block coordinates") model_state = getattr(runner, "model_state", None) - if model_state is not None and not hasattr(model_state, "kv_guard"): - raise ValueError("KV guard: unsupported model state") + if model_state is None or not hasattr(model_state, "kv_guard"): + raise ValueError("KV guard: unsupported V2 model state") self.state.configure(sizes) - if model_state is not None: - model_state.kv_guard = self + model_state.kv_guard = self self.report("guard.enabled", failure_action=self.action) def observe_metadata(self, metadata: Any, request_ids: list[str] | None = None, full_graph: bool = False) -> None: @@ -145,7 +141,7 @@ def _slots(self, metadata: Any, group: int) -> list[int] | None: slots = [slot for row in slots for slot in (row if isinstance(row, list) else [row])] return slots - def _audit(self, runner: Any) -> None: + def _audit(self) -> None: if not self.state.scheduled: for group, metadata in self._metadata.items(): slots = self._slots(metadata, group) @@ -156,13 +152,14 @@ def _audit(self, runner: Any) -> None: self.report("guard.violation", reason="unscheduled_attention_plan", group_id=group) return request_ids = self._request_ids - if request_ids is None: - request_ids = list(getattr(getattr(runner, "input_batch", None), "req_ids", ())) for group in self.state.group_sizes: metadata = self._metadata.get(group) if metadata is None: self.gap("metadata_group_missing", group_id=group) continue + if request_ids is None: + self.gap("batch_request_ids_missing", group_id=group) + continue if len(request_ids) != len(set(request_ids)) or set(request_ids) != set(self.state.scheduled): self.gap("batch_request_mismatch", group_id=group) continue @@ -205,7 +202,7 @@ def execution(self, runner: Any, scheduler_output: Any): self.check_failure() yield try: - self._audit(runner) + self._audit() except Exception as exc: self.gap(f"access_audit_failed:{type(exc).__name__}") self.check_failure() diff --git a/vllm_ascend/worker/v1/model_runner.py b/vllm_ascend/worker/v1/model_runner.py index 99b1335787e5..5584227654e7 100644 --- a/vllm_ascend/worker/v1/model_runner.py +++ b/vllm_ascend/worker/v1/model_runner.py @@ -152,7 +152,6 @@ ) from vllm_ascend.compilation.breakable_aclgraph import BreakableACLGraphWrapper from vllm_ascend.core.kv_cache_interface import is_circular_kv_cache_spec -from vllm_ascend.debug.kv_guard_runtime import GuardRuntime, create_guard from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.distributed.kv_transfer.kv_p2p.sfa_pd_rd2h import get_prebound_copy_sfa_slots from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise.layerwise_cache_layout import ( @@ -379,8 +378,6 @@ class ExecuteModelState(NamedTuple): class NPUModelRunner(GPUModelRunner): - _kv_guard: GuardRuntime | None = None - # vLLM #51718 describes compatible KV cache groups as views over one # standardized backing allocation. The default runner preserves that # contract for its layer/block-compact Attention+Mamba path. @@ -399,8 +396,6 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): with _torch_cuda_wrapper(): super().__init__(vllm_config, device) - - self._kv_guard = create_guard(vllm_config) self.lookback_token_ids: CpuGpuBuffer | None = None self.device_metadata_executor: DeviceMetadataExecutor | None = None @@ -2353,15 +2348,6 @@ def execute_model( self, scheduler_output: "SchedulerOutput", intermediate_tensors: IntermediateTensors | None = None, - ) -> ModelRunnerOutput | IntermediateTensors | None: - guard = self._kv_guard - with guard.execution(self, scheduler_output) if guard is not None else nullcontext(): - return self._execute_model(scheduler_output, intermediate_tensors) - - def _execute_model( - self, - scheduler_output: "SchedulerOutput", - intermediate_tensors: IntermediateTensors | None = None, ) -> ModelRunnerOutput | IntermediateTensors | None: self._cpp_execution_time_ms = None profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config @@ -2896,8 +2882,6 @@ def _skip_drafting( def sample_tokens( self, grammar_output: "GrammarOutput | None" ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: - if self._kv_guard is not None: - self._kv_guard.check_failure() kv_connector_output = self.kv_connector_output self.kv_connector_output = None pp = get_pp_group() @@ -3589,12 +3573,6 @@ def _model_forward( if executor is not None and executor.submission_in_flight: executor.release() - if self._kv_guard is not None: - self._kv_guard.observe_metadata( - forward_context.attn_metadata, - request_ids=list(self.input_batch.req_ids), - full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL, - ) return hidden_states def _prepare_device_metadata_for_forward( @@ -4597,17 +4575,14 @@ def dummy_drafter_compute_logits(hidden_states): if not is_graph_capturing and self.ascend_config.enable_force_eplb \ and self.vllm_config.model_config.is_moe: build_force_eplb_topk(self.device, self.max_num_tokens) - guard = self._kv_guard - phase = "warmup" if is_profile or is_graph_capturing else "dummy" - with guard.dummy_scope(self, phase=phase) if guard is not None else nullcontext(): - outputs = self._model_forward( - num_tokens_padded, - input_ids, - positions, - intermediate_tensors, - inputs_embeds, - decoder_replay_dummy=True, - ) + outputs = self._model_forward( + num_tokens_padded, + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + decoder_replay_dummy=True, + ) if active_device_metadata_executor is not None and active_device_metadata_executor.submission_in_flight: active_device_metadata_executor.release() self.kvpp.complete_forward() @@ -5131,9 +5106,6 @@ def initialize_kv_cache_tensors( kv_caches = self._reshape_kv_cache_tensors( kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes ) - guard = self._kv_guard - if guard is not None: - guard.configure(self) # Set up cross-layer KV cache sharing for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): @@ -6969,8 +6941,6 @@ def _check_and_update_cudagraph_mode( def capture_model(self) -> int: """Capture NPU graphs and return actual graph pool memory bytes consumed.""" - if self._kv_guard is not None: - self._kv_guard.check_capture() parent_module_name = _get_gpu_model_runner_module_name(self) self._engram_capture_active = True try: