diff --git a/docs/source/user_guide/configuration/additional_config.md b/docs/source/user_guide/configuration/additional_config.md index 0d556ee1653b..279b2d467b00 100644 --- a/docs/source/user_guide/configuration/additional_config.md +++ b/docs/source/user_guide/configuration/additional_config.md @@ -38,6 +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 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 new file mode 100644 index 000000000000..30a47d700eec --- /dev/null +++ b/docs/source/user_guide/feature_guide/kv_cache_guard.md @@ -0,0 +1,118 @@ +# KV cache address guard + +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 and requires Model Runner V2. Enable it for an engine with: + +```bash +VLLM_USE_V2_MODEL_RUNNER=1 vllm serve MODEL --additional-config '{"kv_guard_config":{"enabled":true,"action":"raise"}}' +``` + +| Option | Type | Default | Meaning | +| --- | --- | --- | --- | +| `enabled` | bool | `false` | Enable worker metadata-plan checks. | +| `action` | `"log"` or `"raise"` | `"log"` | Report failures or stop returning worker results. | + +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 +failure coordination across workers. + +## Independent execution plan + +The worker consumes the unmodified `SchedulerOutput` before its runner updates: + +- 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 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 + +Run focused tests with: + +```bash +pytest -q --confcutdir=tests/ut/debug tests/ut/debug +``` + +The service regression compares disabled/enabled greedy outputs for Runner V2, +eager/FULL, repeated prompts and concurrent requests: + +```bash +pytest -q tests/e2e/pull_request/one_card/test_kv_cache_guard.py +``` + +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/mkdocs.yml b/mkdocs.yml index 137103f1a3fc..a3a13f810f80 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -296,6 +296,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..3f10abeb3d0e --- /dev/null +++ b/release_notes/fragments/17833.yaml @@ -0,0 +1,13 @@ +pr: 17833 +entries: + - category: feature + text: >- + Added opt-in KV cache metadata-plan diagnostics for dense PA and MLA paths + 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 + 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/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..c68f37af22fe --- /dev/null +++ b/tests/e2e/pull_request/one_card/test_kv_cache_guard.py @@ -0,0 +1,96 @@ +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( + "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(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": "1"} + with ( + RemoteOpenAIServer( + model, + [ + *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", + 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/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 new file mode 100644 index 000000000000..9cd366386e1e --- /dev/null +++ b/tests/ut/debug/test_kv_guard.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shadow bindings use standard scheduler messages, never runner tables.""" + +import pytest + +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_runtime.py b/tests/ut/debug/test_kv_guard_runtime.py new file mode 100644 index 000000000000..8983daff3c29 --- /dev/null +++ b/tests/ut/debug/test_kv_guard_runtime.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One current execution plan per group is compared with an independent mirror.""" + +from types import SimpleNamespace as NS + +import pytest +import torch + +from .helpers import message, metadata, observe, reasons + + +@pytest.mark.parametrize("full_graph", [False, True]) +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 + assert not hasattr(runtime, "write") + assert not hasattr(runtime, "read") + + +@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 test_corrupted_runner_table_cannot_change_shadow_plan(runtime): + 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) + assert "read_block_mapping_mismatch" in reasons(runtime) + assert runtime.state.requests["r"] == ([2],) + + +@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) + + +@pytest.mark.parametrize( + "slots,tables,lengths,reason", + [ + ([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_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_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_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_group_reports_only_that_group(runtime): + runtime.state.configure({0: 4, 1: 8}) + runtime.layer_groups = {"layer0": 0, "layer1": 1} + 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(), 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_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"])): + 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_is_synchronous_and_sticky_without_modifying_output_getter(runtime, subject): + runtime.action = "raise" + 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 new file mode 100644 index 000000000000..004491bf611c --- /dev/null +++ b/tests/ut/debug/test_kv_guard_scope.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-runner ownership, explicit unsupported modes, and capture isolation.""" + +from types import SimpleNamespace as NS + +import pytest + +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"])]), + 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): + 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) + 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=[]))) + + +@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: + 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/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 821deb648037..f8e69b112c44 100644 --- a/vllm_ascend/ascend_config.py +++ b/vllm_ascend/ascend_config.py @@ -32,6 +32,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 @@ -433,6 +434,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, "sfa_dcp_force_tmajor_restore": false, @@ -582,6 +584,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" sfa_dcp_force_tmajor_restore: bool = False @@ -689,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/__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..6cac42e51c93 --- /dev/null +++ b/vllm_ascend/debug/kv_guard.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent request/block mirror built only from standard SchedulerOutput. + +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 + +from dataclasses import dataclass +from typing import Any + + +@dataclass(slots=True) +class AccessPlan: + slots: list[int] + tables: list[list[int]] + lengths: list[int] + + +class KVGuard: + def __init__(self, emit: Any): + self.emit = emit + self.requests: dict[str, tuple[list[int], ...]] = {} + self.scheduled: dict[str, tuple[int, int]] = {} + self.group_sizes: dict[int, int] = {} + + def configure(self, group_sizes: dict[int, int]) -> None: + 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: + 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: + 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 new file mode 100644 index 000000000000..42cb5b42de33 --- /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 options for worker KV metadata-plan diagnostics.""" + +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 new file mode 100644 index 000000000000..ae19b4a6a828 --- /dev/null +++ b/vllm_ascend/debug/kv_guard_runtime.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare one attention execution plan per cache group with SchedulerOutput. + +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 + +from collections import OrderedDict +from contextlib import contextmanager +from typing import Any + +from vllm_ascend.debug.kv_guard import KVGuard +from vllm_ascend.debug.kv_guard_config import KVGuardConfig + +_MAX_REPORT_KEYS = 1024 + + +class KVGuardError(RuntimeError): + """A violation or missing coverage prevents further model output.""" + + +class GuardRuntime: + def __init__(self, vllm_config: Any): + from vllm.logger import logger + + self.logger = logger + self.config = vllm_config + self.action = KVGuardConfig.from_vllm_config(vllm_config).action + self.state = KVGuard(self.report) + self.layer_groups: dict[str, int] = {} + self._failure: str | None = None + self._reported: OrderedDict[tuple, None] = OrderedDict() + self._phase: str | None = None + self._reset_step() + + 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: + 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) + + 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, 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), + ) + ): + raise ValueError("KV guard: unsupported connector, CP, DBO, speculation or encoder/decoder mode") + sizes = {} + self.layer_groups = {} + for group_id, group in enumerate(runner.kv_cache_config.kv_cache_groups): + sizes[group_id] = group.kv_cache_spec.block_size + for name in group.layer_names: + 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 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") + model_state = getattr(runner, "model_state", None) + if model_state is None or not hasattr(model_state, "kv_guard"): + raise ValueError("KV guard: unsupported V2 model state") + self.state.configure(sizes) + 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)): + 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 _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 + 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) -> None: + 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 = self._request_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 + 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) + 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 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 + try: + self._audit() + except Exception as exc: + self.gap(f"access_audit_failed:{type(exc).__name__}") + self.check_failure() + finally: + self._phase = None + self._reset_step() + + def check_capture(self) -> None: + self.check_failure() + if self.state.requests: + raise ValueError("KV guard: capture requires no live requests") + + @contextmanager + 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: + yield + 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 + 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 create_guard(vllm_config: Any) -> GuardRuntime | None: + if not KVGuardConfig.from_vllm_config(vllm_config).enabled: + return None + return GuardRuntime(vllm_config) diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py index 48cc352d7222..225b95ad68f7 100644 --- a/vllm_ascend/worker/v2/model_runner.py +++ b/vllm_ascend/worker/v2/model_runner.py @@ -76,6 +76,7 @@ _finish_profiling_chunk_timing, _start_profiling_chunk_timing, ) +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, @@ -140,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. @@ -288,6 +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 = create_guard(vllm_config) @property def pcp_manager_cls(self) -> type[AscendPCPManager]: @@ -321,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 @@ -518,6 +524,9 @@ def initialize_kv_cache( self.model_state.kvpp_runtime = self.kvpp self._fill_c8_mxfp_v_scale_caches() + guard = self._kv_guard + if guard is not None: + guard.configure(self) def _fill_c8_mxfp_v_scale_caches(self) -> None: """Initialize C8 V scales and retain views for page reassignment.""" @@ -578,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) @@ -636,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