File size: 3,551 Bytes
16f5171 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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
|