chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

diffusion: fix helios accuracy issue (#20036)

Browse Source
This commit is contained in:
Yuhao Yang
2026-03-15 13:55:51 +08:00
committed by GitHub
parent 6c5bf53a36
commit a6ecf050be
9 changed files with 631 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/configs/pipeline_configs/helios.py
View File

@@ -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:

2
python/sglang/multimodal_gen/configs/sample/sampling_params.py
View File

@@ -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

77
python/sglang/multimodal_gen/runtime/models/dits/helios.py
View File

@@ -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)

5
python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py
View File

@@ -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"]

11
python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py
View File

@@ -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):

68
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_decoding.py Normal file
View File

@@ -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

314
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py
View File

@@ -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

257
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -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
}
}
}

37
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -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.