[diffusion] Refactor ComfyUI integrated mode onto native pipelines + adapters

[niehen6174]

Summary

ComfyUI integrated mode is no longer a third copy of each diffusion pipeline. It is a DiT-only runtime profile on the native Flux / Z-Image / Qwen pipelines (--comfyui-mode).

• Drop the three comfyui_*_pipeline.py forks (~1450 lines).
• Keep one shared executor. Model differences live in a checkpoint spec (single-file .safetensors remap) and a step adapter (pack / unpack / fill_req).
• Keep the existing process topology (ComfyUI process + spawned workers). Local hops use CUDA IPC handles instead of pickling GPU tensors through host memory. Multi-rank --comfyui-mode detaches CUDA tensors and broadcasts them over NCCL. Conditioning is cached in a run-scoped worker session.

Adding another image DiT that already has ComfyUI UNET / CLIP / VAE / KSampler should be two small files, not another 400–700 line pipeline.

Motivation

Integrated mode exists so ComfyUI keeps CLIP, VAE, the sampler loop, CFG, and community nodes, while SGLang only accelerates the DiT. That boundary is unchanged.

The old implementation encoded that boundary by copying a pipeline and an executor per model. Flux / Z-Image / Qwen already had native pipelines. The ComfyUI copies existed only because:

1. ComfyUI gives a single .safetensors UNET, not model_index.json.
2. The worker must run one DiT step with a pass-through scheduler.
3. ComfyUI key names do not match SGLang param_names_mapping.

comfyui_mode was already scattered through DenoisingStage / SamplingParams, but it was never a first-class profile, so each new model copied the skeleton again.

Layer	Before	Actually model-specific
ComfyUIXxxPipeline	~1450 lines across 3 files	weight mapping, QKV split, scale/shift order
XxxExecutor	~300 lines across 4 classes	pack / unpack, timestep scale, which Req fields get embeds
generator.py tables	edit two dicts per model	none

A new image DiT needed ~500 lines of glue. Most of it was identical.

Architecture

• Befor (loop, CFG, model_options callables, community nodes) stays in the ComfyUI process.
• After (DiT, SP/TP, TeaCache / BCG / FSDP) stays in the SGLang workers.
• Between them: tensors only (IPC-able). No Python closures across the process boundary.

No new comfyui_*_pipeline.py. pipeline_class_dict / executor_class_dict are generated from adapter.model_types.

Why this is easier to extend

	Before	After
New image DiT	pipeline 400–700 + executor 70–170 + two registries	checkpoint spec + adapter (Qwen spec ~55, Z-Image ~63, typical adapter 40–90)
Weight mapping	ComfyUI copy + native config, easy to diverge	one ComfyUICheckpointSpec; inherit_config_mapping when the two mappings share a source namespace
Loader features	Flux used a 175-line handwritten state_dict loop (no FSDP / quant)	shared TransformerLoader / weights_iterator (same shape as GGUF)
H3 / packed video	executor only understood forward(x, t, context) + 4D image pack	PackedForward.extra_req and nested latents are reserved; H3 adapter is not in this PR
Server mode	untouched	untouched

Concrete correctness win from making the profile real: the old Flux ComfyUI loader merged native FluxConfig.param_names_mapping (targets like to_qkv, ff.linear_in) with the ComfyUI aliases. Mapping is applied to a fixpoint, so already-correct diffusers names were rewritten to names the live module does not have. The handwritten loader built on a real device and left unmatched params randomly initialized with a warning. The new path builds on meta device and errors instead, which is how the bug surfaced. Flux now sets inherit_config_mapping=False. Z-Image keeps True because its native mapping must merge split to_q/k/v back into fused to_qkv.

The adapter interface is sized for H3 (nested latents, structured payload), not for Flux's three-tensor case. This PR does not add MiniMax-H3 support; it stops the next model from needing another pipeline fork.

Transport / performance

These hops were using the per-request distribution path on every sampler step.

Hop	Before	After
ComfyUI ↔ rank 0	pickle.dumps copies CUDA tensors to host. Measured: 460 KB tensor → 461 KB pickle	CudaIpcRef handle. Unit test: handle pickle < 2 KB. Receiver _new_shared_cuda + clone
Multi-rank --comfyui-mode	whole Req over broadcast_pyobj / gloo CPU, then _fix_tensor_device	control envelope still gloo; GPU tensors via existing NCCL broadcast_tensor_dict. General SP / CFG / TP recv is unchanged
Conditioning	CLIP / pooled / image embeds resent every step	comfyui_session_id; later steps send latent + timestep

On current image models this is not a large wall-clock win. Flux 1024² packed latent is ~0.5 MB/step, so host pickle was already cheap next to DiT (~124 ms Flux / ~232 ms Qwen). The leftover ~10 ms/step is mostly scheduler wakeup, not the tensor copy.

The transport work is for correctness and the next model: stay off host pickle / gloo, keep GPU tensors as IPC / NCCL handles, and cache conditioning so later steps do not resend embeds. A video DiT (H3 ~14 MB/step) would otherwise copy gigabytes over a full sampler. Processes stay separate on purpose — folding ComfyUI into rank-0 would shave that last ~10 ms, but mixes its event loop with NCCL and process-global CUDA state.

Test plan

