[diffusion] comfyui: split batched DiT calls into per-row requests by OnePunchMonk · Pull Request #43163 · sgl-project/sglang

OnePunchMonk · 2026-10-08T17:43:13Z

Motivation

ComfyUI puts CFG cond and uncond into one diffusion_model call whenever their shapes allow it and memory permits (for Flux, T5 is fixed length, so that is the normal case), and a latent with batch_size > 1 is batched too. The ComfyUI executor crashed on the first such call:

RuntimeError: a Tensor with 2 elements cannot be converted to Scalar   (should_suppress_logs)

Fixing that line alone would not help, because the worker's comfyui path is per-sample: ComfyUIPassThroughScheduler takes req.timesteps as the step schedule, and Flux's require_text_seq_lens expects one seq len per row while the adapters send one per request.

Depends on #43109, please merge that first (the first two commits here are #43109). Once rows go through the per-run conditioning cache one by one, main's old key makes two rows with the same text but different reference latents (Qwen-Image-Edit) share a key, and the second row silently runs on the first row's reference image. test_batched_rows_with_same_text_keep_their_own_reference fails with this change on top of main's key and passes on #43109.

Modifications

SGLDiffusionExecutor.forward: when x has batch > 1, run each row as its own request through the existing single-sample path and torch.cat the outputs. _row slices tensors whose leading dim equals the batch (recursing into plain lists/tuples such as ref_latents) and passes everything else through (transformer_options, shared refs). MiniMax-H3 passes x as [video, audio] and is not split, matching ComfyUI's native H3, which only supports batch 1.
Tests: one request per row with the right inputs and output order (Flux adapter); only batched values are sliced; batch 1 and list-shaped x are not split; same text with per-row reference latents keeps each row's own reference through the real cache path.

Limitation: a batched call now costs B worker round trips instead of one. True batching on the worker would need per-row timesteps and seq lens in the comfyui request contract; that is a bigger change and could be a follow-up.

Accuracy Tests

CPU end-to-end: real ComfyUI KSampler (ComfyUI master 58b176f) with the plugin's FluxExecutor + SGLDModelPatcher, worker half in-process (real bind_comfyui_session -> FluxPipelineConfig.prepare_pos_cond_kwargs -> the DiT's time_text_embed), and a noise_pred that depends on each row's conditioning. The reference is the same run with ComfyUI batching disabled (or each latent row sampled alone):

before (main e2fbc6a):
cfg=3.5, batch_size=1: CRASH RuntimeError: a Tensor with 2 elements cannot be converted to Scalar
cfg=1.0, batch_size=2: CRASH RuntimeError: a Tensor with 2 elements cannot be converted to Scalar
after:
cfg=3.5, batch_size=1: OK, 8 worker requests, each latents (1, 64, 64); matches unbatched reference: True (max diff 0.00e+00)
cfg=1.0, batch_size=2: OK, 8 worker requests, each latents (1, 64, 64); matches unbatched reference: True (max diff 0.00e+00)

