diff --git a/docs/design/feature/model_local_kv_caches.md b/docs/design/feature/model_local_kv_caches.md new file mode 100644 index 000000000..371b989ff --- /dev/null +++ b/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. diff --git a/docs/design/index.md b/docs/design/index.md index cff0b1eb3..6481aa49b 100644 --- a/docs/design/index.md +++ b/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 diff --git a/tests/model_executor/models/test_model_local_kv.py b/tests/model_executor/models/test_model_local_kv.py new file mode 100644 index 000000000..5e2c9cdca --- /dev/null +++ b/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" diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py index 47c6ac89f..3107900c0 100644 --- a/tests/tools/test_check_tts_adapter.py +++ b/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 diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index bb01e522c..ec4faadb8 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/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 = [] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index d34cce465..f01e20688 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/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 diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 8e3e55611..293ec2fb3 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/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( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 315ae159d..d70bc30d0 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/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, diff --git a/vllm_omni/model_executor/models/model_local_kv.py b/vllm_omni/model_executor/models/model_local_kv.py new file mode 100644 index 000000000..7328e9967 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py index a1eb071ee..3c73f5d13 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index ea11e64a7..9572d9ee2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/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 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 80d55020f..8dd9f5f58 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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):