| """Threaded Fish Speech inference queue for the canonical MXFP8 checkpoint.""" |
|
|
| from __future__ import annotations |
|
|
| import queue |
| import threading |
| import traceback |
| from pathlib import Path |
|
|
| import torch |
| from loguru import logger |
|
|
| from fish_speech.models.text2semantic.inference import ( |
| GenerateRequest, |
| WrappedGenerateResponse, |
| decode_one_token_ar, |
| generate_long, |
| ) |
|
|
| from .modules import MXFP8Linear, load_mxfp8_checkpoint |
|
|
|
|
| def launch_mxfp8_thread_safe_queue( |
| checkpoint_path: str | Path, |
| device: str | torch.device, |
| precision: torch.dtype, |
| compile: bool = False, |
| *, |
| max_length: int = 4096, |
| verify_checksums: bool = False, |
| ) -> queue.Queue: |
| """Load MXFP8 once and serialize Fish Speech generation through a queue.""" |
|
|
| if precision is not torch.bfloat16: |
| raise ValueError("The qualified MXFP8 service requires BF16 exclusions") |
| if compile: |
| raise ValueError("torch.compile is not qualified for the MXFP8 API service") |
| if max_length < 1024: |
| raise ValueError("MXFP8 API cache length must be at least 1024") |
|
|
| input_queue: queue.Queue = queue.Queue() |
| init_event = threading.Event() |
| init_error: list[BaseException] = [] |
|
|
| def worker() -> None: |
| try: |
| model = load_mxfp8_checkpoint( |
| checkpoint_path, |
| device=device, |
| max_length=max_length, |
| verify_checksums=verify_checksums, |
| ) |
| with torch.device(device): |
| model.setup_caches( |
| max_batch_size=1, |
| max_seq_len=model.config.max_seq_len, |
| dtype=torch.bfloat16, |
| ) |
| model._cache_setup_done = True |
| module_count = sum( |
| isinstance(module, MXFP8Linear) for module in model.modules() |
| ) |
| if module_count != 180: |
| raise RuntimeError( |
| f"Expected 180 native MXFP8 projections, found {module_count}" |
| ) |
| logger.info( |
| "Loaded qualified MXFP8 slow-transformer checkpoint with " |
| f"{module_count} native projections on {device}" |
| ) |
| except BaseException as error: |
| init_error.append(error) |
| logger.error(traceback.format_exc()) |
| init_event.set() |
| return |
|
|
| init_event.set() |
| while True: |
| item: GenerateRequest | None = input_queue.get() |
| if item is None: |
| break |
| response_queue = item.response_queue |
| try: |
| for chunk in generate_long( |
| model=model, |
| decode_one_token=decode_one_token_ar, |
| **item.request, |
| ): |
| response_queue.put( |
| WrappedGenerateResponse(status="success", response=chunk) |
| ) |
| except Exception as error: |
| logger.error(traceback.format_exc()) |
| response_queue.put( |
| WrappedGenerateResponse(status="error", response=error) |
| ) |
| finally: |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| threading.Thread( |
| target=worker, |
| name="fish-s2-pro-mxfp8-worker", |
| daemon=True, |
| ).start() |
| init_event.wait() |
| if init_error: |
| raise RuntimeError("MXFP8 model worker failed to initialize") from init_error[0] |
| return input_queue |
|
|
