[Perf][Higgs-Audio-V3] LRU cache for voice-clone ref-audio encode (#4200)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
1 parent
82e74ce972
commit
fce0f51089
2 files changed
+118
-1
No files matched your search
@@ -1935,10 +1935,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
|
||||
apply_delay_pattern,
|
||||
encode_reference_audio,
|
||||
)
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.ref_audio_cache import (
|
||||
cached_encode_reference_audio,
|
||||
)
|
||||
|
||||
wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
|
||||
wav = np.asarray(wav_list, dtype=np.float32)
|
||||
ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, int(sr))
|
||||
# Cache the codec encode by content hash so identical ref_audio
|
||||
# payloads skip the encode and reuse the same [num_frames,
|
||||
# num_codebooks] codes tensor across requests.
|
||||
ref_codes_raw = await asyncio.to_thread(cached_encode_reference_audio, wav, int(sr), encode_reference_audio)
|
||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
||||
|
||||
prompt_ids = adapter.build_prompt(
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bounded LRU cache for Higgs Audio v3 reference-audio codec encoding.
|
||||
|
||||
Voice-clone requests that re-send the same ref_audio are otherwise re-encoded
|
||||
through the DAC encoder on every call. The cache keys by a stable content
|
||||
hash (size + sample rate + head/mid/tail sentinel bytes) so the same audio
|
||||
served via different URLs or repeated base64 payloads also hits warm.
|
||||
Tensors are cloned on get/put so callers and the cache cannot alias storage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
__all__ = ["cached_encode_reference_audio", "cache_stats", "clear_cache"]
|
||||
|
||||
_DEFAULT_MAX_ENTRIES = 128
|
||||
_SENTINEL_BYTES = 4096
|
||||
|
||||
_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
|
||||
_cache_lock = threading.Lock()
|
||||
_hits = 0
|
||||
_misses = 0
|
||||
|
||||
|
||||
def _content_hash(wav: np.ndarray, sr: int) -> str:
|
||||
"""Stable content hash. Head/mid/tail sentinel keeps cost O(1)."""
|
||||
raw = wav.tobytes()
|
||||
n = len(raw)
|
||||
h = hashlib.sha256()
|
||||
h.update(n.to_bytes(8, "little"))
|
||||
h.update(int(sr).to_bytes(4, "little"))
|
||||
if n == 0:
|
||||
return h.hexdigest()
|
||||
if n <= _SENTINEL_BYTES * 3:
|
||||
h.update(raw)
|
||||
return h.hexdigest()
|
||||
h.update(raw[:_SENTINEL_BYTES])
|
||||
mid = (n - _SENTINEL_BYTES) // 2
|
||||
h.update(raw[mid : mid + _SENTINEL_BYTES])
|
||||
h.update(raw[-_SENTINEL_BYTES:])
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def cached_encode_reference_audio(
|
||||
wav: np.ndarray,
|
||||
sr: int,
|
||||
encode_fn: Callable[[np.ndarray, int], torch.Tensor],
|
||||
max_entries: int = _DEFAULT_MAX_ENTRIES,
|
||||
) -> torch.Tensor:
|
||||
"""LRU-cached wrapper around ``encode_reference_audio``.
|
||||
|
||||
``encode_fn`` is the actual codec encoder; passed in so this module
|
||||
stays decoupled from the tokenizer import (avoids circular imports
|
||||
during stage-input-processor discovery).
|
||||
|
||||
Returns a CPU tensor clone — never the cache's interior storage.
|
||||
"""
|
||||
global _hits, _misses
|
||||
key = _content_hash(wav, sr)
|
||||
with _cache_lock:
|
||||
cached = _cache.get(key)
|
||||
if cached is not None:
|
||||
_hits += 1
|
||||
_cache.move_to_end(key)
|
||||
return cached.clone()
|
||||
|
||||
# Miss: encode outside the lock so concurrent misses on different
|
||||
# refs don't serialize on this single mutex.
|
||||
codes = encode_fn(wav, sr)
|
||||
codes_cpu = codes.detach().to("cpu").contiguous()
|
||||
with _cache_lock:
|
||||
if key in _cache:
|
||||
# Concurrent miss winner already stored it; refresh LRU position
|
||||
# and return that copy to avoid divergent tensors per request.
|
||||
_cache.move_to_end(key)
|
||||
_misses += 1 # we still paid the encode cost
|
||||
return _cache[key].clone()
|
||||
_cache[key] = codes_cpu
|
||||
if len(_cache) > max_entries:
|
||||
_cache.popitem(last=False)
|
||||
_misses += 1
|
||||
return codes_cpu.clone()
|
||||
|
||||
|
||||
def cache_stats() -> dict:
|
||||
"""Return (hits, misses, size, hit_rate) — for /metrics-style reporting."""
|
||||
with _cache_lock:
|
||||
total = _hits + _misses
|
||||
return {
|
||||
"hits": _hits,
|
||||
"misses": _misses,
|
||||
"size": len(_cache),
|
||||
"hit_rate": (_hits / total) if total > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def clear_cache() -> None:
|
||||
"""Drop all cached entries. Used by tests."""
|
||||
global _hits, _misses
|
||||
with _cache_lock:
|
||||
_cache.clear()
|
||||
_hits = 0
|
||||
_misses = 0
|
||||
Reference in new issue
Block a user