From ab43c6d351c812348ba8b697cf2742ad315e9d6c Mon Sep 17 00:00:00 2001 From: qiujiawei Date: Mon, 19 Oct 2026 19:35:31 +0800 Subject: [PATCH] [Feat] SenseNova-U1: CFG parallelism on two GPUs Run the conditional branch on CFG rank 0 and the unconditional branch on rank 1, exchange v_pred with all_gather every step and combine it with the unchanged serial code, so the image is bit-identical to one GPU (1.95x at 2048x2048 on 2x H200). A failure flag exchanged after the prefill makes both ranks fail a request that failed on either. Accept num_gpus=2 only with --enable-cfg-parallel (tp/sp/dp 1) and reject dynamic batching with it. Co-Authored-By: Claude Opus 5.5 --- .../SenseNova/SenseNova-U1.5-8B-MoT.mdx | 39 +- .../configs/SenseNova/sensenova-u15.jsx | 14 +- .../configs/pipeline_configs/sensenova_u1.py | 34 +- .../configs/sample/sensenova_u1.py | 5 + .../neo_unify/modeling_neo_chat.py | 451 ++++++++++-------- .../runtime/pipelines/sensenova_u1.py | 9 +- .../sensenova_u1/stages/generation.py | 12 +- .../test/unit/test_sensenova_u1.py | 274 ++++++++++- 8 files changed, 607 insertions(+), 231 deletions(-) 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():