[diffusion] fix: serve requests that turn cfg off on a cfg-parallel server (#39303)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Mick
2026-09-14 11:26:37 +08:00
committed by GitHub
co-authored by Mick Qian
parent 4c71d14fba
commit 3480112a9c
5 changed files with 255 additions and 75 deletions
@@ -394,34 +394,30 @@ class InputValidationStage(PipelineStage):
f"Guidance scale must be positive, but got {batch.guidance_scale}"
)
# Reject requests that do not enable CFG on a server launched with
# --enable-cfg-parallel. CFG-parallel splits cond/uncond across ranks,
# so rank 1 has no work and returns None for noise_pred, which crashes
# scheduler.step() ~30 minutes later under a gloo broadcast timeout.
# Earlier, field-specific checks above (negative_prompt missing,
# guidance_scale < 0) fire first and produce better messages for those
# cases; this is the catch-all for any combination that still leaves
# do_classifier_free_guidance=False under cfg-parallel.
# A request that leaves CFG off is servable under CFG parallelism: the
# dispatcher gives branch 0 to rank 0, and every other rank runs branch 0
# too so the all-gather has shapes to work with. Both ranks then read the
# owner's prediction, so the answer is the single-branch answer and the
# extra ranks are only redundant.
#
# This used to raise. That guard was added for a warmup hang (#23198)
# two weeks BEFORE the multi-branch refactor (#23736) taught the
# dispatcher to handle a single branch, and the warmup path has since
# grown its own fix -- the warmup builder forces CFG on whenever
# cfg-parallel is enabled. What was left was a server refusing traffic
# it could serve, and the runtime AUTO-enables cfg-parallel from the
# model's default sampling params, so `sglang serve --num-gpus 2` on a
# CFG-defaulting model rejected every guidance_scale=1.0 request while
# blaming a flag the user never passed.
if server_args.enable_cfg_parallel and not batch.do_classifier_free_guidance:
neg_prompt_state = (
"not set"
if batch.negative_prompt is None
else "empty"
if batch.negative_prompt == ""
else "set"
)
raise ValueError(
f"Server was launched with --enable-cfg-parallel but this "
f"request does not use classifier-free guidance "
f"(do_classifier_free_guidance={batch.do_classifier_free_guidance}, "
f"guidance_scale={batch.guidance_scale}, "
f"true_cfg_scale={batch.true_cfg_scale}, "
f"negative_prompt={neg_prompt_state}). "
f"CFG-parallel splits cond/uncond across ranks and requires "
f"both to be active. Either disable --enable-cfg-parallel or "
f"ensure the request enables CFG (set guidance_scale > 1.0 or "
f"true_cfg_scale > 1.0, with a non-empty negative_prompt or "
f"negative_prompt_embeds)."
logger.warning_once(
"CFG parallelism is enabled but this request does not use "
"classifier-free guidance (guidance_scale=%s, true_cfg_scale=%s), "
"so it has one branch and the other CFG rank(s) recompute it "
"redundantly. Pass --cfg-parallel-size 1 to spend those GPUs on "
"another parallelism instead.",
batch.guidance_scale,
batch.true_cfg_scale,
)
# for i2v, get image from image_path
@@ -1466,9 +1466,11 @@ class ServerArgs(DisaggServerArgsMixin):
self.enable_cfg_parallel = auto_cfg_parallel_degree > 1
if self.enable_cfg_parallel:
logger.info(
"Automatically enabled CFG parallel at degree %d for %d GPUs. "
"Use --sp-degree / --ulysses-degree to use sequence "
"parallelism instead.",
"Automatically enabled CFG parallel at degree %d for %d GPUs "
"because this model uses classifier-free guidance by default. "
"A request that turns CFG off still runs, but it has one branch, "
"so the other CFG rank(s) recompute it redundantly. Override with "
"--cfg-parallel-size 1, --tp-size, or --sp-degree / --ulysses-degree.",
self.cfg_parallel_degree,
self.num_gpus,
)
@@ -100,16 +100,53 @@ class SortedHelpFormatter(argparse.HelpFormatter):
super().add_arguments(actions)
@lru_cache
def _print_info_once(logger: Logger, msg: str) -> None:
# Set the stacklevel to 2 to print the original caller's line info
logger.info(msg, stacklevel=2)
# `logger.warning_once(msg, *args)` is bound as MethodType(_print_warning_once,
# logger), so there is exactly ONE frame between the caller and logger.warning --
# and stacklevel=2 is part of the observable contract, asserted literally by
# test_diffusion_bcg_padding. Any helper in between pushes the record's filename
# to this file and breaks that assertion, so the dedup cannot be an lru_cache on
# a second function.
#
# It also cannot be an lru_cache on THIS function: keyed on the arguments it would
# hold a strong reference to each one for the life of the process, and callers
# here pass tensors. Hence a set of formatted text, which stores only strings.
#
# The args themselves are new: these helpers used to take the message alone, so a
# caller that formatted lazily -- the way the standard contract implies -- raised
# TypeError instead of logging, always on a branch too rare to have been seen.
_logged_once: set[tuple[str, int, str]] = set()
@lru_cache
def _print_warning_once(logger: Logger, msg: str) -> None:
# Set the stacklevel to 2 to print the original caller's line info
logger.warning(msg, stacklevel=2)
def _log_once_guard(logger: Logger, level: int, msg: str, *args: Any) -> str | None:
"""The text to log, or None when this message has already been logged."""
text = msg % args if args else msg
key = (logger.name, level, text)
if key in _logged_once:
return None
_logged_once.add(key)
return text
def _print_info_once(logger: Logger, msg: str, *args: Any) -> None:
text = _log_once_guard(logger, logging.INFO, msg, *args)
# stacklevel=2 is asserted literally by test_diffusion_bcg_padding, so it is
# contract rather than a tuning knob. It does NOT reach the caller: init_logger
# also patches `warning` into a forwarder to `logger.log`, adding a frame, so
# the record names this module. That was true before these helpers too.
if text is not None:
logger.info(text, stacklevel=2)
def _print_warning_once(logger: Logger, msg: str, *args: Any) -> None:
text = _log_once_guard(logger, logging.WARNING, msg, *args)
if text is not None:
logger.warning(text, stacklevel=2)
# These were lru_cache objects, so `.cache_clear()` was part of their surface and
# a test resets the dedup through it.
_print_info_once.cache_clear = _logged_once.clear
_print_warning_once.cache_clear = _logged_once.clear
def get_is_main_process():
@@ -167,19 +204,19 @@ class _SGLDiffusionLogger(Logger):
`intel_extension_for_pytorch.utils._logger`.
"""
def info_once(self, msg: str) -> None:
def info_once(self, msg: str, *args: Any) -> None:
"""
As :meth:`info`, but subsequent calls with the same message
are silently dropped.
and args are silently dropped.
"""
_print_info_once(self, msg)
_print_info_once(self, msg, *args)
def warning_once(self, msg: str) -> None:
def warning_once(self, msg: str, *args: Any) -> None:
"""
As :meth:`warning`, but subsequent calls with the same message
are silently dropped.
and args are silently dropped.
"""
_print_warning_once(self, msg)
_print_warning_once(self, msg, *args)
def info( # type: ignore[override]
self,
@@ -4,8 +4,8 @@ Covers warmup and cfg-parallel guard paths introduced alongside this file:
- build_warmup_reqs synthesizes warmup Reqs that actually enable
classifier-free guidance when cfg-parallel is on.
- DiffGenerator sends explicit warmup resolutions through the scheduler client.
- InputValidationStage.forward rejects non-CFG requests when the server
has cfg-parallel on.
- InputValidationStage.forward ACCEPTS non-CFG requests when the server
has cfg-parallel on, and the branch dispatcher serves them.
- Server-based warmup can opt into model-default negative prompts so warmup
populates the negative text embedding cache.
- Req-based warmup remains available only through the lazy legacy path.
@@ -1047,26 +1047,28 @@ class TestImageVaeEncodingLatentRetrieval(unittest.TestCase):
)
class TestInputValidationCfgParallelGuard(unittest.TestCase):
"""Commit 2: per-request cfg-parallel check.
class TestInputValidationCfgParallelSingleBranch(unittest.TestCase):
"""A request that turns CFG off must still be served under cfg-parallel.
This used to raise. The guard came from a warmup hang (#23198, 2026-04-23);
two weeks later the multi-branch refactor (#23736) taught the dispatcher to
handle a single branch, and the warmup builder grew its own fix (it forces
CFG on whenever cfg-parallel is enabled). What the guard still did was refuse
live traffic the runtime could serve -- and because cfg-parallel is
AUTO-enabled from the model's default sampling params, a plain
`sglang serve --num-gpus 2` on a CFG-defaulting model rejected every
guidance_scale=1.0 request, citing a flag the user never passed.
Both tests patch _generate_seeds (the first statement of
InputValidationStage.forward, input_validation.py:274) to sidestep
its device-lookup / generator-creation code which pulls in torch
CUDA bindings — keeps the suite strictly CPU-only. We still need
num_inference_steps on the Req because the stage's
"num_inference_steps <= 0" check at L305-308 raises TypeError on
None before the new commit-2 check is reached.
InputValidationStage.forward) to sidestep its device-lookup / generator
creation, keeping the suite CPU-only. num_inference_steps must be set because
the "num_inference_steps <= 0" check raises TypeError on None first.
"""
def test_input_validation_rejects_cfg_parallel_without_cfg(self):
# negative_prompt="" (non-None) ensures the existing
# negative_prompt-is-None check at input_validation.py:295-298
# does NOT fire first — this isolates the new commit-2 check.
# width/height/num_outputs_per_prompt pre-set so the stage's
# default-dimension block at L352-361 doesn't mutate the Req
# in a way that obscures the assertion target.
req = Req(
def _single_branch_req(self) -> Req:
# negative_prompt="" (non-None) keeps the negative_prompt-is-None check
# from firing first, so this isolates the cfg-parallel path.
return Req(
prompt="test",
negative_prompt="",
guidance_scale=1.0,
@@ -1076,29 +1078,26 @@ class TestInputValidationCfgParallelGuard(unittest.TestCase):
width=512,
height=512,
)
def test_input_validation_accepts_cfg_parallel_without_cfg(self):
req = self._single_branch_req()
self.assertIs(
req.do_classifier_free_guidance,
False,
"Sanity: test setup must leave do_cfg=False so the "
"commit-2 check is the one that fires, not an upstream check.",
"Sanity: the setup must leave do_cfg=False, or this tests nothing.",
)
stage = _make_input_validation_stage()
server_args = _make_validation_server_args(enable_cfg_parallel=True)
with patch.object(InputValidationStage, "_generate_seeds"):
with self.assertRaises(ValueError) as ctx:
try:
stage.forward(req, server_args)
msg = str(ctx.exception).lower()
self.assertIn("cfg-parallel", msg)
for field in (
"do_classifier_free_guidance",
"guidance_scale",
"true_cfg_scale",
"negative_prompt",
):
self.assertIn(field, str(ctx.exception))
except ValueError as e:
self.fail(
"forward() rejected a single-branch request under "
f"cfg-parallel; the dispatcher can serve it: {e}"
)
def test_input_validation_passes_cfg_parallel_with_cfg(self):
req = Req(
@@ -1127,5 +1126,78 @@ class TestInputValidationCfgParallelGuard(unittest.TestCase):
self.fail(f"forward() raised ValueError on a valid CFG request: {e}")
class TestCfgParallelServesOneBranch(unittest.TestCase):
"""The property that makes accepting a single-branch request safe.
Dropping the validation guard is only correct because the dispatcher already
handles n_branches=1 on a 2-rank CFG group: rank 0 owns the branch, every
other rank runs it too so the all-gather has shapes, and the reorder step
hands both ranks the owner's prediction. Pin it from the rank that owns
nothing -- that is the rank the old comment said returned None and hung a
gloo broadcast for half an hour.
"""
def _run_on_rank(self, cfg_rank: int, n_branches: int = 1, world_size: int = 2):
from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
CFGBranch,
CFGPolicy,
)
mod = "sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils"
branches = [CFGBranch(f"b{i}", i == 0, {"tag": i}) for i in range(n_branches)]
policy = CFGPolicy(branches=branches)
seen: list[int] = []
def predict_fn(branch):
seen.append(branch.kwargs["tag"])
return torch.full((1, 2), float(branch.kwargs["tag"]))
# A real 2-rank gather returns one tensor per rank. Both ranks ran the
# same branch here, so both contributions carry the same values.
def fake_all_gather(t, dim=0, separate_tensors=False):
return [t.clone() for _ in range(world_size)]
with (
patch(f"{mod}.get_classifier_free_guidance_rank", return_value=cfg_rank),
patch(
f"{mod}.get_classifier_free_guidance_world_size",
return_value=world_size,
),
patch(f"{mod}.get_local_torch_device", return_value=torch.device("cpu")),
patch(f"{mod}.cfg_model_parallel_all_gather", side_effect=fake_all_gather),
):
from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
run_cfg_parallel,
)
return run_cfg_parallel(policy, predict_fn), seen
def test_branch_owner_gets_the_single_prediction(self):
preds, seen = self._run_on_rank(cfg_rank=0)
self.assertEqual(len(preds), 1)
self.assertEqual(seen, [0], "the owning rank runs branch 0 once")
self.assertTrue(torch.equal(preds[0], torch.zeros(1, 2)))
def test_rank_without_a_branch_still_returns_the_owners_prediction(self):
preds, seen = self._run_on_rank(cfg_rank=1)
self.assertEqual(
seen,
[0],
"the rank that owns no branch must still run one, or the "
"all-gather has no shapes to work with",
)
self.assertEqual(len(preds), 1)
self.assertIsNotNone(preds[0])
self.assertTrue(torch.equal(preds[0], torch.zeros(1, 2)))
def test_two_branches_still_split_across_the_ranks(self):
from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
dispatch_branches,
)
self.assertEqual(dispatch_branches(1, 2), [[0], []])
self.assertEqual(dispatch_branches(2, 2), [[0], [1]])
if __name__ == "__main__":
unittest.main()
@@ -1,10 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
import gc
import logging
import unittest
import weakref
from sglang.multimodal_gen.runtime.utils.logging_utils import (
globally_suppress_loggers,
init_logger,
)
@@ -27,5 +30,75 @@ class TestSuppressNoisyDependencyLogs(unittest.TestCase):
)
class TestLogOnceTakesFormatArgs(unittest.TestCase):
"""`warning_once(msg, *args)` must format, not raise.
The helpers took only the message, so every caller that formatted lazily --
the way logger.warning wants -- raised TypeError instead of logging. Three
call sites did, and each sat on a branch that rarely runs, so the bug was
invisible: cfg_parallel_utils only reaches its call when a CFG-parallel
group has more ranks than branches, which input validation used to reject
outright. Removing that rejection turned the latent TypeError into a crash
on the first single-branch request.
"""
def test_warning_once_formats_lazy_args(self):
logger = init_logger("sglang.test.logonce.warning")
with self.assertLogs(logger, level=logging.WARNING) as captured:
logger.warning_once("cfg_parallel_size=%d > n_branches=%d", 2, 1)
self.assertIn("cfg_parallel_size=2 > n_branches=1", captured.output[0])
def test_info_once_formats_lazy_args(self):
logger = init_logger("sglang.test.logonce.info")
with self.assertLogs(logger, level=logging.INFO) as captured:
logger.info_once("degree %d on %d GPUs", 2, 2)
self.assertIn("degree 2 on 2 GPUs", captured.output[0])
def test_record_does_not_name_the_caller(self):
"""Documents a wart, so nobody "fixes" it and breaks the bcg assertion.
init_logger also replaces `logger.warning` with a patched method that
forwards to `logger.log`, so there is one more frame than the stacklevel
accounts for and the record names this module rather than the caller.
That predates these helpers -- the original passed stacklevel=2 through
the same patched method -- and raising the number would contradict
test_diffusion_bcg_padding, which asserts the literal `stacklevel=2`.
"""
logger = init_logger("sglang.test.logonce.stacklevel")
with self.assertLogs(logger, level=logging.WARNING) as captured:
logger.warning_once("from the caller %d", 1)
self.assertEqual(captured.records[0].filename, "logging_utils.py")
def test_arguments_are_not_retained(self):
"""The once-cache must key on text, not on the arguments.
An lru_cache keyed on the arguments holds a strong reference to each of
them for the life of the process, and callers in this package pass
tensors. Formatting first and caching the result keeps only strings.
"""
logger = init_logger("sglang.test.logonce.retain")
class _Heavy:
def __repr__(self):
return "<heavy>"
obj = _Heavy()
ref = weakref.ref(obj)
with self.assertLogs(logger, level=logging.WARNING) as captured:
logger.warning_once("holding %s", obj)
self.assertIn("<heavy>", captured.output[0])
del obj
gc.collect()
self.assertIsNone(ref(), "the once-cache kept the argument alive")
def test_same_message_and_args_logs_once(self):
logger = init_logger("sglang.test.logonce.dedup")
with self.assertLogs(logger, level=logging.WARNING) as captured:
logger.warning_once("idle ranks: %d", 1)
logger.warning_once("idle ranks: %d", 1)
logger.warning_once("idle ranks: %d", 2) # different args, new line
self.assertEqual(len(captured.output), 2, captured.output)
if __name__ == "__main__":
unittest.main()