[Refactor] Declare model-local KV held outside the paged manager

[linyueqian]

Purpose

Part of #4855 (K3). Four models keep attention KV outside the engine's paged manager, as HuggingFace transformers cache objects. This has each of them declare what it holds and sums it at load, so the memory is visible rather than inferred.

Two rounds of external review shaped this. The first found the multiplicity model wrong and two of my claims unsupported; the second found that my replacement fixed those but introduced a regression of its own. Both sets of corrections are described below, including what I got wrong, because some of it was previously stated as fact in this description.

What each model declares

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, twice over 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 (duplex only)	EncoderDecoderCache	embed_positions rows = 1500	140.62 MiB	duplex_max_sessions
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. Sizing any of them by sequence length overstates it by orders of magnitude.

ming's geometry does not exist in this repo. self._llm_config is a Qwen2Config resolved from the checkpoint, so layers, kv-heads, head-dim and dtype are only knowable after load. That is what makes this a post-load query rather than a table, and it is the one row a static table could not have held.

Lifetime is not multiplicity

scope records when memory appears and goes. It does not decide how wide the allocation is. Width is declared by naming the driver: Fixed(n, because=...) for an extent the model controls, MaxNumSeqs() or DuplexMaxSessions() for one the engine controls. Fixed rejects a blank reason or a non-positive count, because an unexplained constant cannot be reviewed.

rows counts row-equivalents at peak and deliberately does not distinguish "one allocation N rows wide" from "N allocations of one row". They cost the same, and the object layout goes in allocation_note where it cannot be mistaken for arithmetic.

A declaration is inert when its path is. Qwen3-TTS holds KV only in async-chunk mode, so stateless decoding declares nothing. Where the model cannot see the deciding setting -- omni_config collapses "duplex off" and "duplex with one session" both to duplex_max_sessions=1 -- it states the condition in only_when rather than guessing.

Corrections, first round

The multiplier was wrong in two directions. Deriving it from scope made every non-MODEL cache scale with max_num_seqs. ming's talker is per-call but serialized -- forward takes runtime_additional_information[0] and runs the whole AR loop inline -- so it was over-reported by that factor. MiniCPM-o's encoder is per duplex session, capped by duplex_max_sessions, which is a different setting.

MiMo's bytes were right and its description was not. base_local_forward builds one DynamicCache whose batch dimension is the group size. One allocation of B rows costs what B allocations of one row cost, so the total was correct while the report said "N instances" of something there was one of.

Qwen3-TTS reported zero under enforce_eager. The declaration hung off the CUDA-graph wrapper, but the eager async-chunk path allocates its cache whether or not graphs were captured. The per-decode entry moved to the decoder; the wrapper declares only the graph-resident copies.

spec_from_hf_config defeated its own required field. The dataclass documented the batch extent as required; the helper every declarer uses defaulted it to 1.

I retract the claim that the collector had a bug I caught. The previous description said walking a CUDAGraphWrapper found nothing and that the fix was what made declarations appear. That is false: the wrapper's __getattr__ forwards to its runnable, so named_modules resolves and the original code worked. I inferred a bug from a stage that legitimately has no caches. The collector now uses the supported unwrap() accessor, which also removes a real double-counting path when a root-level model declares.

The accounting motivation was wrong for the largest figures. Graph capture happens during load_model, and determine_available_memory() then samples per-process memory through NVML, so the resident portion is already charged; subtracting the declared total would double-count it. What is genuinely unaccounted is the per-request half, allocated after profiling -- roughly 1.1 GiB for Qwen3-TTS at max_num_seqs=256. The consumer reports and does not subtract.

Corrections, second round

The fixes above were reviewed again, and the replacement had problems of its own.

I introduced a regression, and it is the reason this went back to draft. Having moved the per-decode entry to the decoder, I made it unconditional. Stateless decoding runs _forward_exact, which calls the transformer without use_cache and allocates no KV at all, so that declared roughly a gigabyte of nothing at max_num_seqs=256 -- worse than the revision before it, which reported the correct zero there. Declarations are now gated on async_chunk, which is the condition that actually decides. The eager path also deep-copies a working cache that is live alongside the retained one, so async-chunk mode declares two entries rather than one.

