[javey-q]

Motivation

Part of #37742 (the "Multi-GPU inference ... including CFG / sequence parallelism" item).
#40663 covers data-parallel replicas and states that CFG parallelism is not supported, so
this PR adds that part.

SenseNova-U1 runs classifier-free guidance by default (guidance_scale=4.0), so every
denoising step does two full backbone forwards: one over the condition prefix KV cache and
one over the empty-prompt unconditional cache. The pipeline config rejected num_gpus > 1
outright, so a second GPU could not help a single request. The official deployment stack
(LightX2V neopp) runs exactly this split, cond on CFG rank 0 and uncond on rank 1, and
reports ~1.97x on 2x H100.

Merging the two branches into one B=2 forward on a single GPU does not help at the
recommended 2K resolution: measured on H200 it is 1.7% slower per image (GPU-bound), so
splitting across GPUs is the lever.

Modifications

Design:

• The stage stays REPLICATED: the scheduler already broadcasts each request to every rank
  and only the output rank replies. t2i_generate receives the CFG GroupCoordinator and
  splits the work by rank_in_group.
• Rank 0 prefills and denoises only the condition branch (including the think_mode
  decode), rank 1 only the unconditional branch, so the two prefix forwards also run
  concurrently.
• Each step: one local forward, then all_gather of the bf16 v_pred (25 MB at 2K, ~0.2 ms
  over NVLink). Every rank combines the two predictions with the unchanged serial code
  (none / global / channel / cfg_zero_star) and steps its own copy of the latent. No
  all-reduce of g*cond + (1-g)*uncond: that changes bf16 rounding and cannot express the
  norm variants.
• Steps outside cfg_interval: rank 1 skips the forward and rank 0 broadcasts the condition
  prediction.
• Noise comes from the per-request seed on each rank, so all ranks hold the same latent
  without a broadcast.
• guidance_scale <= 1: both ranks run the serial single-branch path (the runtime's
  documented behavior for CFG-off requests under CFG parallel).
• Failure handling: only rank 0 does prompt-dependent work (tokenization, condition prefill,
  think decode), so an error there can hit rank 0 alone. Without coordination rank 1 would
  already sit in the step-0 all_gather, the scheduler's next broadcast_pyobj would hang
  until the group timeout, and a request arriving in between could pair with the stale
  collective and return a wrong image. The prefill (through the flash-cache setup)
  therefore ends with a one-element all_reduce of a preallocated failure flag; if either
  rank failed, both raise. In the denoising loop both ranks run the same code on the same
  image tokens, so an ordinary error hits both. An OOM can still hit one rank alone, because
  rank 0 also holds the condition prefix cache. [Fix] diffusion: Fail closed after local worker loss or multi-rank OOM #46025 already handles that for every
  multi-rank forward: the worker that ran out of memory exits, and the server fails closed
  (the request fails, /health and new POSTs return 503) until it is restarted.

Files:

• neo_unify/modeling_neo_chat.py: t2i_generate(cfg_group=...); the branch prefill
  (_t2i_prefill, _t2i_think_prefill), the batch expansion of the prefix cache and the
  CFG combine (_apply_t2i_cfg) move into helpers with their bodies unchanged;
  _cfg_ranks_fail_together wraps the prefill.
• stages/generation.py, pipelines/sensenova_u1.py: pass the CFG group when
  --enable-cfg-parallel is set.
• configs/pipeline_configs/sensenova_u1.py: accept num_gpus=2 with
  --enable-cfg-parallel (degree 2, tp_size == sp_degree == dp_size == 1); keep rejecting
  every other multi-GPU layout. ServerArgs never checks sp_degree against
  num_gpus / cfg degree, so --sp-degree 2 used to pass and then fail at distributed init
  with a world-size error. Reject dynamic batching (--batching-max-size > 1) together with
  CFG parallel (see Notes). Drop the unread supports_cfg_parallel=False field and
  deployment flag.
• configs/sample/sensenova_u1.py: should_use_classifier_free_guidance() follows
  guidance_scale alone (the model builds its own empty-prompt uncond). Without it every
  CFG request logged "CFG parallelism is enabled but this request does not use
  classifier-free guidance ... recompute it redundantly".
• Cookbook page + deployment snippet: two-GPU recipe (§5.2) and benchmark tab; the snippet
  refuses the 2-GPU + dynamic batching combination.

Accuracy Tests

Measured on main aecfbce3fd. The rebase conflicts only with #46046, which reads the
per-step CFG window to the host once (cfg_window) instead of comparing t on every step;
the split loop now uses that list, and otherwise modeling_neo_chat.py equals this PR's
version. (The earlier rebase onto ed0b8f433f overlapped only with #44690's dead-code
removal.)

