diff --git a/tests/model_executor/models/common/test_code_predictor_incremental.py b/tests/model_executor/models/common/test_code_predictor_incremental.py new file mode 100644 index 000000000..14b0e19f3 --- /dev/null +++ b/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 diff --git a/tests/model_executor/models/qwen3_tts/test_cpu_fast_sampler.py b/tests/model_executor/models/qwen3_tts/test_cpu_fast_sampler.py new file mode 100644 index 000000000..3b9139c03 --- /dev/null +++ b/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 diff --git a/vllm_omni/edge/step_stats.py b/vllm_omni/edge/step_stats.py index 09e30a3ab..11f68dbfc 100644 --- a/vllm_omni/edge/step_stats.py +++ b/vllm_omni/edge/step_stats.py @@ -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: diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 4af31a82a..e29840c70 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -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) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 67be40eac..f280b8b0e 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 18a0e9e49..875816326 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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)