diff --git a/docs/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT.mdx b/docs/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT.mdx
index 4d9d915fd1f6..3fce7d98e481 100644
--- a/docs/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT.mdx
+++ b/docs/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT.mdx
@@ -105,7 +105,29 @@ in one batch must use the same resolution, inference steps, guidance settings,
and output options. Requests with multiple outputs or `think_mode=true` are
executed sequentially.
-### 5.2 Ascend NPU optimizations
+### 5.2 CFG parallelism on two GPUs
+
+With classifier-free guidance enabled (the default `guidance_scale` is 4.0),
+every denoising step runs a conditional and an unconditional forward pass.
+`--enable-cfg-parallel` places one branch on each of two GPUs, so each GPU runs
+one forward pass per step and the two exchange the predicted velocity. The
+image is bit-identical to the single-GPU result.
+
+```bash Command
+sglang serve \
+ --model-path sensenova/SenseNova-U1.5-8B-MoT \
+ --port 30000 \
+ --num-gpus 2 \
+ --enable-cfg-parallel
+```
+
+Each GPU holds a full copy of the weights. Requests with `guidance_scale` 1 have
+a single branch, which both GPUs compute. SenseNova-U1 does not support tensor,
+sequence or data parallelism yet, so `--num-gpus 2` requires
+`--enable-cfg-parallel`. CFG parallelism works with `think_mode`, but not yet
+with dynamic request batching (`--batching-max-size` greater than 1).
+
+### 5.3 Ascend NPU optimizations
On Ascend NPU, SenseNova-U1 automatically uses fused inference attention,
RMSNorm, and SwiGLU MLP operators for supported inputs. FIA and SwiGLU fall back
@@ -172,6 +194,21 @@ first use.
------------------------------------------------------------
```
+
+
+ Environment: two NVIDIA H200 GPUs connected by NVLink; workload: the
+ Section 3.1 `sglang generate` command (2048 x 2048, 50 denoising steps,
+ CFG 4, seed 42) with `--warmup-mode request`; repetitions: 2 per row, both
+ within 0.03 s; latency: the warmed-up request time.
+
+ | GPUs | Extra flags | Latency (s) | Speedup | Peak reserved memory on GPU 0 (MB) |
+ |---:|---|---:|---:|---:|
+ | 1 | | 16.26 | 1.00x | 36446 |
+ | 2 | `--num-gpus 2 --enable-cfg-parallel` | 8.32 | 1.95x | 35606 |
+
+ Both configurations produce byte-identical PNG files.
+
+
Environment: one Ascend 910C; workload: four concurrent 2048 x 2048
diff --git a/docs/src/snippets/configs/SenseNova/sensenova-u15.jsx b/docs/src/snippets/configs/SenseNova/sensenova-u15.jsx
index af69e2a82a56..1f0abca937e3 100644
--- a/docs/src/snippets/configs/SenseNova/sensenova-u15.jsx
+++ b/docs/src/snippets/configs/SenseNova/sensenova-u15.jsx
@@ -5,12 +5,18 @@ export const config = (() => {
})[s.checkpoint || "default"];
const plan = (s) => {
const flags = ["--model-path " + modelOf(s), ...["--port {{PORT}}"]];
- let gpus = 1;
+ const gpus = Number(s.gpus_per_node);
let t = { tp_size: 1, ulysses_degree: 1, ring_degree: 1 };
const env = [];
const warnings = [];
const errors = [];
- if (s.batching === "on") flags.push("--batching-max-size 2", "--batching-delay-ms 100");
+ if (gpus === 2) {
+ flags.push("--num-gpus 2", "--enable-cfg-parallel");
+ warnings.push(
+ "CFG parallelism runs one guidance branch per GPU; requests with guidance_scale 1 compute the same branch on both GPUs.",
+ );
+ if (s.batching === "on") errors.push("Dynamic batching does not support CFG parallelism yet; use one GPU.");
+ } else if (s.batching === "on") flags.push("--batching-max-size 2", "--batching-delay-ms 100");
return { flags, gpus, topology: t, env, warnings, errors };
};
const defaults = {
@@ -44,7 +50,7 @@ export const config = (() => {
commandBuilder: {
defaultSelection: defaults,
resource: {
- limits: { nodes: { min: 1, max: 1 }, gpus_per_node: { min: 1, max: 1 } },
+ limits: { nodes: { min: 1, max: 1 }, gpus_per_node: { min: 1, max: 2 } },
verifiedRecipes: ["cuda"].map((hw) => {
const s = { ...defaults, hw };
const p = plan(s);
@@ -64,7 +70,7 @@ export const config = (() => {
if (!config.supportedHardware.includes(s.hw)) errors.push("Choose a listed hardware platform.");
if (Number(s.nodes) !== 1) errors.push("This recipe covers one node.");
const count = Number(s.gpus_per_node);
- if (!Number.isInteger(count) || count < 1 || count > 1)
+ if (!Number.isInteger(count) || count < 1 || count > 2)
errors.push("Choose an integer GPU count within the documented limits.");
if (Number(s.gpus_per_node) !== p.gpus) errors.push("This recipe requires " + p.gpus + " GPUs / node.");
if (["tp_size", "ulysses_degree", "ring_degree"].some((key) => Number(t[key]) !== p.topology[key]))
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/sensenova_u1.py b/python/sglang/multimodal_gen/configs/pipeline_configs/sensenova_u1.py
index f1495eac220b..d0193dc70a9f 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/sensenova_u1.py
+++ b/python/sglang/multimodal_gen/configs/pipeline_configs/sensenova_u1.py
@@ -46,6 +46,32 @@ def _component_residency_requests_offload(value) -> bool:
return False
+def validate_sensenova_u1_parallelism(server_args) -> None:
+ # ServerArgs rejects CFG parallel on one GPU, so one GPU needs no more checks.
+ if server_args.num_gpus == 1:
+ return
+ if not (
+ server_args.num_gpus == 2
+ and server_args.enable_cfg_parallel
+ and server_args.cfg_parallel_degree == 2
+ and server_args.tp_size == 1
+ and server_args.sp_degree == 1
+ and server_args.dp_size == 1
+ ):
+ raise ValueError(
+ "SenseNovaU1Pipeline supports num_gpus=1, or num_gpus=2 with "
+ "--enable-cfg-parallel (one CFG branch per GPU). Native tensor, "
+ "sequence and data parallelism are not implemented yet."
+ )
+ # Each rank times the batching delay on its own clock, so the two CFG
+ # ranks can dispatch different batches and deadlock in a collective.
+ if server_args.batching_max_size > 1:
+ raise ValueError(
+ "SenseNovaU1Pipeline does not support dynamic batching "
+ "(--batching-max-size > 1) with --enable-cfg-parallel yet."
+ )
+
+
def _set_compatible_runtime_defaults(server_args) -> None:
compatible_defaults = {
"component_residency": None,
@@ -80,7 +106,6 @@ class SenseNovaU1PipelineConfig(PipelineConfig):
task_type: ModelTaskType = ModelTaskType.T2I
model_precision: str = "bf16"
should_use_guidance: bool = True
- supports_cfg_parallel: bool = False
def supports_dynamic_batching(self):
return True
@@ -108,11 +133,7 @@ def supports_sequential_multi_output_inference(self):
return True
def validate_server_args(self, server_args) -> None:
- if server_args.num_gpus != 1:
- raise ValueError(
- "SenseNovaU1Pipeline currently supports num_gpus=1. "
- "Native tensor/pipeline parallelism is not implemented yet."
- )
+ validate_sensenova_u1_parallelism(server_args)
if getattr(server_args, "enable_torch_compile", False):
raise ValueError(
"SenseNovaU1Pipeline does not support torch.compile yet. "
@@ -193,7 +214,6 @@ def get_model_deployment_config(self) -> ModelDeploymentConfig:
speed_mode_enable_torch_compile_by_default=False,
keep_resident_min_available_gb=80,
auto_enable_cfg_parallel=False,
- supports_cfg_parallel=False,
)
diff --git a/python/sglang/multimodal_gen/configs/sample/sensenova_u1.py b/python/sglang/multimodal_gen/configs/sample/sensenova_u1.py
index 3dea19b47c80..b112142f9036 100644
--- a/python/sglang/multimodal_gen/configs/sample/sensenova_u1.py
+++ b/python/sglang/multimodal_gen/configs/sample/sensenova_u1.py
@@ -68,6 +68,11 @@ def get_cli_args(cls, args):
if key in _PUBLIC_OVERRIDE_FIELDS
}
+ def should_use_classifier_free_guidance(self) -> bool:
+ # The model builds its own empty-prompt unconditional branch, so CFG
+ # depends on the guidance scale alone, not on negative_prompt.
+ return self.guidance_scale > 1.0
+
def __post_init__(self) -> None:
if isinstance(self.cfg_interval, list):
self.cfg_interval = tuple(float(x) for x in self.cfg_interval)
diff --git a/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_neo_chat.py b/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_neo_chat.py
index 34990bdd18ce..2e26c0482fb0 100644
--- a/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_neo_chat.py
+++ b/python/sglang/multimodal_gen/runtime/models/sensenova_u1/neo_unify/modeling_neo_chat.py
@@ -148,6 +148,69 @@ def prepare_flash_kv_cache(
layer.flash_v_cache = v_cache
+@contextlib.contextmanager
+def _cfg_ranks_fail_together(cfg_group):
+ # Only CFG rank 0 runs the prompt-dependent prefill. A one-sided error must
+ # stop both ranks before the per-step collectives pair up wrongly.
+ if cfg_group is None:
+ yield
+ return
+ # Allocated up front so that reporting an OOM needs no device memory.
+ failed = torch.zeros(1, dtype=torch.int32, device=cfg_group.device)
+ try:
+ yield
+ except Exception:
+ cfg_group.all_reduce(failed.fill_(1))
+ raise
+ if cfg_group.all_reduce(failed).item():
+ raise RuntimeError(
+ "SenseNova-U1 prefill failed on the other CFG rank; see its log."
+ )
+
+
+def _expand_prefix_cache_batch(past_key_values, batch_size: int) -> None:
+ for layer in past_key_values.layers:
+ layer.keys = layer.keys.expand(batch_size, *layer.keys.shape[1:])
+ layer.values = layer.values.expand(batch_size, *layer.values.shape[1:])
+
+
+def _apply_t2i_cfg(
+ v_pred_condition,
+ v_pred_uncondition,
+ *,
+ cfg_scale: float,
+ cfg_norm: str,
+ step_i: int,
+ batch_size: int,
+):
+ if cfg_norm == "cfg_zero_star":
+ positive_flat = v_pred_condition.view(batch_size, -1)
+ negative_flat = v_pred_uncondition.view(batch_size, -1)
+
+ alpha = optimized_scale(positive_flat, negative_flat)
+ alpha = alpha.view(batch_size, *([1] * (len(v_pred_condition.shape) - 1)))
+ alpha = alpha.to(positive_flat.dtype)
+
+ if step_i <= 0:
+ return v_pred_condition * 0.0
+ return v_pred_uncondition * alpha + cfg_scale * (
+ v_pred_condition - v_pred_uncondition * alpha
+ )
+
+ v_pred = v_pred_uncondition + cfg_scale * (v_pred_condition - v_pred_uncondition)
+ if cfg_norm == "global":
+ norm_v_condition = torch.norm(v_pred_condition, dim=(1, 2), keepdim=True)
+ norm_v_cfg = torch.norm(v_pred, dim=(1, 2), keepdim=True)
+ scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+ v_pred = v_pred * scale
+ elif cfg_norm == "channel":
+ norm_v_condition = torch.norm(v_pred_condition, dim=-1, keepdim=True)
+ norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
+ scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(min=0, max=1.0)
+ v_pred = v_pred * scale
+ return v_pred
+
+
def clear_flash_kv_cache(past_key_values):
if past_key_values is None:
return
@@ -516,6 +579,62 @@ def _t2i_prefix_forward(self, input_ids, indexes, attention_mask):
)
return out.past_key_values, out.last_hidden_state
+ def _t2i_prefill(self, tokenizer, query, *, token_h, token_w):
+ (
+ input_ids,
+ indexes,
+ attention_mask,
+ key_valid_mask,
+ prefix_lengths,
+ ) = self._build_t2i_text_inputs(tokenizer, query)
+ indexes_image = self._build_t2i_image_indexes(
+ token_h,
+ token_w,
+ (
+ int(prefix_lengths[0].item())
+ if prefix_lengths.numel() == 1
+ else prefix_lengths
+ ),
+ device=input_ids.device,
+ )
+ past_key_values, _ = self._t2i_prefix_forward(
+ input_ids, indexes, attention_mask
+ )
+ return past_key_values, indexes_image, key_valid_mask, prefix_lengths
+
+ def _t2i_think_prefill(
+ self, tokenizer, query, *, token_h, token_w, IMG_START_TOKEN
+ ):
+ (
+ input_ids,
+ indexes,
+ attention_mask,
+ key_valid_mask,
+ prefix_lengths,
+ ) = self._build_t2i_text_inputs(tokenizer, query)
+ outputs = self._think_prefix_forward(
+ input_ids=input_ids,
+ indexes=indexes,
+ attention_mask=attention_mask,
+ )
+ past_key_values, t_index, think_text = self._generate_think(
+ tokenizer,
+ outputs,
+ outputs.past_key_values,
+ indexes[0].max().item(),
+ IMG_START_TOKEN,
+ )
+ indexes_image = self._build_t2i_image_indexes(
+ token_h, token_w, t_index + 1, device=input_ids.device
+ )
+ return (
+ past_key_values,
+ indexes_image,
+ key_valid_mask,
+ prefix_lengths,
+ think_text,
+ )
+
def _think_prefix_forward(self, **kwargs):
"""Build the Think prefix cache without materialising per-token logits."""
return self.language_model(use_cache=True, logits_to_keep=1, **kwargs)
@@ -696,6 +815,7 @@ def t2i_generate(
t_eps=0.02,
think_mode=False,
seed=0,
+ cfg_group=None,
):
assert self.concat_time_token_num == 0
assert cfg_norm in ["cfg_zero_star", "global", "none", "channel"]
@@ -718,162 +838,99 @@ def t2i_generate(
self.config.t_eps = t_eps
think_text = ""
needs_cfg = cfg_scale > 1
-
- think_content = (
- "\n" if think_mode else "\n\n\n\n" + IMG_START_TOKEN
- )
- query_conditions = [
- self._build_t2i_query(
- item,
- system_message=SYSTEM_MESSAGE_FOR_GEN,
- append_text=think_content,
- )
- for item in prompts
- ]
- query_condition = (
- query_conditions[0] if len(query_conditions) == 1 else query_conditions
- )
- query_uncondition = (
- self._build_t2i_query("", append_text=IMG_START_TOKEN)
- if needs_cfg
- else None
+ cfg_parallel = needs_cfg and cfg_group is not None and cfg_group.world_size > 1
+ # CFG rank 0 owns the condition branch and rank 1 the unconditional one;
+ # each rank prefills and denoises only the branch it owns.
+ run_condition = not cfg_parallel or cfg_group.rank_in_group == 0
+ run_uncondition = needs_cfg and (
+ not cfg_parallel or cfg_group.rank_in_group == 1
)
- (
- input_ids_condition,
- indexes_condition,
- attention_mask_condition_prefix,
- condition_key_valid_mask,
- condition_prefix_lengths,
- ) = self._build_t2i_text_inputs(tokenizer, query_condition)
- if query_uncondition is not None:
- (
- input_ids_uncondition,
- indexes_uncondition,
- attention_mask_uncondition_prefix,
- _,
- uncondition_prefix_lengths,
- ) = self._build_t2i_text_inputs(tokenizer, query_uncondition)
- else:
- input_ids_uncondition = indexes_uncondition = (
- attention_mask_uncondition_prefix
- ) = None
- uncondition_prefix_lengths = None
-
token_h = image_size[1] // (self.patch_size * merge_size)
token_w = image_size[0] // (self.patch_size * merge_size)
- indexes_image_condition = self._build_t2i_image_indexes(
- token_h,
- token_w,
- (
- int(condition_prefix_lengths[0].item())
- if len(prompts) == 1
- else condition_prefix_lengths
- ),
- device=input_ids_condition.device,
- )
- indexes_image_uncondition = (
- self._build_t2i_image_indexes(
- token_h,
- token_w,
- int(uncondition_prefix_lengths[0].item()),
- device=input_ids_uncondition.device,
- )
- if indexes_uncondition is not None
- else None
- )
-
- if think_mode:
- outputs_condition = self._think_prefix_forward(
- input_ids=input_ids_condition,
- indexes=indexes_condition,
- attention_mask=attention_mask_condition_prefix,
- )
- past_key_values_condition = outputs_condition.past_key_values
- device = outputs_condition.logits.device
- dtype = outputs_condition.logits.dtype
- t_index_condition = indexes_condition[0].max().item()
- past_key_values_condition, t_index_condition, think_text = (
- self._generate_think(
- tokenizer,
- outputs_condition,
- past_key_values_condition,
- t_index_condition,
- IMG_START_TOKEN,
+ past_key_values_condition = indexes_image_condition = None
+ condition_key_valid_mask = condition_prefix_lengths = None
+ past_key_values_uncondition = indexes_image_uncondition = None
+ uncondition_prefix_lengths = None
+ with _cfg_ranks_fail_together(cfg_group if cfg_parallel else None):
+ if run_condition:
+ think_content = (
+ "\n"
+ if think_mode
+ else "\n\n\n\n" + IMG_START_TOKEN
)
- )
- indexes_image_condition = self._build_t2i_image_indexes(
- token_h,
- token_w,
- t_index_condition + 1,
- device=input_ids_condition.device,
- )
- del outputs_condition
- else:
- past_key_values_condition, prefix_hidden_states = self._t2i_prefix_forward(
- input_ids_condition, indexes_condition, attention_mask_condition_prefix
- )
- device = prefix_hidden_states.device
- dtype = prefix_hidden_states.dtype
- del prefix_hidden_states
- past_key_values_uncondition = None
- if input_ids_uncondition is not None:
- past_key_values_uncondition, _ = self._t2i_prefix_forward(
- input_ids_uncondition,
- indexes_uncondition,
- attention_mask_uncondition_prefix,
- )
-
- del input_ids_condition, indexes_condition, attention_mask_condition_prefix
- if input_ids_uncondition is not None:
- del (
- input_ids_uncondition,
- indexes_uncondition,
- attention_mask_uncondition_prefix,
- )
- self._notify_layer_offload_phase("denoise")
-
- for layer_idx in range(len(past_key_values_condition.layers)):
- past_key_values_condition.layers[
- layer_idx
- ].keys = past_key_values_condition.layers[layer_idx].keys.expand(
- batch_size, *past_key_values_condition.layers[layer_idx].keys.shape[1:]
- )
- past_key_values_condition.layers[
- layer_idx
- ].values = past_key_values_condition.layers[layer_idx].values.expand(
- batch_size,
- *past_key_values_condition.layers[layer_idx].values.shape[1:],
- )
- if past_key_values_uncondition is not None:
- past_key_values_uncondition.layers[
- layer_idx
- ].keys = past_key_values_uncondition.layers[layer_idx].keys.expand(
- batch_size,
- *past_key_values_uncondition.layers[layer_idx].keys.shape[1:],
+ query_conditions = [
+ self._build_t2i_query(
+ item,
+ system_message=SYSTEM_MESSAGE_FOR_GEN,
+ append_text=think_content,
+ )
+ for item in prompts
+ ]
+ query_condition = (
+ query_conditions[0]
+ if len(query_conditions) == 1
+ else query_conditions
)
- past_key_values_uncondition.layers[
- layer_idx
- ].values = past_key_values_uncondition.layers[layer_idx].values.expand(
- batch_size,
- *past_key_values_uncondition.layers[layer_idx].values.shape[1:],
+ if think_mode:
+ (
+ past_key_values_condition,
+ indexes_image_condition,
+ condition_key_valid_mask,
+ condition_prefix_lengths,
+ think_text,
+ ) = self._t2i_think_prefill(
+ tokenizer,
+ query_condition,
+ token_h=token_h,
+ token_w=token_w,
+ IMG_START_TOKEN=IMG_START_TOKEN,
+ )
+ else:
+ (
+ past_key_values_condition,
+ indexes_image_condition,
+ condition_key_valid_mask,
+ condition_prefix_lengths,
+ ) = self._t2i_prefill(
+ tokenizer, query_condition, token_h=token_h, token_w=token_w
+ )
+ if run_uncondition:
+ (
+ past_key_values_uncondition,
+ indexes_image_uncondition,
+ _,
+ uncondition_prefix_lengths,
+ ) = self._t2i_prefill(
+ tokenizer,
+ self._build_t2i_query("", append_text=IMG_START_TOKEN),
+ token_h=token_h,
+ token_w=token_w,
)
+ self._notify_layer_offload_phase("denoise")
- # prepare flash cache once
- prepare_flash_kv_cache(
- past_key_values_condition,
- current_len=token_h * token_w,
- batch_size=batch_size,
- prefix_lengths=condition_prefix_lengths,
- )
- if past_key_values_uncondition is not None:
- prepare_flash_kv_cache(
- past_key_values_uncondition,
- current_len=token_h * token_w,
- batch_size=batch_size,
- prefix_lengths=uncondition_prefix_lengths,
+ own_cache = (
+ past_key_values_condition
+ if run_condition
+ else past_key_values_uncondition
)
+ device = own_cache.layers[0].keys.device
+ dtype = own_cache.layers[0].keys.dtype
+ # prepare flash cache once
+ for past_key_values, prefix_lengths in (
+ (past_key_values_condition, condition_prefix_lengths),
+ (past_key_values_uncondition, uncondition_prefix_lengths),
+ ):
+ if past_key_values is None:
+ continue
+ _expand_prefix_cache_batch(past_key_values, batch_size)
+ prepare_flash_kv_cache(
+ past_key_values,
+ current_len=token_h * token_w,
+ batch_size=batch_size,
+ prefix_lengths=prefix_lengths,
+ )
# init noise image tokens
grid_h = image_size[1] // self.patch_size
@@ -896,7 +953,12 @@ def t2i_generate(
)
attention_mask_condition = {"full_attention": None}
- if device.type == "npu" and batch_size > 1 and not npu_fia_available():
+ if (
+ run_condition
+ and device.type == "npu"
+ and batch_size > 1
+ and not npu_fia_available()
+ ):
condition_key_valid_mask = condition_key_valid_mask.expand(batch_size, -1)
image_key_valid_mask = torch.ones(
(batch_size, token_h * token_w),
@@ -962,69 +1024,57 @@ def t2i_generate(
).view(batch_size, token_h * token_w, -1)
image_embeds = image_embeds + timestep_embeddings
- v_pred_condition = self._t2i_predict_v(
- image_embeds,
- indexes_image_condition,
- attention_mask_condition,
- past_key_values_condition,
- t,
- z,
- image_token_num=token_h * token_w,
- timestep_embeddings=timestep_embeddings,
- image_size=image_size,
- )
-
- if cfg_window[step_i] and cfg_scale > 1:
- v_pred_uncondition = self._t2i_predict_v(
+ v_pred_condition = None
+ if run_condition:
+ v_pred_condition = self._t2i_predict_v(
image_embeds,
- indexes_image_uncondition,
- attention_mask_uncondition,
- past_key_values_uncondition,
+ indexes_image_condition,
+ attention_mask_condition,
+ past_key_values_condition,
t,
z,
image_token_num=token_h * token_w,
timestep_embeddings=timestep_embeddings,
image_size=image_size,
)
- if cfg_norm == "cfg_zero_star":
- positive_flat = v_pred_condition.view(batch_size, -1)
- negative_flat = v_pred_uncondition.view(batch_size, -1)
- alpha = optimized_scale(positive_flat, negative_flat)
- alpha = alpha.view(
- batch_size, *([1] * (len(v_pred_condition.shape) - 1))
+ if cfg_window[step_i] and cfg_scale > 1:
+ v_pred_uncondition = None
+ if run_uncondition:
+ v_pred_uncondition = self._t2i_predict_v(
+ image_embeds,
+ indexes_image_uncondition,
+ attention_mask_uncondition,
+ past_key_values_uncondition,
+ t,
+ z,
+ image_token_num=token_h * token_w,
+ timestep_embeddings=timestep_embeddings,
+ image_size=image_size,
)
- alpha = alpha.to(positive_flat.dtype)
-
- if step_i <= 0:
- v_pred = v_pred_condition * 0.0
- else:
- v_pred = v_pred_uncondition * alpha + cfg_scale * (
- v_pred_condition - v_pred_uncondition * alpha
- )
- else:
- v_pred = v_pred_uncondition + cfg_scale * (
- v_pred_condition - v_pred_uncondition
+ if cfg_parallel:
+ # Every rank combines both gathered branches with the serial
+ # op order, so z stays bit-identical across ranks.
+ v_pred_condition, v_pred_uncondition = cfg_group.all_gather(
+ v_pred_condition if run_condition else v_pred_uncondition,
+ dim=0,
+ separate_tensors=True,
)
- if cfg_norm == "global":
- norm_v_condition = torch.norm(
- v_pred_condition, dim=(1, 2), keepdim=True
- )
- norm_v_cfg = torch.norm(v_pred, dim=(1, 2), keepdim=True)
- scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(
- min=0, max=1.0
- )
- v_pred = v_pred * scale
- elif cfg_norm == "channel":
- norm_v_condition = torch.norm(
- v_pred_condition, dim=-1, keepdim=True
- )
- norm_v_cfg = torch.norm(v_pred, dim=-1, keepdim=True)
- scale = (norm_v_condition / (norm_v_cfg + 1e-8)).clamp(
- min=0, max=1.0
- )
- v_pred = v_pred * scale
+ v_pred = _apply_t2i_cfg(
+ v_pred_condition,
+ v_pred_uncondition,
+ cfg_scale=cfg_scale,
+ cfg_norm=cfg_norm,
+ step_i=step_i,
+ batch_size=batch_size,
+ )
else:
+ if cfg_parallel:
+ # v_pred has z's shape and dtype.
+ v_pred_condition = cfg_group.broadcast(
+ v_pred_condition if run_condition else torch.empty_like(z),
+ src=0,
+ )
v_pred = v_pred_condition
z = z + (t_next - t) * v_pred
@@ -1034,8 +1084,7 @@ def t2i_generate(
)
clear_flash_kv_cache(past_key_values_condition)
- if past_key_values_uncondition is not None:
- clear_flash_kv_cache(past_key_values_uncondition)
+ clear_flash_kv_cache(past_key_values_uncondition)
self.last_think_content = think_text
if think_mode:
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sensenova_u1.py b/python/sglang/multimodal_gen/runtime/pipelines/sensenova_u1.py
index fbed05c87f19..d634f526edc8 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines/sensenova_u1.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines/sensenova_u1.py
@@ -7,6 +7,7 @@
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
SenseNovaU1PipelineConfig,
+ validate_sensenova_u1_parallelism,
)
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
SenseNovaU1SamplingParams,
@@ -49,22 +50,18 @@ def load_modules(
if loaded_modules is not None and {"model", "tokenizer"} <= set(loaded_modules):
return loaded_modules
- if server_args.num_gpus != 1:
- raise ValueError(
- "SenseNovaU1Pipeline currently supports num_gpus=1. "
- "Native tensor/pipeline parallelism is not implemented yet."
- )
+ validate_sensenova_u1_parallelism(server_args)
modules = load_model_and_tokenizer(self.model_path, server_args)
logger.info("Loaded SenseNova-U1 model from %s", self.model_path)
return modules
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
- del server_args
self.add_stage(InputValidationStage())
self.add_stage(
SenseNovaU1GenerationStage(
model=self.get_module("model"),
tokenizer=self.get_module("tokenizer"),
+ enable_cfg_parallel=server_args.enable_cfg_parallel,
),
"sensenova_u1_generation_stage",
)
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sensenova_u1/stages/generation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sensenova_u1/stages/generation.py
index d1a8c03dcfe7..d478c196689b 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sensenova_u1/stages/generation.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sensenova_u1/stages/generation.py
@@ -16,6 +16,7 @@
SENSENOVA_U1_REQUEST_EXTRA_KEY,
)
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_cfg_group
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
OutputBatch,
Req,
@@ -54,10 +55,18 @@ def from_batch(cls, batch: Req) -> SenseNovaU1GenerationOptions:
class SenseNovaU1GenerationStage(PipelineStage):
- def __init__(self, model: torch.nn.Module, tokenizer: Any):
+ def __init__(
+ self,
+ model: torch.nn.Module,
+ tokenizer: Any,
+ *,
+ enable_cfg_parallel: bool = False,
+ ):
super().__init__()
self.model = model
self.tokenizer = tokenizer
+ # Every CFG rank runs this stage; the model splits the branches by rank.
+ self.enable_cfg_parallel = enable_cfg_parallel
@property
def role_affinity(self) -> RoleType:
@@ -117,6 +126,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
t_eps=options.t_eps,
think_mode=options.think_mode,
seed=seed,
+ cfg_group=get_cfg_group() if self.enable_cfg_parallel else None,
)
think_text = None
if options.think_mode:
diff --git a/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py b/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
index fe3d5a8f42ac..ea47b2a89d9a 100644
--- a/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
+++ b/python/sglang/multimodal_gen/test/unit/test_sensenova_u1.py
@@ -1,7 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
import asyncio
+import copy
import json
import math
+import threading
import time
from collections import deque
from types import SimpleNamespace
@@ -14,6 +16,7 @@
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
SenseNovaU1PipelineConfig,
+ validate_sensenova_u1_parallelism,
)
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
SenseNovaU1SamplingParams,
@@ -37,6 +40,7 @@
modeling_fm_modules,
)
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.configuration_neo_chat import (
+ NEOChatConfig,
NEOLLMConfig,
)
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.configuration_neo_vit import (
@@ -66,6 +70,9 @@
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1 import (
SenseNovaU1GenerationStage,
)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1.stages import (
+ generation as sensenova_generation,
+)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.perf_logger import MemorySnapshot
@@ -484,6 +491,213 @@ def test_sensenova_u1_randn_fallback_preserves_device_rng(monkeypatch):
assert torch.equal(device_module.get_rng_state(device), device_rng_state)
+class _ThreadedCFGGroup:
+ """Two-rank CFG group whose collectives exchange tensors between threads."""
+
+ world_size = 2
+ device = torch.device("cpu")
+
+ def __init__(self, rank_in_group, slots, barrier):
+ self.rank_in_group = rank_in_group
+ self._slots = slots
+ self._barrier = barrier
+
+ def _exchange(self, tensor):
+ self._slots[self.rank_in_group] = tensor.clone()
+ self._barrier.wait()
+ gathered = list(self._slots)
+ self._barrier.wait()
+ return gathered
+
+ def all_gather(self, input_, dim=0, separate_tensors=False):
+ assert dim == 0 and separate_tensors
+ return self._exchange(input_)
+
+ def all_reduce(self, input_):
+ return input_.copy_(sum(self._exchange(input_)))
+
+ def broadcast(self, input_, src=0):
+ return input_.copy_(self._exchange(input_)[src])
+
+
+def _make_tiny_sensenova_u1_model():
+ config = NEOChatConfig(
+ vision_config=dict(
+ patch_size=4, hidden_size=16, llm_hidden_size=32, downsample_ratio=0.5
+ ),
+ llm_config=dict(
+ architectures=["Qwen3ForCausalLM"],
+ vocab_size=512,
+ hidden_size=32,
+ intermediate_size=48,
+ num_hidden_layers=2,
+ num_attention_heads=2,
+ num_key_value_heads=1,
+ head_dim=16,
+ max_position_embeddings=512,
+ max_position_embeddings_hw=64,
+ ),
+ downsample_ratio=0.5,
+ template="neo1_0",
+ fm_head_layers=2,
+ use_pixel_head=False,
+ concat_time_token_num=0,
+ noise_scale=1.0,
+ noise_scale_mode="none",
+ noise_scale_base_image_seq_len=64,
+ add_noise_scale_embedding=False,
+ noise_scale_max_value=8.0,
+ time_schedule="standard",
+ time_shift_type="exponential",
+ base_shift=0.5,
+ max_shift=1.15,
+ base_image_seq_len=64,
+ max_image_seq_len=4096,
+ )
+ torch.manual_seed(0)
+ return NEOChatModel(config).eval()
+
+
+@pytest.mark.parametrize(
+ ("generation_kwargs", "uncondition_steps"),
+ [
+ ({"prompt": "a red apple", "cfg_norm": "none", "seed": 1}, 6),
+ # timestep_shift=3 puts 5 of the 6 steps at t <= 0.5; the last step
+ # runs without CFG, so rank 0 broadcasts its condition prediction.
+ (
+ {
+ "prompt": "a red apple",
+ "cfg_norm": "cfg_zero_star",
+ "cfg_interval": (0.0, 0.5),
+ "seed": 1,
+ },
+ 5,
+ ),
+ (
+ {
+ "prompt": ["a red apple", "two cats on a long sofa"],
+ "batch_size": 2,
+ "cfg_norm": "global",
+ "cfg_interval": (0.2, 1.0),
+ "seed": [3, 4],
+ },
+ 3,
+ ),
+ ],
+)
+def test_sensenova_u1_cfg_parallel_splits_branches_and_matches_serial(
+ generation_kwargs, uncondition_steps
+):
+ """Each CFG rank denoises only its own branch, yet both ranks reproduce
+ the single-GPU image bit for bit."""
+ num_threads = torch.get_num_threads()
+ # Keep CPU kernel partitioning identical between the serial and rank runs.
+ torch.set_num_threads(1)
+ try:
+ model = _make_tiny_sensenova_u1_model()
+ kwargs = dict(
+ cfg_scale=4.0,
+ image_size=(32, 32),
+ num_steps=6,
+ timestep_shift=3.0,
+ **generation_kwargs,
+ )
+ with torch.inference_mode():
+ expected = model.t2i_generate(_FakeTokenizer(), **kwargs)
+
+ slots = [None, None]
+ barrier = threading.Barrier(2, timeout=60)
+ outputs, branch_calls, errors = [None, None], [0, 0], []
+
+ def run_rank(rank):
+ rank_model = copy.deepcopy(model)
+ predict_v = rank_model._t2i_predict_v
+
+ def counted_predict_v(*args, **inner_kwargs):
+ branch_calls[rank] += 1
+ return predict_v(*args, **inner_kwargs)
+
+ rank_model._t2i_predict_v = counted_predict_v
+ try:
+ with torch.inference_mode():
+ outputs[rank] = rank_model.t2i_generate(
+ _FakeTokenizer(),
+ **kwargs,
+ cfg_group=_ThreadedCFGGroup(rank, slots, barrier),
+ )
+ except Exception as error:
+ errors.append(error)
+ barrier.abort()
+
+ threads = [threading.Thread(target=run_rank, args=(rank,)) for rank in (0, 1)]
+ for thread in threads:
+ thread.start()
+ for thread in threads:
+ thread.join()
+ finally:
+ torch.set_num_threads(num_threads)
+
+ assert not errors, errors
+ assert branch_calls == [kwargs["num_steps"], uncondition_steps]
+ for output in outputs:
+ torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
+def test_sensenova_u1_cfg_parallel_prefill_error_fails_both_ranks():
+ """A condition-prefill error, which only CFG rank 0 can hit, must fail the
+ request on both ranks, and the next request must still pair up correctly
+ instead of reusing collectives the failed request left behind."""
+ num_threads = torch.get_num_threads()
+ torch.set_num_threads(1)
+ try:
+ model = _make_tiny_sensenova_u1_model()
+ kwargs = dict(
+ prompt="a red apple", cfg_scale=4.0, image_size=(32, 32), num_steps=2
+ )
+ with torch.inference_mode():
+ expected = model.t2i_generate(_FakeTokenizer(), **kwargs)
+
+ slots = [None, None]
+ barrier = threading.Barrier(2, timeout=10)
+ errors, outputs = [None, None], [None, None]
+
+ def run_rank(rank):
+ rank_model = copy.deepcopy(model)
+ prefill = rank_model._t2i_prefill
+ if rank == 0:
+
+ def failing_prefill(*args, **inner_kwargs):
+ rank_model._t2i_prefill = prefill
+ raise torch.OutOfMemoryError("condition prefill")
+
+ rank_model._t2i_prefill = failing_prefill
+ cfg_group = _ThreadedCFGGroup(rank, slots, barrier)
+ with torch.inference_mode():
+ try:
+ rank_model.t2i_generate(
+ _FakeTokenizer(), **kwargs, cfg_group=cfg_group
+ )
+ except Exception as error:
+ errors[rank] = error
+ outputs[rank] = rank_model.t2i_generate(
+ _FakeTokenizer(), **kwargs, cfg_group=cfg_group
+ )
+
+ threads = [threading.Thread(target=run_rank, args=(rank,)) for rank in (0, 1)]
+ for thread in threads:
+ thread.start()
+ for thread in threads:
+ thread.join()
+ finally:
+ torch.set_num_threads(num_threads)
+
+ assert isinstance(errors[0], torch.OutOfMemoryError)
+ assert isinstance(errors[1], RuntimeError)
+ assert "other CFG rank" in str(errors[1])
+ for output in outputs:
+ torch.testing.assert_close(output, expected, rtol=0, atol=0)
+
+
def test_sensenova_u1_registry_resolves_local_and_hf_paths(tmp_path):
_get_config_info.cache_clear()
get_model_info.cache_clear()
@@ -771,21 +985,49 @@ def test_sensenova_u1_scheduler_merge_and_split_preserve_request_order():
assert scheduler._try_merge_generation_reqs(requests) is None
-def test_sensenova_u1_rejects_multi_gpu_during_arg_validation():
+_TWO_GPU_CFG_PARALLEL_ARGS = dict(
+ num_gpus=2,
+ enable_cfg_parallel=True,
+ cfg_parallel_degree=2,
+ tp_size=1,
+ sp_degree=1,
+ dp_size=1,
+ batching_max_size=1,
+)
+
+
+@pytest.mark.parametrize(
+ "override",
+ [
+ dict(enable_cfg_parallel=False, cfg_parallel_degree=1),
+ dict(num_gpus=4),
+ dict(num_gpus=4, cfg_parallel_degree=4),
+ # ServerArgs never checks sp_degree against num_gpus / cfg degree.
+ dict(sp_degree=2),
+ ],
+)
+def test_sensenova_u1_rejects_multi_gpu_without_two_way_cfg_parallel(override):
config = SenseNovaU1PipelineConfig()
- with pytest.raises(ValueError, match="num_gpus=1"):
+ with pytest.raises(ValueError, match="num_gpus=2 with --enable-cfg-parallel"):
config.validate_server_args(
- SimpleNamespace(
- num_gpus=2,
- enable_torch_compile=False,
- lora_path=None,
- attention_backend=None,
- component_attention_backends={},
- )
+ SimpleNamespace(**{**_TWO_GPU_CFG_PARALLEL_ARGS, **override})
+ )
+
+
+def test_sensenova_u1_rejects_dynamic_batching_with_cfg_parallel():
+ """CFG ranks time the batching delay independently and could dispatch
+ different batches into the same collectives."""
+ with pytest.raises(ValueError, match="dynamic batching"):
+ validate_sensenova_u1_parallelism(
+ SimpleNamespace(**{**_TWO_GPU_CFG_PARALLEL_ARGS, "batching_max_size": 2})
)
+def test_sensenova_u1_accepts_two_gpu_cfg_parallel():
+ validate_sensenova_u1_parallelism(SimpleNamespace(**_TWO_GPU_CFG_PARALLEL_ARGS))
+
+
def test_sensenova_u1_clears_auto_tuned_runtime_defaults():
config = SenseNovaU1PipelineConfig()
args = SimpleNamespace(
@@ -1016,7 +1258,12 @@ def test_sensenova_u1_cli_args_expose_only_sglang_compatible_fields():
assert request.perf_dump_path == "/tmp/sensenova-perf.json"
-def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch():
+@pytest.mark.parametrize("enable_cfg_parallel", [False, True])
+def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch(
+ monkeypatch, enable_cfg_parallel
+):
+ cfg_group = object()
+ monkeypatch.setattr(sensenova_generation, "get_cfg_group", lambda: cfg_group)
sampling = SenseNovaU1SamplingParams(
prompt="a mountain lake",
width=2304,
@@ -1037,7 +1284,9 @@ def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch
metrics=None,
)
model = _FakeSenseNovaModel()
- stage = SenseNovaU1GenerationStage(model=model, tokenizer="tok")
+ stage = SenseNovaU1GenerationStage(
+ model=model, tokenizer="tok", enable_cfg_parallel=enable_cfg_parallel
+ )
output = stage.forward(batch, server_args=SimpleNamespace())
@@ -1059,6 +1308,9 @@ def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch
assert model.call_kwargs["num_steps"] == 30
assert model.call_kwargs["batch_size"] == 1
assert model.call_kwargs["seed"] == 123
+ # Without the group every CFG rank silently computes both branches.
+ expected_group = cfg_group if enable_cfg_parallel else None
+ assert model.call_kwargs["cfg_group"] is expected_group
def test_sensenova_u1_generation_stage_passes_dynamic_batch_inputs():