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

[diffusion] chore: enable sequence shard for wan by default (#19311)

Browse Source
This commit is contained in:
Mick
2026-02-25 18:21:44 +08:00
committed by GitHub
parent 307a2b73ed
commit 9840cd3f68
2 changed files with 67 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -156,7 +156,7 @@ class SamplingParams:
suppress_logs: bool = False
return_file_paths_only: bool = True
enable_sequence_shard: bool = False
enable_sequence_shard: bool | None = None
def _set_output_file_ext(self):
# add extension if needed
@@ -390,10 +390,20 @@ class SamplingParams:
)
logger.warning(error_msg)
pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower()
if "wan" in pipeline_name_lower and (
self.enable_sequence_shard is None or self.enable_sequence_shard
):
self.enable_sequence_shard = True
logger.debug("Automatically enabled enable_sequence_shard")
else:
self.enable_sequence_shard = False
if self.enable_sequence_shard:
self.adjust_frames = False
logger.info(
f"Sequence dimension shard is enabled, disabling frame adjustment"
f"Sequence dimension shard is enabled, disabling frame adjustment for better performance"
)
if pipeline_config.task_type.is_image_gen():
@@ -402,52 +412,66 @@ class SamplingParams:
logger.debug(f"Setting `num_frames` to 1 for image generation model")
self.num_frames = 1
elif self.adjust_frames:
else:
# mandatory frame adjusting logic, mod
# NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.
# If we apply it after, adjust_num_frames might modify the frame count
# and break the divisibility constraint (alignment) required by num_gpus.
original_num_frames = self.num_frames
self.num_frames = server_args.pipeline_config.adjust_num_frames(
self.num_frames
original_num_frames
)
logger.info(
"Adjusting number of frames from %s to %s based on model",
original_num_frames,
self.num_frames,
)
# Adjust number of frames based on number of GPUs for video task
use_temporal_scaling_frames = (
pipeline_config.vae_config.use_temporal_scaling_frames
)
num_frames = self.num_frames
num_gpus = server_args.num_gpus
temporal_scale_factor = (
pipeline_config.vae_config.arch_config.temporal_compression_ratio
)
if use_temporal_scaling_frames:
orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1
if orig_latent_num_frames % server_args.num_gpus != 0:
# Adjust latent frames to be divisible by number of GPUs
if self.num_frames_round_down:
# Ensure we have at least 1 batch per GPU
new_latent_num_frames = (
max(1, (orig_latent_num_frames // num_gpus)) * num_gpus
)
else:
new_latent_num_frames = (
math.ceil(orig_latent_num_frames / num_gpus) * num_gpus
)
if self.adjust_frames:
# Adjust number of frames based on number of GPUs for video task
use_temporal_scaling_frames = (
pipeline_config.vae_config.use_temporal_scaling_frames
)
num_frames = self.num_frames
num_gpus = server_args.num_gpus
temporal_scale_factor = (
pipeline_config.vae_config.arch_config.temporal_compression_ratio
)
if use_temporal_scaling_frames:
# Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
new_num_frames = (
new_latent_num_frames - 1
) * temporal_scale_factor + 1
orig_latent_num_frames = (
num_frames - 1
) // temporal_scale_factor + 1
else:
orig_latent_num_frames = num_frames
logger.info(
"Adjusting number of frames from %s to %s based on number of GPUs (%s)",
self.num_frames,
new_num_frames,
server_args.num_gpus,
)
self.num_frames = new_num_frames
if orig_latent_num_frames % server_args.num_gpus != 0:
# Adjust latent frames to be divisible by number of GPUs
if self.num_frames_round_down:
# Ensure we have at least 1 batch per GPU
new_latent_num_frames = (
max(1, (orig_latent_num_frames // num_gpus)) * num_gpus
)
else:
new_latent_num_frames = (
math.ceil(orig_latent_num_frames / num_gpus) * num_gpus
)
if use_temporal_scaling_frames:
# Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
new_num_frames = (
new_latent_num_frames - 1
) * temporal_scale_factor + 1
else:
new_num_frames = new_latent_num_frames
logger.info(
"Adjusting number of frames from %s to %s based on number of GPUs (%s)",
self.num_frames,
new_num_frames,
server_args.num_gpus,
)
self.num_frames = new_num_frames
if not server_args.comfyui_mode:
self._set_output_file_name()

8
python/sglang/multimodal_gen/configs/sample/wan.py
View File

@@ -239,8 +239,8 @@ class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
guidance_scale_2: float = 3.0 # low_noise
num_inference_steps: int = 40
fps: int = 16
# NOTE(will): default boundary timestep is tracked by PipelineConfig, but
# can be overridden during sampling
num_frames: int = 81
# Wan2.2 T2V A14B supported resolutions
supported_resolutions: list[tuple[int, int]] | None = field(
@@ -259,8 +259,8 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
guidance_scale_2: float = 3.5 # low_noise
num_inference_steps: int = 40
fps: int = 16
# NOTE(will): default boundary timestep is tracked by PipelineConfig, but
# can be overridden during sampling
num_frames: int = 81
# Wan2.2 I2V A14B supported resolutions
supported_resolutions: list[tuple[int, int]] | None = field(