[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>
This commit is contained in:
tzhouamandClaude Fable 5.1 committed
1 parent
936c10be4b
commit
6ce8de1b28
6 files changed
+604
-8
No files matched your search
@@ -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
|
||||
@@ -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
|
||||
@@ -212,6 +212,28 @@ def instrument_engine_core(engine_core: Any) -> None:
|
||||
engine_core.step_fn = stats.wrap("core.step_ms", step_fn)
|
||||
|
||||
|
||||
def instrument_sampler(sampler: Any) -> None:
|
||||
"""Time the phases of vLLM's v1 ``Sampler`` in place (idempotent, no-op when disabled)."""
|
||||
stats = StepStats.get()
|
||||
if not stats.enabled or sampler is None or getattr(sampler, "_omni_step_stats_instrumented", False):
|
||||
return
|
||||
for name in (
|
||||
"apply_logits_processors",
|
||||
"apply_penalties",
|
||||
"apply_temperature",
|
||||
"sample",
|
||||
"compute_logprobs",
|
||||
"gather_logprobs",
|
||||
):
|
||||
fn = getattr(sampler, name, None)
|
||||
if callable(fn):
|
||||
setattr(sampler, name, stats.wrap(f"sampler.{name}_ms", fn))
|
||||
tk = getattr(sampler, "topk_topp_sampler", None)
|
||||
if tk is not None and callable(getattr(tk, "forward", None)):
|
||||
tk.forward = stats.wrap("sampler.topk_topp_ms", tk.forward)
|
||||
sampler._omni_step_stats_instrumented = True
|
||||
|
||||
|
||||
def shm_lockfile_age_ms(key: str, now: float | None = None) -> float | None:
|
||||
"""Age of the SharedMemoryConnector lock file for ``key`` (written at put time)."""
|
||||
try:
|
||||
|
||||
@@ -15,6 +15,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable, Sequence
|
||||
|
||||
@@ -270,11 +271,58 @@ class CodePredictorAttention(nn.Module):
|
||||
sync=True,
|
||||
)[0]
|
||||
|
||||
def forward_static_step(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
kv_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
cache_pos: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""One new token against a fixed-size cache (shape-static, ``torch.compile``-friendly).
|
||||
|
||||
``cache_pos`` is a 1-element int64 tensor (the slot to write), ``attn_mask`` a bool
|
||||
``[B, 1, 1, max_seq]`` mask that is True for the valid slots ``<= cache_pos``.
|
||||
"""
|
||||
bsz, seq_len, _ = hidden_states.shape
|
||||
qkv = self.qkv_proj(hidden_states)
|
||||
q_raw, k_raw, v_raw = self._split_qkv(qkv)
|
||||
q = self.q_norm(q_raw.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2)
|
||||
k = self.k_norm(k_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2)
|
||||
v = v_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
||||
cos, sin = position_embeddings
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
q = (q * cos) + (_rotate_half(q) * sin)
|
||||
k = (k * cos) + (_rotate_half(k) * sin)
|
||||
k_cache, v_cache = kv_cache
|
||||
k_cache.index_copy_(2, cache_pos, k.to(k_cache.dtype))
|
||||
v_cache.index_copy_(2, cache_pos, v.to(v_cache.dtype))
|
||||
attn_out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k_cache[:bsz],
|
||||
v_cache[:bsz],
|
||||
attn_mask=attn_mask,
|
||||
scale=self.scaling,
|
||||
enable_gqa=self.is_gqa,
|
||||
)
|
||||
attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
|
||||
return self.o_proj(attn_out)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
cache_pos: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Causal self-attention.
|
||||
|
||||
With ``kv_cache=(k_cache, v_cache)`` of shape ``[B, kv_heads, max_seq, head_dim]`` the
|
||||
new positions ``cache_pos .. cache_pos+seq_len`` are written into the cache and attended
|
||||
together with the cached prefix (incremental decoding of the residual codebooks); a
|
||||
multi-token call is only allowed from an empty cache (``cache_pos == 0``).
|
||||
"""
|
||||
bsz, seq_len, _ = hidden_states.shape
|
||||
hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
|
||||
hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)
|
||||
@@ -299,12 +347,21 @@ class CodePredictorAttention(nn.Module):
|
||||
else:
|
||||
q = (q * cos) + (_rotate_half(q) * sin)
|
||||
k = (k * cos) + (_rotate_half(k) * sin)
|
||||
if kv_cache is not None:
|
||||
if seq_len > 1 and cache_pos != 0:
|
||||
raise ValueError("incremental code-predictor attention: multi-token append needs an empty cache")
|
||||
k_cache, v_cache = kv_cache
|
||||
end = cache_pos + seq_len
|
||||
k_cache[:bsz, :, cache_pos:end] = k
|
||||
v_cache[:bsz, :, cache_pos:end] = v
|
||||
k = k_cache[:bsz, :, :end]
|
||||
v = v_cache[:bsz, :, :end]
|
||||
attn_out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
scale=self.scaling,
|
||||
is_causal=True,
|
||||
is_causal=seq_len > 1,
|
||||
enable_gqa=self.is_gqa,
|
||||
)
|
||||
|
||||
@@ -355,14 +412,35 @@ class CodePredictorDecoderLayer(nn.Module):
|
||||
self.input_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def forward_static_step(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
kv_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
cache_pos: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states = self.self_attn.forward_static_step(
|
||||
hidden_states, position_embeddings, kv_cache, cache_pos, attn_mask
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
return residual + hidden_states
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
position_embeddings: tuple[torch.Tensor, torch.Tensor],
|
||||
kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
cache_pos: int = 0,
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states = self.self_attn(hidden_states, position_embeddings)
|
||||
hidden_states = self.self_attn(hidden_states, position_embeddings, kv_cache=kv_cache, cache_pos=cache_pos)
|
||||
if current_omni_platform.is_npu():
|
||||
hidden_states, _, residual = torch_npu.npu_add_rms_norm(
|
||||
hidden_states,
|
||||
@@ -448,6 +526,55 @@ class CodePredictorBaseModel(nn.Module):
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states.to(input_dtype)
|
||||
|
||||
def new_kv_caches(self, bsz: int, max_seq: int, device: torch.device, dtype: torch.dtype) -> list:
|
||||
"""Per-layer ``(k_cache, v_cache)`` buffers for :meth:`forward_cached`."""
|
||||
attn0 = self.layers[0].self_attn
|
||||
shape = (bsz, attn0.num_kv_heads, max_seq, attn0.head_dim)
|
||||
return [
|
||||
(torch.zeros(shape, device=device, dtype=dtype), torch.zeros(shape, device=device, dtype=dtype))
|
||||
for _ in self.layers
|
||||
]
|
||||
|
||||
def forward_static_step(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
kv_caches: list,
|
||||
cache_pos: torch.Tensor,
|
||||
attn_mask: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""One-token step with tensor-valued position and mask (no Python ints -> one compiled graph)."""
|
||||
hidden_states = inputs_embeds
|
||||
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||
for layer, cache in zip(self.layers, kv_caches, strict=True):
|
||||
hidden_states = layer.forward_static_step(hidden_states, position_embeddings, cache, cache_pos, attn_mask)
|
||||
return self.norm(hidden_states)
|
||||
|
||||
def forward_cached(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
position_ids: torch.Tensor,
|
||||
kv_caches: list,
|
||||
cache_pos: int,
|
||||
) -> torch.Tensor:
|
||||
"""Incremental forward: process only ``inputs_embeds`` (positions ``cache_pos..``) against the caches.
|
||||
|
||||
Same math as :meth:`forward` over the full prefix, without re-running the
|
||||
cached positions; used for the residual-codebook loop on CPU where each
|
||||
of the 15 sub-steps would otherwise recompute up to 17 positions.
|
||||
"""
|
||||
input_dtype = inputs_embeds.dtype
|
||||
use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type != "cpu"
|
||||
if use_fp32:
|
||||
inputs_embeds = inputs_embeds.float()
|
||||
hidden_states = inputs_embeds
|
||||
with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32):
|
||||
position_embeddings = self.rotary_emb(hidden_states, position_ids)
|
||||
for layer, cache in zip(self.layers, kv_caches, strict=True):
|
||||
hidden_states = layer(hidden_states, position_embeddings, kv_cache=cache, cache_pos=cache_pos)
|
||||
hidden_states = self.norm(hidden_states)
|
||||
return hidden_states.to(input_dtype)
|
||||
|
||||
# Mapping from HF checkpoint shard name to the (fused parameter,
|
||||
# slice index) pair used to assemble the fused tensor. See
|
||||
# ``load_weights`` below.
|
||||
@@ -998,6 +1125,149 @@ class CodePredictorWrapper(nn.Module):
|
||||
]
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Incremental (KV-cached) residual decoding
|
||||
# ------------------------------------------------------------------
|
||||
_INCREMENTAL_ENV = "VLLM_OMNI_PREDICTOR_INCREMENTAL"
|
||||
_INT8_ENV = "VLLM_OMNI_PREDICTOR_INT8" # CPU: dynamic int8 Linear layers (default on; "0" disables)
|
||||
_PARITY_ENV = "VLLM_OMNI_PREDICTOR_PARITY_CHECK" # "1": compare int8 logits with the bf16 model per sub-step
|
||||
|
||||
def _parity_check(self, step: int, proj_row: torch.Tensor, logits_int8: torch.Tensor, ref_caches: list) -> None:
|
||||
"""Record top-1 agreement and max |logit diff| between the int8 and the reference predictor."""
|
||||
from vllm_omni.edge.step_stats import StepStats
|
||||
|
||||
ref_model, ref_heads = self._parity_ref
|
||||
rdtype = next(ref_model.parameters()).dtype
|
||||
bsz = proj_row.shape[0]
|
||||
if step == 1:
|
||||
x = proj_row.to(rdtype) # [B, 2, H] prefill rows
|
||||
h = ref_model.forward_cached(x, self._incremental_pos_ids(bsz, 0, 2, proj_row.device), ref_caches, 0)[:, 1]
|
||||
else:
|
||||
x = proj_row.to(rdtype) # [B, 1, H]
|
||||
h = ref_model.forward_cached(x, self._incremental_pos_ids(bsz, step, 1, proj_row.device), ref_caches, step)[
|
||||
:, 0
|
||||
]
|
||||
ref_logits = ref_heads[step - 1](h).float()
|
||||
stats = StepStats.get()
|
||||
agree = (ref_logits.argmax(-1) == logits_int8.float().argmax(-1)).float().mean().item()
|
||||
stats.add("predictor.parity_top1_agree", agree)
|
||||
stats.add("predictor.parity_logit_maxdiff", (ref_logits - logits_int8.float()).abs().max().item())
|
||||
|
||||
def _setup_cpu_fast_path(self, device: torch.device) -> None:
|
||||
"""CPU: dynamic int8 quantization of every Linear (weights int8 per-channel, activations per call).
|
||||
|
||||
Measured on a Xeon 8480C (12 threads): bf16 M=1 GEMVs carry ~30 us of fixed
|
||||
kernel overhead each and the 15 residual sub-steps cost 70 ms per frame; with
|
||||
int8-dynamic Linear layers (fp32 activations) the same loop costs 36 ms.
|
||||
"""
|
||||
if getattr(self, "_cpu_fast_path_done", False):
|
||||
return
|
||||
self._cpu_fast_path_done = True
|
||||
if device.type != "cpu" or current_omni_platform.is_npu():
|
||||
return
|
||||
if os.environ.get(self._INT8_ENV, "1").strip().lower() in ("0", "false", "off", "no"):
|
||||
return
|
||||
from torch.ao.quantization import quantize_dynamic
|
||||
|
||||
self._parity_ref = None
|
||||
if os.environ.get(self._PARITY_ENV, "0").strip().lower() in ("1", "true", "on", "yes"):
|
||||
# keep an unquantized copy to measure int8 fidelity in situ (see _parity_check)
|
||||
import copy as _copy
|
||||
|
||||
self._parity_ref = (_copy.deepcopy(self.model), _copy.deepcopy(self.lm_head))
|
||||
self.model = quantize_dynamic(self.model.float(), {nn.Linear}, dtype=torch.qint8)
|
||||
self.lm_head = quantize_dynamic(self.lm_head.float(), {nn.Linear}, dtype=torch.qint8)
|
||||
self.small_to_mtp_projection = self.small_to_mtp_projection.float()
|
||||
self._model_dtype = torch.float32
|
||||
self._lm_heads_list = list(self.lm_head)
|
||||
self._codec_embeds_list = list(self.model.codec_embedding)
|
||||
self._compiled_model_fwd = self.model.forward
|
||||
self._proj_buf = None # re-allocated in fp32 by _ensure_buffers
|
||||
logger.info("code_predictor: CPU fast path (KV-cached sub-steps, dynamic int8 Linear, fp32 activations)")
|
||||
|
||||
_THREADS_ENV = "VLLM_OMNI_PREDICTOR_THREADS" # CPU: OpenMP threads for the residual loop (0 = unchanged)
|
||||
_COMPILE_ENV = (
|
||||
"VLLM_OMNI_PREDICTOR_COMPILE" # "1": torch.compile the per-layer MLP/norm blocks (CPU, static shapes)
|
||||
)
|
||||
|
||||
def _incremental_threads(self) -> int:
|
||||
try:
|
||||
return max(0, int(os.environ.get(self._THREADS_ENV, "0")))
|
||||
except ValueError:
|
||||
return 0
|
||||
|
||||
def _setup_incremental_compile(self) -> None:
|
||||
"""Compile the whole shape-static one-token step (``forward_static_step``) for the incremental path.
|
||||
|
||||
The step takes the fixed-size caches, a tensor position and a boolean mask, so a
|
||||
single ``torch.compile`` graph serves all 14 residual sub-steps without
|
||||
recompilation; the 2-token prefill of the first sub-step stays eager.
|
||||
"""
|
||||
if getattr(self, "_incremental_compiled", False):
|
||||
return
|
||||
self._incremental_compiled = True
|
||||
self._static_step_fn = None
|
||||
if os.environ.get(self._COMPILE_ENV, "0").strip().lower() not in ("1", "true", "on", "yes", "step"):
|
||||
return
|
||||
self._static_step_fn = torch.compile(self.model.forward_static_step, dynamic=False)
|
||||
logger.info("code_predictor: incremental path with a compiled static one-token step")
|
||||
|
||||
def _use_incremental(self, device: torch.device) -> bool:
|
||||
"""KV-cached sub-steps: default on CPU (no device graphs there), ``=1`` forces, ``=0`` disables."""
|
||||
mode = os.environ.get(self._INCREMENTAL_ENV, "auto").strip().lower()
|
||||
if mode in ("0", "false", "off", "no"):
|
||||
return False
|
||||
if mode in ("1", "true", "on", "yes"):
|
||||
return not current_omni_platform.is_npu()
|
||||
return device.type == "cpu" and not self._device_graphs
|
||||
|
||||
def _incremental_caches(self, bsz: int, max_seq: int, device: torch.device, dtype: torch.dtype) -> list:
|
||||
key = (bsz, max_seq, str(device), dtype)
|
||||
buffers = getattr(self, "_kv_cache_buffers", None)
|
||||
if buffers is None:
|
||||
buffers = self._kv_cache_buffers = {}
|
||||
caches = buffers.get(key)
|
||||
if caches is None:
|
||||
caches = buffers[key] = self.model.new_kv_caches(bsz, max_seq, device, dtype)
|
||||
return caches
|
||||
|
||||
@staticmethod
|
||||
def _incremental_pos_ids(bsz: int, start: int, n: int, device: torch.device) -> torch.Tensor:
|
||||
return torch.arange(start, start + n, device=device, dtype=torch.long).unsqueeze(0).expand(bsz, -1)
|
||||
|
||||
def _sample_step_code(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sample_generator: _GeneratorLike,
|
||||
*,
|
||||
stored_mode: bool,
|
||||
s_top_k: int,
|
||||
s_top_p: float,
|
||||
use_sampling: bool,
|
||||
inv_temperature: float,
|
||||
top_k: int,
|
||||
) -> torch.Tensor:
|
||||
"""Same sampling as the re-prefill loop body (kept identical so codes match)."""
|
||||
if stored_mode:
|
||||
if s_top_k > 0:
|
||||
topk_vals, _ = logits.topk(s_top_k, dim=-1)
|
||||
logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
|
||||
if s_top_p < 1.0:
|
||||
sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
|
||||
sorted_probs = F.softmax(sorted_logits, dim=-1, dtype=torch.float32)
|
||||
cumulative_probs = sorted_probs.cumsum(dim=-1)
|
||||
remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
|
||||
sorted_logits[remove_mask] = float("-inf")
|
||||
logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
|
||||
return self._sample_codes_gumbel(logits, generator=sample_generator)
|
||||
if use_sampling:
|
||||
scaled = logits * inv_temperature
|
||||
if top_k > 0:
|
||||
topk_vals, _ = scaled.topk(top_k, dim=-1)
|
||||
scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
|
||||
return self._sample_codes_gumbel(scaled, generator=sample_generator)
|
||||
return logits.argmax(dim=-1, keepdim=True)
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(
|
||||
self,
|
||||
@@ -1022,6 +1292,7 @@ class CodePredictorWrapper(nn.Module):
|
||||
# _setup_compile caches _model_dtype on first call; use it for buffers
|
||||
# so they always match model weight precision (#2385).
|
||||
self._setup_compile()
|
||||
self._setup_cpu_fast_path(device)
|
||||
dtype = self._model_dtype
|
||||
|
||||
padded_bsz = self._padded_bsz(bsz)
|
||||
@@ -1067,8 +1338,80 @@ class CodePredictorWrapper(nn.Module):
|
||||
|
||||
_stats = StepStats.get()
|
||||
_t_loop = time.perf_counter() if _stats.enabled else 0.0
|
||||
incremental = self._use_incremental(device)
|
||||
kv_caches = self._incremental_caches(bsz, max_seq, device, dtype) if incremental else None
|
||||
prev_threads = None
|
||||
if incremental:
|
||||
self._setup_incremental_compile()
|
||||
n_threads = self._incremental_threads()
|
||||
if n_threads > 0 and device.type == "cpu" and n_threads != torch.get_num_threads():
|
||||
prev_threads = torch.get_num_threads()
|
||||
torch.set_num_threads(n_threads)
|
||||
sample_kwargs = (
|
||||
dict(
|
||||
stored_mode=True, s_top_k=s_top_k, s_top_p=s_top_p, use_sampling=False, inv_temperature=0.0, top_k=top_k
|
||||
)
|
||||
if stored_mode
|
||||
else dict(
|
||||
stored_mode=False,
|
||||
s_top_k=0,
|
||||
s_top_p=1.0,
|
||||
use_sampling=use_sampling,
|
||||
inv_temperature=inv_temperature,
|
||||
top_k=top_k,
|
||||
)
|
||||
)
|
||||
for step in range(1, num_groups):
|
||||
_t_sub = time.perf_counter() if _stats.enabled else 0.0
|
||||
if incremental:
|
||||
# KV-cached: prefill positions 0..1 at the first sub-step, then one new position per sub-step.
|
||||
if step == 1:
|
||||
hidden_new = self.model.forward_cached(
|
||||
proj_buf[:bsz, :2, :], self._incremental_pos_ids(bsz, 0, 2, device), kv_caches, 0
|
||||
)
|
||||
hidden_step = hidden_new[:, 1, :]
|
||||
elif self._static_step_fn is not None:
|
||||
pos_t = torch.tensor([step], device=device, dtype=torch.long)
|
||||
attn_mask = (
|
||||
(torch.arange(max_seq, device=device) <= step).view(1, 1, 1, max_seq).expand(bsz, 1, 1, max_seq)
|
||||
)
|
||||
hidden_new = self._static_step_fn(
|
||||
proj_buf[:bsz, step : step + 1, :],
|
||||
self._incremental_pos_ids(bsz, step, 1, device),
|
||||
kv_caches,
|
||||
pos_t,
|
||||
attn_mask,
|
||||
)
|
||||
hidden_step = hidden_new[:, 0, :]
|
||||
else:
|
||||
hidden_new = self.model.forward_cached(
|
||||
proj_buf[:bsz, step : step + 1, :],
|
||||
self._incremental_pos_ids(bsz, step, 1, device),
|
||||
kv_caches,
|
||||
step,
|
||||
)
|
||||
hidden_step = hidden_new[:, 0, :]
|
||||
logits = lm_heads[step - 1](hidden_step)
|
||||
if getattr(self, "_parity_ref", None) is not None:
|
||||
if step == 1:
|
||||
ref_model = self._parity_ref[0]
|
||||
self._parity_caches = ref_model.new_kv_caches(
|
||||
bsz, max_seq, device, next(ref_model.parameters()).dtype
|
||||
)
|
||||
self._parity_check(step, proj_buf[:bsz, :2, :], logits, self._parity_caches)
|
||||
else:
|
||||
self._parity_check(step, proj_buf[:bsz, step : step + 1, :], logits, self._parity_caches)
|
||||
code = self._sample_step_code(logits, sample_generator, **sample_kwargs)
|
||||
if self._wrapper_config.return_proj_buf:
|
||||
all_codes[:, step] = code
|
||||
else:
|
||||
all_codes[:, step] = code.reshape(bsz)
|
||||
if step < num_groups - 1 or self._wrapper_config.return_proj_buf:
|
||||
new_embed = codec_embeds[step - 1](code)
|
||||
proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
|
||||
if _stats.enabled:
|
||||
_stats.add("predictor.substep_ms", (time.perf_counter() - _t_sub) * 1000.0)
|
||||
continue
|
||||
graph_key: int | tuple[int, int] = padded_bsz
|
||||
seq_len = max_seq
|
||||
if self._prefix_graphs_enabled:
|
||||
@@ -1144,6 +1487,8 @@ class CodePredictorWrapper(nn.Module):
|
||||
proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1)
|
||||
if _stats.enabled:
|
||||
_stats.add("predictor.substep_ms", (time.perf_counter() - _t_sub) * 1000.0)
|
||||
if prev_threads is not None:
|
||||
torch.set_num_threads(prev_threads)
|
||||
if _stats.enabled:
|
||||
_stats.add("predictor.loop_ms", (time.perf_counter() - _t_loop) * 1000.0)
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import os
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
@@ -23,6 +24,7 @@ from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.data_entry_keys import OmniPayload
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
from vllm_omni.utils.speaker_cache import (
|
||||
get_speaker_cache,
|
||||
iter_custom_voice_profiles,
|
||||
@@ -303,6 +305,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
|
||||
self.have_multimodal_outputs = True
|
||||
self.has_preprocess = True
|
||||
self.has_postprocess = True
|
||||
# CPU: vLLM's stock sampler spends ~8.5 ms per step in a torch.compile'd
|
||||
# random-sample kernel (measured, Xeon 8480C); the model-side sampler
|
||||
# below reproduces the same math eagerly in <0.5 ms. GPU keeps the stock path.
|
||||
self.prefer_model_sampler = current_omni_platform.is_cpu() and os.environ.get(
|
||||
"VLLM_OMNI_CPU_FAST_SAMPLER", "1"
|
||||
).strip().lower() not in ("0", "false", "off", "no")
|
||||
# Qwen3-TTS postprocess() only reads hidden_states[-1, :]. On a prefix-
|
||||
# cache hit, the last hidden state is in the newly computed tail, so
|
||||
# reconstructing the full cached_prefix + new_tail span is wasted work.
|
||||
@@ -641,6 +649,46 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
|
||||
|
||||
# -------------------- Omni multimodal output plumbing --------------------
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# CPU fast sampler (prefer_model_sampler): same math as vLLM's Sampler.forward
|
||||
# (logits processors, penalties, temperature, top-k/top-p, exponential-noise
|
||||
# sampling with per-request generators) without the compiled random-sample
|
||||
# kernel that dominates the CPU step. Returns None to fall back to the stock
|
||||
# sampler whenever logprobs are requested or the tensor is not on CPU.
|
||||
# ------------------------------------------------------------------
|
||||
def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any:
|
||||
if logits is None or logits.device.type != "cpu" or sampling_metadata.max_num_logprobs is not None:
|
||||
return None
|
||||
if getattr(sampling_metadata, "logprob_token_ids", None):
|
||||
return None
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
sampler = getattr(self, "_stock_sampler", None)
|
||||
if sampler is None:
|
||||
sampler = self._stock_sampler = Sampler()
|
||||
logits = logits.to(torch.float32)
|
||||
# apply_logits_processors applies penalties, bad words and min-tokens (vLLM 0.28)
|
||||
logits = sampler.apply_logits_processors(logits, sampling_metadata, False)
|
||||
if sampling_metadata.all_greedy:
|
||||
sampled = logits.argmax(dim=-1)
|
||||
else:
|
||||
greedy = logits.argmax(dim=-1) if not sampling_metadata.all_random else None
|
||||
logits = sampler.apply_temperature(logits, sampling_metadata.temperature, sampling_metadata.all_random)
|
||||
k, p = sampling_metadata.top_k, sampling_metadata.top_p
|
||||
if k is not None or p is not None:
|
||||
logits = apply_top_k_top_p(logits, k, p)
|
||||
probs = logits.softmax(dim=-1, dtype=torch.float32)
|
||||
q = torch.empty_like(probs)
|
||||
q.exponential_()
|
||||
for i, generator in sampling_metadata.generators.items():
|
||||
q[i].exponential_(generator=generator)
|
||||
sampled = probs.div_(q).argmax(dim=-1)
|
||||
if greedy is not None:
|
||||
sampled = torch.where(sampling_metadata.temperature < 1e-5, greedy, sampled)
|
||||
return SamplerOutput(sampled_token_ids=sampled.view(-1, 1), logprobs_tensors=None)
|
||||
|
||||
def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
|
||||
if isinstance(model_outputs, OmniOutput):
|
||||
return model_outputs
|
||||
|
||||
@@ -1394,10 +1394,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
|
||||
fallthrough is load-bearing, not an error path; it is logged once per
|
||||
model for visibility.
|
||||
"""
|
||||
sampling_metadata = self.input_batch.sampling_metadata
|
||||
_stats = _step_stats()
|
||||
with _stats.timed("model.sample_meta_ms"):
|
||||
sampling_metadata = self.input_batch.sampling_metadata
|
||||
if spec_decode_metadata is None:
|
||||
model_sample = getattr(self.model, "sample", None)
|
||||
self.input_batch.update_async_output_token_ids()
|
||||
with _stats.timed("model.sample_async_ids_ms"):
|
||||
self.input_batch.update_async_output_token_ids()
|
||||
if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
|
||||
# Apply logit bias (min_tokens, allowed_token_ids) before
|
||||
# the custom model sampler — the standard GPU sampler does
|
||||
@@ -1422,10 +1425,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
|
||||
"prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.",
|
||||
type(self.model).__name__,
|
||||
)
|
||||
return self.sampler(
|
||||
logits=logits,
|
||||
sampling_metadata=sampling_metadata,
|
||||
)
|
||||
if _stats.enabled:
|
||||
from vllm_omni.edge.step_stats import instrument_sampler
|
||||
|
||||
instrument_sampler(self.sampler)
|
||||
with _stats.timed("model.sampler_call_ms", sync_device=self.device):
|
||||
return self.sampler(
|
||||
logits=logits,
|
||||
sampling_metadata=sampling_metadata,
|
||||
)
|
||||
|
||||
return super()._sample(logits, spec_decode_metadata)
|
||||
|
||||
|
||||
Reference in new issue
Block a user