• test/unit/test_comfyui_profile.py — profile assembly + Flux synthetic single-file load (meta device, fused QKV / adaLN)
• test/unit/test_comfyui_adapters.py — pack / unpack
• test/unit/test_comfyui_session.py — conditioning cache
• test/unit/test_ipc_cuda.py — handle pickle < 2 KB, cross-process roundtrip
• Weight convert equivalence vs the deleted pipelines on synthetic state dicts (Flux / Z-Image / Qwen names, merge info, values)
• Z-Image comfyui_mode one-step (noise_pred [1, 16, 1, 90, 160]) and ComfyUI-format single-file load (6.15B)
• Flux comfyui_mode 2-GPU one-step (noise_pred [1, 3600, 64], ~0.7 s)
• Full ComfyUI UI / KSampler workflow on a machine with matching comfy_aimdo / comfy_kitchen
• Existing plugin e2e (test_flux_pipeline.py, test_zimage_pipeline.py, Qwen) on real checkpoints

---

CI States

Latest PR Test (Base): ❌ Run #36252148993
Latest PR Test (Extra): ❌ Run #36252148992
Latest PR Test (AMD ROCm 10): ❌ Run #36252149165

[mickqian]

/tag-and-rerun-ci

[mickqian]

maybe move it to distributed/utils?

[niehen6174]

Moved it under distributed/ instead of folding it into utils.py.
That file gets star-imported by the package, so dumping the CUDA IPC retain/spill path in there would load it on every distributed import.

[mickqian]

Current-head CI root failure: test_declared_alias_loads_by_exact_key_and_structural_source passes a SimpleNamespace without comfyui_mode to is_comfyui_mode, which directly reads that field: https://github.com/sgl-project/sglang/actions/runs/33407655401/job/99539408863 . Align the fixture with the supported argument contract and validate that real loader test. Component/B200/multi-GPU failures are health-check cascades. The 14:30 UTC snapshot also reported main conflicts (latest query UNKNOWN). I do not have fork push access; please resolve and push before final CI.

[mickqian]

Current-head CI at 3540631 is blocked before tests by check-maintenance: required main base commit 3fc7a66 is not an ancestor (diverged). Exact check: https://github.com/sgl-project/sglang/actions/runs/36115937564/job/108009951064 ; pr-test-finish is only the aggregate failure. Please merge current main normally into this branch, resolve any conflicts, push, and validate the new exact head. Re-running the unchanged head will not satisfy the ancestry gate; no bypass or force-push is needed. The previous-head comfyui_mode fixture failure is not evidence of a test failure on this new head, since this run has not reached the tests. I do not have fork push access.

[mickqian]

Current-head CI follow-up for 359b9a1b31317910601afd77cf68667c95284815, Base run 36207375945:

• Unit job 108306727631 still fails TestComponentLoaderIdentity.test_declared_alias_loads_by_exact_key_and_structural_source: load_modules() calls is_comfyui_mode(server_args), but the native SimpleNamespace fixture lacks comfyui_mode (comfyui_mode.py:54). Final suite: 1 failed, 3343 passed, 132 skipped. Please set the native fixture's explicit comfyui_mode=False and validate this real loader boundary; the nested pytest failure output elsewhere in the log is expected test data, not additional failing tests.
• Dual-GPU partition 1, job 108306767313: four load-latency failures (ms, actual > limit): Wan2.2 I2V 116903.1686 > 98334.75; Wan2.2 TeaCache 129193.2776 > 79922.5375; Wan2.1 I2V 480P 74732.2557 > 58002.775; LTX2.5 decoder 37955.1707 > 35166.55. Initial pass only; remaining 964.5s was below the 1176.5s failed-item retry budget.
• Dual-GPU partition 2, job 108306767357: seven performance failures: MiniMax H3 ref2va E2E 172067.0072 > 84741.475; Wan1.3B Cache-DiT load 46770.6748 > 42957.7375; FSDP load 45665.6623 > 39225.675; LTX2 two-stage load 61270.0318 > 59338.575; Wan14B LoRA E2E 335154.8244 > 122881.8875; Qwen Image E2E 17658.9568 > 12740.6; LTX2.3 one-stage E2E 24015.2872 > 20903.025. Initial pass only; 890.3s remaining < 1224.2s required.

The performance causes are not established by these logs. They do not justify loosening thresholds or attributing the difference to CPU hardware. The finish job is an aggregate failure. No additional rerun requested.

[mickqian]

Exact-head CI follow-up for 020cca500e3545bbab040856d407c917b67b9073: the current unit job is no longer failing the previous native loader fixture. Base run 36252148993 now has one actual failing test job plus the aggregate finish.

Dual-H100 partition 2, job 108432202468, h100-novita6-gpu-01, ends with seven initial performance failures (actual > limit, ms):

• MiniMax H3 ref2va E2E: 183234.8659 > 84741.475.
• LTX2 two-stage load: 63982.4118 > 59338.575.
• LTX2.3 two-stage TI2V average denoise: 309.4081 > 300.192.
• LTX2.3 two-stage T2V refinement: 763.3852 > 736.65.
• Wan2.1 I2V 14B LoRA E2E: 342742.2182 > 122881.8875.
• Qwen Image E2E: 18746.4926 > 12740.6.
• LTX2.3 one-stage TI2V E2E: 25087.4133 > 20903.025.

No failed-item retry was started: 973.2s remaining was below the estimated 1360.3s needed. This is a deadline-budget stop, not six exhausted retries. The diff does touch shared native pipeline/loader entry points as well as ComfyUI; logs alone do not identify which, if any, causes the timings. Please compare this exact head against its base on an independent matching H100 setup with warmed inputs, preserving the gates. No H100 development capacity was listed at this check. I cannot push to the fork and did not request an additional blind rerun or threshold change.

[niehen6174]

This PR was originally intended as the first phase of the work. Since the second-phase changes #35990 have already been merged, this PR can now be closed.