feat(engine): [image, audio] sglang image vllm asr engines - #5121
Conversation
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
…-dev Implement the SGLangImageModel stub with sglang-diffusion offline API (sglang.multimodal_gen.DiffGenerator) so text-to-image models can run on the SGLang engine on Linux with NVIDIA GPUs. Supported models: Qwen-Image, Qwen-Image-2512, Z-Image, Z-Image-Turbo and FLUX.1-dev, matching the sglang-diffusion compatibility matrix. Virtualenv specs install sglang[diffusion] when the SGLang engine is selected. Ref xorbitsai#4896
Add an engine abstraction layer for audio models (mirroring the image module design): AUDIO_ENGINES registry, engine matching and virtualenv marker support. Families without registered engines keep the legacy hardcoded dispatch, so existing behavior is unchanged. Qwen3-ASR-0.6B/1.7B register two engines: the default transformers engine and, on Linux with NVIDIA GPUs, a vLLM engine backed by qwen_asr's Qwen3ASRModel.LLM which shares the same transcribe API. /v1/engines/audio/<model_name> now returns engine params for audio models. Ref xorbitsai#4896
Show the engine select in the launch dialog for audio models that support engine selection (e.g. Qwen3-ASR on transformers/vLLM), and include audio in the model types that query /v1/engines. Audio models without registered engines keep the previous form unchanged. Ref xorbitsai#4896
7e346ee to
77a64ad
Compare
Implement the vLLM image engine placeholder with a real engine backed by vllm-omni's Omni text-to-image entrypoint. The engine matches on Linux with NVIDIA GPUs for the same model set as the SGLang engine: Qwen-Image, Qwen-Image-2512, Z-Image, Z-Image-Turbo and FLUX.1-dev. Add vllm-omni virtualenv markers to the builtin specs so engine-scoped installs work. Ref xorbitsai#4896
vllm-omni does not declare vllm as a dependency but requires a vllm with the same major.minor version (its patch module imports private vllm APIs), so an unpinned marker resolved vllm-omni against whatever vllm the parent environment had and crashed on import. Pin the vllm-omni/vllm pair in the virtualenv packages, surface the original import error with the pairing requirement in the failure message, and assert the pins stay in sync in tests. Ref xorbitsai#4896
The vLLM image engine serialized all requests: VLLMDiffusionModel had no allow_batch flag, so ModelActor wrapped every text_to_image call in its per-model asyncio lock, and each call went through the blocking Omni.generate. The vllm-omni engine itself schedules requests from a queue and pipelines stages, so the serialization was purely on the xinference side. Omni.generate cannot simply be called concurrently: it drains the shared engine output queue and drops messages belonging to other requests, and per-request sampling params (seed, size, steps) cannot be expressed through one batched generate call since sampling_params_list is per-stage, not per-prompt. Set allow_batch = True and submit each request directly via engine.add_request with its own request_id and sampling params. A single dispatcher thread drains the engine output queue and routes messages back to per-request futures; a fatal ErrorMessage fails all in-flight requests, a per-request error fails only its own future. The dispatcher exits when idle and restarts on the next submission. When the running vllm-omni does not expose the needed engine internals, fall back to the previous serial Omni.generate path, now guarded by a model-level lock because the actor no longer serializes calls.
sglang pins its own torch stack (sglang 0.5.7 requires torch==2.9.1 exactly), so the unscoped #system_torch#/#system_numpy# placeholders in the vLLM/SGLang image model specs expand to the host torch (e.g. 2.11.0) and can never co-resolve with sglang[diffusion] in one virtualenv. Scope the system placeholders (and the transformers floor, which sglang also pins itself) to the diffusers and vLLM engines so the SGLang venv resolves sglang[diffusion] self-contained with its own torch. Note the SGLang install path additionally requires the xoscar skip_installed fix for requirements with extras: without it, sglang[diffusion] is treated as satisfied by a bare host sglang and the venv is left empty.
The SGLang/vLLM image engines copied the full ability list from the builtin spec, so running instances reported image2image/inpainting endpoints that these engines do not implement and the image-edit API could auto-select them. Restrict the advertised abilities to text2image on an engine-local copy of the family, keeping the shared builtin spec untouched.
CI's isolated mypy environment types the pydantic copy() as Any, so reassigning the optional model_spec parameter lost the None narrowing and failed with union-attr. Build the restricted copy in a local variable instead.
The worker pickles the constructed model instance into the model actor subprocess before load(); the concurrent-dispatch locks introduced in the vLLM engine broke that with 'cannot pickle _thread.lock object'. Exclude locks and the dispatcher thread from the pickled state and recreate the locks on unpickling, with a regression test.
The model actor injects request_id as uuid.UUID; SGLang's SamplingParams accepts the field but its runtime JSON-serializes the params and crashed with 'Object of type UUID is not JSON serializable'. Convert UUID values to strings when building sampling params.
sglang JIT-compiles kernels through the ninja binary; the ninja wheel installs it into the virtualenv's bin directory, which is not on PATH when the model subprocess is spawned with the venv interpreter, so generation failed with 'No such file or directory: ninja'. Prepend the interpreter's bin directory to PATH at load time.
End-to-end benchmark: Z-Image on diffusers / vLLM / SGLangRan all three image engines for real on a GPU box against this PR's HEAD, same prompt and generation config, measuring client-side wall-clock of Setup
Single-request latency
Concurrency: 4 simultaneous requests (fired at once after warmup)
Reading the numbers:
Cross-request GPU batching: sd3.5-medium (vllm-omni pipeline with Z-Image's vllm-omni pipeline declares
Requests in the same wave complete with identical timestamps (21.478/21.460/21.478 s) -- they went through the DiT as one batched tensor -- and the engine logs All three engines produced valid, prompt-faithful 1024x1024 images (visually verified). Bugs found and fixed while testing (pushed to this branch):
Environment notes (outside this PR's scope)
|
A fatal vLLM-Omni error previously only failed the in-flight waiters; the model stayed open, so the next request passed the closed check, was submitted to the dead backend and could wait forever. Transition to a terminal state atomically with draining the waiters (mark closed, record the fatal error, shut the engine down), and fail subsequent submissions and text_to_image calls immediately with the recorded cause. Dispatcher loop failures take the same path. Regression tests cover both.
sd3.5-medium runs on vllm-omni's SD3 pipeline, which supports request-level batching (supports_request_batch=True), so concurrent requests are batched on the GPU when launched with max_num_seqs > 1. Verified end to end on an RTX 3090 Ti: concurrent requests in the same wave complete with identical timestamps and the engine logs '[RequestBatch] engine init max_num_seqs=4'. Ref xorbitsai#4896
#4896