[Refactor] Declare model-local KV held outside the paged manager (#6171)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

docs/design/feature/model_local_kv_caches.md

@@ -0,0 +1,132 @@
+# Model-local KV caches
+
+Some models keep their attention KV as HuggingFace `transformers` cache objects
+instead of using the engine's paged KV manager. Each such model declares what it
+holds by implementing `model_local_kv_specs()`, and
+`OmniGPUModelRunner.load_model` sums and logs the result after load.
+
+## What is and is not already accounted for
+
+This distinction decides what the protocol is for, and an earlier revision of
+this page got it wrong by claiming none of this memory was accounted for.
+
+CUDA-graph capture happens *during* `load_model` -- MiMo in its constructor,
+Qwen3-TTS in `load_weights`. `determine_available_memory()` then samples
+per-process GPU memory through NVML and subtracts it from the requested budget
+(`vllm_omni/worker/base.py`). So the **graph-resident** portion is already
+charged. Subtracting the declared total from the KV budget would double-count
+it.
+
+What is *not* charged is the **per-request** portion, which is allocated after
+profiling. That half is not small: Qwen3-TTS's codec decoder is 4.44 MiB per
+request, so a server at `max_num_seqs=256` reaches roughly 1.1 GiB that no
+per-stage number predicted.
+
+The consumer therefore reports and does not subtract.
+
+## What the four models declare
+
+| Model | Cache | Bounded by | Per row | Rows |
+|---|---|---|---|---|
+| Qwen3-TTS codec decoder (async-chunk only) | sliding `DynamicCache` | `sliding_window - 1 = 71` | 4.44 MiB | `max_num_seqs`, x2 for the working copy |
+| Qwen3-TTS graph pool | retained captures | same | 4.44 MiB | fixed, from captured shapes |
+| MiMo-Audio local transformer | `DynamicCache` | `group_size + max(delay_pattern) = 11` | 704 KiB | `max_num_seqs` |
+| MiMo-Audio graph pool | captured | same | 704 KiB | fixed 261 (bucket sum) |
+| MiniCPM-o Whisper encoder (streaming only) | `EncoderDecoderCache` | `embed_positions` rows = 1500 | 140.62 MiB | fixed 1, per session |
+| ming_flash_omni talker | `StaticCache` | hardcoded `max_cache_len = 2048` | 24.00 MiB | fixed 1 |
+
+**No cache is bounded by `max_model_len`.** Each has its own mechanism: a
+sliding window, a decode-loop trip count, an encoder frame limit, a constant.
+
+**ming's geometry does not exist in this repo.** `self._llm_config` is a
+`Qwen2Config` from the checkpoint, so layers, kv-heads, head-dim and dtype are
+only knowable after load. That is what makes this a runtime query rather than a
+table.
+
+## Lifetime is not multiplicity
+
+`scope` records when memory appears and goes. It does **not** decide how wide
+the allocation is, and an earlier revision that derived scaling from it was
+wrong in two different directions at once:
+
+- ming's talker is per-call but serialized -- `forward` takes
+  `runtime_additional_information[0]` and runs the whole AR loop inline -- so
+  scaling it by `max_num_seqs` over-reported it by that factor.
+- MiniCPM-o's encoder is per *session*, not per sequence, so `max_num_seqs` is
+  simply the wrong number for it.
+
+So width is declared explicitly, by naming the driver. `RowDriver` has two
+values, because two are all the known caches need:
+
+- `FIXED` -- a count the model controls, with `rows_fixed` and a required
+  `rows_reason`. `rows_fixed=1` on a cache that looks per-request is a claim a
+  reviewer must be able to check.
+- `MAX_NUM_SEQS` -- one row per in-flight sequence; only the engine knows the
+  value.
+
+A third value belongs here when a model actually has a cache the engine widens
+by some other number. An earlier revision added one speculatively, for a single
+declarer, and paid for it with a class hierarchy, an engine-capacity object and
+a conditional-annotation field.
+
+`rows` counts row-equivalents at peak and deliberately does not distinguish
+"one allocation N rows wide" from "N allocations of one row". Those cost the
+same, and a revision that tried to carry both got the object topology wrong in
+three of four declarers: MiniCPM-o keeps one cache object per session while
+asserting batch size 1, and both graph pools keep one distinct object per
+captured bucket. The layout belongs in `allocation_note`, where it cannot be
+mistaken for arithmetic.
+
+## Declaring nothing is a valid declaration
+
+A declaration must be inert when its path is. Qwen3-TTS holds KV only in
+async-chunk mode -- stateless decoding runs `_forward_exact`, which calls the
+transformer without `use_cache` -- so the stateless path declares nothing.
+Getting this wrong is easy in both directions, and both have happened here: one
+revision hung the declaration off the CUDA-graph wrapper and reported zero under
+`enforce_eager`, where the eager async-chunk path does allocate; the next
+declared unconditionally and invented roughly a gigabyte for stateless
+deployments.
+
+Where the model cannot see the deciding setting, it declares the honest unit
+and says so in `allocation_note`. MiniCPM-o's encoder cache exists only on the
+streaming path and is one object per session; the session cap is not visible
+from the model, so it declares one session rather than scaling by a number it
+would have to guess at.
+
+## Finding the declarers
+
+`collect_model_local_kv_specs()` walks `named_modules()`. Two wrinkles:
+
+- The runner's `self.model` may be a `CUDAGraphWrapper`, which is not an
+  `nn.Module` but forwards attribute access to its runnable. The walk therefore
+  works either way, but entering through both the wrapper and its runnable would
+  count a root-level declarer twice. The collector calls the supported
+  `unwrap()` accessor first.
+- An owner that is not a submodule at all (Qwen3-TTS's graph wrapper, ming's
+  `MingAudioGenerator`) is reached by one explicit hop from the module that
+  holds it.
+
+## Adding a model
+
+Return one entry per distinct allocation. Prefer `spec_from_hf_config()` and
+pass the same config object the cache is built from; it raises rather than
+guessing when an attribute is missing. `capacity_source` is diagnostic text --
+never branch on it.
+
+## Known limits
+
+Row counts taken from live state after warmup miss caches `transformers` has
+not yet materialized, which makes Qwen3-TTS's graph-pool figure a floor for the
+xvec path.
+
+MiMo's graph-pool cache has no reachable Python object -- it is a local inside
+the captured frame -- so its width comes from the captured bucket list rather
+than from tensors. Its two entries are also mutually exclusive per call: a call
+either replays a bucket or runs eagerly. The sum is the right ceiling for the
+process and the wrong number for a single request.
+
+MiMo captures every bucket gated only on `torch.cuda.is_available()`
+(`mimo_audio_llm.py:670`), so those 179.44 MiB stay resident even when the
+deploy config sets `enforce_eager: true`. Gating it is a one-line change that
+needs a model boot to verify, so it is noted rather than bundled here.

docs/design/index.md

@@ -26,6 +26,7 @@ implementation contract; it is not, by itself, a general support claim.
 - [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
 - [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md)
 - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
+- [Model-local KV Caches](feature/model_local_kv_caches.md)
 - [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
 
 ### Communication

tests/model_executor/models/test_model_local_kv.py

@@ -0,0 +1,404 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""A declaration must match the cache that actually gets allocated.
+
+Two earlier versions of this file could not fail. The first asserted hand-typed
+constants against each other. The second built real `transformers` caches but
+compared them to specs written in the test, so no production declaration was
+ever executed. What matters here is that a real declarer is called and its
+numbers are checked against a real cache built the way the model builds it.
+"""
+
+import ast
+import pathlib
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers import Qwen2Config
+from transformers.cache_utils import DynamicCache, StaticCache
+
+from vllm_omni.model_executor.models.model_local_kv import (
+    HasModelLocalKV,
+    ModelLocalKVScope,
+    ModelLocalKVSpec,
+    RowDriver,
+    collect_model_local_kv_specs,
+    spec_from_hf_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+MODELS_DIR = pathlib.Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "models"
+
+
+def _total_bytes(model: object, max_num_seqs: int = 1) -> int:
+    """Sum declared peak bytes across a model's declarations."""
+    return sum(spec.peak_bytes(max_num_seqs) for _, spec in collect_model_local_kv_specs(model))
+
+
+def _cache_bytes(cache) -> int:
+    total = 0
+    for layer in cache.layers:
+        for tensor in (layer.keys, layer.values):
+            if tensor is not None:
+                total += tensor.numel() * tensor.element_size()
+    return total
+
+
+def _qwen2_config(layers=4, kv_heads=2, heads=8, hidden=256) -> Qwen2Config:
+    return Qwen2Config(
+        num_hidden_layers=layers,
+        num_attention_heads=heads,
+        num_key_value_heads=kv_heads,
+        hidden_size=hidden,
+        intermediate_size=hidden * 2,
+        vocab_size=1000,
+    )
+
+
+# --------------------------------------------------------------------------
+# A real declarer, called, checked against a real cache.
+# --------------------------------------------------------------------------
+
+
+def test_ming_declaration_matches_the_cache_it_describes():
+    """Calls MingAudioGenerator.model_local_kv_specs and checks the bytes.
+
+    Builds the declarer without its weights -- only `_llm_config` and `_model`
+    are read -- then builds the `StaticCache` that `_init_kv_cache` builds from
+    the same config and compares. A wrong layer count, kv-head count, dtype or
+    capacity in the production declaration fails here.
+    """
+    from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator
+
+    config = _qwen2_config()
+    generator = object.__new__(MingAudioGenerator)
+    generator._llm_config = config
+    generator._model = torch.nn.Linear(1, 1).to(torch.float16)
+
+    (spec,) = generator.model_local_kv_specs()
+
+    cache = StaticCache(
+        config=config,
+        max_batch_size=1,
+        max_cache_len=MingAudioGenerator._STATIC_CACHE_LEN,
+        device="cpu",
+        dtype=torch.float16,
+    )
+    head_dim = config.hidden_size // config.num_attention_heads
+    for layer_idx in range(config.num_hidden_layers):
+        keys = torch.zeros(1, config.num_key_value_heads, 1, head_dim, dtype=torch.float16)
+        cache.update(keys, keys.clone(), layer_idx)
+
+    assert spec.peak_bytes(1) == _cache_bytes(cache)
+
+
+def test_ming_does_not_scale_with_max_num_seqs():
+    """forward() takes runtime_additional_information[0] and loops inline.
+
+    One worker holds one of these no matter how many sequences the scheduler
+    admits. A previous revision derived the multiplier from scope and reported
+    `max_num_seqs` times the real figure.
+    """
+    from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator
+
+    generator = object.__new__(MingAudioGenerator)
+    generator._llm_config = _qwen2_config()
+    generator._model = torch.nn.Linear(1, 1).to(torch.float16)
+
+    (spec,) = generator.model_local_kv_specs()
+    assert spec.peak_bytes(64) == spec.peak_bytes(1)
+
+
+def test_growing_cache_declaration_matches_a_real_dynamic_cache():
+    """The declared positions are the extent the tensor reaches when full."""
+    layers, kv_heads, head_dim, positions = 4, 2, 32, 11
+    config = _qwen2_config(layers=layers, kv_heads=kv_heads, heads=kv_heads, hidden=kv_heads * head_dim)
+    spec = spec_from_hf_config(
+        config,
+        name="local_transformer",
+        dtype=torch.bfloat16,
+        physical_capacity_positions=positions,
+        capacity_source="group_size + max(delay_pattern)",
+        scope=ModelLocalKVScope.INVOCATION,
+        rows=RowDriver.FIXED,
+        rows_reason="test",
+    )
+
+    cache = DynamicCache()
+    for layer_idx in range(layers):
+        keys = torch.zeros(1, kv_heads, positions, head_dim, dtype=torch.bfloat16)
+        cache.update(keys, torch.zeros_like(keys), layer_idx)
+
+    assert spec.peak_bytes() == _cache_bytes(cache)
+
+
+def test_multiplicity_does_not_claim_an_object_layout():
+    """One allocation B rows wide costs the same as B one-row allocations.
+
+    The protocol reports bytes and deliberately does not distinguish the two.
+    An earlier revision carried a separate `allocations` count and got the
+    object topology wrong in three of four declarers, which is worse than not
+    claiming it: the layout belongs in `allocation_note`, where it cannot be
+    mistaken for arithmetic.
+    """
+    config = _qwen2_config(layers=4, kv_heads=2, heads=2, hidden=64)
+    common = dict(
+        config=config,
+        name="c",
+        dtype=torch.bfloat16,
+        physical_capacity_positions=11,
+        capacity_source="test",
+        scope=ModelLocalKVScope.INVOCATION,
+    )
+    batched = spec_from_hf_config(rows=RowDriver.MAX_NUM_SEQS, **common)
+
+    assert batched.row_count(64) == 64
+    assert batched.peak_bytes(64) == batched.bytes_per_row * 64
+    assert not hasattr(batched, "allocations"), "the object-count axis was removed on purpose"
+
+
+def test_a_fixed_driver_ignores_concurrency():
+    """A serialized or capture-pinned cache does not widen with max_num_seqs."""
+    spec = spec_from_hf_config(
+        _qwen2_config(),
+        name="graph_pool",
+        dtype=torch.bfloat16,
+        physical_capacity_positions=11,
+        capacity_source="test",
+        scope=ModelLocalKVScope.MODEL,
+        rows=RowDriver.FIXED,
+        rows_fixed=261,
+        rows_reason="sum of captured buckets",
+    )
+    assert spec.row_count(1) == 261
+    assert spec.row_count(64) == 261
+    assert spec.peak_bytes(64) == spec.bytes_per_row * 261
+
+
+def test_qwen3_tts_declares_nothing_on_the_stateless_path():
+    """Calls the real decoder declarer with no cache-holding mode configured.
+
+    `_forward_exact` runs `pre_transformer` without `use_cache`, so stateless
+    decoding allocates no KV. A revision that declared a per-request cache
+    unconditionally invented over a gigabyte at max_num_seqs=256, and the
+    revision before it declared nothing at all under enforce_eager, where the
+    eager async-chunk path does allocate. Both are regressions this pins.
+    """
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import (
+        modeling_qwen3_tts_tokenizer_v2 as mod,
+    )
+
+    decoder = object.__new__(mod.Qwen3TTSTokenizerV2Decoder)
+    # nn.Module.__setattr__ needs the machinery __init__ would have built, and
+    # this instance deliberately has none of it.
+    setattr_ = object.__setattr__
+    setattr_(
+        decoder,
+        "config",
+        SimpleNamespace(
+            sliding_window=72,
+            num_hidden_layers=8,
+            num_key_value_heads=16,
+            head_dim=64,
+            hidden_size=1024,
+            num_attention_heads=16,
+        ),
+    )
+    setattr_(decoder, "parameters", lambda: iter([torch.zeros(1, dtype=torch.float32)]))
+
+    setattr_(decoder, "_cudagraph_wrapper", None)
+    assert decoder.model_local_kv_specs() == []
+
+    setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=False, model_local_kv_specs=lambda: []))
+    assert decoder.model_local_kv_specs() == []
+
+    setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=True, model_local_kv_specs=lambda: []))
+    names = [s.name for s in decoder.model_local_kv_specs()]
+    assert names == ["codec_decoder_request", "codec_decoder_working_copy"], names
+
+
+def test_fixed_rejects_unreviewable_values():
+    """A bare or unexplained constant cannot be reviewed, so it is rejected."""
+    common = dict(
+        name="x",
+        layers=1,
+        kv_heads=1,
+        head_dim=1,
+        dtype=torch.float16,
+        physical_capacity_positions=1,
+        capacity_source="",
+        scope=ModelLocalKVScope.REQUEST,
+    )
+    with pytest.raises(ValueError, match="rows_reason"):
+        ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=1, rows_reason="   ", **common)
+    with pytest.raises(ValueError, match=">= 1"):
+        ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=0, rows_reason="test", **common)
+
+
+def test_rows_is_required():
+    """Defaulting the batch extent is how the previous guarantee died.
+
+    The dataclass declared it required while the helper every declarer uses
+    defaulted it to 1, so nothing enforced it where it mattered.
+    """
+    with pytest.raises(TypeError):
+        spec_from_hf_config(  # type: ignore[call-arg]
+            _qwen2_config(),
+            name="x",
+            dtype=torch.float16,
+            physical_capacity_positions=1,
+            capacity_source="",
+            scope=ModelLocalKVScope.REQUEST,
+        )
+
+
+def test_spec_from_hf_config_raises_rather_than_guessing():
+    class Empty:
+        pass
+
+    with pytest.raises(ValueError, match="layer count"):
+        spec_from_hf_config(
+            Empty(),
+            name="x",
+            dtype=torch.float16,
+            physical_capacity_positions=1,
+            capacity_source="",
+            scope=ModelLocalKVScope.REQUEST,
+            rows=RowDriver.FIXED,
+            rows_reason="test",
+        )
+
+
+# --------------------------------------------------------------------------
+# Collection
+# --------------------------------------------------------------------------
+
+
+def _spec(**overrides) -> ModelLocalKVSpec:
+    kwargs = dict(
+        name="c",
+        layers=2,
+        kv_heads=2,
+        head_dim=4,
+        dtype=torch.float16,
+        physical_capacity_positions=8,
+        capacity_source="test",
+        scope=ModelLocalKVScope.REQUEST,
+        rows=RowDriver.FIXED,
+        rows_reason="test",
+    )
+    kwargs.update(overrides)
+    return ModelLocalKVSpec(**kwargs)
+
+
+class _Owner(torch.nn.Module):
+    def model_local_kv_specs(self):
+        return [_spec()]
+
+
+class _CUDAGraphWrapperLike:
+    """Mirrors vllm.compilation.cuda_graph.CUDAGraphWrapper.
+
+    Not an nn.Module, and its __getattr__ forwards everything the runnable
+    has. That forwarding is why a previous revision's "unwrap fix" was
+    unnecessary, and why a walk that does not unwrap can enter the tree twice.
+    """
+
+    def __init__(self, runnable):
+        self.runnable = runnable
+
+    def unwrap(self):
+        return self.runnable
+
+    def __getattr__(self, key):
+        return getattr(self.runnable, key)
+
+
+def test_collector_finds_owners_nested_in_the_module_tree():
+    class Mid(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.inner = _Owner()
+
+    class Top(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.mid = Mid()
+
+    model = Top()
+    assert [path for path, _ in collect_model_local_kv_specs(model)] == ["mid.inner"]
+    # 2 layers * 2 kv * 4 head_dim * 8 positions * 2 bytes * 2 (K and V) = 512
+    assert _total_bytes(model) == 512
+
+
+def test_a_root_declarer_behind_a_wrapper_is_counted_once():
+    """The wrapper forwards attributes, so a naive walk sees it twice."""
+
+    class Root(_Owner):
+        pass
+
+    wrapped = _CUDAGraphWrapperLike(Root())
+    collected = collect_model_local_kv_specs(wrapped)
+
+    assert len(collected) == 1, collected
+    assert _total_bytes(wrapped) == 512
+
+
+def test_collector_reaches_owners_through_a_wrapper():
+    class Top(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.owner = _Owner()
+
+    wrapped = _CUDAGraphWrapperLike(Top())
+    assert [path for path, _ in collect_model_local_kv_specs(wrapped)] == ["owner"]
+
+
+def test_collector_survives_an_owner_that_raises():
+    class Broken(torch.nn.Module):
+        def model_local_kv_specs(self):
+            raise RuntimeError("checkpoint field missing")
+
+    class Top(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.broken = Broken()
+            self.ok = _Owner()
+
+    assert _total_bytes(Top()) == 512
+
+
+def test_undeclared_models_report_zero_not_an_error():
+    assert _total_bytes(object()) == 0
+    assert _total_bytes(torch.nn.Linear(2, 2)) == 0
+
+
+def test_protocol_is_structural():
+    assert isinstance(_Owner(), HasModelLocalKV)
+    assert not isinstance(torch.nn.Linear(2, 2), HasModelLocalKV)
+
+
+@pytest.mark.parametrize(
+    ("relative_path", "owner"),
+    [
+        ("qwen3_tts/segmented_graph_wrapper.py", "CUDAGraphDecoderWrapper"),
+        ("qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py", "Qwen3TTSTokenizerV2Decoder"),
+        ("mimo_audio/mimo_audio_llm.py", "MiMoAudioLLMForConditionalGeneration"),
+        ("minicpmo_4_5/minicpmo_4_5_omni_llm.py", "MiniCPMWhisperEncoder"),
+        ("ming_flash_omni/talker_module.py", "MingAudioGenerator"),
+        ("ming_flash_omni/ming_flash_omni_talker.py", "MingFlashOmniTalkerForConditionalGeneration"),
+    ],
+)
+def test_known_cache_owners_still_declare(relative_path: str, owner: str):
+    """Parsed rather than imported: these modules need CUDA and real weights.
+
+    Weak by construction -- it only proves the method exists. The ming tests
+    above are the ones that check a declaration's contents.
+    """
+    tree = ast.parse((MODELS_DIR / relative_path).read_text())
+    classes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == owner]
+    assert classes, f"{owner} not found in {relative_path}"
+    methods = {child.name for cls in classes for child in cls.body if isinstance(child, ast.FunctionDef)}
+    assert "model_local_kv_specs" in methods, f"{owner} no longer declares its model-local KV"

tests/tools/test_check_tts_adapter.py

@@ -63,7 +63,15 @@ def test_gate_passes_on_current_tree():
     assert main([]) == 0
 
 
-def test_budgets_are_the_real_counts():
+def test_budgets_are_not_exceeded():
+    """Budgets are a ceiling, not a target.
+
+    This used to assert equality, which re-imposed through pytest exactly what
+    #6008 removed from the checker: a PR that deleted branches failed CI unless
+    it also hand-edited the constant. That is what broke `main` for 5h19m in
+    #5746. The checker reports an under-budget count as a passing notice
+    (check_tts_adapter.py:140); the test must agree.
+    """
     root = Path(__file__).resolve().parents[2]
     serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py"
-    assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES
+    assert len(_model_type_branches(ast.parse(serving.read_text()))) <= MAX_MODEL_TYPE_BRANCHES

vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py

@@ -58,6 +58,12 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema
 
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import MiMoAudioConfig
+from vllm_omni.model_executor.models.model_local_kv import (
+    ModelLocalKVScope,
+    ModelLocalKVSpec,
+    RowDriver,
+    spec_from_hf_config,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -744,6 +750,67 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
     def get_language_model(self) -> torch.nn.Module:
         return self.model
 
+    def model_local_kv_specs(self) -> list[ModelLocalKVSpec]:
+        """Declare the local transformer's decode cache and its graph copies.
+
+        The cache is created per call inside ``base_local_forward`` and is
+        bounded by the decode loop, not by sequence length: the loop runs
+        ``group_size + max(delay_pattern)`` steps regardless of how long the
+        utterance is.
+
+        The second entry is the one that matters. ``base_local_forward`` runs
+        inside ``torch.cuda.graph`` during capture, so each captured bucket
+        leaves a cache's worth of tensors pinned in the graph pool for the
+        process lifetime. There is no Python object left to inspect -- the
+        cache is a local in the captured frame -- so the width comes from the
+        buckets that captured successfully, which is why it reads
+        ``local_forward_cg_by_bs`` rather than MIMO_CUDAGRAPH_BATCH_SIZES: a
+        bucket whose capture raised is logged and skipped.
+
+        Both entries are declared, and summing them over-states a single step:
+        a given call either replays a captured bucket or runs eagerly, never
+        both. The graph pool is resident regardless, so the sum is the right
+        ceiling for the process and the wrong number for one request.
+        """
+        delay_iters = self.group_size + max(self.delay_pattern)
+        dtype = next(self.local_transformer.parameters()).dtype
+        bound = f"group_size({self.group_size}) + max(delay_pattern)({max(self.delay_pattern)})"
+        specs = [
+            spec_from_hf_config(
+                self.local_config,
+                name="local_transformer",
+                dtype=dtype,
+                physical_capacity_positions=delay_iters,
+                capacity_source=bound,
+                scope=ModelLocalKVScope.INVOCATION,
+                # One cache, B rows wide -- not B caches. base_local_forward
+                # builds a single DynamicCache whose batch dimension is the
+                # number of requests in the group.
+                rows=RowDriver.MAX_NUM_SEQS,
+                allocation_note="one batched allocation per eager call; graph replay uses the pool entry instead",
+            )
+        ]
+        captured_rows = sum(self.local_forward_cg_by_bs)
+        if captured_rows:
+            specs.append(
+                spec_from_hf_config(
+                    self.local_config,
+                    name="local_transformer_graph_pool",
+                    dtype=dtype,
+                    physical_capacity_positions=delay_iters,
+                    capacity_source=bound,
+                    scope=ModelLocalKVScope.MODEL,
+                    rows=RowDriver.FIXED,
+                    rows_fixed=captured_rows,
+                    rows_reason=f"sum of captured buckets {sorted(self.local_forward_cg_by_bs)}",
+                    allocation_note=(
+                        "pinned in the CUDA graph pool; capture is gated only on cuda.is_available(), "
+                        "so this stays resident under enforce_eager too"
+                    ),
+                )
+            )
+        return specs
+
     def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
         if kwargs.get("modality_preprocess") is None:
             mm_dummy_embeddings = []

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -158,6 +158,17 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin)
             patch_size=self.patch_size,
         )
 
+    def model_local_kv_specs(self):
+        """Forward to the generator, which owns the talker's StaticCache.
+
+        ``MingAudioGenerator`` is a plain object rather than a submodule, so
+        the module-tree walk cannot reach it on its own.
+        """
+        generator = getattr(self, "audio_generator", None)
+        if generator is None:
+            return []
+        return generator.model_local_kv_specs()
+
     @property
     def device(self) -> torch.device:
         return next(self.model.parameters()).device

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -41,6 +41,12 @@ from vllm_omni.model_executor.models.common.ming.cfm_solver import (
     integrate_cfm_steps,
 )
 from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer
+from vllm_omni.model_executor.models.model_local_kv import (
+    ModelLocalKVScope,
+    ModelLocalKVSpec,
+    RowDriver,
+    spec_from_hf_config,
+)
 
 logger = init_logger(__name__)
 
@@ -575,6 +581,13 @@ class MingAudioGenerator:
     for a single TTS request. The generator is stateless across requests.
     """
 
+    _STATIC_CACHE_LEN = 2048
+    """Positions the talker's StaticCache is preallocated for.
+
+    Read by both ``_init_kv_cache`` and ``model_local_kv_specs`` so the
+    declaration cannot drift from the allocation.
+    """
+
     def __init__(
         self,
         config,
@@ -609,6 +622,49 @@ class MingAudioGenerator:
         # trailing latent frames prepended on each decode call
         self._vae_decode_pad_frames = 32
 
+    def model_local_kv_specs(self) -> list[ModelLocalKVSpec]:
+        """Declare the talker's preallocated decode cache.
+
+        This is the case that a static table cannot hold. ``self._llm_config``
+        is a ``Qwen2Config`` resolved from the checkpoint at load time, so
+        layers, kv-heads and head-dim do not exist anywhere in this repo, and
+        the dtype is whatever the weights loaded as. Reading them off the same
+        objects ``_init_kv_cache`` builds from is the only way to get a real
+        number.
+
+        Unlike the other three this is not a ceiling that a short utterance
+        stays under. transformers v5 initializes ``StaticLayer`` lazily, but
+        the first write to a layer allocates all ``max_cache_len`` positions at
+        once rather than growing, so a prefill brings the full extent into
+        existence on step 0 of every generation.
+
+        Width is fixed at one row, not ``max_num_seqs``. ``forward`` takes only
+        ``runtime_additional_information[0]`` and runs the whole AR loop inline
+        before returning, so one worker holds one cache no matter how many
+        sequences the scheduler admits. Declaring this per-sequence -- which an
+        earlier revision did, by deriving the multiplier from scope -- reported
+        ``max_num_seqs`` times the real figure.
+
+        The caveat is that nothing enforces the serialization. ``StaticCache``
+        is built outside ``_sampler_pool``'s lock, so overlapping calls to
+        ``generate_latents`` would coexist. If this model ever gains a
+        concurrent entry point, this row count is the line to revisit.
+        """
+        return [
+            spec_from_hf_config(
+                self._llm_config,
+                name="talker_llm_static_cache",
+                dtype=next(self._model.parameters()).dtype,
+                physical_capacity_positions=self._STATIC_CACHE_LEN,
+                capacity_source=f"hardcoded max_cache_len={self._STATIC_CACHE_LEN} in _init_kv_cache",
+                scope=ModelLocalKVScope.INVOCATION,
+                rows=RowDriver.FIXED,
+                rows_fixed=1,
+                rows_reason="forward() handles one request and runs the AR loop inline",
+                allocation_note="full extent on first write per layer, not grown; rebuilt per text segment",
+            )
+        ]
+
     @cached_property
     def _sampler_pool(self) -> CFMGraphExecutorPool | None:
         device = next(self._model.parameters()).device
@@ -790,7 +846,7 @@ class MingAudioGenerator:
     def _init_kv_cache(
         self, use_static_cache: bool, device: torch.device, dtype: torch.dtype
     ) -> tuple[StaticCache | None, int]:
-        max_cache_len = 2048
+        max_cache_len = self._STATIC_CACHE_LEN
         if not use_static_cache:
             return None, max_cache_len
         cache = StaticCache(

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -136,6 +136,13 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]:
 
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
+from vllm_omni.model_executor.models.model_local_kv import (
+    ModelLocalKVScope,
+    ModelLocalKVSpec,
+    RowDriver,
+    spec_from_hf_config,
+)
+
 logger = init_logger(__name__)
 hf_logger = logging.get_logger(__name__)
 
@@ -2621,6 +2628,51 @@ class MiniCPMWhisperEncoder(WhisperEncoder):
             [MiniCPMWhisperEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)]
         )
 
+    def model_local_kv_specs(self) -> list[ModelLocalKVSpec]:
+        """Declare the streaming encoder's self-attention cache.
+
+        Bounded by learned position embeddings, not by ``max_model_len``: the
+        cache cannot outgrow ``embed_positions``, and past that point the
+        forward pass repeats the last position rather than extending.
+
+        Whisper is encoder-only self-attention here, so kv-heads equal
+        attention heads. Layer count is read off the built ``self.layers``
+        rather than the config, so a partially built encoder reports what it
+        actually has.
+
+        Declared per session, not per sequence. Each streaming session state
+        owns its own ``audio_past_key_values`` and the encoder runs at batch
+        size one, so ``max_num_seqs`` is the wrong number here.
+
+        One session is the unit rather than the configured cap, because the cap
+        is not visible from the model and the setting that carries it cannot
+        distinguish "duplex off" from "one session" -- so a cap-scaled figure
+        would report a cache that a non-duplex deployment never allocates. An
+        earlier revision grew a capacity driver and an engine-side field to
+        carry that number for this one declarer; the arithmetic is a
+        multiplication the reader can do when they need the ceiling.
+        """
+        return [
+            spec_from_hf_config(
+                self.config,
+                name="whisper_encoder_self_attn",
+                dtype=self.conv1.weight.dtype,
+                layers=len(self.layers),
+                kv_heads=self.config.encoder_attention_heads,
+                head_dim=self.config.d_model // self.config.encoder_attention_heads,
+                physical_capacity_positions=int(self.embed_positions.weight.shape[0]),
+                capacity_source="embed_positions rows (max_source_positions)",
+                scope=ModelLocalKVScope.SESSION,
+                rows=RowDriver.FIXED,
+                rows_fixed=1,
+                rows_reason="one EncoderDecoderCache per streaming session, encoder runs at batch size 1",
+                allocation_note=(
+                    "only allocated on the streaming path, where use_cache is set; scales with the "
+                    "configured duplex session cap"
+                ),
+            )
+        ]
+
     def forward(
         self,
         input_features,

vllm_omni/model_executor/models/model_local_kv.py

@@ -0,0 +1,347 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Declaration protocol for attention KV kept outside the paged manager.
+
+Several models hold their own attention KV as HuggingFace ``transformers`` cache
+objects rather than using the engine's paged KV manager. That memory is
+allocated after the profiling run that decides how much KV the engine may claim,
+so no per-stage footprint number accounts for it, and nothing catches a cache
+that silently stops working.
+
+This module lets a model *declare* what it holds. It describes; it never
+allocates. Allocation for the diffusion path is owned by the DiT KV manager
+(RFC #5244 / PR #6094) and this protocol must not grow into a second allocator.
+
+Why a runtime query rather than a static table: geometry is frequently not
+knowable until weights are loaded. ``ming_flash_omni``'s talker builds its cache
+from a ``Qwen2Config`` that comes from the checkpoint, so layers, kv-heads,
+head-dim and dtype simply do not exist in this repo. A table would have a hole
+in it; a post-load query does not.
+
+Why capacity is a resolved number rather than a formula: the four known caches
+are each bounded by a different mechanism -- a sliding window, a decode-loop
+trip count, an encoder frame limit, a hardcoded constant -- and not one of them
+is ``max_model_len``. Only the model knows its own bound.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+from typing import Protocol, runtime_checkable
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+__all__ = [
+    "HasModelLocalKV",
+    "ModelLocalKVScope",
+    "ModelLocalKVSpec",
+    "RowDriver",
+    "collect_model_local_kv_specs",
+    "spec_from_hf_config",
+]
+
+
+class ModelLocalKVScope(str, Enum):
+    """How long one allocation lives.
+
+    Diagnostic only. Lifetime and size are independent: a per-call cache can be
+    the widest thing a model owns, and a process-lifetime one can be a single
+    row. ``ModelLocalKVSpec.rows`` carries the size; this carries only the
+    reader's mental model of when the memory appears and goes.
+    """
+
+    INVOCATION = "invocation"
+    """Dies when the call returns (e.g. a per-step working copy)."""
+
+    REQUEST = "request"
+    """Retained across steps of one request, released when it finishes."""
+
+    SESSION = "session"
+    """Outlives a request; belongs to a duplex/streaming session."""
+
+    MODEL = "model"
+    """Lives as long as the model (e.g. captured into a CUDA-graph pool)."""
+
+
+class RowDriver(str, Enum):
+    """What sets how many rows of a cache are live at peak.
+
+    Two values, because two are all the known caches need. A third belongs
+    here only when a model actually has a cache the engine widens by some
+    other number -- adding one speculatively costs a branch everywhere and
+    buys nothing.
+    """
+
+    FIXED = "fixed"
+    """A count the model controls: a captured-bucket sum, a serialized call path."""
+
+    MAX_NUM_SEQS = "max_num_seqs"
+    """One row per in-flight sequence. Only the engine knows the value."""
+
+
+@dataclass(frozen=True)
+class ModelLocalKVSpec:
+    """One declared KV allocation.
+
+    A model returns one entry per *distinct lifetime*, not per cache object: if
+    the same geometry exists as both a retained per-request cache and a
+    short-lived working copy, that is two entries.
+    """
+
+    name: str
+    layers: int
+    kv_heads: int
+    head_dim: int
+    dtype: torch.dtype
+
+    physical_capacity_positions: int
+    """Positions the tensor can physically hold.
+
+    Not the logical sequence position. A sliding-window cache truncates on every
+    write, so a stream of thousands of tokens may only ever occupy
+    ``sliding_window - 1`` slots.
+    """
+
+    capacity_source: str
+    """Free-text note on where the bound comes from, for diagnostics only.
+
+    Never branch on this. It exists so a reader can tell "71 because sliding
+    window" from "2048 because someone typed 2048", which is otherwise
+    invisible.
+    """
+
+    scope: ModelLocalKVScope
+    """How long one allocation lives. Diagnostic only.
+
+    Deliberately not the multiplier. An earlier revision derived scaling from
+    scope, which made every non-``MODEL`` cache multiply by ``max_num_seqs``.
+    That over-reported ming by that factor (its calls are serialized) and used
+    the wrong driver entirely for MiniCPM-o (duplex sessions, not sequences).
+    Lifetime does not imply replication topology; ``rows`` says what does.
+    """
+
+    rows: RowDriver
+    """What sets how many rows of this shape are live at peak.
+
+    A row is one batch entry's worth of the geometry above. Deliberately does
+    not distinguish "one allocation N rows wide" from "N allocations of one
+    row": those cost the same, and an earlier revision that tried to carry both
+    got the object topology wrong in three of four declarers. Bytes are what
+    this protocol reports; describe the layout in ``allocation_note``.
+
+    Required, because defaulting it silently understates every cache whose
+    multiplicity is not 1.
+    """
+
+    rows_fixed: int = 1
+    """The count, when ``rows`` is ``FIXED``. Ignored otherwise."""
+
+    rows_reason: str = ""
+    """Why ``rows_fixed`` is that number. Required when ``rows`` is ``FIXED``.
+
+    A bare constant cannot be reviewed. ``rows_fixed=1`` on a cache that looks
+    per-request is exactly the claim a reader has to be able to check.
+    """
+
+    allocation_note: str | None = None
+    """Optional detail such as "rebuilt per text segment". Diagnostic only."""
+
+    def __post_init__(self) -> None:
+        if self.rows is RowDriver.FIXED:
+            if self.rows_fixed < 1:
+                raise ValueError(f"rows_fixed must be >= 1, got {self.rows_fixed}")
+            if not self.rows_reason.strip():
+                raise ValueError("a FIXED row count needs rows_reason; an unexplained constant cannot be reviewed")
+
+    @property
+    def bytes_per_row(self) -> int:
+        """Bytes for a single batch row of this cache."""
+        return (
+            2  # K and V
+            * self.layers
+            * self.kv_heads
+            * self.head_dim
+            * self.physical_capacity_positions
+            * torch.empty((), dtype=self.dtype).element_size()
+        )
+
+    def row_count(self, max_num_seqs: int = 1) -> int:
+        """Rows live at peak, given the engine's concurrency."""
+        if self.rows is RowDriver.FIXED:
+            return self.rows_fixed
+        return max(1, max_num_seqs)
+
+    def peak_bytes(self, max_num_seqs: int = 1) -> int:
+        """Peak bytes once the engine's concurrency is known.
+
+        The model supplies geometry and names the driver; only the engine knows
+        the number behind it.
+        """
+        return self.bytes_per_row * self.row_count(max_num_seqs)
+
+
+def _resolve(config: object, candidates: Sequence[str], what: str) -> int:
+    for attr in candidates:
+        value = getattr(config, attr, None)
+        if value is not None:
+            return int(value)
+    raise ValueError(
+        f"Cannot determine {what} from {type(config).__name__}: tried "
+        f"{', '.join(candidates)}. Pass it explicitly to spec_from_hf_config()."
+    )
+
+
+def spec_from_hf_config(
+    config: object,
+    *,
+    name: str,
+    physical_capacity_positions: int,
+    capacity_source: str,
+    scope: ModelLocalKVScope,
+    dtype: torch.dtype,
+    rows: RowDriver,
+    rows_fixed: int = 1,
+    rows_reason: str = "",
+    allocation_note: str | None = None,
+    layers: int | None = None,
+    kv_heads: int | None = None,
+    head_dim: int | None = None,
+) -> ModelLocalKVSpec:
+    """Build a spec from the same config object the cache is built from.
+
+    Every known model-local cache is constructed from an HF config, so the
+    geometry half of a declaration is the same three lookups each time. Pass
+    ``config`` and this derives them; pass ``layers``/``kv_heads``/``head_dim``
+    explicitly to override when the config uses encoder-decoder naming or when
+    the built module is a more truthful source than the config.
+
+    Raises rather than guessing when an attribute is absent: a silently wrong
+    geometry would under-report, which is the failure this protocol exists to
+    prevent. ``rows`` is keyword-required for the same reason -- an earlier
+    revision defaulted the batch extent to 1 here while documenting it as
+    required on the dataclass, so every declarer that used this helper got the
+    default and the guarantee was decorative.
+    """
+    if layers is None:
+        layers = _resolve(config, ("num_hidden_layers", "encoder_layers", "num_layers"), "layer count")
+    if kv_heads is None:
+        kv_heads = _resolve(
+            config,
+            ("num_key_value_heads", "num_attention_heads", "encoder_attention_heads"),
+            "kv head count",
+        )
+    if head_dim is None:
+        explicit = getattr(config, "head_dim", None)
+        if explicit is not None:
+            head_dim = int(explicit)
+        else:
+            hidden = _resolve(config, ("hidden_size", "d_model"), "hidden size")
+            heads = _resolve(config, ("num_attention_heads", "encoder_attention_heads"), "attention head count")
+            head_dim = hidden // heads
+
+    return ModelLocalKVSpec(
+        name=name,
+        layers=layers,
+        kv_heads=kv_heads,
+        head_dim=head_dim,
+        dtype=dtype,
+        physical_capacity_positions=physical_capacity_positions,
+        capacity_source=capacity_source,
+        scope=scope,
+        rows=rows,
+        rows_fixed=rows_fixed,
+        rows_reason=rows_reason,
+        allocation_note=allocation_note,
+    )
+
+
+@runtime_checkable
+class HasModelLocalKV(Protocol):
+    """Implemented by whichever object owns the cache.
+
+    Implement it on the owner, not on the registered model. The owner is
+    usually several levels down (a codec decoder, a talker backbone), and
+    ``collect_model_local_kv_specs`` finds it by walking the module tree, so
+    no intermediate class has to forward anything.
+    """
+
+    def model_local_kv_specs(self) -> Sequence[ModelLocalKVSpec]: ...
+
+
+def _unwrap_to_module(model: object) -> object:
+    """Unwrap runner-side wrappers down to the real module tree.
+
+    ``vllm.compilation.cuda_graph.CUDAGraphWrapper`` is a plain callable, not
+    an ``nn.Module``, but its ``__getattr__`` forwards anything the runnable
+    has -- including ``named_modules``. So a check for ``named_modules`` does
+    not detect the wrapper, and the wrapper and its runnable both answer to the
+    same attributes. That matters twice over: it means walking the wrapper
+    happens to work, and it means a naive walk can enter the tree as both the
+    wrapper and the root module and count a root-level declarer twice.
+
+    Prefer the supported ``unwrap()`` accessor, which returns the runnable
+    directly.
+    """
+    seen: set[int] = set()
+    current = model
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        unwrap = getattr(current, "unwrap", None)
+        if not callable(unwrap):
+            return current
+        try:
+            nxt = unwrap()
+        except Exception:
+            return current
+        if nxt is None or nxt is current:
+            return current
+        current = nxt
+    return model
+
+
+def collect_model_local_kv_specs(model: object) -> list[tuple[str, ModelLocalKVSpec]]:
+    """Collect every declaration in a loaded model, with the owner's path.
+
+    Walks ``named_modules()`` because the cache owner is an inner module in all
+    known cases. Requiring each registered model to forward the call would put
+    the burden on classes that do not own a cache and would silently report
+    zero the moment someone forgot -- exactly the failure mode this is meant to
+    surface.
+
+    A raising declaration is logged and skipped rather than propagated:
+    reporting memory must not be able to break model load.
+    """
+    root = _unwrap_to_module(model)
+
+    owners: list[tuple[str, object]] = []
+    seen_ids: set[int] = set()
+    if isinstance(root, HasModelLocalKV):
+        owners.append(("", root))
+        seen_ids.add(id(root))
+
+    named_modules = getattr(root, "named_modules", None)
+    if callable(named_modules):
+        for path, module in named_modules():
+            if id(module) in seen_ids or not isinstance(module, HasModelLocalKV):
+                continue
+            seen_ids.add(id(module))
+            owners.append((path, module))
+
+    collected: list[tuple[str, ModelLocalKVSpec]] = []
+    for path, owner in owners:
+        try:
+            specs = owner.model_local_kv_specs()
+        except Exception:
+            logger.warning(
+                "model_local_kv_specs() raised on %s; skipping its declaration",
+                type(owner).__name__,
+                exc_info=True,
+            )
+            continue
+        collected.extend((path, spec) for spec in specs)
+    return collected

vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py

@@ -17,6 +17,13 @@ from transformers.cache_utils import DynamicCache
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 
+from vllm_omni.model_executor.models.model_local_kv import (
+    ModelLocalKVScope,
+    ModelLocalKVSpec,
+    RowDriver,
+    spec_from_hf_config,
+)
+
 logger = init_logger(__name__)
 
 
@@ -166,6 +173,85 @@ class CUDAGraphDecoderWrapper:
         if self._stats_log_every > 0 and self._stats_total_requests % self._stats_log_every == 0:
             self.log_decode_stats()
 
+    @staticmethod
+    def _retained_kv_caches(states: dict) -> list[tuple[int, DynamicCache]]:
+        """Find the KV caches a capture kept alive, with their batch rows.
+
+        Walks the stored ``cache`` payload instead of assuming which capture
+        path populates it: ``_decode_icl_first_chunk`` fills its dict on the
+        decoder side, so the only reliable count is the one taken from the
+        objects themselves after warmup.
+        """
+        found: list[tuple[int, DynamicCache]] = []
+        for state in states.values():
+            payload = state.get("cache")
+            if not isinstance(payload, dict):
+                continue
+            for value in payload.values():
+                if not isinstance(value, DynamicCache):
+                    continue
+                layers = getattr(value, "layers", None)
+                keys = getattr(layers[0], "keys", None) if layers else None
+                if keys is None:
+                    continue
+                found.append((int(keys.shape[0]), value))
+        return found
+
+    def model_local_kv_specs(self) -> list[ModelLocalKVSpec]:
+        """Declare only the graph-resident copies this wrapper retains.
+
+        The per-decode cache is declared by the decoder itself, because it
+        exists whether or not graphs were captured. What is specific to this
+        wrapper is that every captured shape keeps its dummy ``DynamicCache``
+        alive for replay (``combined_states``, ``icl_prefix_states`` and
+        ``xvec_prefix_states`` all store ``cache``), for the process lifetime.
+
+        Rows come from the state dicts rather than the configured capture
+        lists: a capture that raised is logged and skipped, so the configured
+        shape count would over-report. The trade is that a capture whose
+        tensors transformers has not yet materialized contributes nothing, so
+        this is a floor rather than a contract.
+        """
+        config = self.decoder.config
+        try:
+            dtype = next(self.decoder.parameters()).dtype
+        except StopIteration:  # parameterless stub, only reachable in tests
+            return []
+
+        # Physical, not logical: the sliding window truncates on every write,
+        # so a stream of thousands of frames never occupies more than this.
+        physical = self.prefix_length - 1
+        retained = [
+            *self._retained_kv_caches(self.combined_states),
+            *self._retained_kv_caches(self.icl_prefix_states),
+            *self._retained_kv_caches(self.xvec_prefix_states),
+        ]
+        batched_captures = sum(rows for rows, _ in retained)
+
+        if not batched_captures:
+            return []
+        return [
+            spec_from_hf_config(
+                config,
+                name="codec_decoder_graph_pool",
+                dtype=dtype,
+                physical_capacity_positions=physical,
+                capacity_source=f"sliding_window({self.prefix_length}) - 1",
+                scope=ModelLocalKVScope.MODEL,
+                rows=RowDriver.FIXED,
+                rows_fixed=batched_captures,
+                rows_reason=(
+                    f"rows across {len(retained)} retained caches in "
+                    f"{len(self.combined_states)} suffix / {len(self.icl_prefix_states)} icl-prefix / "
+                    f"{len(self.xvec_prefix_states)} xvec-prefix captures"
+                ),
+                allocation_note=(
+                    "counted from caches transformers has materialized; captures whose tensors are still "
+                    "unallocated report nothing, so this is a floor for the xvec path"
+                ),
+            )
+        ]
+
     def log_decode_stats(self) -> None:
         if not getattr(self, "_stats_enabled", False) or self._stats_total_requests == 0:
             return

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -854,6 +854,75 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         self._cudagraph_enabled = False
         self._cudagraph_wrapper = None
 
+    def model_local_kv_specs(self):
+        """Declare this decoder's attention KV.
+
+        Only the async-chunk path holds KV. In stateless mode the decoder runs
+        ``_forward_exact``, which calls ``pre_transformer`` without
+        ``use_cache`` and allocates no cache at all, so declaring a per-request
+        cache there invents memory that does not exist. The previous revision
+        did exactly that and reported over a gigabyte for stateless
+        deployments; the one before it hung the declaration off the CUDA-graph
+        wrapper and so reported zero under ``enforce_eager``, where the eager
+        async-chunk path does allocate. Both mistakes come from tying this to
+        the wrong condition.
+
+        Two entries when async-chunk is on: the retained per-request cache, and
+        the working copy ``_decode_icl_first_chunk`` deep-copies from it, which
+        is live at the same time.
+
+        The wrapper, when present, adds the graph-resident copies. It is a
+        plain object rather than a submodule, so the module-tree walk cannot
+        reach it; this is the hop it needs.
+        """
+        from vllm_omni.model_executor.models.model_local_kv import (
+            ModelLocalKVScope,
+            RowDriver,
+            spec_from_hf_config,
+        )
+
+        config = self.config
+        try:
+            dtype = next(self.parameters()).dtype
+        except StopIteration:  # parameterless stub, only reachable in tests
+            return []
+
+        window = int(getattr(config, "sliding_window", 0) or 0)
+        if window <= 1:
+            return []
+
+        specs = []
+        wrapper = getattr(self, "_cudagraph_wrapper", None)
+        # async_chunk is the wrapper's record of the mode; without a wrapper the
+        # decoder has not been told, and only the stateless path is reachable.
+        if wrapper is not None and getattr(wrapper, "async_chunk", False):
+            common = dict(
+                config=config,
+                dtype=dtype,
+                physical_capacity_positions=window - 1,
+                capacity_source=f"sliding_window({window}) - 1",
+                rows=RowDriver.MAX_NUM_SEQS,
+            )
+            specs.append(
+                spec_from_hf_config(
+                    name="codec_decoder_request",
+                    scope=ModelLocalKVScope.REQUEST,
+                    allocation_note="retained across chunks of one request",
+                    **common,
+                )
+            )
+            specs.append(
+                spec_from_hf_config(
+                    name="codec_decoder_working_copy",
+                    scope=ModelLocalKVScope.INVOCATION,
+                    allocation_note="deep-copied from the retained cache each chunk; live alongside it",
+                    **common,
+                )
+            )
+        if wrapper is not None:
+            specs.extend(wrapper.model_local_kv_specs())
+        return specs
+
     def precompute_snake_caches(self):
         """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules."""
         count = 0

vllm_omni/worker/gpu_model_runner.py

@@ -36,6 +36,7 @@ from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding
+from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.platforms import current_omni_platform
 
@@ -194,6 +195,52 @@ class OmniGPUModelRunner(GPUModelRunner):
         self._maybe_enable_output_token_ids_for_model_sampler()
         self._init_talker_mtp()
         self._prewarm_attention_capture_workspaces()
+        self._report_model_local_kv()
+
+    def _report_model_local_kv(self) -> None:
+        """Log attention KV this model holds outside the paged manager.
+
+        Reporting only. Two of these caches are captured into CUDA graphs
+        during ``load_model``, so NVML already charges them to the process
+        before ``determine_available_memory()`` samples it; subtracting the
+        declared total from the KV budget would double-count them. What is not
+        charged is the per-request half, which appears after profiling and
+        scales with ``max_num_seqs``.
+
+        Wrapped whole: a memory report has no business breaking model load,
+        and this runs on platforms these declarations have never been
+        exercised on.
+        """
+        try:
+            self._log_model_local_kv()
+        except Exception:
+            logger.warning("Model-local KV reporting failed; continuing", exc_info=True)
+
+    def _log_model_local_kv(self) -> None:
+        specs = collect_model_local_kv_specs(getattr(self, "model", None))
+        if not specs:
+            return
+        max_num_seqs = max(1, int(self.scheduler_config.max_num_seqs))
+        total = sum(spec.peak_bytes(max_num_seqs) for _, spec in specs)
+        logger.info(
+            "Model-local KV (outside the paged manager): %.2f MiB across %d declaration(s) at max_num_seqs=%d",
+            total / (1 << 20),
+            len(specs),
+            max_num_seqs,
+        )
+        for path, spec in specs:
+            logger.info(
+                "  %s.%s: %.2f MiB (%d rows from %s x %.2f MiB/row, scope=%s, %d positions from %s)",
+                path or type(self.model).__name__,
+                spec.name,
+                spec.peak_bytes(max_num_seqs) / (1 << 20),
+                spec.row_count(max_num_seqs),
+                spec.rows.value if spec.rows.value != "fixed" else f"fixed({spec.rows_fixed}; {spec.rows_reason})",
+                spec.bytes_per_row / (1 << 20),
+                spec.scope.value,
+                spec.physical_capacity_positions,
+                spec.capacity_source,
+            )
 
     def _maybe_enable_output_token_ids_for_model_sampler(self) -> None:
         if getattr(self.model, "logitsprocs_need_output_token_ids", False):