tzhouamandClaude Fable 5.1
37c0ea9a6b
- vllm_omni/edge/step_stats.py: process-local counters enabled by VLLM_OMNI_STEP_STATS_DIR (scheduler schedule/update_from_output, engine-core step, chunk hop enqueue->put / put->get / get, orchestrator dispatch), dumped per process and merged into an attribution table (share of the 80 ms frame). Hooks: stage_engine_core_proc (instrument after construction, dump at exit), chunk_transfer_adapter (put/get timing, put->get via the shm lock-file mtime), orchestrator dispatch loop. Bench: stream_latency_bench.py --step-stats. - vllm_omni/edge/scheduling.py: playback slack, prefill-throttle decision and the no-preemption KV admission rule (exact bytes/token from the talker KV geometry); adapt.derive_overrides now sizes max_num_seqs with it and auto_profile passes the geometry read from the model config. - Engine re-reads parallel_stage_init after apply_deploy_profile so the edge profile's orchestrator default takes effect (was read before the profile). - benchmarks/tts/cpu_thread_sweep.py: CPU-only 8/16/32-thread sweep (+inductor) with per-step attribution. - Tests: tests/core/test_edge_scheduling.py, tests/edge/test_step_stats.py, tests/benchmarks/test_cpu_thread_sweep.py, profile default test. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
98 lines
4.1 KiB
Python
98 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Edge scheduling helpers (WP8): playback slack, prefill throttling, no-preemption admission.
|
|
|
|
Pure functions so the policy can be unit-tested without an engine. They are
|
|
consumed by the hardware adaptation layer (``vllm_omni.edge.adapt``) and are the
|
|
reference for the measurement-gated scheduler changes described in
|
|
``analysis/edge_branch_plan.md``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
_DTYPE_BYTES = {"float32": 4, "float16": 2, "bfloat16": 2, "fp8": 1, "float8_e4m3fn": 1, "auto": 2}
|
|
|
|
|
|
# --------------------------------------------------------------- playback slack
|
|
@dataclass(frozen=True)
|
|
class ChunkArrival:
|
|
t_ms: float # arrival time on the consumer clock
|
|
audio_ms: float # audio duration carried by the chunk
|
|
|
|
|
|
def playback_slack_ms(chunks: list[ChunkArrival], now_ms: float, start_ms: float | None = None) -> float:
|
|
"""Audio still buffered ahead of a player that started at ``start_ms``.
|
|
|
|
The player starts at the first chunk's arrival unless ``start_ms`` is
|
|
given. Slack < 0 means the player has already underrun.
|
|
"""
|
|
if not chunks:
|
|
return 0.0
|
|
start = chunks[0].t_ms if start_ms is None else start_ms
|
|
buffered = sum(c.audio_ms for c in chunks if c.t_ms <= now_ms)
|
|
return buffered - max(0.0, now_ms - start)
|
|
|
|
|
|
def should_throttle_prefills(slack_ms: float, next_chunk_eta_ms: float, safety_margin_ms: float) -> bool:
|
|
"""Throttle other requests' prefills when this request's next chunk would land too late.
|
|
|
|
``next_chunk_eta_ms`` is the expected time until the next chunk is
|
|
delivered (steps x step_ms + hop). Throttling is worthwhile only when the
|
|
buffered audio would not cover that wait plus the safety margin.
|
|
"""
|
|
return slack_ms - next_chunk_eta_ms < safety_margin_ms
|
|
|
|
|
|
# --------------------------------------------------------------- KV admission
|
|
def kv_bytes_per_token(
|
|
num_layers: int, num_kv_heads: int, head_dim: int, dtype: str = "bfloat16", kv_cache_dtype: str = "auto"
|
|
) -> int:
|
|
"""Bytes of KV cache one token occupies (K and V, all layers)."""
|
|
d = kv_cache_dtype if kv_cache_dtype != "auto" else dtype
|
|
b = _DTYPE_BYTES.get(d, 2)
|
|
return 2 * num_layers * num_kv_heads * head_dim * b
|
|
|
|
|
|
def kv_bytes_for_seqs(max_num_seqs: int, max_model_len: int, bytes_per_token: int, block_size: int = 16) -> int:
|
|
"""KV bytes needed so ``max_num_seqs`` sequences of ``max_model_len`` never preempt."""
|
|
blocks_per_seq = -(-max_model_len // block_size)
|
|
return max_num_seqs * blocks_per_seq * block_size * bytes_per_token
|
|
|
|
|
|
def max_seqs_without_preemption(
|
|
kv_cache_bytes: int, max_model_len: int, bytes_per_token: int, block_size: int = 16
|
|
) -> int:
|
|
"""Largest ``max_num_seqs`` the KV budget admits with no recompute preemption (>= 0)."""
|
|
per_seq = kv_bytes_for_seqs(1, max_model_len, bytes_per_token, block_size)
|
|
return int(kv_cache_bytes // per_seq) if per_seq > 0 else 0
|
|
|
|
|
|
def admission_allows_no_preemption(
|
|
kv_cache_bytes: int, max_num_seqs: int, max_model_len: int, bytes_per_token: int, block_size: int = 16
|
|
) -> bool:
|
|
return max_seqs_without_preemption(kv_cache_bytes, max_model_len, bytes_per_token, block_size) >= max_num_seqs
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class KVGeometry:
|
|
num_layers: int
|
|
num_kv_heads: int
|
|
head_dim: int
|
|
|
|
@classmethod
|
|
def from_hf_config(cls, cfg: dict) -> KVGeometry | None:
|
|
"""Talker geometry from a Qwen3-TTS (or plain decoder) HF config dict."""
|
|
t = cfg.get("talker_config") if isinstance(cfg.get("talker_config"), dict) else cfg
|
|
try:
|
|
layers = int(t["num_hidden_layers"])
|
|
kv_heads = int(t.get("num_key_value_heads") or t["num_attention_heads"])
|
|
head_dim = int(t.get("head_dim") or (int(t["hidden_size"]) // int(t["num_attention_heads"])))
|
|
except (KeyError, TypeError, ValueError):
|
|
return None
|
|
return cls(layers, kv_heads, head_dim)
|
|
|
|
def bytes_per_token(self, dtype: str = "bfloat16") -> int:
|
|
return kv_bytes_per_token(self.num_layers, self.num_kv_heads, self.head_dim, dtype)
|