[Edge][CPU] KV-cached code predictor, int8-dynamic CPU fast path, model-side CPU sampler

Measured on a Xeon 8480C (12+4 threads, 0.6B): the 15 residual-codebook sub-steps
re-ran the 5-layer predictor over the whole prefix each time (97 ms per frame of
the 134 ms talker step) and vLLM's stock sampler spent 9 ms per step in its
torch.compile'd random-sample kernel on CPU.

- code predictor: forward_cached (KV cache per layer, 2-token prefill then one
  token per sub-step) and a shape-static forward_static_step (tensor position +
  boolean mask) that torch.compile handles as one graph; incremental decoding
  is the default on CPU (VLLM_OMNI_PREDICTOR_INCREMENTAL), optional compile
  (VLLM_OMNI_PREDICTOR_COMPILE=step), optional thread cap.
- CPU fast path: dynamic int8 Linear layers (fp32 activations) for the predictor
  and its lm heads (VLLM_OMNI_PREDICTOR_INT8, default on; micro-benchmark 70 ->
  36 ms per frame for the 15 sub-steps); VLLM_OMNI_PREDICTOR_PARITY_CHECK=1
  records top-1 agreement / max logit diff against the bf16 model in situ.
- Qwen3-TTS talker: prefer_model_sampler on CPU with a sample() that mirrors
  vLLM's Sampler (logits processors + penalties, temperature, top-k/top-p,
  exponential-noise sampling with per-request generators) without the compiled
  kernel (VLLM_OMNI_CPU_FAST_SAMPLER, default on).
- step_stats: sampler phase timers and sample-metadata timers.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

tests/model_executor/models/common/test_code_predictor_incremental.py

