[rwang5203]

Motivation

Add native SGLang Diffusion serving for inclusionAI/LLaDA-Image and inclusionAI/LLaDA-Image-Turbo in BF16 on a single GPU. The integration supports text-to-image generation and single-image editing through the OpenAI-compatible Images API. Sequence parallelism and the FP8 checkpoints follow in a separate PR.

Contributors

• InclusionAI, Ant Group. Kun Chen (@kunchengit), Chengzhang Wu (@Anmuliar), Chaoyang Wang (@wzy-ustc)
• SGLang community. Kaige Li (@stephenkgli), Yuhan Zhou (@Dayuxiaoshui), Yihao Wang (@AgainstEntropy), Mick Qian (@mickqian), Richard Wang (@rwang5203)

Modifications

• Pipeline. Add the LLaDA-Image DiT, pipeline configuration, sampling parameters, flow-match scheduler support for its sigma schedule, an embedded LLaDA2 MoE text encoder, source-image conditioning, and its registration in the model's own config module.
• Serving. Support /v1/images/generations and /v1/images/edits, with request validation for image dimensions and text length. Request validation errors from the Images API now return 400 with their message instead of 500 for every image model, as the video API already does.
• Text encoder. The LLaDA2 MoE text encoder runs as an embedded srt model with its own runtime context, its own srt process groups, and a FlashInfer attention backend that applies the conditioning mask. While it builds its groups it sets aside the TP, attention TP, and MoE groups that the diffusion runtime mirrors into srt, restores them afterwards, and uses its own groups for each encode. The srt changes below act only when the LLaDA-Image pipeline builds this encoder, except the LLaDA2 loading step, which leaves existing LLaDA2 checkpoints loading as before.
• srt attention backend. llada2_cfg_flashinfer subclasses the FlashInfer backend. When a forward batch carries LLaDA-Image conditioning text lengths, it plans one ragged prefill with a per-request mask in which text tokens attend only to text tokens and query tokens attend to the whole sequence, and it refuses decode, speculative, cached-prefix, and multi-wrapper batches. Every other batch runs the stock FlashInfer path.
• srt forward batch. ForwardBatch.llada_image_conditioning_text_lens_cpu carries those lengths through batch rebuilds. It defaults to None, so other batches are unchanged, and the encoder raises an error if a forward ran without the mask.
• srt runtime context. create_context builds a runtime context for the embedded encoder without replacing the diffusion worker's published config, and use_context activates it around encoder setup and each encode, restoring the previous context on exit and on error. publish now shares its configuration step with create_context and behaves as before.
• srt LLaDA2 model. LLaDA2MoeModelLM gains get_input_embeddings, and its weight loading maps the encoder checkpoint's model.language_model.* and model.lm_head.* names and its fused expert tensors to the native layout.
• Safety and lifecycle. Propagate the checkpoint revision and remote-code trust into the embedded encoder, keep the encoder resident, and reject multi-GPU, quantized, and cache-dit launches at startup with a clear error.
• Defaults. inclusionAI/LLaDA-Image uses the model card recipe of 50 steps and guidance 5.0, and inclusionAI/LLaDA-Image-Turbo registers its own 4 step defaults without CFG.
• Input limits. Generation sizes must be multiples of 16 and editing sizes multiples of 32, as in the reference pipeline, checked when the request is validated. An editing size whose image tokens alone exceed the 32768-row sequence RoPE table is refused there too, and the DiT checks the exact length with the prompt before the RoPE lookup, so an oversized edit no longer breaks the worker for later requests.
• Editing source. The source image is resized to the full output frame instead of center cropped, matching the reference pipeline for sources whose aspect ratio differs from the output.

Reproduction

Start Turbo.

sglang serve --model-path inclusionAI/LLaDA-Image-Turbo --revision f4afc52d925bbac4e22a1c947111fc1f127e37e5 --trust-remote-code --port 30000

Start Base.

sglang serve --model-path inclusionAI/LLaDA-Image --revision e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039 --trust-remote-code --port 30000

Generate an image with Turbo.

curl -sS -X POST http://127.0.0.1:30000/v1/images/generations -H "Content-Type: application/json" -d '{"prompt":"A cinematic photograph of a red fox standing in fresh snow at golden hour, detailed fur, shallow depth of field","size":"1024x1024","n":1,"response_format":"b64_json","num_inference_steps":4,"guidance_scale":1.0,"seed":42}'

Edit an image with Turbo.

curl -sS -X POST http://127.0.0.1:30000/v1/images/edits -F "image=@input.png" -F "prompt=Place the subject on a beach" -F "size=1024x1024" -F "n=1" -F "response_format=b64_json" -F "num_inference_steps=4" -F "guidance_scale=1.0" -F "seed=42"

For Base, use "num_inference_steps":50 and "guidance_scale":5.0.

Accuracy Tests

These quality results were measured with the BF16 single-GPU path at an earlier revision of this branch. A re-run on the final commit with the same prompts, seeds, and evaluators is in progress, and this section will be updated with its results.

The benchmark generated and decoded all 6704 expected artifacts, 2212 GenEval images, 1280 LongText images, 1212 GEdit images, and 2000 CVTG images.

Configuration	GenEval macro	LongText en	LongText zh	CVTG Word Accuracy
Base BF16	0.8404	0.9120	0.9198	0.7792
Turbo 4-step	0.8190	0.8959	0.9123	0.8267

Base used 20 steps for GenEval and CVTG and 50 steps for LongText, all with CFG 5. Turbo used 4 steps with CFG 1.

Turbo GEdit metric	Score
en semantic / quality / overall	7.9242 / 7.3682 / 7.5728
zh semantic / quality / overall	7.7375 / 7.3867 / 7.4347

The GEdit scores predate the editing source change above, which resizes the source to the full output frame instead of center cropping it.

The final commit also passes these checks.

• Both registered LLaDA-Image server cases, text to image and editing, pass in CI, including their performance baseline checks.
• The srt unit tests touched here pass in the CPU CI job.
• The new LLaDA2 loading step yields every tensor of LLaDA2.0-mini, LLaDA2.1-mini, and LLaDA2.2-mini with its name and data unchanged, so these checkpoints load as they do on main.

Speed Tests and Profiling

The following values are median client-wall latency measured on this branch before its latest merges of main, on one NVIDIA H200 at 1024x1024 with one request at a time. Each cell pools 40 measured requests from four fresh server processes, each running two warmups and then ten measured requests.

Task	Turbo 4-step	Base 50-step
T2I	0.5822s	13.1110s
Edit	1.4372s	31.4010s

Known Limitations

• The embedded text encoder requires reviewed remote code, so deployments must use --trust-remote-code and should pin the checkpoint revision.
• Sequence parallelism and the FP8 checkpoints are not served by this PR and are refused at startup. They follow in a separate PR.
• The embedded text encoder must remain resident and does not support generic memory sleep or hot weight updates.
• An edit with more than one source image, or an edit whose prompt and image together exceed the 32768-row sequence limit, is refused by the worker without affecting later requests, but the client receives a 500 response instead of a 400. A follow-up PR moves both checks into request validation.

---

CI States

Latest PR Test (Base): ⏳ Run #38063465852
Latest PR Test (Extra): ❌ Run #38063465537
Latest PR Test (AMD ROCm 10): ❌ Run #38063465628

[rwang5203]

/rerun-failed-ci