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>
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
"""WP8 edge scheduling policy: playback slack, prefill throttling, no-preemption admission (CPU)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.edge import adapt
|
|
from vllm_omni.edge.hardware_probe import HardwareProfile
|
|
from vllm_omni.edge.scheduling import (
|
|
ChunkArrival,
|
|
KVGeometry,
|
|
admission_allows_no_preemption,
|
|
kv_bytes_for_seqs,
|
|
kv_bytes_per_token,
|
|
max_seqs_without_preemption,
|
|
playback_slack_ms,
|
|
should_throttle_prefills,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
QWEN3_TTS_TALKER = KVGeometry(num_layers=28, num_kv_heads=8, head_dim=128)
|
|
|
|
|
|
def test_playback_slack_tracks_buffered_audio_minus_elapsed():
|
|
# measured default schedule: 1 frame (80 ms) at t=30, 25 frames (2000 ms) at t=190
|
|
chunks = [ChunkArrival(30.0, 80.0), ChunkArrival(190.0, 2000.0)]
|
|
assert playback_slack_ms(chunks, now_ms=30.0) == pytest.approx(80.0)
|
|
assert playback_slack_ms(chunks, now_ms=110.0) == pytest.approx(0.0) # underrun starts here
|
|
assert playback_slack_ms(chunks, now_ms=150.0) == pytest.approx(-40.0)
|
|
assert playback_slack_ms(chunks, now_ms=190.0) == pytest.approx(2080.0 - 160.0)
|
|
# edge ramp: 2,4 frames at 42 ms / 200 ms never underruns
|
|
ramp = [ChunkArrival(42.0, 160.0), ChunkArrival(200.0, 320.0)]
|
|
assert playback_slack_ms(ramp, now_ms=199.0) > 0
|
|
assert playback_slack_ms([], now_ms=5.0) == 0.0
|
|
|
|
|
|
def test_throttle_decision_uses_eta_and_margin():
|
|
assert should_throttle_prefills(slack_ms=50.0, next_chunk_eta_ms=40.0, safety_margin_ms=20.0)
|
|
assert not should_throttle_prefills(slack_ms=500.0, next_chunk_eta_ms=160.0, safety_margin_ms=20.0)
|
|
assert should_throttle_prefills(slack_ms=-10.0, next_chunk_eta_ms=0.0, safety_margin_ms=0.0)
|
|
|
|
|
|
def test_kv_geometry_and_bytes():
|
|
assert kv_bytes_per_token(28, 8, 128, "bfloat16") == 2 * 28 * 8 * 128 * 2 == 114688
|
|
assert kv_bytes_per_token(28, 8, 128, "float32") == 229376
|
|
assert QWEN3_TTS_TALKER.bytes_per_token("bfloat16") == 114688
|
|
geo = KVGeometry.from_hf_config(
|
|
{
|
|
"talker_config": {
|
|
"num_hidden_layers": 28,
|
|
"num_key_value_heads": 8,
|
|
"head_dim": 128,
|
|
"num_attention_heads": 16,
|
|
}
|
|
}
|
|
)
|
|
assert geo == QWEN3_TTS_TALKER
|
|
assert KVGeometry.from_hf_config({"foo": 1}) is None
|
|
# blocks round up to the block size
|
|
assert kv_bytes_for_seqs(1, 17, 100, block_size=16) == 32 * 100
|
|
|
|
|
|
def test_admission_rule_for_qwen3_tts_edge_profile():
|
|
bpt = QWEN3_TTS_TALKER.bytes_per_token("bfloat16")
|
|
# 4 seqs x 2048 tokens x 112 KiB = 896 MiB needed
|
|
need = kv_bytes_for_seqs(4, 2048, bpt)
|
|
assert need == 4 * 2048 * 114688
|
|
assert admission_allows_no_preemption(need, 4, 2048, bpt)
|
|
assert not admission_allows_no_preemption(need - 1, 4, 2048, bpt)
|
|
assert max_seqs_without_preemption(need - 1, 2048, bpt) == 3
|
|
assert max_seqs_without_preemption(0, 2048, bpt) == 0
|
|
|
|
|
|
ADAPT_PROFILES = {
|
|
"x86_cpu_8g": dict(
|
|
arch="x86_64",
|
|
accelerator="none",
|
|
ram_total_bytes=8 << 30,
|
|
ram_available_bytes=6 << 30,
|
|
cpu_flags=["avx2"],
|
|
n_physical=4,
|
|
n_logical=4,
|
|
),
|
|
"jetson_8g": dict(
|
|
arch="aarch64",
|
|
accelerator="cuda_unified",
|
|
ram_total_bytes=8 << 30,
|
|
ram_available_bytes=6 << 30,
|
|
gpu_mem_bytes=8 << 30,
|
|
cpu_flags=["neon"],
|
|
n_physical=6,
|
|
n_logical=6,
|
|
supports_bf16=True,
|
|
),
|
|
"l20x": dict(
|
|
arch="x86_64",
|
|
accelerator="cuda_discrete",
|
|
ram_total_bytes=256 << 30,
|
|
ram_available_bytes=200 << 30,
|
|
gpu_mem_bytes=48 << 30,
|
|
cpu_flags=["avx2", "avx512"],
|
|
n_physical=32,
|
|
n_logical=64,
|
|
supports_bf16=True,
|
|
),
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("profile_kwargs", list(ADAPT_PROFILES.values()), ids=list(ADAPT_PROFILES.keys()))
|
|
def test_derived_budgets_never_allow_preemption(profile_kwargs):
|
|
"""The derived stage-0 KV budget always holds max_num_seqs x max_model_len talker tokens."""
|
|
profile = HardwareProfile(**profile_kwargs)
|
|
ov = adapt.derive_overrides(profile, weights_bytes=2 << 30, kv_geometry=QWEN3_TTS_TALKER)
|
|
s0 = ov["stages"][0]
|
|
kv0 = s0["engine_args"]["kv_cache_memory_bytes"]
|
|
dtype = s0["engine_args"]["dtype"]
|
|
bpt = QWEN3_TTS_TALKER.bytes_per_token(dtype)
|
|
assert s0["max_num_seqs"] >= 1
|
|
assert admission_allows_no_preemption(kv0, s0["max_num_seqs"], s0["max_model_len"], bpt), (kv0, s0)
|