The object-count axis claimed a topology that was wrong in three of four declarers. MiniCPM-o keeps one cache object per session while asserting batch size 1; both graph pools keep one distinct object per captured bucket. All three folded their objects into a single wide allocation. allocations is gone; rows means row-equivalents, and layout is described in prose where it cannot be mistaken for arithmetic.

The "unresolved driver" mechanism was built on a false premise. OmniModelConfig.duplex_max_sessions exists and is read this way elsewhere in the tree. I concluded it did not because I grepped for it while my checkout had lost the file, and then designed around a gap that was not there. The machinery is deleted and the attribute read directly.

One bug fixed, one superseded

The TTS ratchet test asserted count == budget, re-imposing through pytest exactly what #6008 removed from the checker. A PR removing a branch then fails CI unless it also hand-edits the constant, which is what broke main for 5h19m in #5746. Budgets are a ceiling.

Superseded: MiniCPM-o's dead streaming cache. Measuring turned up that the vendored Whisper layer rebound its cache to None on every layer, so the encoder re-encoded its whole prefix each chunk. main fixed this independently while this PR was in draft, and did it better: it selects the kwarg name through self._past_key_values_kwarg rather than hardcoding the transformers v5 spelling, and guards the read-back with if len(attn_out) > 2 so a future 3-tuple still works, where this branch had deleted the branch outright. The merge takes main's version; nothing of that fix remains mine to claim.

Test Plan

pytest tests/model_executor/models/test_model_local_kv.py
pytest tests/tools/test_check_tts_adapter.py

Test Result

35 passed. The suite now calls the real declarers rather than comparing specs written in the test to caches built in the test. Verified by mutation twice: reintroducing the ming multiplier fails test_ming_does_not_scale_with_max_num_seqs, and removing the async-chunk gate fails test_qwen3_tts_declares_nothing_on_the_stateless_path. Both would have passed under the suite this replaced.

ming, against the real checkpoint config (Jonathan1909/Ming-flash-omni-2.0, talker/llm/config.json):

checkpoint geometry: L=24 kv=2 hidden=896 heads=14 dtype=torch.bfloat16
declared bytes_per_instance=25165824 (24.00 MiB)
real StaticCache bytes=25165824 (24.00 MiB)   MATCH

Qwen3-TTS and MiMo-Audio, real weights on an H200, under the previous revision's multiplicity model. Declared resident bytes equalled measured bytes exactly for Qwen3-TTS (5 x 4,653,056 = 23,265,280), and MiMo's 261 rows were read off the live model. Those measurements constrain geometry, which this revision does not change; the row-count rework is covered by the unit tests above and has not been re-run on GPU.

Not verified. ming and MiniCPM-o have not been booted. ming's checkpoint was not available; MiniCPM-o loads and declares but its pipeline needs cosyvoice2, which is absent from the machine I have.

Deliberately not fixed here. MiMo captures its CUDA graphs gated only on torch.cuda.is_available() (mimo_audio_llm.py:670), so 179.44 MiB stays resident even under enforce_eager: true. One line, but verifying it needs a model boot.

• vLLM version: 0.27.0

Note on splitting

The ratchet fix is ~5 lines and unblocks anyone removing branches from serving_speech.py today. If the abstraction needs more debate, say so and I will pull that fix into its own PR rather than let it wait.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[hsliuustc0106]

As shipped, nothing implements HasModelLocalKV and nothing in the engine calls total_declared_bytes — the protocol's only exercise is the test constructing specs by hand. All four declarations live as prose in the doc and literals in the test, so changing sliding_window in the Qwen3-TTS checkpoint tomorrow breaks nothing here, which is exactly the staleness the doc argues a table suffers from.

Consider landing one live edge in this PR — the Qwen3-TTS codec decoder is the natural first declarer (geometry fully known, numbers already derived) — or a minimal consumer that logs total_declared_bytes at startup, so the API can't rot unexercised until the K3 consumer lands.

One design question worth deciding before the first declarer freezes the API: for REQUEST scope, multiplicity is engine-side concurrency, not something the model knows — the test hardcodes max_live_instances=1 for a per-request cache. If the engine is expected to multiply by live requests, that contract should be explicit in ModelLocalKVSpec.

[linyueqian]

Both of your suggestions landed since this comment, and you were right about the helper.