The output is bit-identical, not merely close. Cookbook §3.1 command (2048x2048, 50 steps,
CFG 4, seed 42), two runs each: main on 1 GPU, this PR on 1 GPU and this PR on 2 GPUs all
write a PNG with md5 cf5d361f1f0f2ae40a28a36da8c58455. So the refactor is a no-op for
single-GPU users and the split reproduces it exactly. The same held on the earlier bases
4ec19ca27f and ed0b8f433f (all md5 b95d2e9f57b48dc4e91fa20de8deca74).

On an earlier base (f7ad4345b8), also byte-identical between 1 and 2 GPUs:

• Python API at 1024x1024: default, cfg_norm=global + cfg_interval=(0.2, 1.0),
  cfg_zero_star + cfg_interval=(0, 0.5), guidance_scale=1, think_mode=True.
• With the layout / fused-MLP PR ([Perf] SenseNova-U1: BSHD q/k/v, fused MLP gate_up + SiLU, RoPE table… #44922) and the FA3 attention change stacked underneath:
  ca5f3dc015ac07a3b31a46cc822a3e42 on 1 and 2 GPUs.

Failure paths, 2-GPU server, with a real CUDA OOM forced on one rank through a
sitecustomize probe (1024x1024, 20 steps; a fresh server per row):

OOM	Failed request	Afterwards
rank 1, denoising step 0	HTTP 500 after 0.8 s	/health and new POSTs 503 ("Scheduler worker exited; restart the server.")
rank 0, denoising step 10	HTTP 500 after 1.3 s	same
rank 0, condition prefill of a 50,000-word prompt (Tried to allocate 150.52 GiB); rank 1 raises "SenseNova-U1 prefill failed on the other CFG rank"	HTTP 500 after 1.1 s	same

No rank is left waiting in a collective, and both GPUs are free after shutdown. Before
#46025 (this PR on ed0b8f433f), the first row left rank 0 in the step-0 all_gather until
the NCCL watchdog fired after 600 s, and the server then crashed. Launching with
--batching-max-size 2 or --sp-degree 2 stops with the SenseNova ValueError.

Unit tests on CPU on aecfbce3fd: test_sensenova_u1.py 82 passed, 1 skipped;
test_sensenova_u1_backbone.py, test_cfg_gating.py, test_cfg_policy.py and
test_scheduler_client.py 70 passed, 2 skipped. New or changed cases:

• test_sensenova_u1_cfg_parallel_splits_branches_and_matches_serial: a tiny random
  NEO-Unify model on CPU, two ranks simulated by threads with a barrier-backed group;
  asserts both ranks reproduce the serial image bit for bit and that rank 0 runs only cond
  forwards and rank 1 only the CFG steps' uncond forwards (covers the broadcast path and
  B=2). Mutation-checked: a rank that also runs the other branch, and a broadcast from the
  wrong rank, both fail it.
• test_sensenova_u1_cfg_parallel_prefill_error_fails_both_ranks: rank 0's condition
  prefill raises OOM; both ranks must fail that request, and the next request must match the
  serial image. Without the flag exchange rank 1 does not fail at all: its denoise loop pairs
  with rank 0's next request. (In a server an OOM additionally stops the worker through
  [Fix] diffusion: Fail closed after local worker loss or multi-rank OOM #46025; the next-request check is what keeps non-OOM errors safe.)
• Parallelism validation: 2 GPUs without CFG parallel, CFG degree 2 on 4 GPUs, degree 4, and
  sp_degree=2 are rejected (main accepted the last one); dynamic batching with CFG
  parallel is rejected; 2 GPUs with degree 2 is accepted.
• The generation stage forwards the CFG group only when enabled.

Speed Tests and Profiling

2x H200 (NVLink), Xeon Platinum 8558, cookbook §3.1 command with --warmup-mode request,
reading the (with warmup excluded) line, 2-3 runs per row:

Tree	1 GPU	2 GPUs --enable-cfg-parallel	Speedup
main aecfbce3fd	16.27 / 16.31 s	—	—
this PR on aecfbce3fd	16.30 / 16.30 s	8.35 / 8.34 s	1.95x
this PR + #44922 + FA3 (on f7ad4345b8)	13.76 / 13.78 / 13.77 s	7.03 / 7.04 s	1.96x

