""" AsyncLLM-based streaming for Qwen3-ASR — concurrency via vLLM continuous batching. Ports qwen-asr 0.0.6's streaming-state algorithm (init_streaming_state / streaming_transcribe / finish_streaming_transcribe) onto vllm.v1.engine.AsyncLLM, so concurrent generate() calls from many WebSocket streams batch on the GPU instead of serializing through the offline LLM. Benchmark showed ~3.8x at 8 concurrent and ~28x KV-cache headroom on one L4. This module owns ONE shared AsyncLLM engine + processor; each utterance is an independent AsyncUtteranceSession (its own buffer/state), so sessions are isolated and only the GPU engine is shared. """ import os import itertools import numpy as np SAMPLE_RATE = 16000 MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen3-ASR-1.7B") GPU_MEMORY_UTILIZATION = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.80")) STREAMING_MAX_NEW_TOKENS = int(os.getenv("STREAMING_MAX_NEW_TOKENS", "1024")) MAX_MODEL_LEN = int(os.getenv("MAX_MODEL_LEN", "0")) or None CHUNK_SIZE_SEC = float(os.getenv("CHUNK_SIZE_SEC", "4.0")) UNFIXED_CHUNK_NUM = int(os.getenv("UNFIXED_CHUNK_NUM", "5")) UNFIXED_TOKEN_NUM = int(os.getenv("UNFIXED_TOKEN_NUM", "15")) import logging log = logging.getLogger("qwen3-asr-async") _engine = None # vllm AsyncLLM _processor = None # HF processor (chat template + tokenizer) _sp = None # SamplingParams _parse = None # qwen_asr.inference.utils.parse_asr_output _req_ids = itertools.count() def is_ready() -> bool: return _engine is not None def _pick_dtype() -> str: """ Choose engine precision per GPU. bfloat16 needs compute capability >= 8.0 (Ampere+, e.g. L4=8.9); Turing (T4=7.5) supports only float16 and vLLM hard-errors on bf16 there. Override with the DTYPE env var if needed. """ override = os.getenv("DTYPE", "").strip().lower() if override: return override try: import torch cap = torch.cuda.get_device_capability() return "bfloat16" if cap[0] >= 8 else "float16" except Exception: return "float16" def _enforce_eager() -> bool: """ Whether to skip vLLM's torch.compile + CUDA-graph capture at startup. Pre-Ampere GPUs (T4 = compute 7.5) crash during that compile step ('Engine core initialization failed' / 'Not enough SMs'), so default to eager there. Ampere+ (L4/A100) keep the faster compiled path. Override with ENFORCE_EAGER=true/false. """ v = os.getenv("ENFORCE_EAGER", "").strip().lower() if v in ("1", "true", "yes"): return True if v in ("0", "false", "no"): return False try: import torch return torch.cuda.get_device_capability()[0] < 8 except Exception: return False async def init_engine(): """Build the shared AsyncLLM engine + processor (call once, inside the loop).""" global _engine, _processor, _sp, _parse if _engine is not None: return from transformers import AutoConfig, AutoModel, AutoProcessor from qwen_asr.core.transformers_backend import ( Qwen3ASRConfig, Qwen3ASRForConditionalGeneration as HFModel, Qwen3ASRProcessor, ) AutoConfig.register("qwen3_asr", Qwen3ASRConfig) AutoModel.register(Qwen3ASRConfig, HFModel) AutoProcessor.register(Qwen3ASRConfig, Qwen3ASRProcessor) from qwen_asr.core.vllm_backend import Qwen3ASRForConditionalGeneration as VLLMModel from qwen_asr.inference.utils import parse_asr_output from vllm import ModelRegistry, AsyncEngineArgs, SamplingParams from vllm.v1.engine.async_llm import AsyncLLM # V1 async engine (confirmed by spike) ModelRegistry.register_model("Qwen3ASRForConditionalGeneration", VLLMModel) dtype = _pick_dtype() eager = _enforce_eager() engine_args = AsyncEngineArgs( model=MODEL_ID, gpu_memory_utilization=GPU_MEMORY_UTILIZATION, dtype=dtype, max_model_len=MAX_MODEL_LEN, enforce_eager=eager, limit_mm_per_prompt={"audio": 1}, ) log.info(f"Building AsyncLLM for {MODEL_ID} (gpu_mem={GPU_MEMORY_UTILIZATION}, dtype={dtype}, enforce_eager={eager})...") _engine = AsyncLLM.from_engine_args(engine_args) _processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) _sp = SamplingParams(temperature=0.0, max_tokens=STREAMING_MAX_NEW_TOKENS) _parse = parse_asr_output log.info("AsyncLLM engine ready") def _build_prompt(context: str, force_language) -> str: """Mirror qwen_asr._build_text_prompt (chat template + optional forced lang).""" msgs = [ {"role": "system", "content": context or ""}, {"role": "user", "content": [{"type": "audio", "audio": ""}]}, ] base = _processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False) if force_language: base = base + f"language {force_language}" return base async def transcribe_audio(audio: np.ndarray, context: str, force_language) -> str: """ One-shot (non-streaming) transcription of a full clip via the shared AsyncLLM engine — used by the batch /v1/audio/transcriptions endpoint. Returns the parsed transcript text (raw, un-romanized). """ prompt = _build_prompt(context, force_language) inp = {"prompt": prompt, "multi_modal_data": {"audio": [np.asarray(audio, dtype=np.float32)]}} rid = f"b{next(_req_ids)}" out = None async for o in _engine.generate(prompt=inp, sampling_params=_sp, request_id=rid): out = o gen = out.outputs[0].text if (out and out.outputs) else "" _lang, txt = _parse(gen, user_language=force_language) return txt class AsyncUtteranceSession: """ One streaming utterance. Faithful async port of qwen-asr's streaming state: buffer audio into chunks, re-feed accumulated audio each step, roll back the last UNFIXED_TOKEN_NUM tokens for the prefix prompt after UNFIXED_CHUNK_NUM chunks, decode via AsyncLLM, parse to (language, text). """ def __init__(self, context: str, force_language): self.context = context self.force_language = force_language # canonical name or None self.prompt_raw = _build_prompt(context, force_language) self.chunk_samples = max(1, int(round(CHUNK_SIZE_SEC * SAMPLE_RATE))) self.buffer = np.zeros((0,), dtype=np.float32) self.audio_accum = np.zeros((0,), dtype=np.float32) self.chunk_id = 0 self._raw_decoded = "" self.text = "" self.language = "" async def feed(self, audio: np.ndarray) -> str: x = np.asarray(audio, dtype=np.float32) if x.ndim != 1: x = x.reshape(-1) if x.shape[0] > 0: self.buffer = np.concatenate([self.buffer, x]) while self.buffer.shape[0] >= self.chunk_samples: chunk = self.buffer[:self.chunk_samples] self.buffer = self.buffer[self.chunk_samples:] self._accumulate(chunk) await self._decode_step(final=False) return self.text async def finish(self) -> str: if self.buffer.shape[0] > 0: tail = self.buffer self.buffer = np.zeros((0,), dtype=np.float32) self._accumulate(tail) await self._decode_step(final=True) return self.text def _accumulate(self, chunk: np.ndarray): if self.audio_accum.shape[0] == 0: self.audio_accum = chunk else: self.audio_accum = np.concatenate([self.audio_accum, chunk], axis=0) def _compute_prefix(self, final: bool) -> str: if self.chunk_id < UNFIXED_CHUNK_NUM: return "" tok = _processor.tokenizer cur_ids = tok.encode(self._raw_decoded) k = int(UNFIXED_TOKEN_NUM) if final: end_idx = max(1, len(cur_ids) - k) return tok.decode(cur_ids[:end_idx]) # per-chunk: grow rollback until the decoded prefix has no broken char while True: end_idx = max(0, len(cur_ids) - k) prefix = tok.decode(cur_ids[:end_idx]) if end_idx > 0 else "" if "�" not in prefix: return prefix if end_idx == 0: return "" k += 1 async def _decode_step(self, final: bool): prefix = self._compute_prefix(final) prompt = self.prompt_raw + prefix inp = {"prompt": prompt, "multi_modal_data": {"audio": [self.audio_accum]}} rid = f"u{next(_req_ids)}" out = None async for o in _engine.generate(prompt=inp, sampling_params=_sp, request_id=rid): out = o gen_text = out.outputs[0].text if (out and out.outputs) else "" self._raw_decoded = prefix + gen_text self.language, self.text = _parse(self._raw_decoded, user_language=self.force_language) self.chunk_id += 1