Live declarers now exist rather than prose: qwen3_tts/segmented_graph_wrapper.py, qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py, mimo_audio/mimo_audio_llm.py, ming_flash_omni/talker_module.py with its forwarder in ming_flash_omni_talker.py, and minicpmo_4_5/minicpmo_4_5_omni_llm.py.

The engine consumer is OmniGPUModelRunner._log_model_local_kv (gpu_model_runner.py:215): it collects at startup and logs the total plus one line per declaration, wrapped so a reporting failure cannot break model load. So a sliding_window change in the Qwen3-TTS checkpoint now moves a logged number, and the test asserts the declaration matches the cache it describes.

On total_declared_bytes specifically, it stayed dead even after that consumer landed: the consumer needs the individual specs for its per-declaration lines and sums them inline, so the helper never acquired a production caller and the design doc never referenced it. Removed in 1f971f6. collect_model_local_kv_specs() is the single public entry point now.

On the REQUEST scope question, that contract is now explicit instead of a hardcoded test literal. Multiplicity moved onto ModelLocalKVSpec.rows as RowDriver: MAX_NUM_SEQS means one row per in-flight sequence with the value supplied by the engine, and FIXED carries rows_fixed plus a required rows_reason, validated in __post_init__. peak_bytes(max_num_seqs) and row_count(max_num_seqs) take the engine number as an argument. Scope stayed deliberately independent of size and diagnostic only, so a REQUEST-scoped cache the engine does not widen per sequence declares FIXED and has to say why.

[hsliuustc0106]

This page isn't linked from docs/design/index.md, so it stays out of the design-doc navigation even though RTD builds it. A one-line entry ("Runtime and stage execution" seems the closest fit) would fix it.

[linyueqian]

Fixed. docs/design/index.md links it under "Runtime and stage execution", which was the fit you suggested.

[hsliuustc0106]

@FayeSpica PTAL NPU ci failures

[hsliuustc0106]

NPU failure:

https://buildkite.com/vllm/vllm-omni-npu-ci/builds/5580/canvas?sid=01a01c6e-e49a-4f70-9d67-26e1d2922442&tab=output

=================================================================== short test summary info ====================================================================
--
  | FAILED tests/platforms/npu/test_minicpmo_code2wav_npugraph.py::test_minicpmo_cfm_estimator_npugraph_matches_eager - AttributeError: '_Flow' object has no attribute 'encoder'. Did you mean: 'decoder'?
  | ================================================================ 1 failed, 14 warnings in 9.26s ================================================================
  | 🚨 Error: The command exited with status 1

[hsliuustc0106]

why not move to the omni_config folder?

[linyueqian]

There is no vllm_omni/omni_config/ in the tree, so I read this as vllm_omni/config/. I would rather leave it under model_executor/models/, though I will move it if you still prefer that after this.

The module is a model-side declaration protocol rather than configuration. Models implement model_local_kv_specs() on themselves, and the helpers walk a live model object through named_modules(), unwrapping CUDA-graph wrappers and deduplicating a root declarer seen twice. Nothing in it is parsed from or serialized to a config file, and none of the values come from user input: they are read off the loaded checkpoint config and the model's own captured buckets.

docs/design/module/vllm_omni_config.md scopes that contract to vllm_omni/config/**, vllm_omni/deploy/** and pipeline.py, which is deploy and stage schema plus construction. Putting a Protocol implemented by model classes behind that boundary would separate it from the five declarers that implement it and give the config contract ownership of runtime model introspection.

[hsliuustc0106]

does it only occupy two class args and one of them related to duplex? this abstraction looks over desgined

[linyueqian]

Agreed, and addressed in b55d434.

EngineCapacity is gone. Once the duplex driver was removed it held a single field, so peak_bytes() takes max_num_seqs directly instead of a wrapper whose only job was to carry it.

The four row classes (ModelLocalKVRows, Fixed, MaxNumSeqs, DuplexMaxSessions) collapsed into one RowDriver enum with two values, plus rows_fixed and rows_reason on the spec. Validation moved into __post_init__, so a FIXED count still cannot ship without a positive value and a stated reason. total_declared_bytes has since gone as well (1f971f6), leaving collect_model_local_kv_specs() as the only public entry point.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit d8f31eb79f3a produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.