GPU: not run. I set up a 2-GPU sanity run (FLUX.1-dev in comfyui mode: a 2-row batched call through the real FluxExecutor, compared against each row run alone, plus the plugin's test_flux_pipeline.py), but no 2x L40S container was scheduled within 10 minutes, so I stopped it. Happy to run it if someone can trigger GPU CI.

Unit tests (CPU):

before fix: 2 failed, 4 passed   test_comfyui_adapters.py
after fix:  100 passed, 11 skipped (CUDA only), 3 xfailed   test_comfyui_{session,adapters,h3,generator,profile}.py, test_conditioning_cache.py
            12 passed   apps/ComfyUI_SGLDiffusion/test/test_h3_request.py

Speed Tests and Profiling

Batch-1 calls are unchanged. Batched calls (CFG > 1, batch_size > 1) used to crash; they now make one worker request per row.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations. (not needed)
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. (CPU, see above)
Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #37953314216
Latest PR Test (Extra): ❌ Run #37953313899
Latest PR Test (AMD ROCm 10): ❌ Run #37953314129

_cond_key used the shape plus the first and last scalar of prompt_embeds[0]. For Flux that is the pooled CLIP-L vector, so prompts that only differ in T5 got the same key and later steps reused the first prompt's T5 context. Hash every field drop_cached_fields removes instead. Fixes sgl-project#43107

Adapters now declare cached_extra_keys; drop_cached_fields and the cond key both read it, so H3's h3_* fields and sigmas are part of the key. The worker stored packed extras in one shared slot, so on a cache hit the positive cond could get the negative's h3_context/payload back. Store them per comfyui_cond_key instead. Hash nested dicts/lists by content with framing instead of repr().

ComfyUI batches CFG cond/uncond (and batch_size > 1) into one diffusion_model call. The executor crashed on it in should_suppress_logs, and the worker's comfyui path is per-sample anyway (req.timesteps is the schedule, seq lens are per request). Send one request per row and concatenate the outputs. Fixes sgl-project#43111

niehen6174 · 2026-10-09T03:05:07Z

+        # ComfyUI batches CFG cond/uncond (and batch_size > 1) into one call,
+        # but the worker's comfyui path is per-sample: req.timesteps is its
+        # schedule and seq lens are per request. Send one request per row.
+        return torch.cat(


Native worker batching would be a useful follow-up for throughput.

will open a follow up pr for this. keeping this unresolved.

Native batching depends on #43109 and #43163. I'll pick this up once these two PRs are wrapped up, and I'll open an RFC to agree on per-row timesteps and sequence lengths in the request contract.

hi @niehen6174 , opened the rfc here: #43411

niehen6174 · 2026-10-09T03:05:08Z

        return self.adapter.unpack(output_batch.noise_pred, packed, x)

    def forward(self, x, timestep, context, **kwargs):
+        batch = int(x.shape[0]) if torch.is_tensor(x) else 1


is checking shape[0] == batch_size sufficient to distinguish batched inputs from shared conditions?

Checked this. I think this is sufficient for now, since we do not support models like Trellis2 (those which share tensors). Pls tell me your thoughts on long-term maintenance, or a better way to generalise this.

Rename _row to _slice_batch_row. Recurse into dicts so ControlNet's control residuals are split per row, and reject tensors whose leading dim is neither 1 nor the batch instead of passing them through. transformer_options stays shared except its per-chunk entries (cond_or_uncond, uuids) and sigmas, which ComfyUI sets to one chunk's timestep, so row i takes sigmas[i % len].

OnePunchMonk · 2026-10-10T04:13:20Z

Superseded by #43492 (rebuilt on current main; same 4 commits). Closing.

OnePunchMonk added 3 commits October 8, 2026 21:06

OnePunchMonk requested review from AgainstEntropy, BBuf, HaiShaw, kevin-mii, mickqian, niehen6174, ping1jing2 and yichiche as code owners October 8, 2026 17:43

github-actions Bot added the diffusion SGLang Diffusion label Oct 8, 2026

niehen6174 reviewed Oct 9, 2026

View reviewed changes

Comment thread python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py Resolved

niehen6174 reviewed Oct 9, 2026

View reviewed changes

niehen6174 mentioned this pull request Oct 9, 2026

[diffusion] comfyui: zero-fill missing Flux pooled vector #43137

Closed

3 of 5 tasks

OnePunchMonk force-pushed the fix/comfyui-batched-forward branch 2 times, most recently from bd935f2 to f467ad9 Compare October 9, 2026 09:43

OnePunchMonk force-pushed the fix/comfyui-batched-forward branch from f467ad9 to a7d520e Compare October 9, 2026 15:39

OnePunchMonk mentioned this pull request Oct 9, 2026

[RFC] Native worker batching for ComfyUI diffusion requests #43411

Open

OnePunchMonk closed this Oct 10, 2026

This was referenced Oct 10, 2026

[diffusion] comfyui: resolve Flux guidance_embeds from checkpoint keys instead of assuming True #43394

Closed

ComfyUI ideas OnePunchMonk/sglang#4

Open

Navigation Menu

Repository navigation

[diffusion] comfyui: split batched DiT calls into per-row requests - #43163

[diffusion] comfyui: split batched DiT calls into per-row requests#43163
OnePunchMonk wants to merge 4 commits into
sgl-project:mainfrom
OnePunchMonk:fix/comfyui-batched-forward

OnePunchMonk commented Oct 8, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Uh oh!

niehen6174 Oct 9, 2026

Uh oh!

OnePunchMonk Oct 9, 2026 •

edited

Loading

Uh oh!

OnePunchMonk Oct 9, 2026

Uh oh!

OnePunchMonk Oct 9, 2026

Uh oh!

niehen6174 Oct 9, 2026

Uh oh!

OnePunchMonk Oct 9, 2026

Uh oh!

OnePunchMonk commented Oct 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

OnePunchMonk commented Oct 8, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

Uh oh!

niehen6174 Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

OnePunchMonk Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

OnePunchMonk Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

OnePunchMonk Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

niehen6174 Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

OnePunchMonk Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

OnePunchMonk commented Oct 10, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

OnePunchMonk commented Oct 8, 2026 •

edited by github-actions Bot

Loading

OnePunchMonk Oct 9, 2026 •

edited

Loading