@@ -0,0 +1,79 @@
+"""KV-cached incremental residual decoding of the Qwen3 code predictor equals the re-prefill forward (CPU)."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorBaseModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _tiny_config():
+    return SimpleNamespace(
+        hidden_size=64,
+        num_hidden_layers=3,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        intermediate_size=96,
+        rms_norm_eps=1e-6,
+        vocab_size=32,
+        num_code_groups=16,
+        rope_theta=10000.0,
+        attention_bias=False,
+    )
+
+
+@torch.inference_mode()
+def test_forward_cached_matches_full_prefix():
+    torch.manual_seed(0)
+    model = CodePredictorBaseModel(_tiny_config()).eval()
+    for p in model.parameters():
+        p.normal_(0, 0.2)
+    bsz, max_seq, hidden = 2, 17, 64
+    x = torch.randn(bsz, max_seq, hidden)
+    pos = torch.arange(max_seq).unsqueeze(0).expand(bsz, -1)
+    full = model(x, pos)  # [B, 17, H], causal over the whole prefix
+
+    caches = model.new_kv_caches(bsz, max_seq, x.device, x.dtype)
+    out_prefill = model.forward_cached(x[:, :2], pos[:, :2], caches, 0)
+    assert torch.allclose(out_prefill, full[:, :2], atol=1e-4, rtol=1e-4)
+    for step in range(2, max_seq):
+        out_step = model.forward_cached(x[:, step : step + 1], pos[:, step : step + 1], caches, step)
+        assert torch.allclose(out_step[:, 0], full[:, step], atol=1e-4, rtol=1e-4), step
+    # caches hold exactly the prefix K/V that a fresh full pass would produce
+    assert caches[0][0].shape == (bsz, 2, max_seq, 16)
+
+
+@torch.inference_mode()
+def test_forward_cached_rejects_multi_token_append():
+    model = CodePredictorBaseModel(_tiny_config()).eval()
+    x = torch.randn(1, 2, 64)
+    pos = torch.arange(2).unsqueeze(0)
+    caches = model.new_kv_caches(1, 17, x.device, x.dtype)
+    model.forward_cached(x, pos, caches, 0)
+    with pytest.raises(ValueError):
+        model.forward_cached(x, pos + 2, caches, 2)
+
+
+@torch.inference_mode()
+def test_forward_static_step_matches_full_prefix():
+    torch.manual_seed(1)
+    model = CodePredictorBaseModel(_tiny_config()).eval()
+    for p in model.parameters():
+        p.normal_(0, 0.2)
+    bsz, max_seq = 2, 17
+    x = torch.randn(bsz, max_seq, 64)
+    pos = torch.arange(max_seq).unsqueeze(0).expand(bsz, -1)
+    full = model(x, pos)
+    caches = model.new_kv_caches(bsz, max_seq, x.device, x.dtype)
+    model.forward_cached(x[:, :2], pos[:, :2], caches, 0)  # eager 2-token prefill
+    for step in range(2, max_seq):
+        pos_t = torch.tensor([step])
+        mask = (torch.arange(max_seq) <= step).view(1, 1, 1, max_seq).expand(bsz, 1, 1, max_seq)
+        out = model.forward_static_step(x[:, step : step + 1], pos[:, step : step + 1], caches, pos_t, mask)
+        assert torch.allclose(out[:, 0], full[:, step], atol=1e-4, rtol=1e-4), step

tests/model_executor/models/qwen3_tts/test_cpu_fast_sampler.py

@@ -0,0 +1,94 @@
+"""Qwen3-TTS talker model-side CPU sampler: same contract as vLLM's Sampler, no compiled kernel (CPU)."""
+
+from __future__ import annotations
+
+import dataclasses
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+V, B = 3072, 2
+
+
+def _metadata(temperature=0.9, top_k=50, rep=1.05, generators=None, max_num_logprobs=None):
+    fields = {f.name: f for f in dataclasses.fields(SamplingMetadata)}
+    kw = {}
+    temp = torch.full((B,), float(temperature))
+    for name, f in fields.items():
+        if name == "temperature":
+            kw[name] = temp
+        elif name == "all_greedy":
+            kw[name] = temperature == 0
+        elif name == "all_random":
+            kw[name] = temperature > 0
+        elif name == "top_p":
+            kw[name] = None
+        elif name == "top_k":
+            kw[name] = None if top_k <= 0 else torch.full((B,), top_k, dtype=torch.int64)
+        elif name == "generators":
+            kw[name] = generators or {}
+        elif name == "max_num_logprobs":
+            kw[name] = max_num_logprobs
+        elif name == "no_penalties":
+            kw[name] = rep == 1.0
+        elif name == "prompt_token_ids":
+            kw[name] = None if rep == 1.0 else torch.randint(0, 2048, (B, 8))
+        elif name in ("frequency_penalties", "presence_penalties"):
+            kw[name] = torch.zeros(B)
+        elif name == "repetition_penalties":
+            kw[name] = torch.full((B,), float(rep))
+        elif name == "output_token_ids":
+            kw[name] = [[1, 2, 3], [4, 5]]
+        elif name == "allowed_token_ids_mask":
+            kw[name] = None
+        elif name == "bad_words_token_ids":
+            kw[name] = {}
+        elif name == "logitsprocs":
+            from vllm.v1.sample.logits_processor import LogitsProcessors
+
+            kw[name] = LogitsProcessors()
+        elif name == "logprob_token_ids":
+            kw[name] = None
+        elif f.default is not dataclasses.MISSING:
+            kw[name] = f.default
+        elif getattr(f, "default_factory", dataclasses.MISSING) is not dataclasses.MISSING:
+            kw[name] = f.default_factory()
+        else:
+            kw[name] = None
+    return SamplingMetadata(**kw)
+
+
+def _sample(logits, md):
+    fake = SimpleNamespace()
+    return Qwen3TTSTalkerForConditionalGeneration.sample(fake, logits, md)
+
+
+def test_greedy_matches_argmax():
+    logits = torch.randn(B, V)
+    out = _sample(logits.clone(), _metadata(temperature=0.0, top_k=0, rep=1.0))
+    assert out.sampled_token_ids.shape == (B, 1) and out.logprobs_tensors is None
+    assert torch.equal(out.sampled_token_ids.view(-1), logits.argmax(-1))
+
+
+def test_seeded_generators_are_deterministic_and_top_k_respected():
+    logits = torch.randn(B, V)
+    outs = []
+    for _ in range(2):
+        gens = {i: torch.Generator().manual_seed(1234 + i) for i in range(B)}
+        outs.append(_sample(logits.clone(), _metadata(top_k=5, rep=1.0, generators=gens)).sampled_token_ids)
+    assert torch.equal(outs[0], outs[1])
+    top5 = logits.topk(5, dim=-1).indices
+    for i in range(B):
+        assert int(outs[0][i, 0]) in top5[i].tolist()
+
+
+def test_falls_back_when_logprobs_requested_or_not_cpu():
+    logits = torch.randn(B, V)
+    assert _sample(logits, _metadata(max_num_logprobs=1)) is None
+    assert _sample(None, _metadata()) is None