Peak reserved memory on GPU 0 (Max peak): main 36542 MB, this PR 36446 MB on 1 GPU and
35606 MB on 2 GPUs (each rank keeps only its own branch's KV cache).

Notes

• Conflicts with diffusion: support SenseNova DP/TP/SP parallelism #40663 (DP replicas) in the parallelism validator, load_modules, the
  validator tests and cookbook §5.2: both replace the num_gpus != 1 check. The combined
  rule is one replica per dp_size, each with 1 GPU or 2 GPUs with CFG parallel; whichever
  lands second merges them. Merges cleanly with [Perf] SenseNova-U1: BSHD q/k/v, fused MLP gate_up + SiLU, RoPE table… #44922.
• Dynamic batching with CFG parallel is rejected for now. Each rank applies the batching
  delay to its own enqueue timestamps, so the two ranks can dispatch different batches into
  the same collectives (reproduced on CPU with the real scheduler loop: about 1 hang in 380
  delay-triggered dispatches, deterministic with a 1 ms stall). The fix belongs in the
  generic scheduler: let the output rank decide and broadcast the batch. That logic is
  shared by every dynamic-batching pipeline running on more than one GPU, so it is a
  separate PR.
• An OOM under CFG parallel follows [Fix] diffusion: Fail closed after local worker loss or multi-rank OOM #46025, like every multi-rank pipeline: the server must be
  restarted. DP replicas (diffusion: support SenseNova DP/TP/SP parallelism #40663) keep request-level OOM recovery. A model-level alternative
  was prototyped and verified on 2x H200: the failed rank keeps joining the per-step
  collectives with zeros, and both ranks raise after the loop, so the server keeps serving.
  It is not included because it would have to exempt this model from [Fix] diffusion: Fail closed after local worker loss or multi-rank OOM #46025's policy.
• Ascend NPU: the code path is backend-agnostic (HCCL collectives through the same
  GroupCoordinator), but it was not run on NPU.
• --num-gpus 2 does not auto-enable CFG parallel: _model_default_uses_cfg() keys on
  negative_prompt, which SenseNova does not use. Pass --enable-cfg-parallel explicitly.
• guidance_scale <= 1 requests compute the single branch on both GPUs; the rank-1 result
  is discarded.
• think_mode has a separate pre-existing bug on main (the image tokens cannot attend to
  the reasoning, so the image is noise); a fix follows in its own PR. Whichever lands second
  moves the one-line change between t2i_generate and _t2i_think_prefill.

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.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #41632388147
Latest PR Test (Extra): ❌ Run #41632388148
Latest PR Test (AMD ROCm 10): ❌ Run #41632388026

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[AuFlow]

Thanks for the heads-up! The CFG split looks like a useful complement to the DP replicas in #40663. We can reconcile the validation rules once either PR lands.

I had one concern while reviewing: the failure coordination currently only covers prefill. If one rank OOMs during denoising, the other could get stuck in all_gather() or broadcast() until timeout. The conditional branch keeps a larger KV cache, so this seems possible even with the same image-token count. Could we cover that path too, ideally with a test that injects a denoising failure on one rank?
@javey-q

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit ab43c6d.

Approve preview

[javey-q]

Thanks for the heads-up! The CFG split looks like a useful complement to the DP replicas in #40663. We can reconcile the validation rules once either PR lands.

I had one concern while reviewing: the failure coordination currently only covers prefill. If one rank OOMs during denoising, the other could get stuck in all_gather() or broadcast() until timeout. The conditional branch keeps a larger KV cache, so this seems possible even with the same image-token count. Could we cover that path too, ideally with a test that injects a denoising failure on one rank? @javey-q

Thanks @AuFlow, good catch. You're right, and the PR's reasoning ("same image tokens, so both ranks fail together") was wrong: rank 0 also holds the condition prefix KV cache, so a denoising OOM can hit one rank alone. I reproduced it on 2x H200 by forcing a real CUDA OOM in rank 1's first denoising step on the PR's original base: rank 0 sat in the step-0 all_gather until the NCCL watchdog fired after 600 s, and the server then crashed.

On current main this is covered by #46025: an OOM during any multi-rank forward terminates that worker and the server fails closed instead of leaving the peer in a collective. I rebased #45161 onto main (aecfbce) and checked it on 2x H200 with a real OOM forced on one rank:

OOM	Failed request	Afterwards
rank 1, denoising step 0	HTTP 500 after 0.8 s	/health and new POSTs 503
rank 0, denoising step 10	HTTP 500 after 1.3 s	same
rank 0, prefill of a 50k-word prompt	HTTP 500 after 1.1 s	same

So the denoising loop stays as is. The prefill flag exchange stays too, since #46025 only covers OOM and only rank 0 does prompt-dependent work. The cookbook command still gives a bit-identical image on 1 and 2 GPUs (1.95x).

I also prototyped model-level recovery, with a test that injects a failure on either rank before and after the all_gather and broadcast steps: the failed rank keeps joining the per-step collectives with zeros and both ranks raise after the loop, so the server keeps serving. It works on GPU, but it would have to exempt this model from #46025's fail-closed policy, so I left it out. Happy to send it as a follow-up if you or think request-level OOM recovery is worth having for CFG parallel