diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py b/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py index d14e92769..fcd0ebf8d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py @@ -19,7 +19,8 @@ logger = init_logger(__name__) # Helios UMT5 max sequence length (used for both tokenizer and post-processing padding) -HELIOS_MAX_SEQUENCE_LENGTH = 226 +# Matches diffusers HeliosPipeline.__call__ default max_sequence_length=512 +HELIOS_MAX_SEQUENCE_LENGTH = 512 def umt5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 42eba884d..bd8256de5 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -431,7 +431,7 @@ class SamplingParams: pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower() - if "wan" in pipeline_name_lower and ( + if ("wan" in pipeline_name_lower or "helios" in pipeline_name_lower) and ( self.enable_sequence_shard is None or self.enable_sequence_shard ): self.enable_sequence_shard = True diff --git a/python/sglang/multimodal_gen/runtime/models/dits/helios.py b/python/sglang/multimodal_gen/runtime/models/dits/helios.py index f4eecd42e..ce62e521f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/helios.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/helios.py @@ -19,8 +19,13 @@ import torch.nn.functional as F from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.runtime.distributed import ( divide, + get_sp_world_size, get_tp_world_size, ) +from sglang.multimodal_gen.runtime.distributed.communication_op import ( + sequence_model_parallel_all_gather, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group from sglang.multimodal_gen.runtime.layers.attention import USPAttention from sglang.multimodal_gen.runtime.layers.layernorm import ( FP32LayerNorm, @@ -40,6 +45,7 @@ from sglang.multimodal_gen.runtime.layers.visual_embedding import ( PatchEmbed, TimestepEmbedder, ) +from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -295,8 +301,13 @@ class HeliosSelfAttention(nn.Module): q = apply_rotary_emb_transposed(q, rotary_emb) k = apply_rotary_emb_transposed(k, rotary_emb) + history_seq_len = ( + hidden_states.shape[1] - original_context_length + if original_context_length is not None + else 0 + ) + if self.is_amplify_history and original_context_length is not None: - history_seq_len = hidden_states.shape[1] - original_context_length if history_seq_len > 0: scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * ( self.max_scale - 1.0 @@ -308,7 +319,7 @@ class HeliosSelfAttention(nn.Module): dim=1, ) - x = self.attn(q, k, v) + x = self.attn(q, k, v, num_replicated_prefix=history_seq_len) x = x.flatten(2) x, _ = self.to_out(x) return x @@ -356,7 +367,7 @@ class HeliosCrossAttention(nn.Module): num_heads=self.local_num_heads, head_size=self.head_dim, causal=False, - is_cross_attention=True, + skip_sequence_parallel=True, ) def forward(self, hidden_states, encoder_hidden_states): @@ -624,6 +635,7 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): self.cnt = 0 self.__post_init__() self.layer_names = ["blocks"] + self.sp_size = get_sp_world_size() def forward( self, @@ -644,6 +656,15 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] + # Check if sequence parallelism is enabled + forward_batch = get_forward_context().forward_batch + if forward_batch is not None: + sequence_shard_enabled = ( + forward_batch.enable_sequence_shard and self.sp_size > 1 + ) + else: + sequence_shard_enabled = False + batch_size = hidden_states.shape[0] p_t, p_h, p_w = self.patch_size @@ -672,6 +693,40 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): rotary_emb = rotary_emb.flatten(2).transpose(1, 2) original_context_length = hidden_states.shape[1] + # Sequence parallelism: shard current tokens and RoPE across SP ranks + seq_shard_pad = 0 + if sequence_shard_enabled: + sp_rank = get_sp_group().rank_in_group + seq_len = hidden_states.shape[1] + if seq_len % self.sp_size != 0: + seq_shard_pad = self.sp_size - (seq_len % self.sp_size) + hs_pad = torch.zeros( + batch_size, + seq_shard_pad, + hidden_states.shape[2], + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + re_pad = torch.zeros( + batch_size, + seq_shard_pad, + rotary_emb.shape[2], + dtype=rotary_emb.dtype, + device=rotary_emb.device, + ) + hidden_states = torch.cat([hidden_states, hs_pad], dim=1) + rotary_emb = torch.cat([rotary_emb, re_pad], dim=1) + local_seq_len = hidden_states.shape[1] // self.sp_size + hidden_states = hidden_states.view( + batch_size, self.sp_size, local_seq_len, -1 + )[:, sp_rank, :, :].contiguous() + rotary_emb = rotary_emb.view(batch_size, self.sp_size, local_seq_len, -1)[ + :, sp_rank, :, : + ].contiguous() + effective_context_length = local_seq_len + else: + effective_context_length = original_context_length + # 3. Process short history if ( latents_history_short is not None @@ -743,7 +798,7 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) - history_context_length = hidden_states.shape[1] - original_context_length + history_context_length = hidden_states.shape[1] - effective_context_length # 6. Compute condition embeddings if indices_hidden_states is not None and self.zero_history_timestep: @@ -772,7 +827,7 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): if indices_hidden_states is not None and not self.zero_history_timestep: main_repeat_size = hidden_states.shape[1] else: - main_repeat_size = original_context_length + main_repeat_size = effective_context_length temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand( batch_size, 6, main_repeat_size, -1 @@ -796,11 +851,21 @@ class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): encoder_hidden_states, timestep_proj, rotary_emb, - original_context_length, + effective_context_length, ) self.cnt += 1 + # SP: all-gather current tokens before output + if sequence_shard_enabled: + current_tokens = hidden_states[:, -local_seq_len:, :].contiguous() + current_tokens = sequence_model_parallel_all_gather(current_tokens, dim=1) + if seq_shard_pad > 0: + current_tokens = current_tokens[:, :original_context_length, :] + hidden_states = current_tokens + # Re-create temb for norm_out (all current tokens share same timestep) + temb = temb[:, :1, :].expand(batch_size, original_context_length, -1) + # 8. Output norm & projection hidden_states = self.norm_out(hidden_states, temb, original_context_length) hidden_states, _ = self.proj_out(hidden_states) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py index 3b3f64e92..bdb1adec6 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py @@ -731,4 +731,7 @@ class HeliosScheduler: return self.config.num_train_timesteps -EntryClass = HeliosScheduler +# Alias for Helios-Distilled which uses "HeliosDMDScheduler" in scheduler_config.json +HeliosDMDScheduler = HeliosScheduler + +EntryClass = [HeliosScheduler, "HeliosDMDScheduler"] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py index eb6a9e398..ff6de5661 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py @@ -13,6 +13,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipel from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_decoding import ( + HeliosDecodingStage, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import ( HeliosChunkedDenoisingStage, ) @@ -69,8 +72,12 @@ class HeliosPipeline(LoRAPipeline, ComposedPipelineBase): ), "helios_chunked_denoising_stage", ) - # Standard DecodingStage handles VAE decode of the denoised latents - self.add_standard_decoding_stage() + # Helios-specific decoding: decode each chunk's latents separately + # to avoid temporal artifacts from Wan VAE causal convolutions + self.add_stage( + HeliosDecodingStage(vae=self.get_module("vae"), pipeline=self), + "helios_decoding_stage", + ) class HeliosPyramidPipeline(HeliosPipeline): diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py new file mode 100644 index 000000000..899a18937 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Helios-specific decoding stage. + +Decodes latent chunks one at a time (matching diffusers HeliosPipeline behavior) +to avoid temporal artifacts at chunk boundaries caused by Wan VAE's causal convolutions. +""" + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import ( + DecodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class HeliosDecodingStage(DecodingStage): + """ + Helios-specific decoding stage that decodes latent chunks independently. + + The Wan VAE uses causal 3D convolutions with feature caching. When decoding + the full latent sequence at once, the causal conv processes all frames with + continuous context, producing a different number of output frames per latent + frame compared to chunk-by-chunk decoding. This causes temporal misalignment + and visible seams at chunk boundaries. + + This stage decodes each chunk's latents separately (matching diffusers' + HeliosPipeline behavior) and concatenates the results in pixel space. + """ + + @torch.no_grad() + def forward( + self, + batch: Req, + server_args: ServerArgs, + ) -> OutputBatch: + latent_chunks = getattr(batch, "latent_chunks", None) + + if latent_chunks is None or len(latent_chunks) <= 1: + # No chunked latents or single chunk — use standard decode + return super().forward(batch, server_args) + + # Load VAE if needed + self.load_model() + + # Decode each chunk separately and concatenate in pixel space + video_chunks = [] + for chunk_latents in latent_chunks: + chunk_video = self.decode(chunk_latents, server_args) + video_chunks.append(chunk_video) + + frames = torch.cat(video_chunks, dim=2) + frames = server_args.pipeline_config.post_decoding(frames, server_args) + + output_batch = OutputBatch( + output=frames, + trajectory_timesteps=batch.trajectory_timesteps, + trajectory_latents=batch.trajectory_latents, + trajectory_decoded=None, + metrics=batch.metrics, + ) + + self.offload_model() + return output_batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py index 119f1a202..29b7a3b13 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.utils import PRECISION_TO_TYPE logger = init_logger(__name__) @@ -56,11 +57,13 @@ def sample_block_noise( _, ph, pw = patch_size block_size = ph * pw + # Explicitly use CPU to avoid requiring MAGMA for cholesky on ROCm/CUDA cov = ( - torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma + torch.eye(block_size, device="cpu") * (1 + gamma) + - torch.ones(block_size, block_size, device="cpu") * gamma ) dist = torch.distributions.MultivariateNormal( - torch.zeros(block_size, device=cov.device), covariance_matrix=cov + torch.zeros(block_size, device="cpu"), covariance_matrix=cov ) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) @@ -113,55 +116,33 @@ class HeliosChunkedDenoisingStage(PipelineStage): zero_steps=1, batch=None, server_args=None, + global_step_offset=0, ): """Denoise a single chunk with full timestep loop.""" batch_size = latents.shape[0] do_cfg = guidance_scale > 1.0 for i, t in enumerate(timesteps): - timestep = t.expand(batch_size) - latent_model_input = latents.to(target_dtype) - - with set_forward_context( - current_timestep=t, - forward_batch=None, - attn_metadata=None, + with StageProfiler( + f"denoising_step_{global_step_offset + i}", + logger=logger, + metrics=batch.metrics if batch is not None else None, + perf_dump_path_provided=( + batch.perf_dump_path is not None if batch is not None else False + ), ): - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - indices_hidden_states=indices_hidden_states, - indices_latents_history_short=indices_latents_history_short, - indices_latents_history_mid=indices_latents_history_mid, - indices_latents_history_long=indices_latents_history_long, - latents_history_short=( - latents_history_short.to(target_dtype) - if latents_history_short is not None - else None - ), - latents_history_mid=( - latents_history_mid.to(target_dtype) - if latents_history_mid is not None - else None - ), - latents_history_long=( - latents_history_long.to(target_dtype) - if latents_history_long is not None - else None - ), - ) + timestep = t.expand(batch_size) + latent_model_input = latents.to(target_dtype) - if do_cfg: with set_forward_context( current_timestep=t, - forward_batch=None, + forward_batch=batch, attn_metadata=None, ): - noise_uncond = self.transformer( + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, @@ -183,29 +164,62 @@ class HeliosChunkedDenoisingStage(PipelineStage): ), ) - if is_cfg_zero_star: - noise_pred_text = noise_pred - positive_flat = noise_pred_text.reshape(batch_size, -1) - negative_flat = noise_uncond.reshape(batch_size, -1) - - alpha = optimized_scale(positive_flat, negative_flat) - alpha = alpha.view( - batch_size, *([1] * (len(noise_pred_text.shape) - 1)) - ) - alpha = alpha.to(noise_pred_text.dtype) - - if (i <= zero_steps) and use_zero_init: - noise_pred = noise_pred_text * 0.0 - else: - noise_pred = noise_uncond * alpha + guidance_scale * ( - noise_pred_text - noise_uncond * alpha + if do_cfg: + with set_forward_context( + current_timestep=t, + forward_batch=batch, + attn_metadata=None, + ): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), ) - else: - noise_pred = noise_uncond + guidance_scale * ( - noise_pred - noise_uncond - ) - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if is_cfg_zero_star: + noise_pred_text = noise_pred + positive_flat = noise_pred_text.reshape(batch_size, -1) + negative_flat = noise_uncond.reshape(batch_size, -1) + + alpha = optimized_scale(positive_flat, negative_flat) + alpha = alpha.view( + batch_size, *([1] * (len(noise_pred_text.shape) - 1)) + ) + alpha = alpha.to(noise_pred_text.dtype) + + if (i <= zero_steps) and use_zero_init: + noise_pred = noise_pred_text * 0.0 + else: + noise_pred = noise_uncond * alpha + guidance_scale * ( + noise_pred_text - noise_uncond * alpha + ) + else: + noise_pred = noise_uncond + guidance_scale * ( + noise_pred - noise_uncond + ) + + latents = self.scheduler.step( + noise_pred, t, latents, return_dict=False + )[0] return latents @@ -234,6 +248,7 @@ class HeliosChunkedDenoisingStage(PipelineStage): zero_steps=1, batch=None, server_args=None, + global_step_offset=0, ): """Denoise a single chunk using pyramid super-resolution (Stage 2).""" batch_size, num_channel, num_frames, height, width = latents.shape @@ -256,6 +271,7 @@ class HeliosChunkedDenoisingStage(PipelineStage): start_point_list = [latents] do_cfg = guidance_scale > 1.0 + step_counter = global_step_offset for i_s in range(pyramid_num_stages): # Compute mu for current resolution @@ -306,49 +322,26 @@ class HeliosChunkedDenoisingStage(PipelineStage): # Denoising loop for this pyramid stage for idx, t in enumerate(timesteps): - timestep = t.expand(batch_size) - latent_model_input = latents.to(target_dtype) - - with set_forward_context( - current_timestep=t, - forward_batch=None, - attn_metadata=None, + with StageProfiler( + f"denoising_step_{step_counter}", + logger=logger, + metrics=batch.metrics if batch is not None else None, + perf_dump_path_provided=( + batch.perf_dump_path is not None if batch is not None else False + ), ): - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - indices_hidden_states=indices_hidden_states, - indices_latents_history_short=indices_latents_history_short, - indices_latents_history_mid=indices_latents_history_mid, - indices_latents_history_long=indices_latents_history_long, - latents_history_short=( - latents_history_short.to(target_dtype) - if latents_history_short is not None - else None - ), - latents_history_mid=( - latents_history_mid.to(target_dtype) - if latents_history_mid is not None - else None - ), - latents_history_long=( - latents_history_long.to(target_dtype) - if latents_history_long is not None - else None - ), - ) + timestep = t.expand(batch_size) + latent_model_input = latents.to(target_dtype) - if do_cfg: with set_forward_context( current_timestep=t, - forward_batch=None, + forward_batch=batch, attn_metadata=None, ): - noise_uncond = self.transformer( + noise_pred = self.transformer( hidden_states=latent_model_input, timestep=timestep, - encoder_hidden_states=negative_prompt_embeds, + encoder_hidden_states=prompt_embeds, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, @@ -370,44 +363,81 @@ class HeliosChunkedDenoisingStage(PipelineStage): ), ) - if is_cfg_zero_star: - noise_pred_text = noise_pred - positive_flat = noise_pred_text.reshape(batch_size, -1) - negative_flat = noise_uncond.reshape(batch_size, -1) - - alpha_cfg = optimized_scale(positive_flat, negative_flat) - alpha_cfg = alpha_cfg.view( - batch_size, - *([1] * (len(noise_pred_text.shape) - 1)), - ) - alpha_cfg = alpha_cfg.to(noise_pred_text.dtype) - - if (i_s == 0 and idx <= zero_steps) and use_zero_init: - noise_pred = noise_pred_text * 0.0 - else: - noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( - noise_pred_text - noise_uncond * alpha_cfg + if do_cfg: + with set_forward_context( + current_timestep=t, + forward_batch=batch, + attn_metadata=None, + ): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), ) - else: - noise_pred = noise_uncond + guidance_scale * ( - noise_pred - noise_uncond - ) - latents = self.scheduler.step( - noise_pred, - t, - latents, - return_dict=False, - cur_sampling_step=idx, - dmd_noisy_tensor=( - start_point_list[i_s] if start_point_list is not None else None - ), - dmd_sigmas=self.scheduler.sigmas, - dmd_timesteps=self.scheduler.timesteps, - all_timesteps=timesteps, - )[0] + if is_cfg_zero_star: + noise_pred_text = noise_pred + positive_flat = noise_pred_text.reshape(batch_size, -1) + negative_flat = noise_uncond.reshape(batch_size, -1) - return latents + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view( + batch_size, + *([1] * (len(noise_pred_text.shape) - 1)), + ) + alpha_cfg = alpha_cfg.to(noise_pred_text.dtype) + + if (i_s == 0 and idx <= zero_steps) and use_zero_init: + noise_pred = noise_pred_text * 0.0 + else: + noise_pred = ( + noise_uncond * alpha_cfg + + guidance_scale + * (noise_pred_text - noise_uncond * alpha_cfg) + ) + else: + noise_pred = noise_uncond + guidance_scale * ( + noise_pred - noise_uncond + ) + + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + cur_sampling_step=idx, + dmd_noisy_tensor=( + start_point_list[i_s] + if start_point_list is not None + else None + ), + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + + step_counter += 1 + + return latents, step_counter def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Run the Helios chunked denoising loop.""" @@ -537,6 +567,8 @@ class HeliosChunkedDenoisingStage(PipelineStage): # Chunk loop image_latents = None total_generated_latent_frames = 0 + chunk_latents_list = [] # Store per-chunk latents for chunk-by-chunk decode + global_step_offset = 0 # Track step index across chunks for perf logging self.log_info( f"Starting chunked denoising: {num_latent_chunk} chunks, " @@ -582,19 +614,30 @@ class HeliosChunkedDenoisingStage(PipelineStage): ) # Generate noise latents for this chunk - latents = torch.randn( + # Use batch.generator to ensure identical noise across SP ranks + latent_shape = ( batch_size, num_channels_latents, (window_num_frames - 1) // vae_scale_factor_temporal + 1, height // vae_scale_factor_spatial, width // vae_scale_factor_spatial, - device=device, + ) + generator = batch.generator + if isinstance(generator, list): + generator = generator[0] if len(generator) > 0 else None + gen_device = generator.device if generator is not None else device + latents = torch.randn( + latent_shape, + generator=generator, + device=gen_device, dtype=torch.float32, ) + if latents.device != device: + latents = latents.to(device) if is_enable_stage2: # Stage 2: Pyramid SR denoising (handles scheduler internally) - latents = self._denoise_one_chunk_stage2( + latents, global_step_offset = self._denoise_one_chunk_stage2( latents=latents, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -618,6 +661,7 @@ class HeliosChunkedDenoisingStage(PipelineStage): zero_steps=zero_steps, batch=batch, server_args=server_args, + global_step_offset=global_step_offset, ) else: # Stage 1: Standard flat denoising @@ -646,7 +690,9 @@ class HeliosChunkedDenoisingStage(PipelineStage): zero_steps=zero_steps, batch=batch, server_args=server_args, + global_step_offset=global_step_offset, ) + global_step_offset += num_inference_steps # Extract first frame as image_latents for subsequent chunks if keep_first_frame and is_first_chunk and image_latents is None: @@ -655,6 +701,7 @@ class HeliosChunkedDenoisingStage(PipelineStage): # Update history total_generated_latent_frames += latents.shape[2] history_latents = torch.cat([history_latents, latents], dim=2) + chunk_latents_list.append(latents) # Move transformer back to CPU after denoising if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: @@ -662,7 +709,10 @@ class HeliosChunkedDenoisingStage(PipelineStage): self.transformer.to("cpu") torch.cuda.empty_cache() - # Store denoised latents for the standard DecodingStage to decode + # Store per-chunk latents for chunk-by-chunk VAE decode (matches diffusers behavior). + # The standard DecodingStage will check for this attribute and decode each chunk + # separately to avoid temporal artifacts at chunk boundaries. + batch.latent_chunks = chunk_latents_list batch.latents = history_latents[:, :, -total_generated_latent_frames:] return batch diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index c2cf0857d..a0b5b5576 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2332,6 +2332,263 @@ "expected_e2e_ms": 10425.77, "expected_avg_denoise_ms": 172.28, "expected_median_denoise_ms": 173.29 + }, + "helios_base_t2v": { + "stages_ms": { + "InputValidationStage": 0.07, + "TextEncodingStage": 1103.97, + "LatentPreparationStage": 0.24, + "HeliosChunkedDenoisingStage": 118580.37, + "HeliosDecodingStage": 664.79, + "per_frame_generation": null + }, + "denoise_step_ms": {}, + "expected_e2e_ms": 120413.51, + "expected_avg_denoise_ms": 0.0, + "expected_median_denoise_ms": 0.0 + }, + "helios_distilled_t2v": { + "stages_ms": { + "InputValidationStage": 0.13, + "TextEncodingStage": 581.79, + "LatentPreparationStage": 0.18, + "HeliosChunkedDenoisingStage": 49752.88, + "HeliosDecodingStage": 666.69, + "per_frame_generation": null + }, + "denoise_step_ms": {}, + "expected_e2e_ms": 51038.66, + "expected_avg_denoise_ms": 0.0, + "expected_median_denoise_ms": 0.0 + }, + "helios_mid_t2v": { + "stages_ms": { + "InputValidationStage": 0.05, + "TextEncodingStage": 1101.99, + "LatentPreparationStage": 0.16, + "HeliosChunkedDenoisingStage": 77728.72, + "HeliosDecodingStage": 661.23, + "per_frame_generation": null + }, + "denoise_step_ms": {}, + "expected_e2e_ms": 79600.62, + "expected_avg_denoise_ms": 0.0, + "expected_median_denoise_ms": 0.0 + }, + "helios_base_t2v": { + "stages_ms": { + "InputValidationStage": 0.04, + "TextEncodingStage": 1102.45, + "LatentPreparationStage": 0.14, + "HeliosChunkedDenoisingStage": 116964.69, + "HeliosDecodingStage": 664.76, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 1893.3, + "1": 1900.93, + "2": 1934.08, + "3": 1897.65, + "4": 1907.59, + "5": 1909.1, + "6": 1911.51, + "7": 1909.25, + "8": 1911.69, + "9": 1911.77, + "10": 1913.35, + "11": 1915.44, + "12": 1912.11, + "13": 1910.08, + "14": 1911.77, + "15": 1908.22, + "16": 1908.83, + "17": 1910.11, + "18": 1908.19, + "19": 1911.99, + "20": 1909.96, + "21": 1910.32, + "22": 1911.76, + "23": 1911.87, + "24": 1908.91, + "25": 1912.41, + "26": 1913.15, + "27": 1908.34, + "28": 1913.21, + "29": 1911.98, + "30": 1912.16, + "31": 1914.17, + "32": 1911.45, + "33": 1912.5, + "34": 1914.48, + "35": 1912.64, + "36": 1912.24, + "37": 1914.48, + "38": 1911.06, + "39": 1915.45, + "40": 1914.0, + "41": 1912.99, + "42": 1913.68, + "43": 1914.09, + "44": 1915.83, + "45": 1913.36, + "46": 1914.84, + "47": 1915.31, + "48": 1915.58, + "49": 1912.63 + }, + "expected_e2e_ms": 118821.41, + "expected_avg_denoise_ms": 1911.64, + "expected_median_denoise_ms": 1912.05 + }, + "helios_mid_t2v": { + "stages_ms": { + "InputValidationStage": 0.09, + "TextEncodingStage": 1102.28, + "LatentPreparationStage": 0.23, + "HeliosChunkedDenoisingStage": 77947.9, + "HeliosDecodingStage": 664.96, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 404.46, + "1": 404.88, + "2": 405.35, + "3": 406.01, + "4": 404.97, + "5": 405.07, + "6": 405.06, + "7": 404.98, + "8": 405.39, + "9": 405.52, + "10": 405.76, + "11": 405.53, + "12": 405.16, + "13": 405.46, + "14": 405.75, + "15": 405.69, + "16": 405.26, + "17": 405.23, + "18": 405.42, + "19": 405.99, + "20": 663.39, + "21": 666.6, + "22": 665.73, + "23": 666.37, + "24": 667.43, + "25": 668.28, + "26": 667.96, + "27": 668.93, + "28": 667.78, + "29": 668.15, + "30": 668.91, + "31": 667.22, + "32": 669.31, + "33": 666.57, + "34": 669.78, + "35": 668.38, + "36": 669.95, + "37": 668.76, + "38": 667.82, + "39": 668.98, + "40": 1891.05, + "41": 1893.52, + "42": 1893.48, + "43": 1892.79, + "44": 1892.03, + "45": 1892.87, + "46": 1895.55, + "47": 1892.19, + "48": 1892.89, + "49": 1892.32, + "50": 1890.25, + "51": 1894.1, + "52": 1890.67, + "53": 1892.09, + "54": 1892.64, + "55": 1891.91, + "56": 1894.27, + "57": 1893.62, + "58": 1892.65, + "59": 1891.9 + }, + "expected_e2e_ms": 79824.32, + "expected_avg_denoise_ms": 988.6, + "expected_median_denoise_ms": 668.05 + }, + "helios_distilled_t2v": { + "stages_ms": { + "InputValidationStage": 0.05, + "TextEncodingStage": 552.02, + "LatentPreparationStage": 0.13, + "HeliosChunkedDenoisingStage": 57879.88, + "HeliosDecodingStage": 663.31, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 207.03, + "1": 204.36, + "2": 203.87, + "3": 204.51, + "4": 206.21, + "5": 205.54, + "6": 205.06, + "7": 205.45, + "8": 205.96, + "9": 205.95, + "10": 205.22, + "11": 204.43, + "12": 205.14, + "13": 205.06, + "14": 205.11, + "15": 206.09, + "16": 205.1, + "17": 204.99, + "18": 204.55, + "19": 205.14, + "20": 337.47, + "21": 337.06, + "22": 337.68, + "23": 336.58, + "24": 335.98, + "25": 335.84, + "26": 336.01, + "27": 335.61, + "28": 335.79, + "29": 335.62, + "30": 336.69, + "31": 335.98, + "32": 336.15, + "33": 336.55, + "34": 336.98, + "35": 337.33, + "36": 336.34, + "37": 335.94, + "38": 336.69, + "39": 336.14, + "40": 954.88, + "41": 956.2, + "42": 953.9, + "43": 953.49, + "44": 957.1, + "45": 956.95, + "46": 955.02, + "47": 954.98, + "48": 956.0, + "49": 956.63, + "50": 958.66, + "51": 957.26, + "52": 956.73, + "53": 955.06, + "54": 957.04, + "55": 958.07, + "56": 958.28, + "57": 957.99, + "58": 957.61, + "59": 956.98 + }, + "expected_e2e_ms": 59168.9, + "expected_avg_denoise_ms": 499.37, + "expected_median_denoise_ms": 336.25 } } } diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 8cea589b6..fe9dac4e8 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -679,6 +679,43 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ ), TI2V_sampling_params, ), + # === Helios T2V === + DiffusionTestCase( + "helios_base_t2v", + DiffusionServerArgs( + model_path="BestWishYsh/Helios-Base", + modality="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + output_size="640x384", + num_frames=33, + ), + ), + DiffusionTestCase( + "helios_mid_t2v", + DiffusionServerArgs( + model_path="BestWishYsh/Helios-Mid", + modality="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + output_size="640x384", + num_frames=33, + ), + ), + DiffusionTestCase( + "helios_distilled_t2v", + DiffusionServerArgs( + model_path="BestWishYsh/Helios-Distilled", + modality="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + output_size="640x384", + num_frames=33, + ), + ), ] # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm.