From e24cb66a104f793aa809aa9dc0b4ad4a65786a7b Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Sat, 22 Aug 2026 11:51:15 +0000 Subject: [PATCH 01/18] Support LLaDA-Image generation and editing on SGLang main Port the LLaDA-Image text-to-image and editing pipeline from LLMInfra sglang PR 6 at dit-main f0548af30 onto current SGLang main. Add native SRT LLaDA2 MoE conditioning, SigVQ and QueryFormer components, the llada2_cfg_flashinfer backend, request validation, lifecycle guards, and SP1 and SP2 Ulysses execution. Support distilled checkpoints with uniform scheduler sigmas, official edit cropping, per-sample generators, and global stochastic noise draws sliced by SP rank. For SP2 conditioning, run each embedded LLaDA2 encoder as a full data replica with singleton attention and MoE tensor groups. Validate the exact topology, keep sparse and shared experts local, and avoid duplicate post-expert reduction. Tests cover request bounds, attention masks, component residency, scheduler parity, parallel group restoration, replicated encoder metadata, communicator widths, and LLaDA2 reduction behavior. --- .../configs/models/dits/__init__.py | 2 + .../configs/models/dits/llada_image.py | 52 + .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/base.py | 29 + .../configs/pipeline_configs/llada_image.py | 376 +++ .../configs/sample/llada_image.py | 20 + .../configs/sample/sampling_params.py | 3 + python/sglang/multimodal_gen/registry.py | 13 + .../runtime/entrypoints/openai/image_api.py | 66 +- .../component_loaders/llada_image_loader.py | 57 + .../runtime/loader/transformer_load_utils.py | 11 +- .../memory_occupation_controller.py | 8 + .../runtime/models/dits/llada_image.py | 2500 +++++++++++++++++ .../scheduling_flow_match_euler_discrete.py | 104 +- .../runtime/pipelines/llada_image.py | 142 + .../runtime/pipelines_core/schedule_batch.py | 1 + .../stages/llada_image_conditioning.py | 427 +++ .../stages/llada_image_source.py | 122 + .../stages/timestep_preparation.py | 8 +- .../runtime/post_training/weights_updater.py | 18 + .../runtime/server_args/server_args.py | 32 + .../multimodal_gen/test/test_llada_image.py | 693 +++++ .../test/unit/test_llada_image_api.py | 168 ++ .../unit/test_llada_image_conditioning.py | 527 ++++ .../unit/test_llada_image_pipeline_config.py | 326 +++ .../test_llada_image_source_conditioning.py | 174 ++ .../test/unit/test_transformer_quant.py | 59 + python/sglang/srt/arg_groups/parallel_hook.py | 124 +- .../layers/attention/attention_registry.py | 14 + .../attention/llada2_attention_utils.py | 31 + .../llada2_cfg_flashinfer_backend.py | 115 + python/sglang/srt/layers/communicator.py | 30 +- .../srt/model_executor/forward_batch_info.py | 4 + python/sglang/srt/models/llada2.py | 20 +- .../sglang/srt/models/llada2_weight_utils.py | 44 + python/sglang/srt/server_args.py | 1 + .../test_layer_communicator_fusion_gate.py | 117 +- .../unit/models/test_llada2_full_replica.py | 125 + .../unit/server_args/test_server_args.py | 54 + test/srt/test_attention_backend_registry.py | 24 + test/srt/test_llada2_attention_utils.py | 38 + .../srt/test_llada2_cfg_flashinfer_backend.py | 160 ++ test/srt/test_llada2_weight_utils.py | 51 + 43 files changed, 6808 insertions(+), 86 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/llada_image.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py create mode 100644 python/sglang/multimodal_gen/configs/sample/llada_image.py create mode 100644 python/sglang/multimodal_gen/runtime/loader/component_loaders/llada_image_loader.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/llada_image.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/llada_image.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_source.py create mode 100644 python/sglang/multimodal_gen/test/test_llada_image.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_api.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py create mode 100644 python/sglang/srt/layers/attention/llada2_attention_utils.py create mode 100644 python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py create mode 100644 python/sglang/srt/models/llada2_weight_utils.py create mode 100644 test/registered/unit/models/test_llada2_full_replica.py create mode 100644 test/srt/test_attention_backend_registry.py create mode 100644 test/srt/test_llada2_attention_utils.py create mode 100644 test/srt/test_llada2_cfg_flashinfer_backend.py create mode 100644 test/srt/test_llada2_weight_utils.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 7b4d0fe3fa21..8549db92c095 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.models.dits.lingbot_world import ( LingBotWorldVideoConfig, ) +from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig from sglang.multimodal_gen.configs.models.dits.longlive2 import LongLive2VideoConfig from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig @@ -32,6 +33,7 @@ "Ideogram4DistilledDiTConfig", "LingBotWorldVideoConfig", "LingBotVideoMoEConfig", + "LLaDAImageDitConfig", "LongLive2VideoConfig", "MiniMaxH3DiTConfig", "WanVideoConfig", diff --git a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py new file mode 100644 index 000000000000..1ef683e1dfe1 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + + +@dataclass +class LLaDAImageArchConfig(DiTArchConfig): + all_patch_size: tuple[int, ...] = (1,) + all_f_patch_size: tuple[int, ...] = (1,) + axes_dims: tuple[int, ...] = (32, 48, 48) + axes_lens: tuple[int, ...] = (32768, 1024, 1024) + cap_feat_dim: int = 2560 + dim: int = 3840 + in_channels: int = 128 + n_heads: int = 30 + n_layers: int = 30 + n_refiner_layers: int = 2 + norm_eps: float = 1e-5 + qk_norm: bool = True + rope_theta: float = 256.0 + semantic_feat_dim: int = 4096 + t_scale: float = 1000.0 + _supported_attention_backends: set[AttentionBackendEnum] = field( + default_factory=lambda: { + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + } + ) + param_names_mapping: dict = field( + default_factory=lambda: { + r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3), + r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3), + r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3), + r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2), + r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2), + } + ) + + def __post_init__(self) -> None: + super().__post_init__() + self.hidden_size = self.dim + self.num_attention_heads = self.n_heads + self.num_channels_latents = self.in_channels + + +@dataclass +class LLaDAImageDitConfig(DiTConfig): + arch_config: LLaDAImageArchConfig = field(default_factory=LLaDAImageArchConfig) + prefix: str = "llada_image" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 9f35fa9c79a4..9f6be9d780a6 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -39,6 +39,9 @@ LingBotWorldCausalDMDConfig, LingBotWorldV2CausalDMDConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( LTX2PipelineConfig, LTX23PipelineConfig, @@ -76,6 +79,7 @@ "Hunyuan3D2PipelineConfig", "Ideogram4PipelineConfig", "Ideogram4DistilledPipelineConfig", + "LLaDAImagePipelineConfig", "FluxPipelineConfig", "Flux2PipelineConfig", "Flux2KleinPipelineConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index e86ad7bde729..3b42943a056f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -436,6 +436,35 @@ def supports_sequential_multi_output_inference(self): """Return whether one request's outputs run through DiT/VAE sequentially.""" return False + def validate_num_outputs_per_prompt( + self, num_outputs_per_prompt: int, server_args + ) -> None: + """Validate a request's per-prompt output count for this pipeline.""" + del num_outputs_per_prompt, server_args + + def supports_memory_release(self) -> bool: + """Return whether release_memory_occupation can offload this pipeline.""" + return True + + def supports_hot_weight_updates(self) -> bool: + """Return whether runtime weight updates cover this pipeline fully.""" + return True + + def validate_request_sampling_params(self, sampling_params, server_args) -> None: + """Validate a request's sampling params for this pipeline. + + May also run before full request materialization with a lightweight + object that only carries width, height, max_sequence_length, + num_outputs_per_prompt, diffusers_kwargs, enable_cache_dit, + cache_dit_params, attention_backend_override, and cfg_gate_step, + so overrides must not require other sampling fields. + """ + del sampling_params, server_args + + def validate_edit_source_count(self, source_count: int, server_args) -> None: + """Validate the number of source images before they are materialized.""" + del source_count, server_args + def estimate_request_cost(self, batch) -> float: """Return the relative cost used for batching admission caps. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py new file mode 100644 index 000000000000..e57c9177522d --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import numpy as np +import torch + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig +from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + SpatialImagePipelineConfig, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +@dataclass +class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): + task_type: ModelTaskType = ModelTaskType.TI2I + should_use_guidance: bool = False + + dit_config: LLaDAImageDitConfig = field(default_factory=LLaDAImageDitConfig) + dit_precision: str = "bf16" + vae_config: Flux2VAEConfig = field(default_factory=Flux2VAEConfig) + vae_precision: str = "bf16" + vae_tiling: bool = False + vae_sp: bool = False + + text_encoder_configs: tuple = () + text_encoder_precisions: tuple[str, ...] = () + preprocess_text_funcs: tuple = () + postprocess_text_funcs: tuple = () + + latent_scale_factor: int = 16 + text_encoder_mem_fraction_static: float = 0.35 + max_request_pixel_area: int = 2048 * 2048 + # The embedded text worker admits 8192 prefill tokens shared by the two + # CFG sequences and each sequence appends 256 query tokens. + max_request_text_tokens: int = 3584 + max_request_total_pixel_area: int = 10 * 1024 * 1024 + + def prepare_sigmas(self, sigmas, num_inference_steps): + if sigmas is not None: + return sigmas + schedule = np.linspace(0.001, 1.0, num_inference_steps + 1)[:-1] + schedule = (1 - (1 - schedule**1.17) ** 0.8) ** 1.1 + return (1 - schedule).tolist() + + def prepare_latent_shape(self, batch, batch_size, num_frames): + del num_frames + if batch.height % self.latent_scale_factor != 0: + raise ValueError("LLaDA-Image height must be divisible by 16") + if batch.width % self.latent_scale_factor != 0: + raise ValueError("LLaDA-Image width must be divisible by 16") + return ( + batch_size, + self.dit_config.num_channels_latents, + batch.height // self.latent_scale_factor, + batch.width // self.latent_scale_factor, + ) + + def prepare_calculated_size(self, image): + del image + return None + + def calculate_condition_image_size(self, image, width, height): + del image, width, height + return None + + def shard_latents_for_sp(self, batch, latents): + sp_degree = get_sp_world_size() + if latents.dim() == 4 and latents.shape[2] % sp_degree != 0: + raise ValueError( + f"LLaDA-Image latent height {latents.shape[2]} must be divisible " + f"by SP degree {sp_degree}. Choose a compatible output height" + ) + return super().shard_latents_for_sp(batch, latents) + + def validate_server_args(self, server_args) -> None: + super().validate_server_args(server_args) + if envs.SGLANG_CACHE_DIT_ENABLED: + raise ValueError("LLaDA-Image does not support SGLANG_CACHE_DIT_ENABLED") + # Replicated condition suffixes are currently de-duplicated only by the + # Ulysses attention path. + if server_args.ring_degree != 1: + raise ValueError("LLaDA-Image sequence parallelism requires ring_degree=1") + if server_args.ulysses_degree != server_args.sp_degree: + raise ValueError( + "LLaDA-Image sequence parallelism requires " + "ulysses_degree == sp_degree" + ) + if server_args.sp_degree not in (1, 2): + raise ValueError("LLaDA-Image currently supports only SP degrees 1 and 2") + if ( + server_args.tp_size != 1 + or server_args.dp_size != 1 + or server_args.cfg_parallel_degree != 1 + ): + raise ValueError("LLaDA-Image requires diffusion parallelism TP=DP=CFG=1") + if server_args.num_gpus != server_args.sp_degree: + raise ValueError( + "LLaDA-Image requires num_gpus == sp_degree so every GPU belongs " + "to the sequence-parallel group" + ) + if server_args.text_encoder_cpu_offload: + if server_args.is_arg_explicitly_set("text_encoder_cpu_offload"): + raise ValueError( + "LLaDA-Image runs its text encoder inside an embedded srt " + "worker that the component residency manager cannot " + "offload. Remove --text-encoder-cpu-offload" + ) + # Auto-tuned default, corrected here because the tuner does not + # know about the embedded worker. + server_args.text_encoder_cpu_offload = False + logger.info( + "LLaDA-Image keeps the embedded text encoder resident and " + "ignores the auto-tuned text_encoder_cpu_offload" + ) + if server_args.residency_mode("text_encoder") != "resident": + if ( + server_args.explicit_residency_mode("text_encoder") is not None + or server_args.is_arg_explicitly_set("layerwise_offload_components") + or server_args.is_arg_explicitly_set("cpu_offload_components") + ): + raise ValueError( + "LLaDA-Image requires its embedded text encoder to " + "remain resident" + ) + # Strip the encoder from auto-tuned offload selections the tuner + # applies without knowing about the embedded worker. + if server_args.layerwise_offload_components is not None: + server_args.layerwise_offload_components = [ + name + for name in server_args.layerwise_offload_components + if name != "text_encoder" + ] + if server_args.cpu_offload_components is not None: + server_args.cpu_offload_components = [ + name + for name in server_args.cpu_offload_components + if name != "text_encoder" + ] + logger.info( + "LLaDA-Image keeps the embedded text encoder resident and " + "removed it from auto-tuned offload selections" + ) + if server_args.residency_mode("text_encoder") != "resident": + raise ValueError( + "LLaDA-Image requires its embedded text encoder to " + "remain resident" + ) + + area_override = server_args.llada_image_max_pixel_area + if area_override is not None and area_override <= 0: + raise ValueError("llada_image_max_pixel_area must be positive") + total_area_override = server_args.llada_image_max_total_pixel_area + if total_area_override is not None and total_area_override <= 0: + raise ValueError("llada_image_max_total_pixel_area must be positive") + tokens_override = server_args.llada_image_max_text_tokens + if tokens_override is not None and not ( + 0 < tokens_override <= self.max_request_text_tokens + ): + raise ValueError( + "llada_image_max_text_tokens must be positive and at most " + f"{self.max_request_text_tokens}, the embedded worker budget" + ) + + # The embedded srt worker and its KV pool are invisible to generic module + # discovery, so partial sleep or update would misreport. + def supports_memory_release(self) -> bool: + return False + + def supports_hot_weight_updates(self) -> bool: + return False + + def validate_num_outputs_per_prompt( + self, num_outputs_per_prompt: int, server_args + ) -> None: + if server_args.sp_degree > 1 and num_outputs_per_prompt != 1: + raise ValueError( + "LLaDA-Image sequence parallelism supports only n=1. " + "submit separate requests for multiple images" + ) + + def validate_edit_source_count(self, source_count: int, server_args) -> None: + del server_args + if source_count != 1: + raise ValueError("LLaDA-Image editing requires exactly one source image") + + def validate_request_sampling_params(self, sampling_params, server_args) -> None: + enable_cache_dit = sampling_params.enable_cache_dit + if enable_cache_dit is not None and enable_cache_dit is not False: + raise ValueError("LLaDA-Image does not support enable_cache_dit") + if sampling_params.cache_dit_params is not None: + raise ValueError("LLaDA-Image does not support cache_dit_params") + + attention_backend = sampling_params.attention_backend_override + supported_attention_backends = {"FA", "TORCH_SDPA"} + if attention_backend is not None and ( + not isinstance(attention_backend, str) + or attention_backend.upper() not in supported_attention_backends + ): + raise ValueError( + "LLaDA-Image attention_backend_override must be fa or torch_sdpa" + ) + + cfg_gate_step = sampling_params.cfg_gate_step + if cfg_gate_step is not None and ( + isinstance(cfg_gate_step, bool) + or not isinstance(cfg_gate_step, (int, float)) + or not 0.0 <= cfg_gate_step <= 1.0 + or not np.isfinite(cfg_gate_step) + ): + raise ValueError("LLaDA-Image cfg_gate_step must be between 0.0 and 1.0") + + width = int(sampling_params.width or 0) + height = int(sampling_params.height or 0) + if width <= 0 or height <= 0: + raise ValueError("LLaDA-Image requires positive width and height") + if width % self.latent_scale_factor or height % self.latent_scale_factor: + raise ValueError( + f"LLaDA-Image width and height must be divisible by " + f"{self.latent_scale_factor}" + ) + sp_height_multiple = self.latent_scale_factor * server_args.sp_degree + if height % sp_height_multiple: + raise ValueError( + f"LLaDA-Image output height must be divisible by " + f"{sp_height_multiple} at SP degree {server_args.sp_degree}" + ) + area_cap = ( + server_args.llada_image_max_pixel_area + if server_args.llada_image_max_pixel_area is not None + else self.max_request_pixel_area + ) + if width * height > area_cap: + raise ValueError( + f"LLaDA-Image output area {width}x{height} exceeds the " + f"supported maximum of {area_cap} pixels" + ) + total_area_cap = ( + server_args.llada_image_max_total_pixel_area + if server_args.llada_image_max_total_pixel_area is not None + else self.max_request_total_pixel_area + ) + num_outputs = sampling_params.num_outputs_per_prompt + if num_outputs is not None: + total_area = width * height * max(1, int(num_outputs)) + if total_area > total_area_cap: + raise ValueError( + f"LLaDA-Image total output area {total_area} pixels across " + f"{num_outputs} outputs exceeds the supported maximum of " + f"{total_area_cap} pixels per request" + ) + tokens_cap = ( + server_args.llada_image_max_text_tokens + if server_args.llada_image_max_text_tokens is not None + else self.max_request_text_tokens + ) + self._validate_max_sequence_length( + sampling_params.max_sequence_length, tokens_cap + ) + extra_kwargs = sampling_params.diffusers_kwargs + if isinstance(extra_kwargs, dict) and "max_sequence_length" in extra_kwargs: + # prepare_request later overrides the request value from this channel. + self._validate_max_sequence_length( + extra_kwargs["max_sequence_length"], tokens_cap + ) + + @staticmethod + def _validate_max_sequence_length(max_sequence_length, tokens_cap: int) -> None: + if max_sequence_length is None: + return + if ( + isinstance(max_sequence_length, bool) + or not isinstance(max_sequence_length, int) + or max_sequence_length <= 0 + ): + raise ValueError("max_sequence_length must be a positive integer") + if max_sequence_length > tokens_cap: + raise ValueError( + f"max_sequence_length {max_sequence_length} exceeds the " + f"embedded text encoder budget of {tokens_cap} tokens" + ) + + @staticmethod + def _prepare_condition_list(values, name: str, expected_size: int, device, dtype): + if values is None: + return None + if len(values) != expected_size: + raise ValueError( + f"LLaDA-Image {name} has {len(values)} entries, " + f"expected {expected_size}" + ) + return [value.to(device=device, dtype=dtype) for value in values] + + def _prepare_source_latents(self, batch, device, dtype): + source_latents = self._prepare_condition_list( + batch.source_latents, + "source_latents", + batch.batch_size, + device, + dtype, + ) + if source_latents and getattr(batch, "did_sp_shard_latents", False): + source_latents = [ + self.shard_latents_for_sp(batch, latent)[0] for latent in source_latents + ] + return source_latents + + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + del rotary_emb + image_embeds = batch.image_embeds + if image_embeds: + image_embeds = self._prepare_condition_list( + image_embeds, + "image_embeds", + batch.batch_size, + device, + dtype, + ) + source_latents = self._prepare_source_latents(batch, device, dtype) + return { + "encoder_hidden_states_image": image_embeds, + "source_latents": source_latents, + } + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + del rotary_emb + image_embeds = batch.image_embeds + if image_embeds: + image_embeds = self._prepare_condition_list( + image_embeds, + "image_embeds", + batch.batch_size, + device, + dtype, + ) + empty_image_embed = image_embeds[0].new_zeros( + (0, image_embeds[0].shape[-1]) + ) + image_embeds = [empty_image_embed] * batch.batch_size + source_latents = self._prepare_source_latents(batch, device, dtype) + return { + "encoder_hidden_states_image": image_embeds, + "source_latents": source_latents, + } + + def get_decode_scale_and_shift(self, device, dtype, vae): + del device, dtype, vae + return 1.0, None + + def preprocess_decoding(self, latents, server_args=None, vae=None): + del server_args + if vae is None or not hasattr(vae, "bn"): + raise ValueError("LLaDA-Image decoding requires the Flux2 VAE BN state") + vae_parameter = next(vae.parameters()) + latents = latents.to(device=vae_parameter.device, dtype=vae_parameter.dtype) + vae_config = getattr(vae.config, "arch_config", vae.config) + latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) + latent_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae_config.batch_norm_eps + ).to(latents) + latents = latents * latent_std + latent_mean + batch_size, channels, height, width = latents.shape + latents = latents.reshape(batch_size, channels // 4, 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + return latents.reshape(batch_size, channels // 4, height * 2, width * 2) + + def post_denoising_loop(self, latents, batch): + del batch + return latents diff --git a/python/sglang/multimodal_gen/configs/sample/llada_image.py b/python/sglang/multimodal_gen/configs/sample/llada_image.py new file mode 100644 index 000000000000..594e4ea89649 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) + + +@dataclass +class LLaDAImageSamplingParams(SamplingParams): + data_type: DataType = DataType.IMAGE + negative_prompt: str | None = None + num_inference_steps: int = 20 + num_frames: int = 1 + height: int = 1024 + width: int = 1024 + guidance_scale: float = 4.5 + max_sequence_length: int = 2048 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 9dfc3462e8a0..89e17fee987e 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -902,6 +902,9 @@ def from_user_sampling_params_args( sampling_params._adjust(server_args) sampling_params._validate_with_pipeline_config(server_args.pipeline_config) + server_args.pipeline_config.validate_request_sampling_params( + sampling_params, server_args + ) return sampling_params diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index a31842f165cd..db3f67cc7e5c 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -38,6 +38,7 @@ HunyuanConfig, LingBotWorldCausalDMDConfig, LingBotWorldV2CausalDMDConfig, + LLaDAImagePipelineConfig, MiniMaxH3PipelineConfig, WanI2V480PConfig, WanI2V720PConfig, @@ -151,6 +152,7 @@ from sglang.multimodal_gen.configs.sample.lingbot_world import ( LingBotWorldSamplingParams, ) +from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams from sglang.multimodal_gen.configs.sample.longcat_image import ( LongCatImageEditSamplingParams, LongCatImageEditTurboSamplingParams, @@ -1065,7 +1067,18 @@ def _register_configs(): hf_model_paths=["krea/Krea-2"], model_detectors=[lambda hf_id: "krea-2" in hf_id.lower()], ) + # LLaDA-Image + register_configs( + sampling_param_cls=LLaDAImageSamplingParams, + pipeline_config_cls=LLaDAImagePipelineConfig, + hf_model_paths=["inclusionAI/LLaDA-Image"], + model_detectors=[ + lambda hf_id: "lladaimagepipeline" in hf_id.lower() + or "llada-image" in hf_id.lower() + ], + ) # Qwen-Image + register_configs( sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImagePipelineConfig, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 88b40d0e7c6f..05412762ba79 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -5,6 +5,7 @@ import contextlib import os import time +from types import SimpleNamespace from typing import Any, List, Optional from fastapi import ( @@ -48,6 +49,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args +from sglang.multimodal_gen.runtime.utils.common import parse_size from sglang.srt.observability.trace import extract_trace_headers router = APIRouter(prefix="/v1/images", tags=["images"]) @@ -83,6 +85,50 @@ def _runtime_sampling_quality(quality: str | None) -> str | None: return None if quality in (None, "auto") else quality +def _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt: int) -> None: + try: + server_args.pipeline_config.validate_num_outputs_per_prompt( + num_outputs_per_prompt, server_args + ) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + +def _build_sampling_params_or_400(request_id: str, **kwargs): + try: + return build_sampling_params(request_id, **kwargs) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + +def _early_validate_edit_bounds( + server_args, size: Optional[str], num_outputs_per_prompt: int = 1 +) -> None: + """Best-effort size rejection before any source image is handled.""" + if not size: + return + width, height = parse_size(size) + if not width or not height: + return + lightweight = SimpleNamespace( + width=width, + height=height, + max_sequence_length=None, + num_outputs_per_prompt=num_outputs_per_prompt, + diffusers_kwargs=None, + enable_cache_dit=None, + cache_dit_params=None, + attention_backend_override=None, + cfg_gate_step=None, + ) + try: + server_args.pipeline_config.validate_request_sampling_params( + lightweight, server_args + ) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + + def _read_b64_for_paths(paths: list[str]) -> list[str]: """Read and base64-encode each file. Must be called before cloud upload deletes them.""" result = [] @@ -266,6 +312,8 @@ async def generations( ): request_id = generate_request_id() server_args = get_global_server_args() + num_outputs_per_prompt = max(1, min(int(request.n or 1), 10)) + _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt) sampling_params_cls = resolve_sampling_params_cls(server_args) model_kwargs = _image_request_model_kwargs(request, sampling_params_cls) output_format = ( @@ -276,13 +324,13 @@ async def generations( ext = choose_output_image_ext(output_format, request.background) with temp_dir_if_disabled(server_args.output_path) as output_dir: - sampling = build_sampling_params( + sampling = _build_sampling_params_or_400( request_id, prompt=request.prompt, size=request.size, width=request.width, height=request.height, - num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)), + num_outputs_per_prompt=num_outputs_per_prompt, output_file_name=f"{request_id}.{ext}", output_path=output_dir, num_frames=1, @@ -437,7 +485,17 @@ async def edits( status_code=422, detail="Field 'image' or 'url' is required" ) + num_outputs_per_prompt = max(1, min(int(n or 1), 10)) + _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt) + _early_validate_edit_bounds(server_args, size, num_outputs_per_prompt) + image_list = merge_image_input_list(images, urls) + try: + server_args.pipeline_config.validate_edit_source_count( + len(image_list), server_args + ) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error with contextlib.ExitStack() as stack: uploads_dir = stack.enter_context( @@ -462,11 +520,11 @@ async def edits( ) ext = choose_output_image_ext(output_format, background) - sampling = build_sampling_params( + sampling = _build_sampling_params_or_400( request_id, prompt=prompt, size=size, - num_outputs_per_prompt=max(1, min(int(n or 1), 10)), + num_outputs_per_prompt=num_outputs_per_prompt, output_file_name=f"{request_id}.{ext}", output_path=output_dir, image_path=input_paths, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/llada_image_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/llada_image_loader.py new file mode 100644 index 000000000000..da9630d33f13 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/llada_image_loader.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import ClassVar + +from safetensors.torch import load_file as safetensors_load_file + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentLoader, +) +from sglang.multimodal_gen.runtime.loader.utils import ( + _list_safetensors_files, + set_default_torch_dtype, +) +from sglang.multimodal_gen.runtime.models.registry import ModelRegistry +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + get_diffusers_component_config, +) +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + + +class LLaDAImageAuxiliaryLoader(ComponentLoader): + """Load the QueryFormer, text projection, and SigVQ modules.""" + + component_names: ClassVar[list[str]] = ["queryformer", "text_projection", "sigvq"] + expected_library = "diffusers" + + def load_customized( + self, + component_model_path: str, + server_args: ServerArgs, + component_name: str, + ): + config = get_diffusers_component_config(component_path=component_model_path) + class_name = config.pop("_class_name", None) + config.pop("_diffusers_version", None) + if class_name is None: + raise ValueError( + f"{component_name} config does not contain a _class_name attribute" + ) + + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] + with set_default_torch_dtype(dtype): + model = model_cls(**config) + + safetensors_list = _list_safetensors_files(component_model_path) + if len(safetensors_list) != 1: + raise ValueError( + f"Expected one safetensors file for {component_name}, " + f"found {len(safetensors_list)}" + ) + state_dict = safetensors_load_file(safetensors_list[0]) + model.load_state_dict(state_dict, strict=True) + server_args.model_paths[component_name] = component_model_path + return model.to(device=get_local_torch_device(), dtype=dtype).eval() diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 78b6df23da82..4a75b7634f1f 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -381,7 +381,7 @@ def prepare(self) -> None: class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter): - """Adapter for diffusion ModelOpt FP8 checkpoints.""" + """Disable whole-DiT CPU offload for FP8 modes that require CUDA setup.""" def __init__( self, @@ -403,7 +403,11 @@ def _maybe_disable_incompatible_dit_offload_modes( if quant_config is None: return - if _get_quant_config_name(quant_config) != "modelopt_fp8": + quant_name = _get_quant_config_name(quant_config) + is_online_fp8 = quant_name == "fp8" and not getattr( + quant_config, "is_checkpoint_fp8_serialized", False + ) + if quant_name != "modelopt_fp8" and not is_online_fp8: return component_offload = _uses_component_offload( @@ -425,8 +429,9 @@ def _maybe_disable_incompatible_dit_offload_modes( feature_name="ModelOpt FP8 diffusion checkpoints", ) logger.warning( - "ModelOpt FP8 diffusion checkpoints keep the DiT resident instead " + "%s diffusion quantization keeps the DiT resident instead " "of using component offload. Layerwise offload remains supported.", + quant_name, ) def prepare(self) -> None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py index 612445b8e669..48b579db5c03 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py @@ -188,6 +188,14 @@ def release_memory_occupation(self) -> dict[str, bool | str]: success=False, message="pipeline not initialized", ) + if not self.pipeline.server_args.pipeline_config.supports_memory_release(): + return self._memory_occupation_result( + success=False, + message=( + "pipeline holds components outside generic module " + "discovery and does not support memory release" + ), + ) self._sleep_restore_map = self._offload_active_modules_to_cpu() self._sleeping = True diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py new file mode 100644 index 000000000000..2c9c4d8b0445 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -0,0 +1,2500 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward +from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import RMSNorm +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import maybe_allow_in_graph +from torch import nn +from torch.nn.utils.rnn import pad_sequence + +from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig +from sglang.multimodal_gen.runtime.distributed import ( + get_sp_parallel_rank, + get_sp_world_size, + get_tp_world_size, +) +from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm as SGLangRMSNorm +from sglang.multimodal_gen.runtime.layers.layernorm import ( + apply_qk_norm_with_optional_rope, +) +from sglang.multimodal_gen.runtime.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + apply_flashinfer_rope_qk_inplace, +) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +ADALN_EMBED_DIM = 256 + + +def apply_rmsnorm_tanh_mul_add( + x: torch.Tensor, + gate: torch.Tensor, + residual: torch.Tensor, + norm: SGLangRMSNorm, +) -> torch.Tensor: + """residual + tanh(gate) * rmsnorm(x), with fp32-upcast norm semantics.""" + return residual + torch.tanh(gate) * norm(x) + + +SEQUENCE_MULTIPLE = 32 +LLADA_IMAGE_ATTENTION_BACKENDS = { + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, +} + + +class LLaDAImageRMSNorm(SGLangRMSNorm): + """RMSNorm with a synthesized unit weight for affine-free checkpoints.""" + + def __init__(self, hidden_size: int, eps: float): + super().__init__(hidden_size, eps=eps) + self.weight.missing_param_init = "ones" + + +@dataclass +class _LLaDAImageSequence: + features: list[torch.Tensor] + position_ids: list[torch.Tensor] + padding_masks: list[torch.Tensor] + noise_masks: list[list[int]] | None = None + + +class LLaDAImageTimestepEmbedder(nn.Module): + def __init__( + self, + output_dim: int, + hidden_dim: int = 1024, + frequency_embedding_dim: int = 256, + ): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_dim, hidden_dim, bias=True), + nn.SiLU(), + nn.Linear(hidden_dim, output_dim, bias=True), + ) + self.frequency_embedding_dim = frequency_embedding_dim + + def forward( + self, timestep: torch.Tensor, hidden_dtype: torch.dtype + ) -> torch.Tensor: + half_dim = self.frequency_embedding_dim // 2 + frequencies = torch.exp( + -math.log(10000) + * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) + / half_dim + ) + arguments = timestep[:, None].float() * frequencies[None] + embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1) + if self.frequency_embedding_dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + return self.mlp(embedding.to(dtype=hidden_dtype)) + + +class LLaDAImageRopeEmbedder(nn.Module): + def __init__( + self, theta: float, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...] + ): + super().__init__() + self.theta = theta + self.axes_dims = axes_dims + self.axes_lens = axes_lens + self.freqs_cis = None + + def _create_frequencies(self, device: torch.device) -> list[torch.Tensor]: + frequencies = [] + for axis_dim, axis_len in zip(self.axes_dims, self.axes_lens): + inverse_frequencies = 1.0 / ( + self.theta + ** ( + torch.arange(0, axis_dim, 2, dtype=torch.float32, device=device) + / axis_dim + ) + ) + positions = torch.arange(axis_len, dtype=torch.float32, device=device) + angles = torch.outer(positions, inverse_frequencies) + frequencies.append(torch.complex(torch.cos(angles), torch.sin(angles))) + return frequencies + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + if self.freqs_cis is None or self.freqs_cis[0].device != position_ids.device: + self.freqs_cis = self._create_frequencies(position_ids.device) + + frequencies = [] + for axis, axis_frequencies in enumerate(self.freqs_cis): + frequencies.append(axis_frequencies[position_ids[:, axis]]) + return torch.cat(frequencies, dim=-1) + + +class LLaDAImageAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "LLaDAImageAttention", + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + freqs_cis: torch.Tensor | None = None, + num_replicated_suffix: int = 0, + skip_sequence_parallel_override: bool = False, + ) -> torch.Tensor: + qkv, _ = attn.to_qkv(hidden_states) + query, key, value = qkv.split( + [attn.inner_dim, attn.inner_dim, attn.inner_dim], dim=-1 + ) + query = query.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() + key = key.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() + value = value.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() + + if freqs_cis is not None: + cos_sin_cache = torch.cat( + [freqs_cis.real.float(), freqs_cis.imag.float()], dim=-1 + ).reshape(-1, attn.head_dim) + positions = torch.arange( + cos_sin_cache.shape[0], device=query.device, dtype=torch.long + ) + if attn.norm_q is not None: + query, key = apply_qk_norm_with_optional_rope( + q=query, + k=key, + q_norm=attn.norm_q, + k_norm=attn.norm_k, + head_dim=attn.head_dim, + cos_sin_cache=cos_sin_cache.contiguous(), + positions=positions, + is_neox=False, + allow_inplace=True, + ) + else: + query, key = apply_flashinfer_rope_qk_inplace( + query, + key, + cos_sin_cache.contiguous(), + head_size=attn.head_dim, + positions=positions, + is_neox=False, + ) + elif attn.norm_q is not None: + query, key = apply_qk_norm_with_optional_rope( + q=query, + k=key, + q_norm=attn.norm_q, + k_norm=attn.norm_k, + head_dim=attn.head_dim, + allow_inplace=True, + ) + + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + + hidden_states = attn.sgl_attention( + query, + key, + value, + attn_mask=attention_mask, + num_replicated_suffix=num_replicated_suffix, + skip_sequence_parallel_override=skip_sequence_parallel_override, + ) + hidden_states = hidden_states.flatten(2, 3) + hidden_states, _ = attn.to_out[0](hidden_states) + return hidden_states + + +class LLaDAImageAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = LLaDAImageAttnProcessor + _available_processors: ClassVar[list[type]] = [LLaDAImageAttnProcessor] + _supports_qkv_fusion = False + + def __init__( + self, + dim: int, + num_heads: int, + norm_eps: float, + qk_norm: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + tp_size = get_tp_world_size() + if num_heads % tp_size != 0: + raise ValueError( + f"num_heads ({num_heads}) must be divisible by TP size ({tp_size})" + ) + self.heads = num_heads + self.local_heads = num_heads // tp_size + self.head_dim = dim // num_heads + self.inner_dim = self.local_heads * self.head_dim + self.to_qkv = MergedColumnParallelLinear( + dim, + [dim, dim, dim], + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", + ) + self.sgl_attention = USPAttention( + num_heads=self.local_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=LLADA_IMAGE_ATTENTION_BACKENDS, + prefix="llada_image", + ) + self.norm_q = ( + LLaDAImageRMSNorm(self.head_dim, eps=norm_eps) if qk_norm else None + ) + self.norm_k = ( + LLaDAImageRMSNorm(self.head_dim, eps=norm_eps) if qk_norm else None + ) + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + dim, + dim, + bias=False, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", + ), + nn.Dropout(0.0), + ] + ) + self.set_processor(self._default_processor_cls()) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + freqs_cis: torch.Tensor, + num_replicated_suffix: int = 0, + skip_sequence_parallel_override: bool = False, + ) -> torch.Tensor: + return self.processor( + self, + hidden_states, + attention_mask, + freqs_cis, + num_replicated_suffix, + skip_sequence_parallel_override, + ) + + +class LLaDAImageFeedForward(nn.Module): + def __init__( + self, + dim: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + hidden_dim = int(dim / 3 * 8) + self.w13 = MergedColumnParallelLinear( + dim, + [hidden_dim, hidden_dim], + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.w13", + ) + self.w2 = RowParallelLinear( + hidden_dim, + dim, + bias=False, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.w2", + ) + self.act = SiluAndMul() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.w13(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.w2(hidden_states) + return hidden_states + + +def _select_per_token( + noisy_value: torch.Tensor, + clean_value: torch.Tensor, + noise_mask: torch.Tensor, + sequence_length: int, +) -> torch.Tensor: + noise_mask = noise_mask.unsqueeze(-1) + return torch.where( + noise_mask == 1, + noisy_value.unsqueeze(1).expand(-1, sequence_length, -1), + clean_value.unsqueeze(1).expand(-1, sequence_length, -1), + ) + + +@maybe_allow_in_graph +class LLaDAImageTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + norm_eps: float, + qk_norm: bool, + modulation: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.modulation = modulation + self.attention = LLaDAImageAttention( + dim, + num_heads, + norm_eps, + qk_norm, + quant_config=quant_config, + prefix=f"{prefix}.attention", + ) + self.feed_forward = LLaDAImageFeedForward( + dim, quant_config=quant_config, prefix=f"{prefix}.feed_forward" + ) + self.attention_norm1 = LLaDAImageRMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = LLaDAImageRMSNorm(dim, eps=norm_eps) + self.attention_norm2 = LLaDAImageRMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = LLaDAImageRMSNorm(dim, eps=norm_eps) + if modulation: + self.adaLN_modulation = nn.Sequential( + nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True) + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + freqs_cis: torch.Tensor, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, + num_replicated_suffix: int = 0, + skip_sequence_parallel_override: bool = False, + ) -> torch.Tensor: + if self.modulation: + sequence_length = hidden_states.shape[1] + if noise_mask is None: + scale_msa, gate_msa, scale_mlp, gate_mlp = ( + self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) + ) + scale_msa = 1.0 + scale_msa + else: + noisy_modulation = self.adaLN_modulation(adaln_noisy) + clean_modulation = self.adaLN_modulation(adaln_clean) + noisy_scale_msa, noisy_gate_msa, noisy_scale_mlp, noisy_gate_mlp = ( + noisy_modulation.chunk(4, dim=1) + ) + clean_scale_msa, clean_gate_msa, clean_scale_mlp, clean_gate_mlp = ( + clean_modulation.chunk(4, dim=1) + ) + scale_msa = _select_per_token( + 1.0 + noisy_scale_msa, + 1.0 + clean_scale_msa, + noise_mask, + sequence_length, + ) + scale_mlp = _select_per_token( + noisy_scale_mlp, + clean_scale_mlp, + noise_mask, + sequence_length, + ) + gate_msa = _select_per_token( + noisy_gate_msa, + clean_gate_msa, + noise_mask, + sequence_length, + ) + gate_mlp = _select_per_token( + noisy_gate_mlp, + clean_gate_mlp, + noise_mask, + sequence_length, + ) + + attention_output = self.attention( + self.attention_norm1(hidden_states) * scale_msa, + attention_mask, + freqs_cis, + num_replicated_suffix, + skip_sequence_parallel_override, + ) + hidden_states = apply_rmsnorm_tanh_mul_add( + attention_output, gate_msa, hidden_states, self.attention_norm2 + ) + ffn_input = self.ffn_norm1(hidden_states) * (1.0 + scale_mlp) + + ffn_output = self.feed_forward(ffn_input) + hidden_states = apply_rmsnorm_tanh_mul_add( + ffn_output, gate_mlp, hidden_states, self.ffn_norm2 + ) + else: + attention_output = self.attention( + self.attention_norm1(hidden_states), + attention_mask, + freqs_cis, + num_replicated_suffix, + skip_sequence_parallel_override, + ) + hidden_states = hidden_states + self.attention_norm2(attention_output) + hidden_states = hidden_states + self.ffn_norm2( + self.feed_forward(self.ffn_norm1(hidden_states)) + ) + return hidden_states + + +class LLaDAImageFinalLayer(nn.Module): + def __init__(self, dim: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(dim, out_channels, bias=True) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(min(dim, ADALN_EMBED_DIM), dim, bias=True), + ) + + def forward( + self, + hidden_states: torch.Tensor, + adaln_input: torch.Tensor | None = None, + noise_mask: torch.Tensor | None = None, + adaln_noisy: torch.Tensor | None = None, + adaln_clean: torch.Tensor | None = None, + ) -> torch.Tensor: + if noise_mask is None: + scale = 1.0 + self.adaLN_modulation(adaln_input) + scale = scale.unsqueeze(1) + else: + sequence_length = hidden_states.shape[1] + noisy_scale = 1.0 + self.adaLN_modulation(adaln_noisy) + clean_scale = 1.0 + self.adaLN_modulation(adaln_clean) + scale = _select_per_token( + noisy_scale, clean_scale, noise_mask, sequence_length + ) + hidden_states = self.norm_final(hidden_states) * scale + return self.linear(hidden_states) + + +class _LLaDAImageTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + The denoising transformer used by LLaDAImage for text-to-image generation and single-image editing. + + This component consumes caption features that have already passed through LLaDAImage's QueryFormer, connector, and + projector. For editing, it additionally consumes GLM/SigVQ features and source-image latents. + + Args: + all_patch_size (`tuple[int, ...]`, defaults to `(1,)`): + Supported spatial patch sizes. + all_f_patch_size (`tuple[int, ...]`, defaults to `(1,)`): + Supported temporal patch sizes paired with `all_patch_size`. + in_channels (`int`, defaults to `128`): + Number of channels in the patchified Flux2 VAE latents. + dim (`int`, defaults to `3840`): + Transformer hidden dimension. + n_layers (`int`, defaults to `30`): + Number of main transformer blocks. + n_refiner_layers (`int`, defaults to `2`): + Number of noise, caption, and SigVQ refiner blocks. + n_heads (`int`, defaults to `30`): + Number of attention heads. + norm_eps (`float`, defaults to `1e-5`): + Epsilon used by RMS normalization layers. + qk_norm (`bool`, defaults to `True`): + Whether to apply RMS normalization to query and key tensors. + cap_feat_dim (`int`, defaults to `2560`): + Dimension of projected QueryFormer caption features. + semantic_feat_dim (`int`, defaults to `4096`): + Dimension of GLM/SigVQ semantic features. + rope_theta (`float`, defaults to `256.0`): + RoPE frequency base. + t_scale (`float`, defaults to `1000.0`): + Scale applied to diffusion timesteps. + axes_dims (`tuple[int, ...]`, defaults to `(32, 48, 48)`): + RoPE dimensions for sequence, height, and width axes. + axes_lens (`tuple[int, ...]`, defaults to `(32768, 1024, 1024)`): + Maximum RoPE positions for sequence, height, and width axes. + """ + + _supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] + _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] + _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ + "t_embedder", + "cap_embedder", + "semantic_embedder", + "sigvq_embedder", + ] + + @register_to_config + def __init__( + self, + all_patch_size: tuple[int, ...] = (1,), + all_f_patch_size: tuple[int, ...] = (1,), + in_channels: int = 128, + dim: int = 3840, + n_layers: int = 30, + n_refiner_layers: int = 2, + n_heads: int = 30, + norm_eps: float = 1e-5, + qk_norm: bool = True, + cap_feat_dim: int = 2560, + semantic_feat_dim: int = 4096, + rope_theta: float = 256.0, + t_scale: float = 1000.0, + axes_dims: tuple[int, ...] = (32, 48, 48), + axes_lens: tuple[int, ...] = (32768, 1024, 1024), + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + if len(all_patch_size) != len(all_f_patch_size): + raise ValueError( + "`all_patch_size` and `all_f_patch_size` must have the same length." + ) + if dim % n_heads != 0: + raise ValueError( + f"`dim` ({dim}) must be divisible by `n_heads` ({n_heads})." + ) + if dim // n_heads != sum(axes_dims): + raise ValueError( + "The attention head dimension must equal the sum of `axes_dims`." + ) + + self.in_channels = in_channels + self.out_channels = in_channels + self.all_patch_size = all_patch_size + self.all_f_patch_size = all_f_patch_size + self.t_scale = t_scale + self.gradient_checkpointing = False + + self.all_x_embedder = nn.ModuleDict() + self.all_final_layer = nn.ModuleDict() + for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size): + patch_key = f"{patch_size}-{f_patch_size}" + patch_dim = f_patch_size * patch_size * patch_size * in_channels + self.all_x_embedder[patch_key] = nn.Linear(patch_dim, dim, bias=True) + self.all_final_layer[patch_key] = LLaDAImageFinalLayer(dim, patch_dim) + + self.noise_refiner = nn.ModuleList( + [ + LLaDAImageTransformerBlock( + dim, + n_heads, + norm_eps, + qk_norm, + modulation=True, + quant_config=quant_config, + prefix=f"noise_refiner.{layer_id}", + ) + for layer_id in range(n_refiner_layers) + ] + ) + self.context_refiner = nn.ModuleList( + [ + LLaDAImageTransformerBlock( + dim, + n_heads, + norm_eps, + qk_norm, + modulation=False, + quant_config=quant_config, + prefix=f"context_refiner.{layer_id}", + ) + for layer_id in range(n_refiner_layers) + ] + ) + self.sigvq_refiner = nn.ModuleList( + [ + LLaDAImageTransformerBlock( + dim, + n_heads, + norm_eps, + qk_norm, + modulation=False, + quant_config=quant_config, + prefix=f"sigvq_refiner.{layer_id}", + ) + for layer_id in range(n_refiner_layers) + ] + ) + self.layers = nn.ModuleList( + [ + LLaDAImageTransformerBlock( + dim, + n_heads, + norm_eps, + qk_norm, + modulation=True, + quant_config=quant_config, + prefix=f"layers.{layer_id}", + ) + for layer_id in range(n_layers) + ] + ) + + self.t_embedder = LLaDAImageTimestepEmbedder(min(dim, ADALN_EMBED_DIM)) + self.cap_embedder = nn.Sequential( + RMSNorm(cap_feat_dim, eps=norm_eps, elementwise_affine=False), + nn.Linear(cap_feat_dim, dim, bias=True), + ) + self.semantic_embedder = nn.Sequential( + RMSNorm(semantic_feat_dim, eps=norm_eps, elementwise_affine=False), + nn.Linear(semantic_feat_dim, dim, bias=True), + ) + self.sigvq_embedder = nn.Sequential( + RMSNorm(semantic_feat_dim, eps=norm_eps, elementwise_affine=False), + nn.Linear(semantic_feat_dim, dim, bias=True), + ) + + nn.init.normal_(self.semantic_embedder[1].weight, mean=0.0, std=0.02) + nn.init.zeros_(self.semantic_embedder[1].bias) + nn.init.normal_(self.sigvq_embedder[1].weight, mean=0.0, std=0.02) + nn.init.zeros_(self.sigvq_embedder[1].bias) + + self.x_pad_token = nn.Parameter(torch.zeros(1, dim)) + self.cap_pad_token = nn.Parameter(torch.zeros(1, dim)) + self.sigvq_pad_token = nn.Parameter(torch.zeros(1, dim)) + nn.init.normal_(self.sigvq_pad_token, mean=0.0, std=0.02) + + self.rope_embedder = LLaDAImageRopeEmbedder(rope_theta, axes_dims, axes_lens) + + @staticmethod + def _create_coordinate_grid( + size: tuple[int, int, int], + start: tuple[int, int, int], + device: torch.device, + ) -> torch.Tensor: + axes = [ + torch.arange( + start_value, start_value + span, dtype=torch.int32, device=device + ) + for start_value, span in zip(start, size) + ] + return torch.stack(torch.meshgrid(axes, indexing="ij"), dim=-1) + + def _patchify_image( + self, + image: torch.Tensor, + patch_size: int, + f_patch_size: int, + ) -> tuple[torch.Tensor, tuple[int, int, int], tuple[int, int, int]]: + channels, frames, height, width = image.shape + frame_tokens = frames // f_patch_size + height_tokens = height // patch_size + width_tokens = width // patch_size + image = image.view( + channels, + frame_tokens, + f_patch_size, + height_tokens, + patch_size, + width_tokens, + patch_size, + ) + image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape( + frame_tokens * height_tokens * width_tokens, + f_patch_size * patch_size * patch_size * channels, + ) + return ( + image, + (frames, height, width), + (frame_tokens, height_tokens, width_tokens), + ) + + def _pad_with_ids( + self, + features: torch.Tensor, + position_grid_size: tuple[int, int, int], + position_start: tuple[int, int, int], + noise_value: int | None = None, + sequence_multiple: int = SEQUENCE_MULTIPLE, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, list[int] | None]: + original_length = len(features) + padding_length = (-original_length) % sequence_multiple + padded_length = original_length + padding_length + device = features.device + + position_ids = self._create_coordinate_grid( + position_grid_size, + position_start, + device, + ).flatten(0, 2) + if padding_length > 0: + padding_position_ids = ( + self._create_coordinate_grid( + (1, 1, 1), + (0, 0, 0), + device, + ) + .flatten(0, 2) + .repeat(padding_length, 1) + ) + position_ids = torch.cat([position_ids, padding_position_ids], dim=0) + features = torch.cat( + [features, features[-1:].repeat(padding_length, 1)], dim=0 + ) + padding_mask = torch.cat( + [ + torch.zeros(original_length, dtype=torch.bool, device=device), + torch.ones(padding_length, dtype=torch.bool, device=device), + ] + ) + else: + padding_mask = torch.zeros(original_length, dtype=torch.bool, device=device) + + noise_mask = [noise_value] * padded_length if noise_value is not None else None + return features, position_ids, padding_mask, padded_length, noise_mask + + @staticmethod + def _batch_sequences( + features: list[torch.Tensor], + frequencies: list[torch.Tensor], + inner_padding_masks: list[torch.Tensor], + pad_token: torch.Tensor, + noise_masks: list[list[int]] | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor | None, list[int], torch.Tensor | None + ]: + sequence_lengths = [len(item) for item in features] + max_sequence_length = max(sequence_lengths) + features = torch.cat(features, dim=0) + inner_padding_mask = torch.cat(inner_padding_masks).unsqueeze(-1) + features = torch.where( + inner_padding_mask.to(device=features.device), + pad_token.to(device=features.device, dtype=features.dtype), + features, + ) + features = list(features.split(sequence_lengths, dim=0)) + + features = pad_sequence(features, batch_first=True, padding_value=0.0) + frequencies = pad_sequence(frequencies, batch_first=True, padding_value=0.0)[ + :, : features.shape[1] + ] + + attention_mask = None + if not all(length == max_sequence_length for length in sequence_lengths): + attention_mask = torch.zeros( + (len(sequence_lengths), max_sequence_length), + dtype=torch.bool, + device=features.device, + ) + for batch_index, sequence_length in enumerate(sequence_lengths): + attention_mask[batch_index, :sequence_length] = True + + noise_mask = None + if noise_masks is not None: + noise_mask = pad_sequence( + [ + torch.tensor(mask, dtype=torch.long, device=features.device) + for mask in noise_masks + ], + batch_first=True, + padding_value=0, + )[:, : features.shape[1]] + + return features, frequencies, attention_mask, sequence_lengths, noise_mask + + def _unpatchify( + self, + hidden_states: list[torch.Tensor], + sizes: list[tuple[int, int, int] | list[tuple[int, int, int]]], + patch_size: int, + f_patch_size: int, + image_offsets: list[tuple[int, int]] | None = None, + ) -> list[torch.Tensor]: + outputs = [] + sequence_multiple = SEQUENCE_MULTIPLE // get_sp_world_size() + for batch_index, batch_hidden_states in enumerate(hidden_states): + if image_offsets is None: + batch_sizes = [sizes[batch_index]] + image_hidden_states = batch_hidden_states + else: + batch_sizes = sizes[batch_index] + start, end = image_offsets[batch_index] + image_hidden_states = batch_hidden_states[start:end] + + current_offset = 0 + output = None + for frames, height, width in batch_sizes: + original_length = ( + (frames // f_patch_size) + * (height // patch_size) + * (width // patch_size) + ) + padding_length = (-original_length) % sequence_multiple + output = ( + image_hidden_states[ + current_offset : current_offset + original_length + ] + .view( + frames // f_patch_size, + height // patch_size, + width // patch_size, + f_patch_size, + patch_size, + patch_size, + self.out_channels, + ) + .permute(6, 0, 3, 1, 4, 2, 5) + .reshape(self.out_channels, frames, height, width) + ) + current_offset += original_length + padding_length + outputs.append(output) + return outputs + + def _prepare_t2i_sequences( + self, + x: list[torch.Tensor], + cap_feats: list[torch.Tensor] | None, + glm_features: list[torch.Tensor] | None, + patch_size: int, + f_patch_size: int, + ) -> tuple[ + _LLaDAImageSequence, + _LLaDAImageSequence | None, + _LLaDAImageSequence | None, + list[tuple[int, int, int]], + ]: + image_sequence = _LLaDAImageSequence([], [], []) + cap_sequence = ( + _LLaDAImageSequence([], [], []) if cap_feats is not None else None + ) + glm_sequence = ( + _LLaDAImageSequence([], [], []) if glm_features is not None else None + ) + image_sizes = [] + + for batch_index, latent in enumerate(x): + position_cursor = 1 + if cap_sequence is not None: + padded_features, position_ids, padding_mask, sequence_length, _ = ( + self._pad_with_ids( + cap_feats[batch_index], + (len(cap_feats[batch_index]), 1, 1), + (position_cursor, 0, 0), + ) + ) + cap_sequence.features.append(padded_features) + cap_sequence.position_ids.append(position_ids) + cap_sequence.padding_masks.append(padding_mask) + position_cursor += sequence_length + + if glm_sequence is not None: + padded_features, position_ids, padding_mask, sequence_length, _ = ( + self._pad_with_ids( + glm_features[batch_index], + (len(glm_features[batch_index]), 1, 1), + (position_cursor, 0, 0), + ) + ) + glm_sequence.features.append(padded_features) + glm_sequence.position_ids.append(position_ids) + glm_sequence.padding_masks.append(padding_mask) + position_cursor += sequence_length + + patches, image_size, token_grid_size = self._patchify_image( + latent, patch_size, f_patch_size + ) + image_height_start = 0 + if get_sp_world_size() > 1: + image_height_start = get_sp_parallel_rank() * token_grid_size[1] + padded_features, position_ids, padding_mask, _, _ = self._pad_with_ids( + patches, + token_grid_size, + (position_cursor, image_height_start, 0), + sequence_multiple=SEQUENCE_MULTIPLE // get_sp_world_size(), + ) + image_sequence.features.append(padded_features) + image_sequence.position_ids.append(position_ids) + image_sequence.padding_masks.append(padding_mask) + image_sizes.append(image_size) + + return image_sequence, cap_sequence, glm_sequence, image_sizes + + def _prepare_editing_sequences( + self, + x: list[torch.Tensor], + cap_feats: list[torch.Tensor], + glm_cap_feats: list[torch.Tensor], + source_latents: list[torch.Tensor], + patch_size: int, + f_patch_size: int, + ) -> tuple[ + _LLaDAImageSequence, + _LLaDAImageSequence, + _LLaDAImageSequence, + list[list[tuple[int, int, int]]], + list[tuple[int, int]], + ]: + image_sequence = _LLaDAImageSequence([], [], [], []) + cap_sequence = _LLaDAImageSequence([], [], [], []) + sigvq_sequence = _LLaDAImageSequence([], [], [], []) + image_sizes = [] + image_offsets = [] + + for batch_index, latent in enumerate(x): + cap_end_positions = [] + position_cursor = 1 + batch_cap_features = [] + batch_cap_positions = [] + batch_cap_padding = [] + batch_cap_noise = [] + for noise_value in (0, 1): + padded_features, position_ids, padding_mask, _, noise_mask = ( + self._pad_with_ids( + cap_feats[batch_index], + (len(cap_feats[batch_index]), 1, 1), + (position_cursor, 0, 0), + noise_value, + ) + ) + batch_cap_features.append(padded_features) + batch_cap_positions.append(position_ids) + batch_cap_padding.append(padding_mask) + batch_cap_noise.extend(noise_mask) + position_cursor += len(cap_feats[batch_index]) + cap_end_positions.append(position_cursor) + position_cursor += 2 + + batch_image_features = [] + batch_image_sizes = [] + batch_image_positions = [] + batch_image_padding = [] + batch_image_noise = [] + for image, position_start, noise_value in zip( + (source_latents[batch_index], latent), + cap_end_positions, + (0, 1), + ): + patches, image_size, token_grid_size = self._patchify_image( + image, patch_size, f_patch_size + ) + image_height_start = 0 + if get_sp_world_size() > 1: + image_height_start = get_sp_parallel_rank() * token_grid_size[1] + padded_features, position_ids, padding_mask, _, noise_mask = ( + self._pad_with_ids( + patches, + token_grid_size, + (position_start, image_height_start, 0), + noise_value, + sequence_multiple=SEQUENCE_MULTIPLE // get_sp_world_size(), + ) + ) + batch_image_features.append(padded_features) + batch_image_sizes.append(image_size) + batch_image_positions.append(position_ids) + batch_image_padding.append(padding_mask) + batch_image_noise.extend(noise_mask) + + batch_cap_features = torch.cat(batch_cap_features, dim=0) + batch_image_features = torch.cat(batch_image_features, dim=0) + cap_sequence.features.append(batch_cap_features) + cap_sequence.position_ids.append(torch.cat(batch_cap_positions, dim=0)) + cap_sequence.padding_masks.append(torch.cat(batch_cap_padding, dim=0)) + cap_sequence.noise_masks.append(batch_cap_noise) + image_sequence.features.append(batch_image_features) + image_sequence.position_ids.append(torch.cat(batch_image_positions, dim=0)) + image_sequence.padding_masks.append(torch.cat(batch_image_padding, dim=0)) + image_sequence.noise_masks.append(batch_image_noise) + image_sizes.append(batch_image_sizes) + image_offsets.append( + ( + len(batch_cap_features), + len(batch_cap_features) + len(batch_image_features), + ) + ) + + padded_features, position_ids, padding_mask, _, noise_mask = ( + self._pad_with_ids( + glm_cap_feats[batch_index], + (len(glm_cap_feats[batch_index]), 1, 1), + ( + len(batch_cap_features) + + len(batch_image_features) * get_sp_world_size() + + 1, + 0, + 0, + ), + 0, + ) + ) + sigvq_sequence.features.append(padded_features) + sigvq_sequence.position_ids.append(position_ids) + sigvq_sequence.padding_masks.append(padding_mask) + sigvq_sequence.noise_masks.append(noise_mask) + + return image_sequence, cap_sequence, sigvq_sequence, image_sizes, image_offsets + + @staticmethod + def _merge_padded_sequences( + feature_groups: tuple[torch.Tensor, ...], + frequency_groups: tuple[torch.Tensor, ...], + length_groups: tuple[list[int], ...], + noise_mask_groups: tuple[torch.Tensor, ...] | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor | None, list[int], torch.Tensor | None + ]: + batch_size = feature_groups[0].shape[0] + merged_features = [] + merged_frequencies = [] + merged_noise_masks = [] if noise_mask_groups is not None else None + + for batch_index in range(batch_size): + device = feature_groups[0].device + merged_features.append( + torch.cat( + [ + features[batch_index, : lengths[batch_index]].to(device) + for features, lengths in zip(feature_groups, length_groups) + ], + dim=0, + ) + ) + merged_frequencies.append( + torch.cat( + [ + frequencies[batch_index, : lengths[batch_index]].to(device) + for frequencies, lengths in zip(frequency_groups, length_groups) + ], + dim=0, + ) + ) + if merged_noise_masks is not None: + merged_noise_masks.append( + torch.cat( + [ + noise_masks[batch_index, : lengths[batch_index]].to(device) + for noise_masks, lengths in zip( + noise_mask_groups, length_groups + ) + ], + dim=0, + ) + ) + + merged_lengths = [len(features) for features in merged_features] + merged_features = pad_sequence( + merged_features, batch_first=True, padding_value=0.0 + ) + merged_frequencies = pad_sequence( + merged_frequencies, batch_first=True, padding_value=0.0 + ) + + attention_mask = None + max_length = max(merged_lengths) + if not all(length == max_length for length in merged_lengths): + attention_mask = torch.zeros( + (batch_size, max_length), + dtype=torch.bool, + device=merged_features.device, + ) + for batch_index, sequence_length in enumerate(merged_lengths): + attention_mask[batch_index, :sequence_length] = True + + noise_mask = None + if merged_noise_masks is not None: + noise_mask = pad_sequence( + merged_noise_masks, batch_first=True, padding_value=0 + )[:, : merged_features.shape[1]] + + return ( + merged_features, + merged_frequencies, + attention_mask, + merged_lengths, + noise_mask, + ) + + def forward( + self, + x: list[torch.Tensor], + t: torch.Tensor, + cap_feats: list[torch.Tensor] | None, + glm_cap_feats: list[torch.Tensor] | None = None, + source_latents: list[torch.Tensor] | None = None, + patch_size: int = 1, + f_patch_size: int = 1, + return_dict: bool = True, + ) -> Transformer2DModelOutput | tuple[list[torch.Tensor]]: + r""" + Args: + x (`list[torch.Tensor]`): + Target latents. Each tensor has shape `(channels, frames, height, width)`. + t (`torch.Tensor`): + Denoising timestep for each batch item. + cap_feats (`list[torch.Tensor]`, *optional*): + Projected QueryFormer features, each with shape `(sequence_length, cap_feat_dim)`. + glm_cap_feats (`list[torch.Tensor]`, *optional*): + GLM/SigVQ features, each with shape `(sequence_length, semantic_feat_dim)`. + source_latents (`list[torch.Tensor]`, *optional*): + Source-image latents for editing. When provided, `cap_feats` and `glm_cap_feats` are required. + patch_size (`int`, defaults to `1`): + Spatial patch size. + f_patch_size (`int`, defaults to `1`): + Temporal patch size. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`]. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + The denoised target latents. + """ + patch_key = f"{patch_size}-{f_patch_size}" + if patch_key not in self.all_x_embedder: + raise ValueError( + f"Unsupported patch sizes: patch_size={patch_size}, f_patch_size={f_patch_size}." + ) + if source_latents is None and cap_feats is None and glm_cap_feats is None: + raise ValueError( + "Text-to-image inference requires `cap_feats` or `glm_cap_feats`." + ) + if source_latents is not None and (cap_feats is None or glm_cap_feats is None): + raise ValueError( + "Editing requires `cap_feats`, `glm_cap_feats`, and `source_latents`." + ) + + batch_size = len(x) + is_editing = source_latents is not None + adaln_input = None + noisy_embedding = None + clean_embedding = None + image_offsets = None + num_replicated_suffix = 0 + + if is_editing: + if t.shape[0] == 1: + t = t.repeat(batch_size) + dual_timestep = torch.cat([t, torch.zeros_like(t)], dim=0) + dual_embedding = self.t_embedder( + dual_timestep.abs() * self.t_scale, x[0].dtype + ) + noisy_embedding = dual_embedding[:batch_size] + clean_embedding = dual_embedding[batch_size:] + image_sequence, cap_sequence, sigvq_sequence, image_sizes, image_offsets = ( + self._prepare_editing_sequences( + x, + cap_feats, + glm_cap_feats, + source_latents, + patch_size, + f_patch_size, + ) + ) + else: + adaln_input = self.t_embedder(t * self.t_scale, x[0].dtype) + glm_features = ( + [ + self.semantic_embedder(batch_features) + for batch_features in glm_cap_feats + ] + if glm_cap_feats is not None + else None + ) + image_sequence, cap_sequence, glm_sequence, image_sizes = ( + self._prepare_t2i_sequences( + x, + cap_feats, + glm_features, + patch_size, + f_patch_size, + ) + ) + + image_lengths = [len(features) for features in image_sequence.features] + image_features = self.all_x_embedder[patch_key]( + torch.cat(image_sequence.features, dim=0) + ) + image_frequencies = list( + self.rope_embedder(torch.cat(image_sequence.position_ids, dim=0)).split( + [len(position_ids) for position_ids in image_sequence.position_ids], + dim=0, + ) + ) + ( + image_features, + image_frequencies, + image_attention_mask, + image_lengths, + image_noise_mask, + ) = self._batch_sequences( + list(image_features.split(image_lengths, dim=0)), + image_frequencies, + image_sequence.padding_masks, + self.x_pad_token, + image_sequence.noise_masks, + ) + + for layer in self.noise_refiner: + if torch.is_grad_enabled() and self.gradient_checkpointing: + if is_editing: + image_features = self._gradient_checkpointing_func( + layer, + image_features, + image_attention_mask, + image_frequencies, + None, + image_noise_mask, + noisy_embedding, + clean_embedding, + ) + else: + image_features = self._gradient_checkpointing_func( + layer, + image_features, + image_attention_mask, + image_frequencies, + adaln_input, + ) + elif is_editing: + image_features = layer( + image_features, + image_attention_mask, + image_frequencies, + noise_mask=image_noise_mask, + adaln_noisy=noisy_embedding, + adaln_clean=clean_embedding, + ) + else: + image_features = layer( + image_features, + image_attention_mask, + image_frequencies, + adaln_input, + ) + + if is_editing: + cap_lengths = [len(features) for features in cap_sequence.features] + cap_features = self.cap_embedder(torch.cat(cap_sequence.features, dim=0)) + cap_frequencies = list( + self.rope_embedder(torch.cat(cap_sequence.position_ids, dim=0)).split( + [len(position_ids) for position_ids in cap_sequence.position_ids], + dim=0, + ) + ) + ( + cap_features, + cap_frequencies, + cap_attention_mask, + cap_lengths, + cap_noise_mask, + ) = self._batch_sequences( + list(cap_features.split(cap_lengths, dim=0)), + cap_frequencies, + cap_sequence.padding_masks, + self.cap_pad_token, + cap_sequence.noise_masks, + ) + + for layer in self.context_refiner: + if torch.is_grad_enabled() and self.gradient_checkpointing: + cap_features = self._gradient_checkpointing_func( + layer, + cap_features, + cap_attention_mask, + cap_frequencies, + None, + None, + None, + None, + 0, + True, + ) + else: + cap_features = layer( + cap_features, + cap_attention_mask, + cap_frequencies, + skip_sequence_parallel_override=True, + ) + + sigvq_lengths = [len(features) for features in sigvq_sequence.features] + sigvq_features = self.sigvq_embedder( + torch.cat(sigvq_sequence.features, dim=0) + ) + sigvq_frequencies = list( + self.rope_embedder(torch.cat(sigvq_sequence.position_ids, dim=0)).split( + [len(position_ids) for position_ids in sigvq_sequence.position_ids], + dim=0, + ) + ) + ( + sigvq_features, + sigvq_frequencies, + sigvq_attention_mask, + sigvq_lengths, + sigvq_noise_mask, + ) = self._batch_sequences( + list(sigvq_features.split(sigvq_lengths, dim=0)), + sigvq_frequencies, + sigvq_sequence.padding_masks, + self.sigvq_pad_token, + sigvq_sequence.noise_masks, + ) + + if any(sigvq_lengths): + for layer in self.sigvq_refiner: + if torch.is_grad_enabled() and self.gradient_checkpointing: + sigvq_features = self._gradient_checkpointing_func( + layer, + sigvq_features, + sigvq_attention_mask, + sigvq_frequencies, + None, + None, + None, + None, + 0, + True, + ) + else: + sigvq_features = layer( + sigvq_features, + sigvq_attention_mask, + sigvq_frequencies, + skip_sequence_parallel_override=True, + ) + + if get_sp_world_size() > 1: + if batch_size != 1: + raise NotImplementedError( + "LLaDA-Image sequence parallelism currently supports batch size 1." + ) + feature_groups = (image_features, cap_features, sigvq_features) + frequency_groups = ( + image_frequencies, + cap_frequencies, + sigvq_frequencies, + ) + length_groups = (image_lengths, cap_lengths, sigvq_lengths) + noise_mask_groups = ( + image_noise_mask, + cap_noise_mask, + sigvq_noise_mask, + ) + num_replicated_suffix = cap_lengths[0] + sigvq_lengths[0] + image_offsets = [(0, image_lengths[0])] + else: + feature_groups = (cap_features, image_features, sigvq_features) + frequency_groups = ( + cap_frequencies, + image_frequencies, + sigvq_frequencies, + ) + length_groups = (cap_lengths, image_lengths, sigvq_lengths) + noise_mask_groups = ( + cap_noise_mask, + image_noise_mask, + sigvq_noise_mask, + ) + + ( + unified_features, + unified_frequencies, + unified_attention_mask, + _, + unified_noise_mask, + ) = self._merge_padded_sequences( + feature_groups, + frequency_groups, + length_groups, + noise_mask_groups, + ) + else: + condition_feature_groups = [] + condition_frequency_groups = [] + condition_length_groups = [] + + if cap_sequence is not None: + cap_lengths = [len(features) for features in cap_sequence.features] + cap_features = self.cap_embedder( + torch.cat(cap_sequence.features, dim=0) + ) + cap_padding_mask = ( + torch.cat(cap_sequence.padding_masks) + .unsqueeze(-1) + .to(cap_features.device) + ) + cap_features = torch.where( + cap_padding_mask, + self.cap_pad_token.to( + device=cap_features.device, dtype=cap_features.dtype + ), + cap_features, + ) + cap_features = pad_sequence( + list(cap_features.split(cap_lengths, dim=0)), + batch_first=True, + padding_value=0.0, + ) + cap_frequencies = list( + self.rope_embedder( + torch.cat(cap_sequence.position_ids, dim=0) + ).split( + [ + len(position_ids) + for position_ids in cap_sequence.position_ids + ], + dim=0, + ) + ) + cap_frequencies = pad_sequence( + cap_frequencies, batch_first=True, padding_value=0.0 + ) + condition_feature_groups.append(cap_features) + condition_frequency_groups.append(cap_frequencies) + condition_length_groups.append(cap_lengths) + + if glm_sequence is not None: + glm_lengths = [len(features) for features in glm_sequence.features] + glm_features = torch.cat(glm_sequence.features, dim=0) + glm_padding_mask = ( + torch.cat(glm_sequence.padding_masks) + .unsqueeze(-1) + .to(glm_features.device) + ) + glm_features = torch.where( + glm_padding_mask, + self.cap_pad_token.to( + device=glm_features.device, dtype=glm_features.dtype + ), + glm_features, + ) + glm_features = pad_sequence( + list(glm_features.split(glm_lengths, dim=0)), + batch_first=True, + padding_value=0.0, + ) + glm_frequencies = list( + self.rope_embedder( + torch.cat(glm_sequence.position_ids, dim=0) + ).split( + [ + len(position_ids) + for position_ids in glm_sequence.position_ids + ], + dim=0, + ) + ) + glm_frequencies = pad_sequence( + glm_frequencies, batch_first=True, padding_value=0.0 + ) + condition_feature_groups.append(glm_features) + condition_frequency_groups.append(glm_frequencies) + condition_length_groups.append(glm_lengths) + + ( + condition_features, + condition_frequencies, + condition_attention_mask, + condition_lengths, + _, + ) = self._merge_padded_sequences( + tuple(condition_feature_groups), + tuple(condition_frequency_groups), + tuple(condition_length_groups), + ) + + for layer in self.context_refiner: + if torch.is_grad_enabled() and self.gradient_checkpointing: + condition_features = self._gradient_checkpointing_func( + layer, + condition_features, + condition_attention_mask, + condition_frequencies, + None, + None, + None, + None, + 0, + True, + ) + else: + condition_features = layer( + condition_features, + condition_attention_mask, + condition_frequencies, + skip_sequence_parallel_override=True, + ) + + ( + unified_features, + unified_frequencies, + unified_attention_mask, + _, + unified_noise_mask, + ) = self._merge_padded_sequences( + (image_features, condition_features), + (image_frequencies, condition_frequencies), + (image_lengths, condition_lengths), + ) + if get_sp_world_size() > 1: + if batch_size != 1: + raise NotImplementedError( + "LLaDA-Image sequence parallelism currently supports batch size 1." + ) + num_replicated_suffix = condition_lengths[0] + + for layer in self.layers: + if torch.is_grad_enabled() and self.gradient_checkpointing: + if is_editing: + unified_features = self._gradient_checkpointing_func( + layer, + unified_features, + unified_attention_mask, + unified_frequencies, + None, + unified_noise_mask, + noisy_embedding, + clean_embedding, + num_replicated_suffix, + ) + else: + unified_features = self._gradient_checkpointing_func( + layer, + unified_features, + unified_attention_mask, + unified_frequencies, + adaln_input, + None, + None, + None, + num_replicated_suffix, + ) + elif is_editing: + unified_features = layer( + unified_features, + unified_attention_mask, + unified_frequencies, + noise_mask=unified_noise_mask, + adaln_noisy=noisy_embedding, + adaln_clean=clean_embedding, + num_replicated_suffix=num_replicated_suffix, + ) + else: + unified_features = layer( + unified_features, + unified_attention_mask, + unified_frequencies, + adaln_input, + num_replicated_suffix=num_replicated_suffix, + ) + + if is_editing: + unified_features = self.all_final_layer[patch_key]( + unified_features, + noise_mask=unified_noise_mask, + adaln_noisy=noisy_embedding, + adaln_clean=clean_embedding, + ) + else: + unified_features = self.all_final_layer[patch_key]( + unified_features, + adaln_input=adaln_input, + ) + + output = self._unpatchify( + list(unified_features.unbind(dim=0)), + image_sizes, + patch_size, + f_patch_size, + image_offsets, + ) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + +@dataclass +class LLaDAImageQueryFormerOutput(BaseOutput): + query_embeds: torch.Tensor + + +class LLaDAImageQueryAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "LLaDAImageQueryAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + query = F.linear( + hidden_states, + attn.in_proj_weight[: attn.inner_dim], + attn.in_proj_bias[: attn.inner_dim], + ) + key = F.linear( + encoder_hidden_states, + attn.in_proj_weight[attn.inner_dim : 2 * attn.inner_dim], + attn.in_proj_bias[attn.inner_dim : 2 * attn.inner_dim], + ) + value = F.linear( + encoder_hidden_states, + attn.in_proj_weight[2 * attn.inner_dim :], + attn.in_proj_bias[2 * attn.inner_dim :], + ) + + query = query.unflatten(-1, (attn.heads, attn.head_dim)) + key = key.unflatten(-1, (attn.heads, attn.head_dim)) + value = value.unflatten(-1, (attn.heads, attn.head_dim)) + + if attention_mask is not None: + attention_mask = attention_mask[:, None, None, :] + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=attn.dropout if attn.training else 0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + return attn.out_proj(hidden_states) + + +class LLaDAImageQueryAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = LLaDAImageQueryAttnProcessor + _available_processors: ClassVar[list[type]] = [LLaDAImageQueryAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, hidden_size: int, num_heads: int, dropout: float): + super().__init__() + self.inner_dim = hidden_size + self.heads = num_heads + self.head_dim = hidden_size // num_heads + self.dropout = dropout + + self.in_proj_weight = nn.Parameter(torch.zeros(3 * hidden_size, hidden_size)) + self.in_proj_bias = nn.Parameter(torch.zeros(3 * hidden_size)) + self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True) + self.set_processor(self._default_processor_cls()) + + nn.init.xavier_uniform_(self.in_proj_weight) + nn.init.zeros_(self.in_proj_bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.processor( + self, hidden_states, encoder_hidden_states, attention_mask + ) + + +@maybe_allow_in_graph +class LLaDAImageQueryFormerBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + intermediate_size: int, + dropout: float, + norm_eps: float, + ): + super().__init__() + self.norm_q = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) + self.norm_k = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) + self.cross_attn = LLaDAImageQueryAttention(hidden_size, num_heads, dropout) + self.dropout = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) + self.mlp = nn.Module() + self.mlp.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) + self.mlp.fc2 = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward( + self, + query_embeds: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + query_embeds = self.norm_q(query_embeds) + encoder_hidden_states = self.norm_k(encoder_hidden_states) + attention_output = self.cross_attn( + query_embeds, encoder_hidden_states, attention_mask + ) + query_embeds = query_embeds + self.dropout(attention_output) + query_embeds = self.norm1(query_embeds) + mlp_output = self.mlp.fc2( + F.gelu(self.mlp.fc1(query_embeds), approximate="tanh") + ) + return query_embeds + self.dropout(mlp_output) + + +class LLaDAImageQueryFormerModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + QueryFormer used by LLaDA-Image to derive learnable image-generation queries from LLaDA token embeddings. + + This model is independent from the LLaDA text encoder. It returns refined query embeddings. The pipeline appends + them to the text embeddings and invokes the text encoder backbone. + + Args: + num_queries (`int`, defaults to `256`): + Number of learnable query tokens. + hidden_size (`int`, defaults to `2048`): + Query and LLaDA token embedding dimension. + num_hidden_layers (`int`, defaults to `1`): + Number of QueryFormer blocks. + num_attention_heads (`int`, defaults to `16`): + Number of cross-attention heads. + intermediate_size (`int`, defaults to `8192`): + Hidden dimension of the QueryFormer MLP. + dropout (`float`, defaults to `0.0`): + Dropout probability. + norm_eps (`float`, defaults to `1e-6`): + Epsilon used by parameter-free layer normalization. + """ + + _supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] + _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] + _skip_layerwise_casting_patterns: ClassVar[list[str]] = ["norm"] + + @register_to_config + def __init__( + self, + num_queries: int = 256, + hidden_size: int = 2048, + num_hidden_layers: int = 1, + num_attention_heads: int = 16, + intermediate_size: int = 8192, + dropout: float = 0.0, + norm_eps: float = 1e-6, + ): + super().__init__() + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." + ) + + self.meta_queries = nn.Parameter(torch.zeros(num_queries, hidden_size)) + nn.init.normal_(self.meta_queries, std=1 / math.sqrt(hidden_size)) + self.query_blocks = nn.ModuleList( + [ + LLaDAImageQueryFormerBlock( + hidden_size, + num_attention_heads, + intermediate_size, + dropout, + norm_eps, + ) + for _ in range(num_hidden_layers) + ] + ) + self.gradient_checkpointing = False + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor, + return_dict: bool = True, + ) -> LLaDAImageQueryFormerOutput | tuple[torch.Tensor]: + r""" + Args: + inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + LLaDA input token embeddings. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`): + Mask whose nonzero entries identify valid text tokens. + return_dict (`bool`, defaults to `True`): + Whether to return [`LLaDAImageQueryFormerOutput`] instead of a tuple. + + Returns: + [`LLaDAImageQueryFormerOutput`] or `tuple`: + The refined query embeddings. + """ + batch_size = inputs_embeds.shape[0] + query_embeds = self.meta_queries.unsqueeze(0).expand(batch_size, -1, -1) + attention_mask = attention_mask.bool() + + for query_block in self.query_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + query_embeds = self._gradient_checkpointing_func( + query_block, + query_embeds, + inputs_embeds, + attention_mask, + ) + else: + query_embeds = query_block(query_embeds, inputs_embeds, attention_mask) + + if not return_dict: + return (query_embeds,) + return LLaDAImageQueryFormerOutput(query_embeds=query_embeds) + + +@dataclass +class LLaDAImageTextProjectionOutput(BaseOutput): + hidden_states: torch.Tensor + + +class LLaDAImageTextProjectionAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "LLaDAImageTextProjectionAttention", + hidden_states: torch.Tensor, + ) -> torch.Tensor: + query = attn.q_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) + key = attn.k_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) + value = attn.v_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) + + query = attn.q_norm(query) + key = attn.k_norm(key) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=attn.dropout if attn.training else 0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + return attn.out_proj(hidden_states) + + +class LLaDAImageTextProjectionAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = LLaDAImageTextProjectionAttnProcessor + _available_processors: ClassVar[list[type]] = [ + LLaDAImageTextProjectionAttnProcessor + ] + _supports_qkv_fusion = False + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_dropout: float, + norm_eps: float, + ): + super().__init__() + self.heads = num_attention_heads + self.head_dim = hidden_size // num_attention_heads + self.dropout = attention_dropout + + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=True) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=True) + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=True) + self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True) + self.q_norm = RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=False) + self.k_norm = RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=False) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states) + + +class LLaDAImageTextProjectionMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) + self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = F.gelu(hidden_states, approximate="tanh") + return self.fc2(hidden_states) + + +@maybe_allow_in_graph +class LLaDAImageTextProjectionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + attention_dropout: float, + norm_eps: float, + ): + super().__init__() + self.self_attn = LLaDAImageTextProjectionAttention( + hidden_size, + num_attention_heads, + attention_dropout, + norm_eps, + ) + self.layer_norm1 = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) + self.mlp = LLaDAImageTextProjectionMLP(hidden_size, intermediate_size) + self.layer_norm2 = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.self_attn(self.layer_norm1(hidden_states)) + hidden_states = hidden_states + self.mlp(self.layer_norm2(hidden_states)) + return hidden_states + + +class LLaDAImageTextProjectionModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + Connector and output projection used to map LLaDA hidden states to the LLaDA-Image denoiser context dimension. + + Args: + hidden_size (`int`, defaults to `2048`): + Input and connector hidden dimension. + intermediate_size (`int`, defaults to `8960`): + Connector MLP hidden dimension. + num_hidden_layers (`int`, defaults to `6`): + Number of connector layers. + num_attention_heads (`int`, defaults to `32`): + Number of connector self-attention heads. + projection_dim (`int`, defaults to `2560`): + Output dimension expected by the denoising transformer. + attention_dropout (`float`, defaults to `0.0`): + Attention dropout probability. + norm_eps (`float`, defaults to `1e-6`): + Epsilon used by parameter-free RMS normalization. + """ + + _supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] + _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] + _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ + "layer_norm", + "q_norm", + "k_norm", + ] + + @register_to_config + def __init__( + self, + hidden_size: int = 2048, + intermediate_size: int = 8960, + num_hidden_layers: int = 6, + num_attention_heads: int = 32, + projection_dim: int = 2560, + attention_dropout: float = 0.0, + norm_eps: float = 1e-6, + ): + super().__init__() + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." + ) + + self.layers = nn.ModuleList( + [ + LLaDAImageTextProjectionBlock( + hidden_size, + intermediate_size, + num_attention_heads, + attention_dropout, + norm_eps, + ) + for _ in range(num_hidden_layers) + ] + ) + self.projector = nn.Linear(hidden_size, projection_dim, bias=True) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + return_dict: bool = True, + ) -> LLaDAImageTextProjectionOutput | tuple[torch.Tensor]: + r""" + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Hidden states produced by the LLaDA text backbone. + return_dict (`bool`, defaults to `True`): + Whether to return [`LLaDAImageTextProjectionOutput`] instead of a tuple. + + Returns: + [`LLaDAImageTextProjectionOutput`] or `tuple`: + Hidden states projected to the denoiser caption dimension. + """ + for layer in self.layers: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func(layer, hidden_states) + else: + hidden_states = layer(hidden_states) + + hidden_states = self.projector(hidden_states) + if not return_dict: + return (hidden_states,) + return LLaDAImageTextProjectionOutput(hidden_states=hidden_states) + + +@dataclass +class LLaDAImageSigVQOutput(BaseOutput): + semantic_features: torch.Tensor + token_ids: torch.Tensor + + +class LLaDAImageSigVQAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__( + self, attn: "LLaDAImageSigVQAttention", hidden_states: torch.Tensor + ) -> torch.Tensor: + query, key, value = attn.qkv(hidden_states).chunk(3, dim=-1) + query = query.unflatten(-1, (attn.heads, attn.head_dim)) + key = key.unflatten(-1, (attn.heads, attn.head_dim)) + value = value.unflatten(-1, (attn.heads, attn.head_dim)) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + dropout_p=attn.dropout if attn.training else 0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + return attn.proj(hidden_states) + + +class LLaDAImageSigVQAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = LLaDAImageSigVQAttnProcessor + _available_processors: ClassVar[list[type]] = [LLaDAImageSigVQAttnProcessor] + _supports_qkv_fusion = False + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + attention_bias: bool, + attention_dropout: float, + ): + super().__init__() + self.heads = num_attention_heads + self.head_dim = hidden_size // num_attention_heads + self.dropout = attention_dropout + self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=attention_bias) + self.proj = nn.Linear(hidden_size, hidden_size, bias=attention_bias) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states) + + +class LLaDAImageSigVQMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) + self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.fc2(F.gelu(self.fc1(hidden_states))) + + +@maybe_allow_in_graph +class LLaDAImageSigVQVisionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + attention_bias: bool, + attention_dropout: float, + norm_eps: float, + ): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, eps=norm_eps) + self.norm2 = nn.LayerNorm(hidden_size, eps=norm_eps) + self.attn = LLaDAImageSigVQAttention( + hidden_size, + num_attention_heads, + attention_bias, + attention_dropout, + ) + self.mlp = LLaDAImageSigVQMLP(hidden_size, intermediate_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class LLaDAImageSigVQPatchEmbed(nn.Module): + def __init__(self, in_channels: int, hidden_size: int, patch_size: int): + super().__init__() + self.in_channels = in_channels + self.patch_size = patch_size + self.proj = nn.Conv2d( + in_channels, hidden_size, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size, channels, height, width = pixel_values.shape + grid_height = height // self.patch_size + grid_width = width // self.patch_size + patches = pixel_values.reshape( + batch_size, + channels, + grid_height, + self.patch_size, + grid_width, + self.patch_size, + ) + patches = patches.permute(0, 2, 4, 1, 3, 5).reshape( + batch_size * grid_height * grid_width, + channels, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(patches).flatten(1) + return hidden_states.reshape(batch_size, grid_height * grid_width, -1) + + +class LLaDAImageSigVQEmbeddings(nn.Module): + def __init__(self, image_size: int, patch_size: int, hidden_size: int): + super().__init__() + num_positions = (image_size // patch_size) ** 2 + self.position_embedding = nn.Embedding(num_positions, hidden_size) + + def forward( + self, hidden_states: torch.Tensor, grid_height: int, grid_width: int + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + position_embedding = self.position_embedding.weight + hidden_size = position_embedding.shape[1] + original_size = int(position_embedding.shape[0] ** 0.5) + position_embedding = position_embedding.reshape( + original_size, original_size, hidden_size + ) + position_embedding = position_embedding.permute(2, 0, 1).unsqueeze(0).float() + + height_coordinates = torch.arange( + grid_height, device=hidden_states.device, dtype=torch.float32 + ) + width_coordinates = torch.arange( + grid_width, device=hidden_states.device, dtype=torch.float32 + ) + height_coordinates, width_coordinates = torch.meshgrid( + height_coordinates, + width_coordinates, + indexing="ij", + ) + normalized_width = ((width_coordinates.flatten() + 0.5) / grid_width) * 2 - 1 + normalized_height = ((height_coordinates.flatten() + 0.5) / grid_height) * 2 - 1 + grid = torch.stack((normalized_width, normalized_height), dim=-1) + grid = grid.reshape(1, grid_height * grid_width, 1, 2).expand( + batch_size, -1, -1, -1 + ) + + position_embedding = F.grid_sample( + position_embedding.expand(batch_size, -1, -1, -1), + grid, + mode="bilinear", + align_corners=False, + padding_mode="border", + ) + position_embedding = ( + position_embedding.squeeze(-1).transpose(1, 2).to(hidden_states.dtype) + ) + return hidden_states + position_embedding + + +class LLaDAImageSigVQQuantizer(nn.Module): + def __init__(self, num_embeddings: int, embedding_dim: int): + super().__init__() + self.embedding = nn.Embedding(num_embeddings, embedding_dim) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.permute(0, 2, 3, 1).contiguous() + hidden_states = F.normalize( + hidden_states.reshape(-1, hidden_states.shape[-1]), p=2, dim=-1 + ) + embedding = F.normalize(self.embedding.weight, p=2, dim=-1) + distances = ( + torch.sum(hidden_states**2, dim=1, keepdim=True) + + torch.sum(embedding**2, dim=1) + - 2 * torch.matmul(hidden_states, embedding.t()) + ) + return torch.argmin(distances, dim=1) + + +class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + Minimal GLM SigVQ image encoder used by LLaDA-Image editing. + + The model contains only the GLM vision encoder, VQ quantizer, and prior token projection used during inference. + Input images must already be RGB tensors normalized to `[-1, 1]`, have one common size, and be divisible by + `patch_size`. + """ + + _supports_gradient_checkpointing = True + _no_split_modules: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] + _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] + _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ + "patch_embed", + "position_embedding", + "norm", + "quantize", + ] + + @register_to_config + def __init__( + self, + image_size: int = 2048, + patch_size: int = 16, + in_channels: int = 3, + hidden_size: int = 1536, + intermediate_size: int = 6144, + num_hidden_layers: int = 40, + num_attention_heads: int = 16, + attention_bias: bool = True, + attention_dropout: float = 0.0, + norm_eps: float = 1e-6, + codebook_size: int = 16384, + codebook_embed_dim: int = 2048, + semantic_embed_dim: int = 4096, + ): + super().__init__() + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." + ) + + self.visual = nn.Module() + self.visual.patch_embed = LLaDAImageSigVQPatchEmbed( + in_channels, hidden_size, patch_size + ) + self.visual.embeddings = LLaDAImageSigVQEmbeddings( + image_size, patch_size, hidden_size + ) + self.visual.blocks = nn.ModuleList( + [ + LLaDAImageSigVQVisionBlock( + hidden_size, + intermediate_size, + num_attention_heads, + attention_bias, + attention_dropout, + norm_eps, + ) + for _ in range(num_hidden_layers) + ] + ) + + self.vqmodel = nn.Module() + self.vqmodel.quant_conv = nn.Conv2d( + hidden_size, codebook_embed_dim, kernel_size=1 + ) + self.vqmodel.quantize = LLaDAImageSigVQQuantizer( + codebook_size, codebook_embed_dim + ) + + self.prior_token_embedding = nn.Embedding(codebook_size, semantic_embed_dim) + self.prior_projector = FeedForward( + semantic_embed_dim, + semantic_embed_dim, + inner_dim=semantic_embed_dim, + activation_fn="linear-silu", + ) + self.gradient_checkpointing = False + + def forward( + self, + pixel_values: torch.Tensor | None = None, + token_ids: torch.Tensor | None = None, + return_dict: bool = True, + ) -> LLaDAImageSigVQOutput | tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + pixel_values (`torch.Tensor` of shape `(batch_size, 3, height, width)`, *optional*): + RGB images normalized to `[-1, 1]`. Mutually exclusive with `token_ids`. + token_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Precomputed VQ codebook IDs. Mutually exclusive with `pixel_values`. + return_dict (`bool`, defaults to `True`): + Whether to return [`LLaDAImageSigVQOutput`] instead of a tuple. + + Returns: + [`LLaDAImageSigVQOutput`] or `tuple`: + The projected semantic features and their discrete token IDs. + """ + if (pixel_values is None) == (token_ids is None): + raise ValueError("Provide exactly one of `pixel_values` or `token_ids`.") + + if pixel_values is not None: + if pixel_values.ndim != 4: + raise ValueError( + f"`pixel_values` must have 4 dimensions, got shape {tuple(pixel_values.shape)}." + ) + height, width = pixel_values.shape[-2:] + if ( + height % self.config.patch_size != 0 + or width % self.config.patch_size != 0 + ): + raise ValueError( + f"Image height and width must be divisible by {self.config.patch_size}, got {height}x{width}." + ) + + grid_height = height // self.config.patch_size + grid_width = width // self.config.patch_size + hidden_states = self.visual.patch_embed(pixel_values) + hidden_states = self.visual.embeddings( + hidden_states, grid_height, grid_width + ) + + for block in self.visual.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states + ) + else: + hidden_states = block(hidden_states) + + hidden_states = hidden_states.transpose(1, 2).reshape( + pixel_values.shape[0], + self.config.hidden_size, + grid_height, + grid_width, + ) + hidden_states = self.vqmodel.quant_conv(hidden_states) + token_ids = self.vqmodel.quantize(hidden_states).reshape( + pixel_values.shape[0], -1 + ) + elif token_ids.ndim != 2: + raise ValueError( + f"`token_ids` must have 2 dimensions, got shape {tuple(token_ids.shape)}." + ) + + semantic_features = self.prior_projector(self.prior_token_embedding(token_ids)) + + if not return_dict: + return semantic_features, token_ids + return LLaDAImageSigVQOutput( + semantic_features=semantic_features, token_ids=token_ids + ) + + +class LLaDAImageTransformer2DModel(_LLaDAImageTransformer2DModel): + """SGLang diffusion adapter that preserves the converted checkpoint layout.""" + + _fsdp_shard_conditions: ClassVar[list] = [ + lambda name, module: isinstance(module, LLaDAImageTransformerBlock) + ] + _compile_conditions: ClassVar[list] = [] + param_names_mapping: ClassVar[dict] = ( + LLaDAImageDitConfig().arch_config.param_names_mapping + ) + reverse_param_names_mapping: ClassVar[dict] = {} + + def __init__(self, config, hf_config: dict, quant_config=None): + init_kwargs = { + key: value for key, value in hf_config.items() if not key.startswith("_") + } + super().__init__(quant_config=quant_config, **init_kwargs) + self.sgl_config = config + self.hidden_size = int(self.config.dim) + self.num_attention_heads = int(self.config.n_heads) + self.num_channels_latents = int(self.config.in_channels) + + def post_load_weights(self) -> None: + """Run model-specific post-load fixups (none are required).""" + return + + @staticmethod + def _as_feature_list( + features: torch.Tensor | list[torch.Tensor], + attention_mask: torch.Tensor | list[torch.Tensor] | None = None, + ) -> list[torch.Tensor]: + if isinstance(features, list): + return features + if features.ndim == 2: + return [features] + if features.ndim != 3: + raise ValueError( + "LLaDA-Image condition features must have two or three dimensions" + ) + if attention_mask is None: + return list(features.unbind(dim=0)) + if isinstance(attention_mask, list): + attention_mask = torch.stack(attention_mask) + return [sample[mask.bool()] for sample, mask in zip(features, attention_mask)] + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + encoder_attention_mask: torch.Tensor | list[torch.Tensor] | None = None, + source_latents: list[torch.Tensor] | None = None, + guidance=None, + **kwargs, + ) -> torch.Tensor: + del guidance, kwargs + if hidden_states.ndim != 4: + raise ValueError( + "LLaDA-Image latents must have shape [batch, channels, height, width]" + ) + + model_dtype = next(self.parameters()).dtype + latent_list = [latent.unsqueeze(1).to(model_dtype) for latent in hidden_states] + cap_feats = self._as_feature_list(encoder_hidden_states, encoder_attention_mask) + glm_cap_feats = ( + None + if encoder_hidden_states_image is None + or ( + isinstance(encoder_hidden_states_image, list) + and not encoder_hidden_states_image + ) + else self._as_feature_list(encoder_hidden_states_image) + ) + output = ( + super() + .forward( + x=latent_list, + t=(timestep / 1000.0).to(model_dtype), + cap_feats=cap_feats, + glm_cap_feats=glm_cap_feats, + source_latents=source_latents, + ) + .sample + ) + return -torch.stack(output, dim=0).squeeze(2).float() + + +EntryClass = [ + LLaDAImageTransformer2DModel, + LLaDAImageQueryFormerModel, + LLaDAImageTextProjectionModel, + LLaDAImageSigVQModel, +] diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index e38d4d792eb5..493ecd7ee96a 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -30,6 +30,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.schedulers.scheduling_utils import SchedulerMixin from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler from sglang.multimodal_gen.runtime.post_training.scheduler_rl_mixin import ( @@ -38,6 +39,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +SP_STOCHASTIC_NOISE_KEY = "sp_stochastic_noise" @dataclass @@ -90,6 +92,9 @@ class FlowMatchEulerDiscreteScheduler( Whether to use exponential sigmas for step sizes in the noise schedule during sampling. use_beta_sigmas (`bool`, defaults to False): Whether to use beta sigmas for step sizes in the noise schedule during sampling. + use_uniform_sigmas (`bool`, defaults to False): + Whether to use a uniform pre-shift sigma grid that excludes the terminal zero. The terminal zero is + appended after shifting. time_shift_type (`str`, defaults to "exponential"): The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". stochastic_sampling (`bool`, defaults to False): @@ -114,6 +119,7 @@ def __init__( use_karras_sigmas: bool | None = False, use_exponential_sigmas: bool | None = False, use_beta_sigmas: bool | None = False, + use_uniform_sigmas: bool | None = False, time_shift_type: str = "exponential", stochastic_sampling: bool = False, ): @@ -123,12 +129,14 @@ def __init__( self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas, + self.config.use_uniform_sigmas, ] ) > 1 ): raise ValueError( - "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, " + "`config.use_karras_sigmas`, `config.use_uniform_sigmas` can be used." ) if time_shift_type not in {"exponential", "linear"}: raise ValueError( @@ -339,13 +347,18 @@ def set_timesteps( sigmas_array: np.ndarray if sigmas is None: - if timesteps_array is None: - timesteps_array = np.linspace( - self._sigma_to_t(self.sigma_max), - self._sigma_to_t(self.sigma_min), - num_inference_steps, - ) - sigmas_array = timesteps_array / self.config.num_train_timesteps + if timesteps_array is None and self.config.use_uniform_sigmas: + sigmas_array = np.linspace( + 1.0, 0.0, num_inference_steps + 1, dtype=np.float32 + )[:-1] + else: + if timesteps_array is None: + timesteps_array = np.linspace( + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, + ) + sigmas_array = timesteps_array / self.config.num_train_timesteps else: sigmas_array = np.array(sigmas).astype(np.float32) num_inference_steps = len(sigmas_array) @@ -443,6 +456,68 @@ def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: else: self._step_index = self._begin_index + @staticmethod + def _draw_stochastic_noise( + sample: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None, + batch, + ) -> torch.Tensor: + batch_extra = getattr(batch, "extra", None) + metadata = batch_extra.get(SP_STOCHASTIC_NOISE_KEY) if batch_extra else None + if metadata is None: + return randn_tensor( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) + + required_keys = {"full_shape", "dim", "start", "length"} + if not isinstance(metadata, dict) or not required_keys.issubset(metadata): + raise ValueError("Invalid SP stochastic noise metadata fields") + if not getattr(batch, "did_sp_shard_latents", False): + raise ValueError("Invalid SP stochastic noise metadata without sharding") + + full_shape = metadata["full_shape"] + dim = metadata["dim"] + start = metadata["start"] + length = metadata["length"] + if ( + not isinstance(full_shape, tuple) + or len(full_shape) != sample.ndim + or any(type(value) is not int or value <= 0 for value in full_shape) + ): + raise ValueError("Invalid SP stochastic noise metadata full shape") + if type(dim) is not int or dim < 1 or dim >= sample.ndim: + raise ValueError("Invalid SP stochastic noise metadata dimension") + if type(start) is not int or type(length) is not int: + raise ValueError("Invalid SP stochastic noise metadata slice") + if length != sample.shape[dim] or start < 0: + raise ValueError("Invalid SP stochastic noise metadata local slice") + if start + length > full_shape[dim]: + raise ValueError("Invalid SP stochastic noise metadata slice bounds") + raw_latent_shape = getattr(batch, "raw_latent_shape", None) + if raw_latent_shape is None or tuple(raw_latent_shape) != full_shape: + raise ValueError("Invalid SP stochastic noise metadata raw shape") + + expected_local_shape = list(full_shape) + expected_local_shape[dim] = length + if tuple(expected_local_shape) != tuple(sample.shape): + raise ValueError("Invalid SP stochastic noise metadata sample shape") + if isinstance(generator, list) and len(generator) != full_shape[0]: + raise ValueError("Invalid SP stochastic noise metadata generator count") + + full_noise = randn_tensor( + full_shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) + noise = full_noise.narrow(dim, start, length).contiguous() + if noise.shape != sample.shape: + raise ValueError("Invalid SP stochastic noise metadata result shape") + return noise + def step( self, model_output: torch.FloatTensor, @@ -452,7 +527,7 @@ def step( s_tmin: float = 0.0, s_tmax: float = float("inf"), s_noise: float = 1.0, - generator: torch.Generator | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, per_token_timesteps: torch.Tensor | None = None, batch=None, return_dict: bool = True, @@ -473,8 +548,8 @@ def step( s_tmax (`float`): s_noise (`float`, defaults to 1.0): Scaling factor for noise added to the sample. - generator (`torch.Generator`, *optional*): - A random number generator. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + A random number generator, or one generator per batch item. per_token_timesteps (`torch.Tensor`, *optional*): The timesteps for each token in the sample. return_dict (`bool`): @@ -533,12 +608,7 @@ def step( else: if self.config.stochastic_sampling: x0 = sample - current_sigma * model_output - noise = torch.randn( - sample.shape, - generator=generator, - device=sample.device, - dtype=sample.dtype, - ) + noise = self._draw_stochastic_noise(sample, generator, batch) prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise else: prev_sample = sample + dt * model_output diff --git a/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py b/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py new file mode 100644 index 000000000000..daf1a186961f --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import ClassVar + +from diffusers.image_processor import VaeImageProcessor +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_sp_parallel_rank, + get_sp_world_size, +) +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + SP_STOCHASTIC_NOISE_KEY, +) +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + DecodingStage, + DenoisingStage, + InputValidationStage, + LatentPreparationStage, + TimestepPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning import ( + LLaDAImageTextConditioningStage, + LLaDAImageTextEncoderRunner, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_source import ( + LLaDAImageSourceImageConditioningStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +class LLaDAImageLatentPreparationStage(LatentPreparationStage): + def forward(self, batch, server_args): + if batch.latents is None: + shape = server_args.pipeline_config.prepare_latent_shape( + batch, batch.batch_size, batch.num_frames + ) + model_dtype = next(self.transformer.parameters()).dtype + batch.latents = randn_tensor( + shape, + generator=batch.generator, + device=get_local_torch_device(), + dtype=model_dtype, + ).float() + batch = super().forward(batch, server_args) + batch.extra.pop(SP_STOCHASTIC_NOISE_KEY, None) + + sp_size = get_sp_world_size() + if self.scheduler.config.stochastic_sampling and sp_size > 1: + full_shape = tuple(int(value) for value in batch.raw_latent_shape) + if len(full_shape) != 4: + raise ValueError( + "LLaDA-Image SP noise requires four-dimensional latents" + ) + latent_height = full_shape[2] + if latent_height % sp_size != 0: + raise ValueError( + "LLaDA-Image SP noise requires height divisible by SP size" + ) + local_height = latent_height // sp_size + sp_rank = get_sp_parallel_rank() + batch.extra[SP_STOCHASTIC_NOISE_KEY] = { + "full_shape": full_shape, + "dim": 2, + "start": sp_rank * local_height, + "length": local_height, + } + + return batch + + +class LLaDAImagePipeline(ComposedPipelineBase): + pipeline_name = "LLaDAImagePipeline" + + _required_config_modules: ClassVar[list[str]] = [ + "queryformer", + "text_projection", + "sigvq", + "tokenizer", + "transformer", + "scheduler", + "vae", + ] + + def create_pipeline_stages(self, server_args: ServerArgs) -> None: + text_runner = LLaDAImageTextEncoderRunner( + model_root=self.model_path, + queryformer=self.get_module("queryformer"), + text_projection=self.get_module("text_projection"), + tokenizer=self.get_module("tokenizer"), + server_args=server_args, + ) + self.add_module("text_encoder", text_runner) + + self.add_stage(InputValidationStage(), "input_validation_stage") + self.add_stage( + LLaDAImageTextConditioningStage(text_runner), + "text_conditioning_stage", + ) + self.add_stage( + LLaDAImageSourceImageConditioningStage( + sigvq=self.get_module("sigvq"), + vae=self.get_module("vae"), + image_processor=VaeImageProcessor( + vae_scale_factor=server_args.pipeline_config.latent_scale_factor + ), + ), + "source_image_conditioning_stage", + ) + self.add_stage( + TimestepPreparationStage(self.get_module("scheduler")), + "timestep_preparation_stage", + ) + self.add_stage( + LLaDAImageLatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + ), + "latent_preparation_stage", + ) + self.add_stage( + DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + pipeline=self, + ), + "denoising_stage", + ) + self.add_stage( + DecodingStage( + vae=self.get_module("vae"), pipeline=self, component_name="vae" + ), + "decoding_stage", + ) + + +EntryClass = [LLaDAImagePipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 0e217caeb338..b027cab1c11a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -130,6 +130,7 @@ class Req: # Latent tensors latents: torch.Tensor | None = None + source_latents: list[torch.Tensor] | None = None y: torch.Tensor | None = None # Flux-2 latent_ids: torch.Tensor | None = None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py new file mode 100644 index 000000000000..69549bbf60de --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py @@ -0,0 +1,427 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +from array import array +from contextlib import nullcontext +from pathlib import Path + +import torch + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_sp_parallel_rank, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +def ensure_conditioning_mask_active(attn_backend) -> None: + """Fail closed if the conditioning block mask did not reach the backend.""" + if not attn_backend.conditioning_mask_active: + raise RuntimeError( + "LLaDA-Image conditioning ran without its block attention mask " + "because the forward batch lost the conditioning text lengths" + ) + + +def format_llada_image_prompt(prompt: str | None) -> str: + if prompt is None: + return "HUMAN Generate an image.\nASSISTANT\n" + return ( + f"HUMAN Generate an image: {prompt.strip()}\n" + "ASSISTANT\n" + ) + + +class LLaDAImageTextEncoderRunner: + """One-shot LLaDA2 prefill runner used to produce diffusion cap features.""" + + def __init__( + self, + model_root: str, + queryformer, + text_projection, + tokenizer, + server_args: ServerArgs, + ) -> None: + import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state + import sglang.srt.distributed.parallel_state as srt_parallel_state + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.chunk_cache import ChunkCache + from sglang.srt.runtime_context import publish as srt_publish + from sglang.srt.server_args import ServerArgs as SRTServerArgs + + # Keep the SRT world aligned with the pure-Ulysses ranks. Within that + # world, SP ranks become data replicas so conditioning matches SP1. + text_tp_size = int(server_args.sp_degree) + text_tp_rank = get_sp_parallel_rank() + text_dp_attention = text_tp_size > 1 + + self.queryformer = queryformer + self.text_projection = text_projection + self.tokenizer = tokenizer + text_encoder_path = str(Path(model_root) / "text_encoder") + device = get_local_torch_device() + gpu_id = device.index + if gpu_id is None: + gpu_id = int(os.environ.get("LOCAL_RANK", "0")) + srt_args = SRTServerArgs( + model_path=text_encoder_path, + tokenizer_path=str(Path(model_root) / "tokenizer"), + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + skip_tokenizer_init=True, + dtype="bfloat16", + tp_size=text_tp_size, + dp_size=text_tp_size if text_dp_attention else 1, + enable_dp_attention=text_dp_attention, + enable_dp_lm_head=text_dp_attention, + dcp_size=1, + attn_cp_size=1, + ep_size=1, + moe_dp_size=text_tp_size if text_dp_attention else 1, + moe_dense_tp_size=1 if text_dp_attention else None, + moe_a2a_backend="none", + pp_size=1, + attention_backend="llada2_cfg_flashinfer", + disable_cuda_graph=True, + disable_radix_cache=True, + chunked_prefill_size=-1, + max_prefill_tokens=8192, + max_total_tokens=8192, + max_running_requests=4 if text_dp_attention else 2, + mem_fraction_static=( + server_args.pipeline_config.text_encoder_mem_fraction_static + ), + ) + # TpModelWorker reads the runtime_context config bags the srt + # scheduler process publishes before construction. + srt_publish(srt_args, role="scheduler") + # The diffusion runtime mirrors its TP group into the srt globals. + # Clear it so the worker can install the encoder's real srt groups, + # then restore the mirror and keep the encoder group for forwards. + saved_srt_tp = srt_parallel_state._TP + saved_srt_attn_tp = srt_parallel_state._ATTN_TP + try: + mm_parallel_state._clear_srt_tp_group() + self.worker = TpModelWorker( + server_args=srt_args, + gpu_id=gpu_id, + ps=ParallelState.trivial( + tp_rank=text_tp_rank, + tp_size=text_tp_size, + dp_rank=text_tp_rank if text_dp_attention else 0, + dp_size=text_tp_size if text_dp_attention else 1, + attn_tp_rank=0 if text_dp_attention else text_tp_rank, + attn_tp_size=1 if text_dp_attention else text_tp_size, + attn_cp_rank=0, + attn_cp_size=1, + attn_dcp_rank=0, + attn_dcp_size=1, + attn_dp_rank=text_tp_rank if text_dp_attention else 0, + attn_dp_size=text_tp_size if text_dp_attention else 1, + moe_ep_rank=0, + moe_ep_size=1, + moe_dp_rank=text_tp_rank if text_dp_attention else 0, + moe_dp_size=text_tp_size if text_dp_attention else 1, + gpu_id=gpu_id, + ), + nccl_port=server_args.nccl_port or 29500, + ) + # Run the post-construction init phases the srt scheduler drives. + self.worker.alloc_memory_pool() + self.worker.init_attention_backends() + self.worker.init_cuda_graphs() + self.encoder_tp_group = srt_parallel_state._TP + self.encoder_attn_tp_group = srt_parallel_state._ATTN_TP + finally: + srt_parallel_state._TP = saved_srt_tp + srt_parallel_state._ATTN_TP = saved_srt_attn_tp + self.server_args = srt_args + self.text_dp_attention = text_dp_attention + self.model_runner = self.worker.model_runner + self.page_size = self.model_runner.page_size + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + self.worker.get_memory_pool() + ) + self.tree_cache = ChunkCache( + CacheInitParams( + disable=True, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + page_size=self.page_size, + ) + ) + + def _prepare_input_embeds( + self, + input_ids: torch.Tensor, + sequence_lengths: list[int], + text_lengths: list[int], + component_context=None, + ) -> torch.Tensor: + model = self.model_runner.model + embeddings = model.get_input_embeddings() + inputs = [] + offset = 0 + context = ( + nullcontext(self.queryformer) + if component_context is None + else component_context( + component_name="queryformer", module=self.queryformer + ) + ) + with context as queryformer: + assert queryformer is not None + queryformer_dtype = queryformer.dtype + for sequence_length, text_length in zip( + sequence_lengths, text_lengths, strict=True + ): + text_ids = input_ids[offset : offset + text_length] + text_embeds = embeddings(text_ids) + text_mask = torch.ones( + (1, text_length), dtype=torch.bool, device=text_embeds.device + ) + query_embeds = queryformer( + text_embeds.unsqueeze(0).to(dtype=queryformer_dtype), + text_mask, + ).query_embeds + query_count = sequence_length - text_length + if query_embeds.shape[1] != query_count: + raise RuntimeError( + f"QueryFormer returned {query_embeds.shape[1]} queries, " + f"expected {query_count}" + ) + inputs.append( + torch.cat( + [text_embeds, query_embeds.squeeze(0).to(text_embeds.dtype)], + dim=0, + ) + ) + offset += sequence_length + if offset != input_ids.numel(): + raise ValueError( + f"Conditioning spans {offset} tokens, got {input_ids.numel()}" + ) + return torch.cat(inputs, dim=0) + + @torch.no_grad() + def encode( + self, prompts: list[str], max_sequence_length: int, component_context=None + ): + import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state + import sglang.srt.distributed.parallel_state as srt_parallel_state + + # srt modules resolve their TP group through the srt globals, which + # hold the diffusion mirror between encodes. + with mm_parallel_state.use_tensor_parallel_group(self.encoder_tp_group): + saved_attn_tp = srt_parallel_state._ATTN_TP + try: + srt_parallel_state._ATTN_TP = self.encoder_attn_tp_group + return self._encode_impl( + prompts, max_sequence_length, component_context=component_context + ) + finally: + srt_parallel_state._ATTN_TP = saved_attn_tp + + def _encode_impl( + self, prompts: list[str], max_sequence_length: int, component_context=None + ): + from sglang.srt.managers.overlap_utils import resolve_forward_inputs + from sglang.srt.managers.schedule_batch import Req as SRTReq + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder + from sglang.srt.mem_cache.common import release_kv_cache + from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + ForwardMode, + ) + from sglang.srt.sampling.sampling_params import ( + SamplingParams as SRTSamplingParams, + ) + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + tokenized = self.tokenizer( + prompts, + add_special_tokens=True, + padding=False, + truncation=True, + max_length=max_sequence_length, + ).input_ids + query_count = int(self.queryformer.config.num_queries) + reqs = [] + text_lengths = [] + sequence_lengths = [] + for index, text_input_ids in enumerate(tokenized): + sampling_params = SRTSamplingParams(max_new_tokens=1, temperature=0.0) + sampling_params.normalize(None) + sequence_ids = list(text_input_ids) + [0] * query_count + req = SRTReq( + rid=f"llada-image-cap-{index}", + origin_input_text="", + origin_input_ids=array("q", sequence_ids), + sampling_params=sampling_params, + vocab_size=self.worker.model_config.vocab_size, + eos_token_ids=set(), + return_hidden_states=True, + dllm_config=None, + ) + req.skip_radix_cache_insert = True + reqs.append(req) + text_lengths.append(len(text_input_ids)) + sequence_lengths.append(len(sequence_ids)) + + chunked_prefill_size = self.server_args.chunked_prefill_size + if chunked_prefill_size <= 0: + chunked_prefill_size = None + + adder = PrefillAdder( + self.page_size, + self.tree_cache, + self.token_to_kv_pool_allocator, + None, + 0.5, + self.server_args.max_prefill_tokens, + chunked_prefill_size, + prefill_max_requests=len(reqs), + dllm_config=None, + ) + for req in reqs: + req.init_next_round_input(self.tree_cache) + result = adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + if req not in adder.can_run_list: + raise RuntimeError( + "Insufficient prefill capacity for LLaDA-Image conditioning" + ) + if result != AddReqResult.CONTINUE: + break + + if len(adder.can_run_list) != len(reqs): + raise RuntimeError("Could not batch all LLaDA-Image conditioning prompts") + + batch = ScheduleBatch.init_new( + reqs=adder.can_run_list, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + tree_cache=self.tree_cache, + model_config=self.worker.model_config, + enable_overlap=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + dllm_config=None, + ) + if self.text_dp_attention: + token_count = sum(sequence_lengths) + batch.global_num_tokens = [token_count] * self.server_args.dp_size + batch.global_num_tokens_for_logprob = [len(reqs)] * ( + self.server_args.dp_size + ) + try: + batch.prepare_for_extend() + batch.forward_mode = ForwardMode.EXTEND + # The FutureMap argument is only consulted on decode/spec paths. + resolve_forward_inputs(batch, None) + forward_batch = ForwardBatch.init_new( + batch, + self.model_runner, + return_hidden_states_before_norm=False, + ) + forward_batch.input_embeds = self._prepare_input_embeds( + forward_batch.input_ids, + sequence_lengths, + text_lengths, + component_context=component_context, + ) + forward_batch.llada_image_conditioning_text_lens_cpu = text_lengths + model_output = self.model_runner.forward(forward_batch=forward_batch) + ensure_conditioning_mask_active(self.model_runner.attn_backend) + hidden_states = model_output.logits_output.hidden_states + if not isinstance(hidden_states, torch.Tensor): + raise TypeError( + "SGLang did not return full LLaDA-Image conditioning states" + ) + + projected = [] + offset = 0 + context = ( + nullcontext(self.text_projection) + if component_context is None + else component_context( + component_name="text_projection", module=self.text_projection + ) + ) + with context as text_projection: + assert text_projection is not None + projection_dtype = text_projection.dtype + for sequence_length in sequence_lengths: + output = text_projection( + hidden_states[offset : offset + sequence_length] + .unsqueeze(0) + .to(dtype=projection_dtype) + ).hidden_states + projected.append(output.squeeze(0)) + offset += sequence_length + return projected + finally: + for req in reqs: + if req.kv is not None: + release_kv_cache(req, self.tree_cache, is_insert=False) + elif req.req_pool_idx is not None: + # KV alloc failed after the slot grab. Freeing the bare + # slot avoids the release_kv_cache lifecycle assert. + self.req_to_token_pool.free(req) + + +class LLaDAImageTextConditioningStage(PipelineStage): + def __init__(self, runner: LLaDAImageTextEncoderRunner) -> None: + super().__init__() + self.runner = runner + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + del server_args + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse(stage_name, "queryformer"), + ComponentUse(stage_name, "text_projection"), + ] + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if not isinstance(batch.prompt, str): + raise TypeError("LLaDA-Image currently supports one prompt per request") + do_cfg = batch.guidance_scale > 1.0 + prompts = [format_llada_image_prompt(batch.prompt)] + if do_cfg: + prompts.append(format_llada_image_prompt(batch.negative_prompt)) + + outputs = self.runner.encode( + prompts, + max_sequence_length=batch.max_sequence_length or 2048, + component_context=self.use_declared_component, + ) + masks = [ + torch.ones(output.shape[0], dtype=torch.bool, device=output.device) + for output in outputs + ] + output_count = batch.num_outputs_per_prompt + batch.prompt_embeds = [outputs[0]] * output_count + batch.prompt_attention_mask = [masks[0]] * output_count + batch.do_classifier_free_guidance = do_cfg + if do_cfg: + batch.negative_prompt_embeds = [outputs[1]] * output_count + batch.negative_attention_mask = [masks[1]] * output_count + else: + batch.negative_prompt_embeds = [] + batch.negative_attention_mask = [] + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_source.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_source.py new file mode 100644 index 000000000000..138c854ed029 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_source.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.server_args import ServerArgs + + +def _module_device_dtype(module: torch.nn.Module) -> tuple[torch.device, torch.dtype]: + parameter = next(module.parameters()) + return parameter.device, parameter.dtype + + +class LLaDAImageSourceImageConditioningStage(PipelineStage): + def __init__(self, sigvq, vae, image_processor) -> None: + super().__init__() + self.sigvq = sigvq + self.vae = vae + self.image_processor = image_processor + + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + del server_args + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse(stage_name, "sigvq"), + ComponentUse(stage_name, "vae"), + ] + + @staticmethod + def _patchify_latents(latents: torch.Tensor) -> torch.Tensor: + batch_size, channels, height, width = latents.shape + if height % 2 != 0 or width % 2 != 0: + raise ValueError("LLaDA-Image source VAE latents must have even dimensions") + latents = latents.reshape(batch_size, channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + return latents.reshape(batch_size, channels * 4, height // 2, width // 2) + + @staticmethod + def _normalize_latents(latents: torch.Tensor, vae) -> torch.Tensor: + if not hasattr(vae, "bn"): + raise ValueError("LLaDA-Image editing requires the Flux2 VAE BN state") + vae_config = getattr(vae.config, "arch_config", vae.config) + eps = vae_config.batch_norm_eps + latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) + latent_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + eps).to(latents) + return (latents - latent_mean) / latent_std + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + del server_args + if batch.condition_image is None: + batch.source_latents = None + return batch + + source_images = ( + batch.condition_image + if isinstance(batch.condition_image, list) + else [batch.condition_image] + ) + if len(source_images) != 1: + raise ValueError("LLaDA-Image editing requires exactly one source image") + + image = self.image_processor.preprocess( + source_images[0], + height=batch.height, + width=batch.width, + resize_mode="crop", + ) + if image.shape[0] == 1 and batch.batch_size > 1: + image = image.repeat(batch.batch_size, 1, 1, 1) + if image.shape[0] != batch.batch_size: + raise ValueError( + "LLaDA-Image source image batch must contain one image or match " + f"the effective batch size {batch.batch_size}, got {image.shape[0]}" + ) + + sigvq_pixels = F.interpolate( + image.float(), + size=(batch.height // 2, batch.width // 2), + mode="bilinear", + align_corners=False, + ) + with self.use_declared_component( + component_name="sigvq", module=self.sigvq + ) as sigvq: + assert sigvq is not None + patch_size = sigvq.config.patch_size + pad_height = (-sigvq_pixels.shape[-2]) % patch_size + pad_width = (-sigvq_pixels.shape[-1]) % patch_size + if pad_height or pad_width: + sigvq_pixels = F.pad( + sigvq_pixels, + (0, pad_width, 0, pad_height), + mode="replicate", + ) + sigvq_device, sigvq_dtype = _module_device_dtype(sigvq) + semantic_features = sigvq( + sigvq_pixels.to(device=sigvq_device, dtype=sigvq_dtype) + ).semantic_features + + with self.use_declared_component(component_name="vae", module=self.vae) as vae: + assert vae is not None + vae_device, vae_dtype = _module_device_dtype(vae) + posterior = vae.encode(image.to(device=vae_device, dtype=vae_dtype)) + latent_dist = getattr(posterior, "latent_dist", posterior) + source_latents = self._patchify_latents(latent_dist.mode()) + source_latents = self._normalize_latents(source_latents, vae) + + batch.image_embeds = list(semantic_features.unbind(dim=0)) + batch.source_latents = [ + latent.unsqueeze(1) for latent in source_latents.unbind(dim=0) + ] + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index 21430f51c825..378ba2540fb7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py @@ -102,7 +102,13 @@ def forward( sigmas = batch.sigmas n_tokens = batch.n_tokens - sigmas = server_args.pipeline_config.prepare_sigmas(sigmas, num_inference_steps) + use_uniform_sigmas = getattr( + getattr(scheduler, "config", None), "use_uniform_sigmas", False + ) + if sigmas is not None or not use_uniform_sigmas: + sigmas = server_args.pipeline_config.prepare_sigmas( + sigmas, num_inference_steps + ) batch.sigmas = sigmas # Prepare extra kwargs for set_timesteps diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 2111beb2ccd6..7475dd39ab76 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -367,6 +367,10 @@ def update_weights_from_disk( """Update model weights from disk without restarting the server.""" logger.info(f"Updating weights from disk: {model_path}") + rejected, message = self._reject_if_partial_update() + if rejected: + return False, message + try: modules_to_update = self._collect_modules(target_modules) except ValueError as e: @@ -440,6 +444,17 @@ def update_weights_from_disk( logger.info(message) return success, message + def _reject_if_partial_update(self) -> tuple[bool, str]: + config = self.pipeline.server_args.pipeline_config + if config.supports_hot_weight_updates(): + return False, "" + message = ( + "pipeline holds components outside generic module discovery, a " + "runtime weight update would leave them at the old checkpoint" + ) + logger.error(message) + return True, message + def _collect_modules( self, target_modules: list[str] | None ) -> list[tuple[str, torch.nn.Module]]: @@ -527,6 +542,9 @@ def update_weights_from_tensor( lora_alpha: int | None = None, lora_rank: int | None = None, ) -> tuple[bool, str]: + rejected, message = self._reject_if_partial_update() + if rejected: + return False, message if weight_update_mode == LORA_MERGE_WEIGHT_UPDATE_MODE: return self._update_lora_from_tensor( named_tensors=named_tensors, diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index b5fe29fe31cb..23f536fb0101 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -399,6 +399,9 @@ class ServerArgs(DisaggServerArgsMixin): use_fsdp_inference: bool | None = None pin_cpu_memory: bool = True ltx2_two_stage_device_mode: str | None = None + llada_image_max_pixel_area: int | None = None + llada_image_max_text_tokens: int | None = None + llada_image_max_total_pixel_area: int | None = None _explicit_arg_names: set[str] = field(default_factory=set, repr=False) _automatic_component_attention_backend_keys: set[str] = field( default_factory=set, init=False, repr=False @@ -2519,6 +2522,35 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "Default is auto: resident on H200/high-memory CUDA GPUs, otherwise original." ), ) + parser.add_argument( + "--llada-image-max-pixel-area", + type=int, + default=ServerArgs.llada_image_max_pixel_area, + help=( + "LLaDA-Image only. Largest accepted output pixel area per " + "request. Defaults to the pipeline config value of 2048*2048." + ), + ) + parser.add_argument( + "--llada-image-max-text-tokens", + type=int, + default=ServerArgs.llada_image_max_text_tokens, + help=( + "LLaDA-Image only. Largest accepted request " + "max_sequence_length. Defaults to the pipeline config value " + "of 3584 and cannot exceed it." + ), + ) + parser.add_argument( + "--llada-image-max-total-pixel-area", + type=int, + default=ServerArgs.llada_image_max_total_pixel_area, + help=( + "LLaDA-Image only. Largest accepted output pixel area times " + "the per-prompt output count. Defaults to the pipeline " + "config value of ten 1024x1024 outputs." + ), + ) parser.add_argument( "--disable-autocast", action=StoreBoolean, diff --git a/python/sglang/multimodal_gen/test/test_llada_image.py b/python/sglang/multimodal_gen/test/test_llada_image.py new file mode 100644 index 000000000000..a7fa54808b23 --- /dev/null +++ b/python/sglang/multimodal_gen/test/test_llada_image.py @@ -0,0 +1,693 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from dataclasses import fields +from itertools import pairwise +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) +from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams +from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.models.dits.llada_image import ( + LLaDAImageTransformerBlock, + _LLaDAImageTransformer2DModel, +) +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( + SP_STOCHASTIC_NOISE_KEY, + FlowMatchEulerDiscreteScheduler, +) +from sglang.multimodal_gen.runtime.pipelines.llada_image import ( + LLaDAImageLatentPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning import ( + format_llada_image_prompt, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( + TimestepPreparationStage, +) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.server_args import ( + get_global_server_args, + set_global_server_args, +) + + +class _CaptureSPBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.skip_values = [] + self.replicated_suffixes = [] + + def forward(self, hidden_states, *args, **kwargs): + self.skip_values.append(kwargs.get("skip_sequence_parallel_override", False)) + self.replicated_suffixes.append(kwargs.get("num_replicated_suffix", 0)) + return hidden_states + + +class TestLLaDAImage(unittest.TestCase): + @classmethod + def setUpClass(cls): + try: + cls.previous_server_args = get_global_server_args() + except ValueError: + cls.previous_server_args = None + set_global_server_args(SimpleNamespace(kv_gather_degree=1, sp_split_auto=False)) + + @classmethod + def tearDownClass(cls): + set_global_server_args(cls.previous_server_args) + + def test_dit_supports_sglang_flash_attention_and_sdpa(self): + backends = ( + LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends + ) + self.assertEqual( + backends, + {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + def test_dit_weight_mapping_fuses_qkv_and_swiglu_inputs(self): + mapping = get_param_names_mapping( + LLaDAImagePipelineConfig().dit_config.arch_config.param_names_mapping + ) + + self.assertEqual( + mapping("layers.0.attention.to_k.weight"), + ("layers.0.attention.to_qkv.weight", 1, 3), + ) + self.assertEqual( + mapping("layers.0.feed_forward.w1.weight"), + ("layers.0.feed_forward.w13.weight", 0, 2), + ) + self.assertEqual( + mapping("layers.0.feed_forward.w3.weight"), + ("layers.0.feed_forward.w13.weight", 1, 2), + ) + + def test_prompt_format_matches_official_pipeline(self): + self.assertEqual( + format_llada_image_prompt("a red car"), + "HUMAN Generate an image: a red car\n" + "ASSISTANT\n", + ) + self.assertEqual( + format_llada_image_prompt(None), + "HUMAN Generate an image.\nASSISTANT\n", + ) + + def test_pipeline_config_uses_llada_image_schedule_and_shape(self): + config = LLaDAImagePipelineConfig() + sigmas = config.prepare_sigmas(None, num_inference_steps=8) + + self.assertEqual(len(sigmas), 8) + self.assertTrue(all(left > right for left, right in pairwise(sigmas))) + self.assertEqual( + config.prepare_latent_shape( + SimpleNamespace(height=1024, width=768), + batch_size=1, + num_frames=1, + ), + (1, 128, 64, 48), + ) + + def test_timestep_stage_uses_uniform_scheduler_schedule(self): + scheduler = FlowMatchEulerDiscreteScheduler( + shift=3.0, + use_uniform_sigmas=True, + ) + stage = TimestepPreparationStage(scheduler) + batch = SimpleNamespace( + scheduler=None, + timesteps=None, + sigmas=None, + num_inference_steps=4, + n_tokens=None, + extra={}, + is_warmup=True, + rollout=False, + ) + server_args = SimpleNamespace(pipeline_config=LLaDAImagePipelineConfig()) + module = ( + "sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation" + ) + + with ( + patch(f"{module}.get_local_torch_device", return_value=torch.device("cpu")), + patch(f"{module}.get_or_create_request_scheduler", return_value=scheduler), + ): + stage.forward(batch, server_args) + + self.assertIsNone(batch.sigmas) + torch.testing.assert_close( + scheduler.sigmas, + torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0]), + ) + + def test_stochastic_scheduler_accepts_per_sample_generators(self): + for seeds in ((11,), (11, 29)): + with self.subTest(batch_size=len(seeds)): + scheduler = FlowMatchEulerDiscreteScheduler( + shift=3.0, + use_uniform_sigmas=True, + stochastic_sampling=True, + ) + scheduler.set_timesteps(4, device="cpu") + + sample = torch.arange(len(seeds) * 8, dtype=torch.float32).reshape( + len(seeds), 2, 2, 2 + ) + model_output = sample / 16 + generators = [ + torch.Generator(device="cpu").manual_seed(seed) for seed in seeds + ] + reference_noise = torch.cat( + [ + torch.randn( + (1, *sample.shape[1:]), + generator=torch.Generator(device="cpu").manual_seed(seed), + dtype=sample.dtype, + ) + for seed in seeds + ] + ) + current_sigma, next_sigma = scheduler.sigmas[:2] + expected = (1.0 - next_sigma) * ( + sample - current_sigma * model_output + ) + next_sigma * reference_noise + + actual = scheduler.step( + model_output=model_output, + timestep=scheduler.timesteps[0], + sample=sample, + generator=generators, + return_dict=False, + )[0] + + self.assertEqual(actual.shape, sample.shape) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_latent_stage_records_sp_stochastic_noise_slice(self): + scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True) + stage = LLaDAImageLatentPreparationStage( + scheduler=scheduler, + transformer=None, + ) + batch = SimpleNamespace( + latents=torch.zeros(1, 2, 4, 3), + raw_latent_shape=torch.Size((1, 2, 4, 3)), + extra={}, + ) + parent_forward = ( + "sglang.multimodal_gen.runtime.pipelines.llada_image." + "LatentPreparationStage.forward" + ) + module = "sglang.multimodal_gen.runtime.pipelines.llada_image" + + with ( + patch(parent_forward, return_value=batch), + patch(f"{module}.get_sp_world_size", return_value=2), + patch(f"{module}.get_sp_parallel_rank", return_value=1), + ): + actual = stage.forward(batch, server_args=None) + + self.assertIs(actual, batch) + self.assertEqual( + batch.extra[SP_STOCHASTIC_NOISE_KEY], + { + "full_shape": (1, 2, 4, 3), + "dim": 2, + "start": 2, + "length": 2, + }, + ) + + with ( + patch(parent_forward, return_value=batch), + patch(f"{module}.get_sp_world_size", return_value=1), + ): + stage.forward(batch, server_args=None) + + self.assertNotIn(SP_STOCHASTIC_NOISE_KEY, batch.extra) + + batch.extra[SP_STOCHASTIC_NOISE_KEY] = {"stale": True} + stage.scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=False) + with ( + patch(parent_forward, return_value=batch), + patch(f"{module}.get_sp_world_size", return_value=2), + ): + stage.forward(batch, server_args=None) + + self.assertNotIn(SP_STOCHASTIC_NOISE_KEY, batch.extra) + + def test_stochastic_scheduler_matches_spatial_sp_slices(self): + full_shape = (1, 2, 4, 3) + + def make_stream(): + generators = [torch.Generator(device="cpu").manual_seed(71)] + initial = torch.randn( + full_shape, + generator=generators[0], + dtype=torch.float32, + ) + return generators, initial + + full_generators, full_sample = make_stream() + rank0_generators, rank0_full = make_stream() + rank1_generators, rank1_full = make_stream() + rank0_sample = rank0_full[:, :, :2, :].contiguous() + rank1_sample = rank1_full[:, :, 2:, :].contiguous() + + schedulers = [ + FlowMatchEulerDiscreteScheduler( + shift=3.0, + use_uniform_sigmas=True, + stochastic_sampling=True, + ) + for _ in range(3) + ] + for scheduler in schedulers: + scheduler.set_timesteps(4, device="cpu") + full_scheduler, rank0_scheduler, rank1_scheduler = schedulers + + def rank_batch(start): + return SimpleNamespace( + rollout=False, + did_sp_shard_latents=True, + raw_latent_shape=torch.Size(full_shape), + extra={ + SP_STOCHASTIC_NOISE_KEY: { + "full_shape": full_shape, + "dim": 2, + "start": start, + "length": 2, + } + }, + ) + + rank0_batch = rank_batch(0) + rank1_batch = rank_batch(2) + + for step_index in range(4): + full_model_output = full_sample / (step_index + 2) + rank0_model_output = full_model_output[:, :, :2, :].contiguous() + rank1_model_output = full_model_output[:, :, 2:, :].contiguous() + + full_sample = full_scheduler.step( + model_output=full_model_output, + timestep=full_scheduler.timesteps[step_index], + sample=full_sample, + generator=full_generators, + return_dict=False, + )[0] + rank0_sample = rank0_scheduler.step( + model_output=rank0_model_output, + timestep=rank0_scheduler.timesteps[step_index], + sample=rank0_sample, + generator=rank0_generators, + batch=rank0_batch, + return_dict=False, + )[0] + rank1_sample = rank1_scheduler.step( + model_output=rank1_model_output, + timestep=rank1_scheduler.timesteps[step_index], + sample=rank1_sample, + generator=rank1_generators, + batch=rank1_batch, + return_dict=False, + )[0] + + torch.testing.assert_close( + torch.cat([rank0_sample, rank1_sample], dim=2), + full_sample, + rtol=0, + atol=0, + ) + self.assertTrue( + torch.equal( + full_generators[0].get_state(), + rank0_generators[0].get_state(), + ) + ) + self.assertTrue( + torch.equal( + full_generators[0].get_state(), + rank1_generators[0].get_state(), + ) + ) + + def test_stochastic_scheduler_without_opt_in_keeps_local_draw(self): + scheduler = FlowMatchEulerDiscreteScheduler( + shift=3.0, + use_uniform_sigmas=True, + stochastic_sampling=True, + ) + scheduler.set_timesteps(4, device="cpu") + sample = torch.arange(12, dtype=torch.float32).reshape(1, 2, 2, 3) + model_output = sample / 16 + generator = [torch.Generator(device="cpu").manual_seed(23)] + reference_noise = torch.randn( + sample.shape, + generator=torch.Generator(device="cpu").manual_seed(23), + ) + current_sigma, next_sigma = scheduler.sigmas[:2] + expected = (1.0 - next_sigma) * ( + sample - current_sigma * model_output + ) + next_sigma * reference_noise + batch = SimpleNamespace( + rollout=False, + did_sp_shard_latents=True, + raw_latent_shape=torch.Size((1, 2, 4, 3)), + extra={}, + ) + + actual = scheduler.step( + model_output=model_output, + timestep=scheduler.timesteps[0], + sample=sample, + generator=generator, + batch=batch, + return_dict=False, + )[0] + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_stochastic_scheduler_rejects_invalid_sp_metadata(self): + full_shape = (1, 2, 4, 3) + valid = { + "full_shape": full_shape, + "dim": 2, + "start": 0, + "length": 2, + } + cases = ( + ("metadata type", "bad", True, full_shape, 1), + ("missing field", {"full_shape": full_shape}, True, full_shape, 1), + ( + "full shape type", + {**valid, "full_shape": list(full_shape)}, + True, + full_shape, + 1, + ), + ( + "full shape ndim", + {**valid, "full_shape": (1, 2, 4)}, + True, + full_shape, + 1, + ), + ( + "full shape value", + {**valid, "full_shape": (1, 2, 0, 3)}, + True, + full_shape, + 1, + ), + ("without sharding", valid, False, full_shape, 1), + ("batch dimension", {**valid, "dim": 0}, True, full_shape, 1), + ("dimension type", {**valid, "dim": 2.0}, True, full_shape, 1), + ("dimension range", {**valid, "dim": 4}, True, full_shape, 1), + ("slice type", {**valid, "start": 0.0}, True, full_shape, 1), + ("negative start", {**valid, "start": -1}, True, full_shape, 1), + ("slice bounds", {**valid, "start": 3}, True, full_shape, 1), + ("local length", {**valid, "length": 1}, True, full_shape, 1), + ("raw shape", valid, True, (1, 2, 6, 3), 1), + ( + "sample shape", + {**valid, "full_shape": (1, 3, 4, 3)}, + True, + (1, 3, 4, 3), + 1, + ), + ("generator count", valid, True, full_shape, 2), + ) + + for name, metadata, did_shard, raw_shape, generator_count in cases: + with self.subTest(name=name): + scheduler = FlowMatchEulerDiscreteScheduler( + shift=3.0, + use_uniform_sigmas=True, + stochastic_sampling=True, + ) + scheduler.set_timesteps(4, device="cpu") + sample = torch.zeros(1, 2, 2, 3) + generators = [ + torch.Generator(device="cpu").manual_seed(seed) + for seed in range(generator_count) + ] + batch = SimpleNamespace( + rollout=False, + did_sp_shard_latents=did_shard, + raw_latent_shape=torch.Size(raw_shape), + extra={SP_STOCHASTIC_NOISE_KEY: metadata}, + ) + + with self.assertRaisesRegex( + ValueError, "Invalid SP stochastic noise metadata" + ): + scheduler.step( + model_output=torch.zeros_like(sample), + timestep=scheduler.timesteps[0], + sample=sample, + generator=generators, + batch=batch, + return_dict=False, + ) + + def test_pipeline_config_rejects_unaligned_resolution(self): + config = LLaDAImagePipelineConfig() + with self.assertRaisesRegex(ValueError, "height must be divisible by 16"): + config.prepare_latent_shape( + SimpleNamespace(height=1023, width=1024), + batch_size=1, + num_frames=1, + ) + + def test_service_sampling_params_do_not_expose_vq_mode(self): + field_names = {field.name for field in fields(LLaDAImageSamplingParams)} + + self.assertNotIn("generation_mode", field_names) + self.assertNotIn("vq_token_ids", field_names) + + def test_dit_block_forwards_sequence_parallel_controls(self): + class CaptureAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.received = None + + def forward( + self, + hidden_states, + attention_mask, + freqs_cis, + num_replicated_suffix=0, + skip_sequence_parallel_override=False, + ): + self.received = ( + num_replicated_suffix, + skip_sequence_parallel_override, + ) + return torch.zeros_like(hidden_states) + + block = object.__new__(LLaDAImageTransformerBlock) + torch.nn.Module.__init__(block) + block.modulation = False + block.attention = CaptureAttention() + block.attention_norm1 = torch.nn.Identity() + block.attention_norm2 = torch.nn.Identity() + block.ffn_norm1 = torch.nn.Identity() + block.ffn_norm2 = torch.nn.Identity() + block.feed_forward = torch.nn.Identity() + + block( + torch.ones(1, 4, 8), + attention_mask=None, + freqs_cis=torch.empty(0), + num_replicated_suffix=32, + skip_sequence_parallel_override=True, + ) + + self.assertEqual(block.attention.received, (32, True)) + + def test_edit_sp_shards_images_and_replicates_conditions(self): + model_module = "sglang.multimodal_gen.runtime.models.dits.llada_image" + linear_module = "sglang.multimodal_gen.runtime.layers.linear" + with ( + patch(f"{model_module}.get_tp_world_size", return_value=1), + patch(f"{linear_module}.get_tp_group", return_value=None), + patch(f"{linear_module}.get_group_size", return_value=1), + patch(f"{linear_module}.get_group_rank", return_value=0), + patch( + "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", + return_value=1, + ), + patch( + "sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args", + return_value=SimpleNamespace(attention_backend="torch_sdpa"), + ), + ): + model = _LLaDAImageTransformer2DModel( + in_channels=4, + dim=64, + n_layers=1, + n_refiner_layers=1, + n_heads=2, + cap_feat_dim=8, + semantic_feat_dim=10, + axes_dims=(8, 12, 12), + axes_lens=(256, 32, 32), + ) + + noise_refiner = _CaptureSPBlock() + context_refiner = _CaptureSPBlock() + sigvq_refiner = _CaptureSPBlock() + main_block = _CaptureSPBlock() + model.noise_refiner = torch.nn.ModuleList([noise_refiner]) + model.context_refiner = torch.nn.ModuleList([context_refiner]) + model.sigvq_refiner = torch.nn.ModuleList([sigvq_refiner]) + model.layers = torch.nn.ModuleList([main_block]) + + with ( + patch(f"{model_module}.get_sp_world_size", return_value=2), + patch(f"{model_module}.get_sp_parallel_rank", return_value=0), + torch.no_grad(), + ): + model( + x=[torch.randn(4, 1, 4, 4)], + t=torch.tensor([0.5]), + cap_feats=[torch.randn(3, 8)], + glm_cap_feats=[torch.randn(5, 10)], + source_latents=[torch.randn(4, 1, 4, 4)], + ) + + self.assertEqual(noise_refiner.skip_values, [False]) + self.assertEqual(context_refiner.skip_values, [True]) + self.assertEqual(sigvq_refiner.skip_values, [True]) + self.assertEqual(main_block.skip_values, [False]) + self.assertEqual(main_block.replicated_suffixes, [96]) + + def test_edit_skips_sigvq_refiner_for_empty_cfg_condition(self): + model_module = "sglang.multimodal_gen.runtime.models.dits.llada_image" + linear_module = "sglang.multimodal_gen.runtime.layers.linear" + with ( + patch(f"{model_module}.get_tp_world_size", return_value=1), + patch(f"{linear_module}.get_tp_group", return_value=None), + patch(f"{linear_module}.get_group_size", return_value=1), + patch(f"{linear_module}.get_group_rank", return_value=0), + patch( + "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", + return_value=1, + ), + patch( + "sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args", + return_value=SimpleNamespace(attention_backend="torch_sdpa"), + ), + ): + model = _LLaDAImageTransformer2DModel( + in_channels=4, + dim=64, + n_layers=1, + n_refiner_layers=1, + n_heads=2, + cap_feat_dim=8, + semantic_feat_dim=10, + axes_dims=(8, 12, 12), + axes_lens=(256, 32, 32), + ) + + noise_refiner = _CaptureSPBlock() + context_refiner = _CaptureSPBlock() + sigvq_refiner = _CaptureSPBlock() + main_block = _CaptureSPBlock() + model.noise_refiner = torch.nn.ModuleList([noise_refiner]) + model.context_refiner = torch.nn.ModuleList([context_refiner]) + model.sigvq_refiner = torch.nn.ModuleList([sigvq_refiner]) + model.layers = torch.nn.ModuleList([main_block]) + + with ( + patch(f"{model_module}.get_sp_world_size", return_value=1), + patch(f"{model_module}.get_sp_parallel_rank", return_value=0), + torch.no_grad(), + ): + model( + x=[torch.randn(4, 1, 4, 4)], + t=torch.tensor([0.5]), + cap_feats=[torch.randn(3, 8)], + glm_cap_feats=[torch.empty(0, 10)], + source_latents=[torch.randn(4, 1, 4, 4)], + ) + + self.assertEqual(noise_refiner.skip_values, [False]) + self.assertEqual(context_refiner.skip_values, [True]) + self.assertEqual(sigvq_refiner.skip_values, []) + self.assertEqual(main_block.skip_values, [False]) + + def test_sp_uses_global_height_coordinates_for_generation_and_edit(self): + calls = [] + + def pad_with_ids( + features, grid_size, start, noise_value=None, sequence_multiple=32 + ): + del sequence_multiple + length = len(features) + calls.append((grid_size, start)) + return ( + features, + torch.zeros((length, 3), dtype=torch.int32), + torch.zeros(length, dtype=torch.bool), + length, + [noise_value] * length if noise_value is not None else None, + ) + + subject = SimpleNamespace( + _patchify_image=lambda _image, _patch, _f_patch: ( + torch.zeros(8, 1), + (1, 2, 4), + (1, 2, 4), + ), + _pad_with_ids=pad_with_ids, + ) + module = "sglang.multimodal_gen.runtime.models.dits.llada_image" + with ( + patch(f"{module}.get_sp_world_size", return_value=2), + patch(f"{module}.get_sp_parallel_rank", return_value=1), + ): + _LLaDAImageTransformer2DModel._prepare_t2i_sequences( + subject, + [torch.zeros(1, 1, 2, 4)], + cap_feats=None, + glm_features=None, + patch_size=1, + f_patch_size=1, + ) + + self.assertEqual(calls[-1], ((1, 2, 4), (1, 2, 0))) + + calls.clear() + with ( + patch(f"{module}.get_sp_world_size", return_value=2), + patch(f"{module}.get_sp_parallel_rank", return_value=1), + ): + _LLaDAImageTransformer2DModel._prepare_editing_sequences( + subject, + [torch.zeros(1, 1, 2, 4)], + cap_feats=[torch.zeros(1, 1)], + glm_cap_feats=[torch.zeros(1, 1)], + source_latents=[torch.zeros(1, 1, 2, 4)], + patch_size=1, + f_patch_size=1, + ) + + image_calls = [call for call in calls if call[0] == (1, 2, 4)] + self.assertEqual(len(image_calls), 2) + self.assertTrue(all(start[1:] == (2, 0) for _, start in image_calls)) + self.assertEqual(calls[-1], ((1, 1, 1), (35, 0, 0))) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py new file mode 100644 index 000000000000..ea864a595e4a --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +from fastapi import HTTPException + +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import ( + _build_sampling_params_or_400, + _early_validate_edit_bounds, + edits, + generations, +) +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + ImageGenerationsRequest, +) + + +class TestLLaDAImageAPI(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.server_args = SimpleNamespace( + pipeline_config=LLaDAImagePipelineConfig(), + sp_degree=2, + input_save_path=None, + output_path=None, + llada_image_max_pixel_area=None, + llada_image_max_text_tokens=None, + llada_image_max_total_pixel_area=None, + ) + self.raw_request = SimpleNamespace(headers={}) + + async def test_generation_sp_rejects_multiple_outputs_before_scheduling(self): + request = ImageGenerationsRequest(prompt="a red car", n=2) + fallback = HTTPException(status_code=418, detail="request was not validated") + + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", + return_value=self.server_args, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", + side_effect=fallback, + ), + self.assertRaises(HTTPException) as context, + ): + await generations(request, self.raw_request) + + self.assertEqual(context.exception.status_code, 400) + self.assertIn("n=1", context.exception.detail) + + async def test_edit_sp_rejects_multiple_outputs_before_saving_input(self): + fallback = HTTPException(status_code=418, detail="request was not validated") + + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", + return_value=self.server_args, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", + new=AsyncMock(side_effect=fallback), + ), + self.assertRaises(HTTPException) as context, + ): + await edits( + raw_request=self.raw_request, + image=None, + image_array=None, + url=["https://example.com/source.png"], + url_array=None, + prompt="make the car blue", + n=2, + ) + + self.assertEqual(context.exception.status_code, 400) + self.assertIn("n=1", context.exception.detail) + + async def test_edit_rejects_oversized_area_before_source_processing(self): + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", + return_value=self.server_args, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", + new=AsyncMock(side_effect=AssertionError("source was processed")), + ), + self.assertRaises(HTTPException) as context, + ): + await edits( + raw_request=self.raw_request, + image=None, + image_array=None, + url=["https://example.com/source.png"], + url_array=None, + prompt="make the car blue", + size="4096x4096", + n=1, + ) + + self.assertEqual(context.exception.status_code, 400) + self.assertIn("exceeds the supported maximum", context.exception.detail) + + async def test_edit_rejects_multiple_sources_before_source_processing(self): + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", + return_value=self.server_args, + ), + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", + new=AsyncMock(side_effect=AssertionError("source was processed")), + ), + self.assertRaises(HTTPException) as context, + ): + await edits( + raw_request=self.raw_request, + image=None, + image_array=None, + url=[ + "https://example.com/source-1.png", + "https://example.com/source-2.png", + ], + url_array=None, + prompt="make the car blue", + n=1, + size=None, + ) + + self.assertEqual(context.exception.status_code, 400) + self.assertIn("exactly one source image", context.exception.detail) + + def test_build_sampling_params_maps_value_error_to_400(self): + with ( + patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", + side_effect=ValueError("bad request field"), + ), + self.assertRaises(HTTPException) as context, + ): + _build_sampling_params_or_400("rid", prompt="x") + + self.assertEqual(context.exception.status_code, 400) + self.assertEqual(context.exception.detail, "bad request field") + + def test_build_sampling_params_passthrough_on_success(self): + marker = object() + with patch( + "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", + return_value=marker, + ): + self.assertIs(_build_sampling_params_or_400("rid", prompt="x"), marker) + + def test_early_validate_edit_bounds_skips_absent_or_unparseable_size(self): + _early_validate_edit_bounds(self.server_args, None) + _early_validate_edit_bounds(self.server_args, "banana") + with self.assertRaises(HTTPException) as context: + _early_validate_edit_bounds(self.server_args, "4096x4096") + self.assertEqual(context.exception.status_code, 400) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py new file mode 100644 index 000000000000..8d0645a09513 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) +from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy +from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning import ( + LLaDAImageTextConditioningStage, + LLaDAImageTextEncoderRunner, +) + +_GLOBAL_ARGS_PATCH = ( + "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" +) + + +class _FakeTextRunner: + def __init__(self): + self.prompts = None + self.component_names = [] + + def encode(self, prompts, max_sequence_length, component_context=None): + self.prompts = (prompts, max_sequence_length) + if component_context is not None: + for component_name in ("queryformer", "text_projection"): + with component_context(component_name=component_name, module=object()): + self.component_names.append(component_name) + return [ + torch.full((index + 2, 4), float(index + 1)) + for index in range(len(prompts)) + ] + + +class TestLLaDAImageTextConditioning(unittest.TestCase): + def setUp(self): + self.runner = _FakeTextRunner() + with patch(_GLOBAL_ARGS_PATCH, return_value=SimpleNamespace()): + self.stage = LLaDAImageTextConditioningStage(self.runner) + + def test_repeats_positive_and_negative_text_for_each_output(self): + batch = SimpleNamespace( + prompt="a red car", + negative_prompt=None, + guidance_scale=5.0, + num_outputs_per_prompt=2, + max_sequence_length=128, + ) + + result = self.stage.forward(batch, server_args=SimpleNamespace()) + + self.assertTrue(result.do_classifier_free_guidance) + self.assertEqual(len(result.prompt_embeds), 2) + self.assertEqual(len(result.negative_prompt_embeds), 2) + self.assertEqual(len(result.prompt_attention_mask), 2) + self.assertEqual(len(result.negative_attention_mask), 2) + self.assertTrue(torch.equal(result.prompt_embeds[0], result.prompt_embeds[1])) + self.assertTrue( + torch.equal( + result.negative_prompt_embeds[0], result.negative_prompt_embeds[1] + ) + ) + + def test_guidance_disabled_has_no_negative_condition(self): + batch = SimpleNamespace( + prompt="a red car", + negative_prompt=None, + guidance_scale=1.0, + num_outputs_per_prompt=2, + max_sequence_length=128, + ) + + result = self.stage.forward(batch, server_args=SimpleNamespace()) + + self.assertFalse(result.do_classifier_free_guidance) + self.assertEqual(len(result.prompt_embeds), 2) + self.assertEqual(result.negative_prompt_embeds, []) + self.assertEqual(result.negative_attention_mask, []) + + def test_conditioning_mask_guard_fails_closed(self): + from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning import ( + ensure_conditioning_mask_active, + ) + + ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=True)) + with self.assertRaisesRegex(RuntimeError, "without its block attention mask"): + ensure_conditioning_mask_active( + SimpleNamespace(conditioning_mask_active=False) + ) + + def test_runtime_lifecycle_apis_reject_partial_coverage(self): + from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + from sglang.multimodal_gen.runtime.managers.memory_managers.memory_occupation_controller import ( + MemoryOccupationController, + ) + from sglang.multimodal_gen.runtime.post_training.weights_updater import ( + WeightsUpdater, + ) + + self.assertTrue(PipelineConfig().supports_memory_release()) + self.assertTrue(PipelineConfig().supports_hot_weight_updates()) + config = LLaDAImagePipelineConfig() + self.assertFalse(config.supports_memory_release()) + self.assertFalse(config.supports_hot_weight_updates()) + + pipeline = SimpleNamespace(server_args=SimpleNamespace(pipeline_config=config)) + controller = MemoryOccupationController( + pipeline=pipeline, rank=0, use_fsdp_inference=False + ) + result = controller.release_memory_occupation() + self.assertFalse(result["success"]) + self.assertIn("does not support memory release", result["message"]) + self.assertFalse(controller.is_sleeping()) + + updater = object.__new__(WeightsUpdater) + updater.pipeline = pipeline + ok, message = updater.update_weights_from_disk("/unused/path") + self.assertFalse(ok) + self.assertIn("old checkpoint", message) + ok, message = updater.update_weights_from_tensor(named_tensors=[]) + self.assertFalse(ok) + self.assertIn("old checkpoint", message) + + def test_forward_batch_declares_conditioning_text_lens_field(self): + import dataclasses + + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + field_names = {field.name for field in dataclasses.fields(ForwardBatch)} + # Batch rebuilds use dataclasses.replace, which drops dynamic attrs. + self.assertIn("llada_image_conditioning_text_lens_cpu", field_names) + + def test_stage_declares_auxiliary_component_residency(self): + uses = self.stage.component_uses(SimpleNamespace(), "conditioning") + + self.assertEqual( + [use.component_name for use in uses], + ["queryformer", "text_projection"], + ) + + batch = SimpleNamespace( + prompt="a red car", + negative_prompt=None, + guidance_scale=1.0, + num_outputs_per_prompt=1, + max_sequence_length=128, + ) + self.stage.forward(batch, server_args=SimpleNamespace()) + + self.assertEqual( + self.runner.component_names, ["queryformer", "text_projection"] + ) + + def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): + resolved_page_size = 16 + fake_worker = SimpleNamespace( + model_runner=SimpleNamespace(page_size=resolved_page_size), + model_config=object(), + get_memory_pool=lambda: (object(), object()), + alloc_memory_pool=lambda: None, + init_attention_backends=lambda: None, + init_cuda_graphs=lambda: None, + ) + srt_args_module = "sglang.srt.server_args.ServerArgs" + worker_module = "sglang.srt.managers.tp_worker.TpModelWorker" + with ( + patch( + worker_module, + return_value=fake_worker, + ) as worker_cls, + patch("sglang.srt.runtime_context.publish"), + patch( + "sglang.srt.mem_cache.cache_init_params.CacheInitParams", + return_value=object(), + ) as cache_init_params_cls, + patch( + "sglang.srt.mem_cache.chunk_cache.ChunkCache", + return_value=object(), + ), + patch( + srt_args_module, + side_effect=lambda **kwargs: SimpleNamespace(page_size=None, **kwargs), + ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.get_local_torch_device", + return_value=torch.device("cpu"), + ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.get_sp_parallel_rank", + return_value=1, + ), + ): + runner = LLaDAImageTextEncoderRunner( + model_root="/unused/model", + queryformer=object(), + text_projection=object(), + tokenizer=object(), + server_args=SimpleNamespace( + sp_degree=2, + nccl_port=29500, + trust_remote_code=True, + revision=None, + pipeline_config=SimpleNamespace( + text_encoder_mem_fraction_static=0.1 + ), + ), + ) + untrusted_runner = LLaDAImageTextEncoderRunner( + model_root="/unused/model", + queryformer=object(), + text_projection=object(), + tokenizer=object(), + server_args=SimpleNamespace( + sp_degree=2, + nccl_port=29500, + trust_remote_code=False, + revision="pinned-rev", + pipeline_config=SimpleNamespace( + text_encoder_mem_fraction_static=0.1 + ), + ), + ) + + self.assertIs(runner.worker, fake_worker) + self.assertIsNone(runner.server_args.page_size) + self.assertEqual(runner.page_size, resolved_page_size) + self.assertEqual( + [call.kwargs["page_size"] for call in cache_init_params_cls.call_args_list], + [resolved_page_size, resolved_page_size], + ) + self.assertEqual(runner.server_args.tp_size, 2) + self.assertEqual(runner.server_args.dp_size, 2) + self.assertTrue(runner.server_args.enable_dp_attention) + self.assertTrue(runner.server_args.enable_dp_lm_head) + self.assertEqual(runner.server_args.attn_cp_size, 1) + self.assertEqual(runner.server_args.ep_size, 1) + self.assertEqual(runner.server_args.moe_dp_size, 2) + self.assertEqual(runner.server_args.moe_dense_tp_size, 1) + self.assertEqual(runner.server_args.moe_a2a_backend, "none") + self.assertEqual(runner.server_args.max_running_requests, 4) + parallel_state = worker_cls.call_args_list[0].kwargs["ps"] + self.assertEqual(parallel_state.tp_rank, 1) + self.assertEqual(parallel_state.tp_size, 2) + self.assertEqual(parallel_state.dp_rank, 1) + self.assertEqual(parallel_state.dp_size, 2) + self.assertEqual(parallel_state.attn_tp_rank, 0) + self.assertEqual(parallel_state.attn_tp_size, 1) + self.assertEqual(parallel_state.attn_cp_rank, 0) + self.assertEqual(parallel_state.attn_cp_size, 1) + self.assertEqual(parallel_state.attn_dp_rank, 1) + self.assertEqual(parallel_state.attn_dp_size, 2) + self.assertEqual(parallel_state.moe_ep_rank, 0) + self.assertEqual(parallel_state.moe_ep_size, 1) + self.assertEqual(parallel_state.moe_dp_rank, 1) + self.assertEqual(parallel_state.moe_dp_size, 2) + self.assertIs(runner.server_args.trust_remote_code, True) + self.assertIs(untrusted_runner.server_args.trust_remote_code, False) + self.assertEqual(untrusted_runner.server_args.revision, "pinned-rev") + + def test_text_runner_scopes_singleton_attention_group_and_restores(self): + import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state + import sglang.srt.distributed.parallel_state as srt_parallel_state + + diffusion_group = object() + encoder_group = SimpleNamespace(world_size=2, rank_in_group=1) + encoder_attention_group = object() + runner = object.__new__(LLaDAImageTextEncoderRunner) + runner.encoder_tp_group = encoder_group + runner.encoder_attn_tp_group = encoder_attention_group + observed = {} + + def fail_encode(*_args, **_kwargs): + observed["tp"] = srt_parallel_state._TP + observed["attn_tp"] = srt_parallel_state._ATTN_TP + raise RuntimeError("encode failed") + + runner._encode_impl = fail_encode + with ( + patch.object(mm_parallel_state, "_TP", diffusion_group), + patch.object(srt_parallel_state, "_TP", diffusion_group), + patch.object(srt_parallel_state, "_ATTN_TP", diffusion_group), + ): + with self.assertRaisesRegex(RuntimeError, "encode failed"): + runner.encode(["hello"], max_sequence_length=16) + + self.assertIs(mm_parallel_state._TP, diffusion_group) + self.assertIs(srt_parallel_state._TP, diffusion_group) + self.assertIs(srt_parallel_state._ATTN_TP, diffusion_group) + + self.assertIs(observed["tp"], encoder_group) + self.assertIs(observed["attn_tp"], encoder_attention_group) + + def test_text_runner_uses_resolved_page_size_for_prefill(self): + class StopAfterPrefillAdder(Exception): + pass + + runner = object.__new__(LLaDAImageTextEncoderRunner) + runner.tokenizer = lambda *args, **kwargs: SimpleNamespace(input_ids=[[1, 2]]) + runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) + runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) + runner.server_args = SimpleNamespace( + page_size=None, + chunked_prefill_size=-1, + max_prefill_tokens=8192, + ) + runner.page_size = 64 + runner.tree_cache = object() + runner.token_to_kv_pool_allocator = object() + with ( + patch( + "sglang.srt.sampling.sampling_params.SamplingParams", + return_value=SimpleNamespace(normalize=lambda _: None), + ), + patch( + "sglang.srt.managers.schedule_batch.Req", + return_value=SimpleNamespace(), + ), + patch( + "sglang.srt.managers.schedule_policy.PrefillAdder", + side_effect=StopAfterPrefillAdder, + ) as prefill_adder_cls, + ): + with self.assertRaises(StopAfterPrefillAdder): + runner._encode_impl(["hello"], max_sequence_length=16) + + self.assertEqual(prefill_adder_cls.call_args.args[0], runner.page_size) + + def test_text_dp_metadata_covers_each_replica(self): + from sglang.srt.managers.schedule_policy import AddReqResult + + class StopAfterMetadata(Exception): + pass + + class FakeReq: + def __init__(self, *_args, **_kwargs): + self.kv = None + self.req_pool_idx = None + + def init_next_round_input(self, _tree_cache): + return None + + class FakePrefillAdder: + def __init__(self, *_args, **_kwargs): + self.can_run_list = [] + + def add_one_req(self, req, **_kwargs): + self.can_run_list.append(req) + return AddReqResult.CONTINUE + + def stop_after_metadata(): + raise StopAfterMetadata + + batch = SimpleNamespace(prepare_for_extend=stop_after_metadata) + runner = object.__new__(LLaDAImageTextEncoderRunner) + runner.tokenizer = lambda *_args, **_kwargs: SimpleNamespace( + input_ids=[[1, 2], [3]] + ) + runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) + runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) + runner.server_args = SimpleNamespace( + chunked_prefill_size=-1, + max_prefill_tokens=8192, + dp_size=2, + ) + runner.text_dp_attention = True + runner.page_size = 1 + runner.tree_cache = object() + runner.token_to_kv_pool_allocator = object() + runner.req_to_token_pool = object() + + with ( + patch( + "sglang.srt.sampling.sampling_params.SamplingParams", + return_value=SimpleNamespace(normalize=lambda _config: None), + ), + patch( + "sglang.srt.managers.schedule_batch.Req", + side_effect=FakeReq, + ), + patch( + "sglang.srt.managers.schedule_policy.PrefillAdder", + side_effect=FakePrefillAdder, + ), + patch( + "sglang.srt.managers.schedule_batch.ScheduleBatch.init_new", + return_value=batch, + ), + ): + with self.assertRaises(StopAfterMetadata): + runner._encode_impl( + ["positive", "negative"], + max_sequence_length=16, + ) + + self.assertEqual(batch.global_num_tokens, [7, 7]) + self.assertEqual(batch.global_num_tokens_for_logprob, [2, 2]) + + def test_text_runner_restores_diffusion_groups_when_worker_init_fails(self): + import sglang.srt.distributed.parallel_state as srt_parallel_state + + diffusion_group = object() + encoder_group = object() + + def fail_after_installing_encoder_groups(**_kwargs): + srt_parallel_state._TP = encoder_group + srt_parallel_state._ATTN_TP = encoder_group + return SimpleNamespace( + alloc_memory_pool=lambda: (_ for _ in ()).throw( + RuntimeError("allocation failed") + ) + ) + + with ( + patch( + "sglang.srt.managers.tp_worker.TpModelWorker", + side_effect=fail_after_installing_encoder_groups, + ), + patch("sglang.srt.runtime_context.publish"), + patch( + "sglang.srt.server_args.ServerArgs", + side_effect=lambda **kwargs: SimpleNamespace(page_size=1, **kwargs), + ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.get_local_torch_device", + return_value=torch.device("cpu"), + ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.get_sp_parallel_rank", + return_value=0, + ), + patch.object(srt_parallel_state, "_TP", diffusion_group), + patch.object(srt_parallel_state, "_ATTN_TP", diffusion_group), + ): + with self.assertRaisesRegex(RuntimeError, "allocation failed"): + LLaDAImageTextEncoderRunner( + model_root="/unused/model", + queryformer=object(), + text_projection=object(), + tokenizer=object(), + server_args=SimpleNamespace( + sp_degree=1, + nccl_port=29500, + trust_remote_code=True, + revision=None, + pipeline_config=SimpleNamespace( + text_encoder_mem_fraction_static=0.1 + ), + ), + ) + + self.assertIs(srt_parallel_state._TP, diffusion_group) + self.assertIs(srt_parallel_state._ATTN_TP, diffusion_group) + + +class TestLLaDAImageConditionKwargs(unittest.TestCase): + def setUp(self): + self.config = LLaDAImagePipelineConfig() + self.semantic = [ + torch.full((3, 5), 1.0), + torch.full((3, 5), 2.0), + ] + self.source = [ + torch.full((8, 1, 2, 3), 3.0), + torch.full((8, 1, 2, 3), 4.0), + ] + self.batch = SimpleNamespace( + batch_size=2, + image_embeds=self.semantic, + source_latents=self.source, + do_classifier_free_guidance=True, + ) + + def test_cfg_uses_semantics_only_on_positive_branch(self): + positive = self.config.prepare_pos_cond_kwargs( + self.batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float64 + ) + negative = self.config.prepare_neg_cond_kwargs( + self.batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float64 + ) + policy = CFGPolicy().build( + self.batch, + {"encoder_hidden_states_image": self.semantic}, + positive, + negative, + ) + + positive_kwargs = policy.branches[0].kwargs + negative_kwargs = policy.branches[1].kwargs + self.assertEqual( + [tuple(x.shape) for x in positive_kwargs["encoder_hidden_states_image"]], + [(3, 5)] * 2, + ) + self.assertEqual( + [tuple(x.shape) for x in negative_kwargs["encoder_hidden_states_image"]], + [(0, 5)] * 2, + ) + self.assertTrue( + all(x.dtype == torch.float64 for x in positive_kwargs["source_latents"]) + ) + for positive_source, negative_source in zip( + positive_kwargs["source_latents"], + negative_kwargs["source_latents"], + strict=True, + ): + torch.testing.assert_close(positive_source, negative_source) + + def test_rejects_condition_batch_length_mismatch(self): + self.batch.image_embeds = self.semantic[:1] + + with self.assertRaisesRegex( + ValueError, "image_embeds has 1 entries, expected 2" + ): + self.config.prepare_pos_cond_kwargs( + self.batch, + torch.device("cpu"), + rotary_emb=None, + dtype=torch.float32, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py new file mode 100644 index 000000000000..c7bc45a98092 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from PIL import Image + +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) + + +class TestLLaDAImagePipelineConfig(unittest.TestCase): + def setUp(self): + self.config = LLaDAImagePipelineConfig() + + def test_edit_keeps_requested_default_output_size(self): + image = Image.new("RGB", (768, 512)) + + self.assertIsNone( + self.config.calculate_condition_image_size(image, image.width, image.height) + ) + self.assertIsNone(self.config.prepare_calculated_size(image)) + + def test_decode_preprocessing_matches_official_vae_dtype_order(self): + class FakeVAE(torch.nn.Module): + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter( + torch.zeros((), dtype=torch.bfloat16), requires_grad=False + ) + self.bn = SimpleNamespace( + running_mean=torch.tensor( + [0.01, -0.02, 0.03, -0.04], dtype=torch.float32 + ), + running_var=torch.full((4,), 0.0129973, dtype=torch.float32), + ) + self.config = SimpleNamespace( + arch_config=SimpleNamespace(batch_norm_eps=0.003) + ) + + vae = FakeVAE() + latents = torch.tensor( + [[[[0.501]], [[-0.249]], [[0.126]], [[-0.751]]]], + dtype=torch.float32, + ) + official_latents = latents.to(torch.bfloat16) + latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(official_latents) + latent_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae.config.arch_config.batch_norm_eps + ).to(official_latents) + expected = official_latents * latent_std + latent_mean + expected = expected.reshape(1, 1, 2, 2, 1, 1) + expected = expected.permute(0, 1, 4, 2, 5, 3).reshape(1, 1, 2, 2) + + actual = self.config.preprocess_decoding(latents, vae=vae) + + self.assertEqual(actual.dtype, torch.bfloat16) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_edit_sp_shards_target_and_source_latents(self): + latents = torch.randn(1, 128, 4, 4) + source_latents = torch.randn(128, 1, 4, 4) + batch = SimpleNamespace( + batch_size=1, + condition_image=Image.new("RGB", (64, 64)), + image_embeds=None, + source_latents=[source_latents], + enable_sequence_shard=False, + ) + base_module = "sglang.multimodal_gen.configs.pipeline_configs.base" + llada_module = "sglang.multimodal_gen.configs.pipeline_configs.llada_image" + with ( + patch(f"{base_module}.get_sp_world_size", return_value=2), + patch(f"{base_module}.get_sp_parallel_rank", return_value=0), + patch(f"{llada_module}.get_sp_world_size", return_value=2), + ): + actual, did_shard = self.config.shard_latents_for_sp(batch, latents) + batch.did_sp_shard_latents = did_shard + source = self.config.prepare_pos_cond_kwargs( + batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float32 + )["source_latents"][0] + + self.assertEqual(actual.shape, (1, 128, 2, 4)) + torch.testing.assert_close(actual, latents[:, :, :2, :]) + self.assertTrue(did_shard) + self.assertEqual(source.shape, (128, 1, 2, 4)) + torch.testing.assert_close(source, source_latents[:, :, :2, :]) + + def test_generation_sp_rejects_latent_height_padding(self): + latents = torch.randn(1, 128, 63, 64) + batch = SimpleNamespace( + condition_image=None, + source_latents=None, + enable_sequence_shard=False, + ) + module = "sglang.multimodal_gen.configs.pipeline_configs.llada_image" + with ( + patch(f"{module}.get_sp_world_size", return_value=2), + self.assertRaisesRegex( + ValueError, "latent height 63 must be divisible by SP degree 2" + ), + ): + self.config.shard_latents_for_sp(batch, latents) + + def test_validates_supported_parallel_topology(self): + defaults = dict( + sp_degree=2, + ulysses_degree=2, + ring_degree=1, + tp_size=1, + dp_size=1, + cfg_parallel_degree=1, + num_gpus=2, + text_encoder_cpu_offload=False, + llada_image_max_pixel_area=None, + llada_image_max_text_tokens=None, + llada_image_max_total_pixel_area=None, + is_arg_explicitly_set=lambda name: False, + residency_mode=lambda name: "resident", + explicit_residency_mode=lambda name: None, + layerwise_offload_components=None, + cpu_offload_components=None, + ) + self.config.validate_server_args(SimpleNamespace(**defaults)) + + with ( + patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}), + self.assertRaisesRegex(ValueError, "does not support"), + ): + self.config.validate_server_args(SimpleNamespace(**defaults)) + + auto_offload = SimpleNamespace( + **(defaults | {"text_encoder_cpu_offload": True}) + ) + self.config.validate_server_args(auto_offload) + self.assertFalse(auto_offload.text_encoder_cpu_offload) + + class _AutoTunedArgs(SimpleNamespace): + def residency_mode(self, name): + components = self.layerwise_offload_components or [] + return "layerwise_offload" if name in components else "resident" + + auto_layerwise = _AutoTunedArgs( + **{key: value for key, value in defaults.items() if key != "residency_mode"} + ) + auto_layerwise.layerwise_offload_components = [ + "text_encoder", + "image_encoder", + ] + self.config.validate_server_args(auto_layerwise) + self.assertEqual(auto_layerwise.layerwise_offload_components, ["image_encoder"]) + + invalid_cases = [ + ({"ulysses_degree": 1, "ring_degree": 2}, "ring_degree=1"), + ( + {"sp_degree": 4, "ulysses_degree": 4, "num_gpus": 4}, + "supports only SP degrees 1 and 2", + ), + ({"cfg_parallel_degree": 2}, "TP=DP=CFG=1"), + ({"sp_degree": 1, "ulysses_degree": 1}, "num_gpus == sp_degree"), + ( + { + "text_encoder_cpu_offload": True, + "is_arg_explicitly_set": lambda name: True, + }, + "residency manager cannot offload", + ), + ( + {"residency_mode": lambda name: "component_offload"}, + "embedded text encoder to remain resident", + ), + ( + { + "residency_mode": lambda name: "layerwise_offload", + "is_arg_explicitly_set": ( + lambda name: name == "layerwise_offload_components" + ), + }, + "embedded text encoder to remain resident", + ), + ({"llada_image_max_pixel_area": 0}, "must be positive"), + ({"llada_image_max_total_pixel_area": 0}, "must be positive"), + ( + {"llada_image_max_text_tokens": 4000}, + "at most 3584", + ), + ] + for overrides, message in invalid_cases: + args = defaults | overrides + with self.subTest(overrides=overrides), self.assertRaisesRegex( + ValueError, message + ): + self.config.validate_server_args(SimpleNamespace(**args)) + + def test_validates_request_sampling_params_bounds(self): + def params(**overrides): + values = dict( + width=1024, + height=1024, + max_sequence_length=2048, + num_outputs_per_prompt=1, + diffusers_kwargs=None, + enable_cache_dit=None, + cache_dit_params=None, + cfg_gate_step=None, + attention_backend_override=None, + ) + values.update(overrides) + return SimpleNamespace(**values) + + def server(**overrides): + values = dict( + sp_degree=1, + llada_image_max_pixel_area=None, + llada_image_max_text_tokens=None, + llada_image_max_total_pixel_area=None, + ) + values.update(overrides) + return SimpleNamespace(**values) + + sp1 = server() + sp2 = server(sp_degree=2) + + self.config.validate_request_sampling_params(params(), sp1) + self.config.validate_request_sampling_params( + params(max_sequence_length=None), sp1 + ) + self.config.validate_request_sampling_params( + params(width=2048, height=2048, max_sequence_length=3584), sp2 + ) + self.config.validate_request_sampling_params( + params(num_outputs_per_prompt=10), sp1 + ) + self.config.validate_request_sampling_params( + params(enable_cache_dit=False), sp1 + ) + self.config.validate_request_sampling_params(params(cfg_gate_step=0.5), sp1) + self.config.validate_request_sampling_params( + params(attention_backend_override="fa"), sp1 + ) + self.config.validate_request_sampling_params( + params(attention_backend_override="TORCH_SDPA"), sp1 + ) + + invalid_cases = [ + (params(width=0), sp1, "positive width and height"), + (params(width=1000), sp1, "divisible by 16"), + (params(height=1040), sp2, "divisible by 32 at SP degree 2"), + (params(width=2064, height=2064), sp1, "exceeds the supported maximum"), + (params(max_sequence_length=0), sp1, "positive integer"), + (params(max_sequence_length=True), sp1, "positive integer"), + (params(max_sequence_length="2048"), sp1, "positive integer"), + ( + params(max_sequence_length=3585), + sp1, + "exceeds the embedded text encoder budget", + ), + ( + params(), + server(llada_image_max_pixel_area=512 * 512), + "exceeds the supported maximum of 262144", + ), + ( + params(), + server(llada_image_max_text_tokens=1024), + "budget of 1024 tokens", + ), + ( + params(width=2048, height=2048, num_outputs_per_prompt=10), + sp1, + "total output area", + ), + ( + params(num_outputs_per_prompt=2), + server(llada_image_max_total_pixel_area=1024 * 1024), + "total output area", + ), + ( + params(diffusers_kwargs={"max_sequence_length": 4096}), + sp1, + "exceeds the embedded text encoder budget", + ), + ( + params(diffusers_kwargs={"max_sequence_length": "x"}), + sp1, + "positive integer", + ), + (params(enable_cache_dit=True), sp1, "does not support enable_cache_dit"), + ( + params(cache_dit_params={}), + sp1, + "does not support cache_dit_params", + ), + ( + params(attention_backend_override="sage_attn"), + sp1, + "must be fa or torch_sdpa", + ), + ( + params(attention_backend_override=1), + sp1, + "must be fa or torch_sdpa", + ), + (params(cfg_gate_step=-0.1), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step=1.1), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step=True), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step=float("nan")), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step=float("inf")), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step=10**1000), sp1, "must be between 0.0 and 1.0"), + (params(cfg_gate_step="0.5"), sp1, "must be between 0.0 and 1.0"), + ] + for sampling_params, server_args, message in invalid_cases: + with self.subTest(message=message), self.assertRaisesRegex( + ValueError, message + ): + self.config.validate_request_sampling_params( + sampling_params, server_args + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py new file mode 100644 index 000000000000..ef9cf0cbef10 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from PIL import Image + +from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_source import ( + LLaDAImageSourceImageConditioningStage, +) + +_GLOBAL_ARGS_PATCH = ( + "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" +) + + +class _FakeImageProcessor: + def __init__(self): + self.calls = [] + self.output = torch.arange(3 * 8 * 12, dtype=torch.float32).reshape(1, 3, 8, 12) + + def preprocess(self, image, height, width, resize_mode): + self.calls.append((image, height, width, resize_mode)) + return self.output.clone() + + +class _FakeSigVQ(torch.nn.Module): + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(()), requires_grad=False) + self.config = SimpleNamespace(patch_size=16) + self.pixel_values = None + + def forward(self, pixel_values): + self.pixel_values = pixel_values + batch_size = pixel_values.shape[0] + features = torch.arange( + batch_size * 3 * 5, + device=pixel_values.device, + dtype=pixel_values.dtype, + ).reshape(batch_size, 3, 5) + return SimpleNamespace(semantic_features=features) + + +class _FakePosterior: + def __init__(self, value): + self.value = value + + def mode(self): + return self.value + + +class _FakeVAE(torch.nn.Module): + def __init__(self): + super().__init__() + self.anchor = torch.nn.Parameter(torch.zeros(()), requires_grad=False) + self.bn = SimpleNamespace( + running_mean=torch.arange(8, dtype=torch.float32), + running_var=torch.full((8,), 3.75, dtype=torch.float32), + ) + self.config = SimpleNamespace(arch_config=SimpleNamespace(batch_norm_eps=0.25)) + self.encoded_image = None + self.raw_latents = None + + def encode(self, image): + self.encoded_image = image + batch_size = image.shape[0] + self.raw_latents = torch.arange( + batch_size * 2 * 4 * 6, + device=image.device, + dtype=image.dtype, + ).reshape(batch_size, 2, 4, 6) + return _FakePosterior(self.raw_latents) + + +class TestLLaDAImageSourceImageConditioningStage(unittest.TestCase): + def setUp(self): + self.processor = _FakeImageProcessor() + self.sigvq = _FakeSigVQ() + self.vae = _FakeVAE() + with patch(_GLOBAL_ARGS_PATCH, return_value=SimpleNamespace()): + self.stage = LLaDAImageSourceImageConditioningStage( + sigvq=self.sigvq, + vae=self.vae, + image_processor=self.processor, + ) + + def test_text_generation_is_a_noop(self): + batch = SimpleNamespace( + condition_image=None, + image_embeds=[], + source_latents=None, + ) + + result = self.stage.forward(batch, server_args=SimpleNamespace()) + + self.assertIs(result, batch) + self.assertEqual(result.image_embeds, []) + self.assertIsNone(result.source_latents) + self.assertEqual(self.processor.calls, []) + + def test_edit_matches_source_preprocessing_and_repeats_outputs(self): + image = Image.new("RGB", (40, 30), color="red") + batch = SimpleNamespace( + condition_image=image, + height=64, + width=80, + batch_size=2, + image_embeds=[], + source_latents=None, + ) + + result = self.stage.forward(batch, server_args=SimpleNamespace()) + + self.assertEqual(self.processor.calls, [(image, 64, 80, "crop")]) + self.assertEqual(tuple(self.sigvq.pixel_values.shape), (2, 3, 32, 48)) + self.assertTrue( + torch.equal(self.vae.encoded_image[0], self.processor.output[0]) + ) + self.assertTrue( + torch.equal(self.vae.encoded_image[0], self.vae.encoded_image[1]) + ) + self.assertEqual([tuple(x.shape) for x in result.image_embeds], [(3, 5)] * 2) + self.assertEqual( + [tuple(x.shape) for x in result.source_latents], [(8, 1, 2, 3)] * 2 + ) + + raw = self.vae.raw_latents + expected = raw.reshape(2, 2, 2, 2, 3, 2) + expected = expected.permute(0, 1, 3, 5, 2, 4).reshape(2, 8, 2, 3) + expected = (expected - self.vae.bn.running_mean.view(1, -1, 1, 1)) / 2 + expected = [sample.unsqueeze(1) for sample in expected] + for actual, wanted in zip(result.source_latents, expected, strict=True): + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + + def test_edit_rejects_more_than_one_source_image(self): + batch = SimpleNamespace( + condition_image=[Image.new("RGB", (32, 32))] * 2, + height=64, + width=64, + batch_size=1, + ) + + with self.assertRaisesRegex(ValueError, "exactly one source image"): + self.stage.forward(batch, server_args=SimpleNamespace()) + + def test_source_normalization_matches_official_std_cast_order(self): + latents = torch.tensor([[[[0.5]]]], dtype=torch.bfloat16) + vae = SimpleNamespace( + bn=SimpleNamespace( + running_mean=torch.tensor([0.01], dtype=torch.float32), + running_var=torch.tensor([0.0129973], dtype=torch.float32), + ), + config=SimpleNamespace(arch_config=SimpleNamespace(batch_norm_eps=0.003)), + ) + expected_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) + expected_std = torch.sqrt( + vae.bn.running_var.view(1, -1, 1, 1) + vae.config.arch_config.batch_norm_eps + ).to(latents) + + actual = self.stage._normalize_latents(latents, vae) + + torch.testing.assert_close( + actual, + (latents - expected_mean) / expected_std, + rtol=0, + atol=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 39ac4992db57..18efcbad11a0 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -98,6 +98,7 @@ def _loads(input_str, _flags=None): from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( TransformerQuantLoadSpec, _Flux2Nvfp4FallbackAdapter, + _ModelOptFp8OffloadAdapter, _needs_device_weight_postprocess, _resolve_quant_config, _resolve_weight_override_quantization, @@ -113,6 +114,7 @@ def _loads(input_str, _flags=None): from sglang.multimodal_gen.runtime.models.dits.flux_2 import ( Flux2Transformer2DModel, ) +from sglang.multimodal_gen.runtime.models.dits.llada_image import LLaDAImageFeedForward from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel from sglang.multimodal_gen.runtime.models.dits.qwen_image import ( QwenImageTransformer2DModel, @@ -162,6 +164,14 @@ def _make_quant_config(name: str, **attrs): return quant_config +class _FakeOnlineFp8Config: + is_checkpoint_fp8_serialized = False + + @classmethod + def get_name(cls): + return "fp8" + + class TestTransformerQuantHelpers(unittest.TestCase): def test_modelopt_fp8_packed_cutlass_preserves_checkpoint_shard_scales(self): method = ModelOptFp8LinearMethod( @@ -618,6 +628,28 @@ def test_modelopt_fp4_preserves_explicit_attention_backend(self): self.assertIsNone(backend) + def test_explicit_fp8_quantization_uses_online_defaults(self): + config = _resolve_quant_config( + hf_config={}, + server_args=self._make_server_args(quantization="fp8"), + safetensors_list=[], + component_model_path="/unused/component/path", + ) + + self.assertEqual(config.get_name(), "fp8") + self.assertFalse(config.is_checkpoint_fp8_serialized) + self.assertEqual(config.activation_scheme, "dynamic") + + def test_online_fp8_disables_whole_dit_cpu_offload(self): + server_args = self._make_server_args(dit_cpu_offload=True) + + _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( + server_args=server_args, + quant_config=_FakeOnlineFp8Config(), + ) + + self.assertFalse(server_args.dit_cpu_offload) + def test_resolve_transformer_checkpoint_files_uses_single_override_file(self): with tempfile.NamedTemporaryFile(suffix=".safetensors") as f: server_args = self._make_server_args(transformer_weights_path=f.name) @@ -1822,6 +1854,33 @@ def test_explicit_modelopt_method_without_algorithm_is_preserved(self): quant_method, ) + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) + @patch( + "sglang.multimodal_gen.runtime.layers.linear.get_tp_group", return_value=None + ) + def test_llada_image_online_fp8_quantizes_fused_linears( + self, + _mock_tp_group, + _mock_group_size, + _mock_group_rank, + ): + fp8_module = "sglang.multimodal_gen.runtime.layers.quantization.fp8" + with patch( + f"{fp8_module}.get_tensor_model_parallel_world_size", + return_value=1, + ): + feed_forward = LLaDAImageFeedForward( + dim=192, + quant_config=Fp8Config(), + prefix="layers.0.feed_forward", + ) + + self.assertIsInstance(feed_forward.w13.quant_method, Fp8LinearMethod) + self.assertIsInstance(feed_forward.w2.quant_method, Fp8LinearMethod) + self.assertEqual(feed_forward.w13.prefix, "layers.0.feed_forward.w13") + self.assertEqual(feed_forward.w2.prefix, "layers.0.feed_forward.w2") + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) @patch( diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 7573a92ed9ec..6d8d15b69723 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -30,6 +30,7 @@ def handle_context_parallelism(server_args: Any): cfg = resolving_view(server_args) + model_arch = None if parse_connector_type(cfg.model_path) != ConnectorType.INSTANCE: from sglang.srt.configs.model_config import is_deepseek_dsa from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES @@ -59,7 +60,8 @@ def handle_context_parallelism(server_args: Any): and not cfg.language_model_only ): raise ValueError( - "MiMo V2 CP-v2 only supports text inference; add --language-only." + "MiMo V2 CP-v2 only supports text inference; add " + "--language-only." ) if cfg.enable_prefill_cp and cfg.cp_strategy is None: @@ -80,39 +82,52 @@ def handle_context_parallelism(server_args: Any): view = resolved_view(server_args) if view.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert cfg.tp_size % view.attn_cp_size == 0, ( - "tp_size must be divisible by attn_cp_size" - ) - assert cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0, ( - "tp_size must be divisible by dp_size * attn_cp_size" - ) + assert ( + cfg.tp_size % view.attn_cp_size == 0 + ), "tp_size must be divisible by attn_cp_size" + assert ( + cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0 + ), "tp_size must be divisible by dp_size * attn_cp_size" - assert not cfg.enable_aiter_allreduce_fusion, ( - "Aiter allreduce fusion is not supported with context parallelism" - ) + assert ( + not cfg.enable_aiter_allreduce_fusion + ), "Aiter allreduce fusion is not supported with context parallelism" if cfg.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert cfg.tp_size % cfg.moe_dp_size == 0, ( - "tp_size must be divisible by moe_dp_size" - ) - assert view.ep_size * cfg.moe_dp_size <= cfg.tp_size, ( - "ep_size * moe_dp_size must be less than or equal to tp_size" - ) + assert ( + cfg.tp_size % cfg.moe_dp_size == 0 + ), "tp_size must be divisible by moe_dp_size" + assert ( + view.ep_size * cfg.moe_dp_size <= cfg.tp_size + ), "ep_size * moe_dp_size must be less than or equal to tp_size" assert cfg.pp_size == 1, "PP is not supported with context parallelism" if view.ep_size > 1: - assert view.ep_size * cfg.moe_dp_size == cfg.tp_size, ( - "ep_size * moe_dp_size must be equal to tp_size" - ) + assert ( + view.ep_size * cfg.moe_dp_size == cfg.tp_size + ), "ep_size * moe_dp_size must be equal to tp_size" - assert not cfg.enable_aiter_allreduce_fusion, ( - "Aiter allreduce fusion is not supported with context parallelism" - ) + assert ( + not cfg.enable_aiter_allreduce_fusion + ), "Aiter allreduce fusion is not supported with context parallelism" if view.attn_cp_size != cfg.moe_dp_size: - assert cfg.moe_dp_size == 1, ( - "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + assert cfg.moe_dp_size == 1 or ( + view.enable_dp_attention + and model_arch == "LLaDA2MoeModelLM" + and view.ep_size == 1 + and view.pp_size == 1 + and cfg.tp_size == cfg.dp_size + and cfg.dp_size == cfg.moe_dp_size + and view.enable_dp_lm_head + and view.attn_cp_size == 1 + and view.dcp_size == 1 + and view.moe_dense_tp_size == 1 + and view.moe_a2a_backend == "none" + ), ( + "attn_cp_size != moe_dp_size requires moe_dp_size == 1 or " + "the LLaDA2 full-replica topology" ) from sglang.srt.layers.cp.base import init_cp_strategy @@ -243,26 +258,26 @@ def handle_dwdp(server_args: Any): if cfg.dwdp_size <= 1: return - assert cfg.dwdp_size >= 2, ( - f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" - ) - assert cfg.dwdp_size == cfg.tp_size, ( - f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" - ) + assert ( + cfg.dwdp_size >= 2 + ), f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" + assert ( + cfg.dwdp_size == cfg.tp_size + ), f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" assert cfg.disaggregation_mode in ( "null", "prefill", ), "DWDP requires --disaggregation-mode null or prefill" - assert not cfg.enable_eplb, ( - "EPLB dynamic migration conflicts with static DWDP partitioning" - ) - assert cfg.speculative_algorithm is None, ( - "DWDP does not support speculative decoding (MTP/draft workers)" - ) + assert ( + not cfg.enable_eplb + ), "EPLB dynamic migration conflicts with static DWDP partitioning" + assert ( + cfg.speculative_algorithm is None + ), "DWDP does not support speculative decoding (MTP/draft workers)" assert cfg.pp_size == 1, "DWDP requires pp_size == 1" - assert not cfg.enable_two_batch_overlap, ( - "DWDP's prefetch event protocol does not support two-batch overlap" - ) + assert ( + not cfg.enable_two_batch_overlap + ), "DWDP's prefetch event protocol does not support two-batch overlap" if cfg.disaggregation_mode == "null": logger.warning( @@ -358,9 +373,7 @@ def handle_elastic_ep(server_args: Any): assert cfg.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", - ], ( - "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." - ) + ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." assert cfg.pp_size == 1, "PP size should be set to 1 under elastic EP" @@ -371,9 +384,9 @@ def handle_elastic_ep(server_args: Any): mooncake_ib_device=validate_ib_devices(cfg.mooncake_ib_device), ) if cfg.ep_join_mode is not None: - assert cfg.elastic_ep_backend is not None, ( - "--elastic-ep-join-mode requires --elastic-ep-backend to be set." - ) + assert ( + cfg.elastic_ep_backend is not None + ), "--elastic-ep-join-mode requires --elastic-ep-backend to be set." if cfg.ep_join_mode == "scale": assert cfg.node_rank == 1, ( "Elastic EP scale-up requires one joining TP group at " @@ -391,9 +404,9 @@ def handle_elastic_ep(server_args: Any): ) assert cfg.ep_join_rank_offset >= 0, "elastic EP join rank offset must be >= 0." if cfg.max_ep_size is not None: - assert cfg.elastic_ep_backend is not None, ( - "--max-ep-size requires --elastic-ep-backend to be set." - ) + assert ( + cfg.elastic_ep_backend is not None + ), "--max-ep-size requires --elastic-ep-backend to be set." assert cfg.max_ep_size > 0, "--max-ep-size must be a positive integer." scaling_active = ( @@ -408,15 +421,16 @@ def handle_elastic_ep(server_args: Any): ) if scaling_active: resolved = resolved_view(server_args) - assert cfg.elastic_ep_scale_timeout > 0, ( - "--elastic-ep-scale-timeout must be greater than zero." - ) + assert ( + cfg.elastic_ep_scale_timeout > 0 + ), "--elastic-ep-scale-timeout must be greater than zero." assert cfg.tokenizer_worker_num == 1, ( - "Elastic EP runtime scale-up currently requires --tokenizer-worker-num 1." - ) - assert not cfg.use_ray, ( - "Elastic EP runtime scale-up does not support --use-ray." + "Elastic EP runtime scale-up currently requires " + "--tokenizer-worker-num 1." ) + assert ( + not cfg.use_ray + ), "Elastic EP runtime scale-up does not support --use-ray." assert not cfg.enable_elastic_expert_backup, ( "Elastic EP runtime scale-up does not support " "--enable-elastic-expert-backup." diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index bbf45f359b71..878517d44f03 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -74,6 +74,20 @@ def create_flashinfer_backend(runner): return FlashInferMLAAttnBackend(runner) +@register_attention_backend("llada2_cfg_flashinfer") +def create_llada2_cfg_flashinfer_backend(runner): + if runner.use_mla_backend: + raise ValueError("LLaDA2 CFG attention does not use an MLA backend") + + from sglang.srt.layers.attention.llada2_cfg_flashinfer_backend import ( + LLaDA2CFGFlashInferAttnBackend, + ) + + return LLaDA2CFGFlashInferAttnBackend( + runner, init_new_workspace=runner.init_new_workspace + ) + + @register_attention_backend("trtllm_mla") def create_trtllm_mla_backend(runner): if not runner.use_mla_backend: diff --git a/python/sglang/srt/layers/attention/llada2_attention_utils.py b/python/sglang/srt/layers/attention/llada2_attention_utils.py new file mode 100644 index 000000000000..5e4f639227e8 --- /dev/null +++ b/python/sglang/srt/layers/attention/llada2_attention_utils.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Mask helpers shared by LLaDA2 attention backends.""" + +from __future__ import annotations + +import torch + + +def build_llada_image_custom_mask( + text_lens: list[int], + sequence_lens: list[int], + device: torch.device | str, +) -> torch.Tensor: + """Build the official one-pass text/query block mask in ragged layout.""" + if len(text_lens) != len(sequence_lens): + raise RuntimeError("LLaDA-Image conditioning metadata batch mismatch") + + flattened_masks: list[torch.Tensor] = [] + for text_len, sequence_len in zip(text_lens, sequence_lens): + text_len = int(text_len) + sequence_len = int(sequence_len) + if text_len <= 0 or text_len >= sequence_len: + raise RuntimeError( + "LLaDA-Image conditioning text length must leave query tokens" + ) + attention_mask = torch.ones( + (sequence_len, sequence_len), dtype=torch.bool, device=device + ) + attention_mask[:text_len, text_len:] = False + flattened_masks.append(attention_mask.flatten()) + return torch.cat(flattened_masks) diff --git a/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py new file mode 100644 index 000000000000..af4d11254b7b --- /dev/null +++ b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FlashInfer backend for LLaDA-Image conditioning masks.""" + +from __future__ import annotations + +import torch +from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper + +from sglang.srt.layers.attention.flashinfer_backend import ( + FlashInferAttnBackend, + PrefillMetadata, +) +from sglang.srt.layers.attention.llada2_attention_utils import ( + build_llada_image_custom_mask, +) + + +class LLaDA2CFGFlashInferAttnBackend(FlashInferAttnBackend): + """Stock FlashInfer plus LLaDA-specific ragged attention masks.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._llada_image_conditioning_mask_active = False + self._cfg_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, "NHD", backend="fa2" + ) + + @property + def conditioning_mask_active(self) -> bool: + return self._llada_image_conditioning_mask_active + + def init_forward_metadata(self, forward_batch): + self._llada_image_conditioning_mask_active = False + text_lens = getattr( + forward_batch, "llada_image_conditioning_text_lens_cpu", None + ) + if text_lens is None: + return super().init_forward_metadata(forward_batch) + + if not forward_batch.forward_mode.is_extend_without_speculative(): + raise RuntimeError( + "LLaDA-Image conditioning mask requires a one-shot extend" + ) + if self.num_wrappers != 1: + raise RuntimeError( + "LLaDA-Image conditioning requires one FlashInfer wrapper" + ) + + seq_lens = forward_batch.seq_lens + prefix_lens = forward_batch.extend_prefix_lens + if bool(torch.any(prefix_lens != 0).item()): + raise RuntimeError( + "LLaDA-Image conditioning does not support cached prefixes" + ) + custom_mask = build_llada_image_custom_mask( + text_lens, + seq_lens.tolist(), + seq_lens.device, + ) + qo_indptr = torch.zeros( + seq_lens.numel() + 1, + dtype=torch.int32, + device=seq_lens.device, + ) + qo_indptr[1:] = torch.cumsum(seq_lens, dim=0) + prefill_indices_updater = self.indices_updater_prefill + self._cfg_prefill_wrapper_ragged.begin_forward( + qo_indptr, + qo_indptr, + prefill_indices_updater.num_qo_heads, + prefill_indices_updater.num_kv_heads, + prefill_indices_updater.head_dim, + custom_mask=custom_mask, + causal=False, + q_data_type=prefill_indices_updater.q_data_type, + kv_data_type=prefill_indices_updater.data_type, + non_blocking=True, + fixed_split_size=self.prefill_split_tile_size, + ) + self._llada_image_conditioning_mask_active = True + self.forward_metadata = PrefillMetadata( + self.prefill_wrappers_paged, + use_ragged=True, + extend_no_prefix=True, + ) + + def forward_extend( + self, + q, + k, + v, + layer, + forward_batch, + save_kv_cache=True, + ): + if not self._llada_image_conditioning_mask_active: + return super().forward_extend( + q, k, v, layer, forward_batch, save_kv_cache=save_kv_cache + ) + + if k is None or v is None: + raise RuntimeError("LLaDA ragged attention requires explicit K/V") + + q_view = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + k_view = k.view(-1, layer.tp_k_head_num, layer.head_dim) + v_view = v.view(-1, layer.tp_v_head_num, layer.head_dim) + attention_output = self._cfg_prefill_wrapper_ragged.forward( + q_view, + k_view, + v_view, + causal=False, + sm_scale=layer.scaling, + logits_soft_cap=layer.logit_cap, + ) + return attention_output.view(-1, layer.tp_q_head_num * layer.head_dim) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index c395f96bf0ee..4dceb7d5b807 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -418,6 +418,8 @@ def _compute_mlp_mode(cls, context: _LayerModeComputationContext): or enable_dwdp() ): return ScatterMode.SCATTERED + if enable_moe_fully_dp(): + return ScatterMode.SCATTERED # DSA CP and MLA CP both don't support MOE_FULL yet; fall back to FULL. if is_enable_moe_cp_allgather() and not ( is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled() @@ -467,6 +469,23 @@ def enable_moe_dense_fully_dp(): return get_parallel().moe_dense_tp_size == 1 +def enable_moe_fully_dp(): + parallel = get_parallel() + return ( + is_dp_attention_enabled() + and parallel.pp_size == 1 + and not parallel.dcp_enabled + and parallel.attn_cp_size == 1 + and parallel.attn_dp_size == parallel.tp_size + and parallel.attn_tp_size == 1 + and parallel.moe_ep_size == 1 + and parallel.moe_tp_size == 1 + and parallel.moe_dp_size == parallel.attn_dp_size + and parallel.moe_dense_tp_size == 1 + and get_moe_a2a_backend().is_none() + ) + + def enable_dwdp(): return get_parallel().dwdp_size > 1 @@ -969,6 +988,15 @@ def init_new(cls): tp_size = get_parallel().tp_size tp_rank = get_parallel().tp_rank moe_cp_size = get_moe_cp_size() + if enable_moe_fully_dp(): + moe_full_size = tp_size // attn_cp_size + else: + if attn_cp_size % moe_cp_size != 0: + raise ValueError( + "Attention CP size must be divisible by the MoE DP group size" + ) + cp_per_moe = attn_cp_size // moe_cp_size + moe_full_size = tp_size // cp_per_moe process_group_sizes = { ScatterMode.SCATTERED: 1, ScatterMode.TP_ATTN_FULL: attn_tp_size, @@ -976,7 +1004,7 @@ def init_new(cls): # With context parallel enabled, we should exclude # the attn_cp_size from the total tp_size ScatterMode.FULL: tp_size // attn_cp_size, - ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size), + ScatterMode.MOE_FULL: moe_full_size, } return cls( process_group_sizes=process_group_sizes, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index f98d5855cd64..6d91ebdce8d0 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -556,6 +556,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Encoder-decoder cross-attention mask cross_attention_custom_mask: Optional[torch.Tensor] = None + # LLaDA-Image conditioning text lengths, declared so batch rebuilds + # carry them to the llada2_cfg_flashinfer backend + llada_image_conditioning_text_lens_cpu: Optional[List[int]] = None + # For DP attention (padding / local info) dp_padding_mode: Optional[DpPaddingMode] = None # for extend, local start pos and num tokens is different in logits processor diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 1e4e804c465c..e0631774ea75 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -40,6 +40,7 @@ LayerCommunicator, LayerScatterModes, enable_moe_dense_fully_dp, + enable_moe_fully_dp, ) from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, @@ -71,6 +72,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llada2_weight_utils import prepare_llada2_language_weights from sglang.srt.models.utils import ( apply_qk_norm, create_fused_set_kv_buffer_arg, @@ -198,6 +200,7 @@ def __init__( self.layer_id = layer_id self.alt_stream = alt_stream self.tp_size = get_parallel().tp_size + self.moe_fully_replicated = enable_moe_fully_dp() self.top_k = config.num_experts_per_tok self.norm_topk_prob = config.norm_topk_prob self.hidden_size = config.hidden_size @@ -299,7 +302,7 @@ def __init__( prefix=add_prefix("shared_experts", prefix), **( dict(tp_rank=0, tp_size=1) - if get_moe_a2a_backend().is_deepep() + if get_moe_a2a_backend().is_deepep() or self.moe_fully_replicated else {} ), ) @@ -386,8 +389,12 @@ def forward_normal( if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if self.tp_size > 1 and not should_skip_post_experts_all_reduce( - is_tp_path=True, + if ( + self.tp_size > 1 + and not self.moe_fully_replicated + and not should_skip_post_experts_all_reduce( + is_tp_path=True, + ) ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -837,6 +844,9 @@ def start_layer(self): def end_layer(self): return self.model.end_layer + def get_input_embeddings(self): + return self.model.word_embeddings + def get_embed_and_head(self): """Used by the eagle_worker.""" return self.model.word_embeddings.weight, self.lm_head.weight @@ -890,7 +900,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) params_dict = dict(self.named_parameters()) - for name, loaded_weight in weights: + for name, loaded_weight in prepare_llada2_language_weights( + weights, num_experts=self.config.num_experts + ): if ( ("v_head" in name) or ("inv_freq" in name) diff --git a/python/sglang/srt/models/llada2_weight_utils.py b/python/sglang/srt/models/llada2_weight_utils.py new file mode 100644 index 000000000000..2e38033ae5d3 --- /dev/null +++ b/python/sglang/srt/models/llada2_weight_utils.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint-layout helpers for LLaDA2 models.""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator + +import torch + +_EXPERT_PROJECTION_NAMES = frozenset({"gate_proj", "up_proj", "down_proj"}) + + +def prepare_llada2_language_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + num_experts: int, +) -> Iterator[tuple[str, torch.Tensor]]: + """Normalize language-model prefixes and expand fused expert tensors.""" + for name, loaded_weight in weights: + if name.startswith("model.language_model."): + name = "model." + name.removeprefix("model.language_model.") + elif name.startswith("model.lm_head."): + name = "lm_head." + name.removeprefix("model.lm_head.") + + expert_prefix, separator, projection_name = name.rpartition(".") + is_fused_expert = ( + separator + and expert_prefix.endswith(".mlp.experts") + and projection_name in _EXPERT_PROJECTION_NAMES + ) + if not is_fused_expert: + yield name, loaded_weight + continue + + if loaded_weight.ndim != 3 or loaded_weight.shape[0] != num_experts: + raise ValueError( + f"Invalid fused expert weight {name!r}: expected first dimension " + f"{num_experts}, got shape={tuple(loaded_weight.shape)}" + ) + for expert_id in range(num_experts): + yield ( + f"{expert_prefix}.{expert_id}.{projection_name}.weight", + loaded_weight[expert_id], + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index afb985c55729..08a1b2158bd7 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -181,6 +181,7 @@ "fa3", "fa4", "flashinfer", + "llada2_cfg_flashinfer", "flashmla", "trtllm_mla", "cutedsl_mla", diff --git a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py index 16f49cf1575f..200bffba3ba4 100644 --- a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py +++ b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py @@ -3,7 +3,13 @@ from unittest.mock import patch from sglang.srt.layers import communicator as comm -from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode +from sglang.srt.layers.communicator import ( + CommunicateContext, + LayerCommunicator, + LayerScatterModes, + ScatterMode, + enable_moe_fully_dp, +) from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -61,5 +67,114 @@ def test_pure_ep_still_fuses(self): self.assertTrue(self._should_fuse(moe_ep_size=4, moe_tp_size=1)) +class TestFullReplicaMoE(CustomTestCase): + def _topology(self, **overrides): + topology = dict( + tp_rank=0, + tp_size=2, + pp_size=1, + dcp_enabled=False, + attn_tp_rank=0, + attn_tp_size=1, + attn_dp_size=2, + attn_cp_rank=0, + attn_cp_size=1, + moe_ep_size=1, + moe_tp_size=1, + moe_dp_size=2, + moe_dense_tp_size=1, + dwdp_size=1, + ) + topology.update(overrides) + return topology + + def test_full_replica_sparse_mlp_stays_local(self): + backend = types.SimpleNamespace(is_none=lambda: True) + mode_args = dict( + num_layers=2, + layer_id=0, + is_layer_sparse=True, + is_previous_layer_sparse=None, + is_next_layer_sparse=True, + ) + with ( + patch.object(comm, "is_dp_attention_enabled", return_value=True), + patch.object(comm, "get_moe_a2a_backend", return_value=backend), + patch.object( + comm, + "should_use_flashinfer_cutlass_moe_fp4_allgather", + return_value=False, + ), + patch.object(comm, "is_enable_moe_cp_allgather", return_value=False), + patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False), + patch.object(comm, "is_mla_prefill_cp_enabled", return_value=False), + patch.object( + comm, + "get_parallel", + return_value=types.SimpleNamespace(**self._topology()), + ), + ): + self.assertTrue(enable_moe_fully_dp()) + self.assertIs( + LayerScatterModes.init_new(**mode_args).mlp_mode, + ScatterMode.SCATTERED, + ) + + with ( + patch.object(comm, "is_dp_attention_enabled", return_value=True), + patch.object(comm, "get_moe_a2a_backend", return_value=backend), + patch.object( + comm, + "should_use_flashinfer_cutlass_moe_fp4_allgather", + return_value=False, + ), + patch.object(comm, "is_enable_moe_cp_allgather", return_value=False), + patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False), + patch.object(comm, "is_mla_prefill_cp_enabled", return_value=False), + patch.object( + comm, + "get_parallel", + return_value=types.SimpleNamespace(**self._topology(moe_tp_size=2)), + ), + ): + self.assertFalse(enable_moe_fully_dp()) + self.assertIs( + LayerScatterModes.init_new(**mode_args).mlp_mode, + ScatterMode.FULL, + ) + + def test_full_replica_moe_full_width_is_defined(self): + backend = types.SimpleNamespace(is_none=lambda: True) + with ( + patch.object(comm, "is_dp_attention_enabled", return_value=True), + patch.object(comm, "get_moe_a2a_backend", return_value=backend), + patch.object(comm, "get_moe_cp_size", return_value=2), + patch.object( + comm, + "get_parallel", + return_value=types.SimpleNamespace(**self._topology()), + ), + ): + context = CommunicateContext.init_new() + + self.assertEqual(context.process_group_sizes[ScatterMode.FULL], 2) + self.assertEqual(context.process_group_sizes[ScatterMode.MOE_FULL], 2) + + def test_invalid_non_replica_moe_width_fails_early(self): + backend = types.SimpleNamespace(is_none=lambda: True) + with ( + patch.object(comm, "is_dp_attention_enabled", return_value=True), + patch.object(comm, "get_moe_a2a_backend", return_value=backend), + patch.object(comm, "get_moe_cp_size", return_value=2), + patch.object( + comm, + "get_parallel", + return_value=types.SimpleNamespace(**self._topology(moe_tp_size=2)), + ), + ): + with self.assertRaisesRegex(ValueError, "must be divisible"): + CommunicateContext.init_new() + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_llada2_full_replica.py b/test/registered/unit/models/test_llada2_full_replica.py new file mode 100644 index 000000000000..f6562c6ae1ca --- /dev/null +++ b/test/registered/unit/models/test_llada2_full_replica.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.models import llada2 # noqa: E402 + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestLLaDA2FullReplica(CustomTestCase): + @staticmethod + def _config(): + return SimpleNamespace( + num_experts_per_tok=1, + norm_topk_prob=True, + hidden_size=4, + num_shared_experts=1, + hidden_act="silu", + num_experts=2, + moe_intermediate_size=8, + ) + + def _shared_expert_kwargs(self, fully_replicated): + backend = SimpleNamespace(is_deepep=lambda: False) + with ( + patch.object( + llada2, + "get_parallel", + return_value=SimpleNamespace(tp_size=2), + ), + patch.object( + llada2, + "get_exec", + return_value=SimpleNamespace( + moe=SimpleNamespace(ep_num_redundant_experts=0) + ), + ), + patch.object( + llada2, + "enable_moe_fully_dp", + return_value=fully_replicated, + ), + patch.object( + llada2, + "get_moe_a2a_backend", + return_value=backend, + ), + patch.object( + llada2, + "LLaDA2MoeGate", + return_value=SimpleNamespace(expert_bias=None), + ), + patch.object(llada2, "TopK", return_value=object()), + patch.object( + llada2, + "get_moe_impl_class", + return_value=lambda **_kwargs: object(), + ), + patch.object( + llada2, + "LLaDA2MoeMLP", + return_value=object(), + ) as shared_mlp, + ): + llada2.LLaDA2MoeSparseMoeBlock( + layer_id=0, + config=self._config(), + ) + + return shared_mlp.call_args.kwargs + + def test_full_replica_shared_expert_is_tp1(self): + replicated_kwargs = self._shared_expert_kwargs(fully_replicated=True) + self.assertEqual(replicated_kwargs["tp_rank"], 0) + self.assertEqual(replicated_kwargs["tp_size"], 1) + + standard_kwargs = self._shared_expert_kwargs(fully_replicated=False) + self.assertNotIn("tp_rank", standard_kwargs) + self.assertNotIn("tp_size", standard_kwargs) + + def test_full_replica_skips_global_post_expert_allreduce(self): + hidden_states = torch.zeros((2, 3)) + + for fully_replicated, expected_calls in ((True, 0), (False, 1)): + with self.subTest(fully_replicated=fully_replicated): + block = SimpleNamespace( + alt_stream=None, + num_shared_experts=1, + tp_size=2, + moe_fully_replicated=fully_replicated, + _forward_shared_experts=lambda tensor: torch.full_like(tensor, 2), + _forward_router_experts=lambda tensor: torch.full_like(tensor, 3), + ) + with ( + patch.object( + llada2, + "should_skip_post_experts_all_reduce", + return_value=False, + ), + patch.object( + llada2, + "tensor_model_parallel_all_reduce", + side_effect=lambda tensor: tensor, + ) as all_reduce, + ): + output = llada2.LLaDA2MoeSparseMoeBlock.forward_normal( + block, + hidden_states, + ) + + self.assertTrue(torch.equal(output, torch.full_like(output, 5))) + self.assertEqual(all_reduce.call_count, expected_calls) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 009e312a6116..29f9659ef998 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1034,8 +1034,13 @@ def _new_cp_args(self, **overrides): tp_size=1, dp_size=1, moe_dp_size=1, + moe_dense_tp_size=None, + moe_a2a_backend="none", ep_size=1, pp_size=1, + enable_dp_attention=False, + enable_dp_lm_head=False, + dcp_size=1, enable_aiter_allreduce_fusion=False, ) defaults.update(overrides) @@ -1117,6 +1122,55 @@ def test_context_parallel_handler_initializes_cp_strategy(self): self.assertTrue(is_cp_enabled()) self.assertTrue(is_interleave()) + def test_llada2_full_replica_is_the_only_mismatched_cp_exception(self): + cases = [ + ("llada2_exact", "LLaDA2MoeModelLM", {}, True), + ("different_architecture", "Qwen3MoeForCausalLM", {}, False), + ("partial_replication", "LLaDA2MoeModelLM", {"tp_size": 4}, False), + ("dcp_enabled", "LLaDA2MoeModelLM", {"dcp_size": 2}, False), + ( + "sharded_lm_head", + "LLaDA2MoeModelLM", + {"enable_dp_lm_head": False}, + False, + ), + ] + + for name, architecture, overrides, should_pass in cases: + with self.subTest(name=name): + topology = dict( + model_path="dummy", + tp_size=2, + dp_size=2, + enable_dp_attention=True, + enable_dp_lm_head=True, + attn_cp_size=1, + moe_dp_size=2, + moe_dense_tp_size=1, + moe_a2a_backend="none", + ep_size=1, + ) + topology.update(overrides) + server_args = self._new_cp_args(**topology) + server_args._model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=[architecture]), + is_multimodal=False, + ) + + with patch.object( + envs.SGLANG_ENABLE_CP_V2, + "is_set", + return_value=True, + ): + if should_pass: + handle_context_parallelism(server_args) + else: + with self.assertRaisesRegex( + AssertionError, + "full-replica topology", + ): + handle_context_parallelism(server_args) + def test_registered_cp_legacy_args_map_to_unified_strategy(self): cases = [ ( diff --git a/test/srt/test_attention_backend_registry.py b/test/srt/test_attention_backend_registry.py new file mode 100644 index 000000000000..9b08ba38aa38 --- /dev/null +++ b/test/srt/test_attention_backend_registry.py @@ -0,0 +1,24 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS +from sglang.srt.server_args import ATTENTION_BACKEND_CHOICES + + +class TestAttentionBackendRegistry(unittest.TestCase): + def test_llada2_cfg_flashinfer_is_registered(self): + backend_name = "llada2_cfg_flashinfer" + + self.assertIn(backend_name, ATTENTION_BACKENDS) + self.assertIn(backend_name, ATTENTION_BACKEND_CHOICES) + + def test_llada2_cfg_flashinfer_rejects_mla(self): + factory = ATTENTION_BACKENDS["llada2_cfg_flashinfer"] + runner = SimpleNamespace(use_mla_backend=True) + + with self.assertRaisesRegex(ValueError, "does not use an MLA backend"): + factory(runner) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/test_llada2_attention_utils.py b/test/srt/test_llada2_attention_utils.py new file mode 100644 index 000000000000..be1079890905 --- /dev/null +++ b/test/srt/test_llada2_attention_utils.py @@ -0,0 +1,38 @@ +import unittest + +import torch + +from sglang.srt.layers.attention.llada2_attention_utils import ( + build_llada_image_custom_mask, +) + + +class TestLLaDA2AttentionUtils(unittest.TestCase): + def test_builds_official_text_query_block_mask(self): + flattened = build_llada_image_custom_mask([2], [4], "cpu") + actual = flattened.view(4, 4) + expected = torch.tensor( + [ + [True, True, False, False], + [True, True, False, False], + [True, True, True, True], + [True, True, True, True], + ] + ) + + torch.testing.assert_close(actual, expected) + + def test_builds_ragged_batch_without_cross_request_entries(self): + flattened = build_llada_image_custom_mask([1, 2], [3, 4], "cpu") + + self.assertEqual(flattened.numel(), 3 * 3 + 4 * 4) + + def test_rejects_invalid_conditioning_spans(self): + with self.assertRaisesRegex(RuntimeError, "metadata batch mismatch"): + build_llada_image_custom_mask([1], [2, 3], "cpu") + with self.assertRaisesRegex(RuntimeError, "must leave query tokens"): + build_llada_image_custom_mask([2], [2], "cpu") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/test_llada2_cfg_flashinfer_backend.py b/test/srt/test_llada2_cfg_flashinfer_backend.py new file mode 100644 index 000000000000..bcb0f4bb5ca5 --- /dev/null +++ b/test/srt/test_llada2_cfg_flashinfer_backend.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend +from sglang.srt.layers.attention.llada2_attention_utils import ( + build_llada_image_custom_mask, +) +from sglang.srt.layers.attention.llada2_cfg_flashinfer_backend import ( + LLaDA2CFGFlashInferAttnBackend, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode + + +class _FakeRaggedWrapper: + def __init__(self): + self.args = None + self.kwargs = None + self.forward_args = None + self.forward_kwargs = None + + def begin_forward(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + def forward(self, *args, **kwargs): + self.forward_args = args + self.forward_kwargs = kwargs + return torch.zeros_like(args[0]) + + +class _FakeTokenToKVPool: + def __init__(self): + self.set_args = None + + def set_kv_buffer(self, *args): + self.set_args = args + + +class _FakePrefillUpdater: + def __init__(self): + self.num_qo_heads = 8 + self.num_kv_heads = 8 + self.head_dim = 64 + self.q_data_type = torch.bfloat16 + self.data_type = torch.bfloat16 + + +def _make_backend(): + backend = object.__new__(LLaDA2CFGFlashInferAttnBackend) + backend._llada_image_conditioning_mask_active = False + backend._cfg_prefill_wrapper_ragged = _FakeRaggedWrapper() + backend.num_wrappers = 1 + backend.prefill_wrappers_paged = [object()] + backend.indices_updater_prefill = _FakePrefillUpdater() + backend.prefill_split_tile_size = None + return backend + + +def _make_forward_batch(prefix_lens=(0, 0)): + seq_lens = torch.tensor([4, 3], dtype=torch.int32) + return SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + seq_lens=seq_lens, + extend_prefix_lens=torch.tensor(prefix_lens, dtype=torch.int32), + llada_image_conditioning_text_lens_cpu=[2, 1], + ) + + +class TestLLaDA2CFGFlashInferBackend(unittest.TestCase): + def test_delegates_batches_without_llada_image_metadata(self): + backend = _make_backend() + backend._cfg_prefill_wrapper_ragged = SimpleNamespace() + forward_batch = _make_forward_batch() + del forward_batch.llada_image_conditioning_text_lens_cpu + + with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: + backend.init_forward_metadata(forward_batch) + + stock_init.assert_called_once_with(forward_batch) + self.assertFalse(backend.conditioning_mask_active) + from sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning import ( + ensure_conditioning_mask_active, + ) + + with self.assertRaisesRegex(RuntimeError, "without its block attention mask"): + ensure_conditioning_mask_active(backend) + + def test_delegates_legacy_dllm_padding_metadata(self): + backend = _make_backend() + forward_batch = _make_forward_batch() + del forward_batch.llada_image_conditioning_text_lens_cpu + forward_batch.forward_mode = ForwardMode.DLLM_EXTEND + forward_batch.dllm_left_pad_lens = torch.tensor([1, 1], dtype=torch.int32) + + with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: + backend.init_forward_metadata(forward_batch) + + stock_init.assert_called_once_with(forward_batch) + + def test_consumes_llada_image_text_lengths_as_custom_mask(self): + backend = _make_backend() + forward_batch = _make_forward_batch() + + with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: + backend.init_forward_metadata(forward_batch) + + stock_init.assert_not_called() + wrapper = backend._cfg_prefill_wrapper_ragged + self.assertEqual(wrapper.args[0].tolist(), [0, 4, 7]) + self.assertEqual(wrapper.args[1].tolist(), [0, 4, 7]) + self.assertFalse(wrapper.kwargs["causal"]) + torch.testing.assert_close( + wrapper.kwargs["custom_mask"], + build_llada_image_custom_mask([2, 1], [4, 3], "cpu"), + ) + self.assertTrue(backend._llada_image_conditioning_mask_active) + self.assertTrue(backend.conditioning_mask_active) + self.assertTrue(backend.forward_metadata.use_ragged) + + def test_masked_forward_uses_ragged_kernel_without_saving_kv(self): + backend = _make_backend() + forward_batch = _make_forward_batch() + forward_batch.token_to_kv_pool = _FakeTokenToKVPool() + backend.init_forward_metadata(forward_batch) + layer = SimpleNamespace( + tp_q_head_num=2, + tp_k_head_num=2, + tp_v_head_num=2, + head_dim=4, + scaling=0.5, + logit_cap=0.0, + is_cross_attention=False, + ) + q = torch.randn(7, 8) + k = torch.randn(7, 8) + v = torch.randn(7, 8) + + output = backend.forward_extend(q, k, v, layer, forward_batch) + + self.assertEqual(output.shape, q.shape) + wrapper = backend._cfg_prefill_wrapper_ragged + torch.testing.assert_close(wrapper.forward_args[0], q.view(7, 2, 4)) + self.assertFalse(wrapper.forward_kwargs["causal"]) + self.assertIsNone(forward_batch.token_to_kv_pool.set_args) + + def test_rejects_cached_prefix_for_one_shot_conditioning(self): + backend = _make_backend() + forward_batch = _make_forward_batch(prefix_lens=(1, 0)) + + with self.assertRaisesRegex(RuntimeError, "does not support cached prefixes"): + backend.init_forward_metadata(forward_batch) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/test_llada2_weight_utils.py b/test/srt/test_llada2_weight_utils.py new file mode 100644 index 000000000000..349b0c5f043e --- /dev/null +++ b/test/srt/test_llada2_weight_utils.py @@ -0,0 +1,51 @@ +import unittest + +import torch + +from sglang.srt.models.llada2_weight_utils import prepare_llada2_language_weights + + +class TestLLaDA2WeightUtils(unittest.TestCase): + def test_prepares_clean_language_checkpoint_layout(self): + fused = torch.arange(24).reshape(2, 3, 4) + lm_head = torch.randn(4, 3) + + expanded = list( + prepare_llada2_language_weights( + [ + ("model.language_model.layers.1.mlp.experts.gate_proj", fused), + ("model.lm_head.weight", lm_head), + ], + num_experts=2, + ) + ) + + self.assertEqual( + [name for name, _ in expanded], + [ + "model.layers.1.mlp.experts.0.gate_proj.weight", + "model.layers.1.mlp.experts.1.gate_proj.weight", + "lm_head.weight", + ], + ) + torch.testing.assert_close(expanded[0][1], fused[0]) + torch.testing.assert_close(expanded[1][1], fused[1]) + torch.testing.assert_close(expanded[2][1], lm_head) + + def test_rejects_wrong_expert_count(self): + with self.assertRaisesRegex(ValueError, "expected first dimension 2"): + list( + prepare_llada2_language_weights( + [ + ( + "model.language_model.layers.1.mlp.experts.down_proj", + torch.empty(3, 4, 5), + ) + ], + num_experts=2, + ) + ) + + +if __name__ == "__main__": + unittest.main() From c4c5567973222407fc1546a450ffad63f66b24d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=81=AA=E9=97=AE?= Date: Wed, 2 Sep 2026 11:02:47 +0800 Subject: [PATCH 02/18] PullRequest: 11 fix(multimodal): use LLaDA-Image deployment defaults MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch kewen/llada-image-few-step-defaults of git@code.alipay.com:LLMInfra/sglang.git into dit-main https://code.alipay.com/LLMInfra/sglang/pull_requests/11 Reviewed-by: 明沄 * fix(multimodal): use LLaDA-Image deployment defaults --- .../configs/pipeline_configs/llada_image.py | 2 +- .../multimodal_gen/configs/sample/llada_image.py | 4 ++-- python/sglang/multimodal_gen/test/test_llada_image.py | 10 +++++++--- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py index e57c9177522d..4fb808faeb3c 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -36,7 +36,7 @@ class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): postprocess_text_funcs: tuple = () latent_scale_factor: int = 16 - text_encoder_mem_fraction_static: float = 0.35 + text_encoder_mem_fraction_static: float | None = None max_request_pixel_area: int = 2048 * 2048 # The embedded text worker admits 8192 prefill tokens shared by the two # CFG sequences and each sequence appends 256 query tokens. diff --git a/python/sglang/multimodal_gen/configs/sample/llada_image.py b/python/sglang/multimodal_gen/configs/sample/llada_image.py index 594e4ea89649..1c0ac99cb1d3 100644 --- a/python/sglang/multimodal_gen/configs/sample/llada_image.py +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -12,9 +12,9 @@ class LLaDAImageSamplingParams(SamplingParams): data_type: DataType = DataType.IMAGE negative_prompt: str | None = None - num_inference_steps: int = 20 + num_inference_steps: int = 4 num_frames: int = 1 height: int = 1024 width: int = 1024 - guidance_scale: float = 4.5 + guidance_scale: float = 1.0 max_sequence_length: int = 2048 diff --git a/python/sglang/multimodal_gen/test/test_llada_image.py b/python/sglang/multimodal_gen/test/test_llada_image.py index a7fa54808b23..c13c7607449c 100644 --- a/python/sglang/multimodal_gen/test/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/test_llada_image.py @@ -100,10 +100,11 @@ def test_prompt_format_matches_official_pipeline(self): "HUMAN Generate an image.\nASSISTANT\n", ) - def test_pipeline_config_uses_llada_image_schedule_and_shape(self): + def test_pipeline_config_defaults_schedule_and_shape(self): config = LLaDAImagePipelineConfig() sigmas = config.prepare_sigmas(None, num_inference_steps=8) + self.assertIsNone(config.text_encoder_mem_fraction_static) self.assertEqual(len(sigmas), 8) self.assertTrue(all(left > right for left, right in pairwise(sigmas))) self.assertEqual( @@ -468,9 +469,12 @@ def test_pipeline_config_rejects_unaligned_resolution(self): num_frames=1, ) - def test_service_sampling_params_do_not_expose_vq_mode(self): - field_names = {field.name for field in fields(LLaDAImageSamplingParams)} + def test_service_sampling_params_use_few_step_defaults_without_vq_mode(self): + params = LLaDAImageSamplingParams() + field_names = {field.name for field in fields(params)} + self.assertEqual(params.num_inference_steps, 4) + self.assertEqual(params.guidance_scale, 1.0) self.assertNotIn("generation_mode", field_names) self.assertNotIn("vq_token_ids", field_names) From ad5f360276b49dcf4294d33757b772f593661210 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=8E=E6=B2=84?= Date: Thu, 3 Sep 2026 14:26:47 +0800 Subject: [PATCH 03/18] PullRequest: 12 feat(llada-image): support dynamic FP8 DiT and MoE text encoder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch mingyun/quant of git@code.alipay.com:LLMInfra/sglang.git into dit-main https://code.alipay.com/LLMInfra/sglang/pull_requests/12 Reviewed-by: 恪问 * feat(llada-image): support dynamic FP8 DiT and MoE text encoder * fix(llada-image): harden FP8 conversion and loading * fix(llada2): isolate experts-only FP8 mode --- .../configs/models/dits/llada_image.py | 6 + .../runtime/models/dits/llada_image.py | 4 +- .../multimodal_gen/test/test_llada_image.py | 8 + .../test_build_modelopt_fp8_transformer.py | 169 ++++++++++++++++++ .../tools/build_modelopt_fp8_transformer.py | 162 +++++++++++++++-- python/sglang/srt/configs/model_config.py | 56 +++++- python/sglang/srt/layers/quantization/fp8.py | 8 + .../sglang/srt/model_loader/weight_utils.py | 6 + python/sglang/srt/models/llada2.py | 52 +++++- .../sglang/srt/models/llada2_weight_utils.py | 15 +- test/srt/test_llada2_fp8.py | 152 ++++++++++++++++ test/srt/test_llada2_weight_utils.py | 86 +++++++-- 12 files changed, 690 insertions(+), 34 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py create mode 100644 test/srt/test_llada2_fp8.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py index 1ef683e1dfe1..92984fdc0a80 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py @@ -34,8 +34,14 @@ class LLaDAImageArchConfig(DiTArchConfig): r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3), r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3), r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3), + r"(.*)\.attention\.to_[qkv]\.weight_scale$": ( + r"\1.attention.to_qkv.weight_scale" + ), r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2), r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2), + r"(.*)\.feed_forward\.w[13]\.weight_scale$": ( + r"\1.feed_forward.w13.weight_scale" + ), } ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index 2c9c4d8b0445..0262e315ff5b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -2418,7 +2418,9 @@ class LLaDAImageTransformer2DModel(_LLaDAImageTransformer2DModel): def __init__(self, config, hf_config: dict, quant_config=None): init_kwargs = { - key: value for key, value in hf_config.items() if not key.startswith("_") + key: value + for key, value in hf_config.items() + if not key.startswith("_") and key != "quantization_config" } super().__init__(quant_config=quant_config, **init_kwargs) self.sgl_config = config diff --git a/python/sglang/multimodal_gen/test/test_llada_image.py b/python/sglang/multimodal_gen/test/test_llada_image.py index c13c7607449c..bbe137c04f2c 100644 --- a/python/sglang/multimodal_gen/test/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/test_llada_image.py @@ -88,6 +88,14 @@ def test_dit_weight_mapping_fuses_qkv_and_swiglu_inputs(self): mapping("layers.0.feed_forward.w3.weight"), ("layers.0.feed_forward.w13.weight", 1, 2), ) + self.assertEqual( + mapping("layers.0.attention.to_q.weight_scale"), + ("layers.0.attention.to_qkv.weight_scale", None, None), + ) + self.assertEqual( + mapping("layers.0.feed_forward.w1.weight_scale"), + ("layers.0.feed_forward.w13.weight_scale", None, None), + ) def test_prompt_format_matches_official_pipeline(self): self.assertEqual( diff --git a/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py new file mode 100644 index 000000000000..c440856edcfe --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest +import torch +from safetensors.torch import load_file, save_file + +from sglang.multimodal_gen.tools.build_modelopt_fp8_transformer import ( + build_fp8_scale_map, + build_modelopt_fp8_transformer, +) + + +def _write_llada_modelopt_export(tmp_path, source_tensors, backbone_state): + source_dir = tmp_path / "source" + output_dir = tmp_path / "output" + source_dir.mkdir() + source_config = { + "_class_name": "LLaDAImageTransformer2DModel", + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + }, + } + (source_dir / "config.json").write_text(json.dumps(source_config)) + save_file(source_tensors, source_dir / "model.safetensors") + + backbone_path = tmp_path / "backbone.pt" + torch.save({"model_state_dict": backbone_state}, backbone_path) + return source_dir, output_dir, backbone_path + + +def test_dynamic_activation_scale_map_only_requires_weight_amax(): + state_dict = { + "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), + } + + scale_map = build_fp8_scale_map( + state_dict, + maxbound=448.0, + require_input_scale=False, + ) + + assert set(scale_map) == {"layers.0.attention.to_qkv.weight"} + assert set(scale_map["layers.0.attention.to_qkv.weight"]) == {"weight_scale"} + torch.testing.assert_close( + scale_map["layers.0.attention.to_qkv.weight"]["weight_scale"], + torch.tensor([0.5]), + ) + + +def test_static_activation_scale_map_still_requires_input_amax(): + weight_amax_only = { + "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), + } + assert build_fp8_scale_map(weight_amax_only, maxbound=448.0) == {} + + state_dict = { + **weight_amax_only, + "layers.0.attention.to_qkv.input_quantizer._amax": torch.tensor(112.0), + } + scale_map = build_fp8_scale_map(state_dict, maxbound=448.0) + + assert set(scale_map["layers.0.attention.to_qkv.weight"]) == { + "weight_scale", + "input_scale", + } + torch.testing.assert_close( + scale_map["layers.0.attention.to_qkv.weight"]["input_scale"], + torch.tensor([0.25]), + ) + + +def test_llada_conversion_writes_dynamic_fp8_without_input_scales(tmp_path): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + { + "layers.0.attention.to_q.weight": torch.ones(2, 2), + "layers.0.attention.to_q.input_scale": torch.tensor([0.25]), + }, + { + "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), + }, + ) + + stats = build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="llada-image", + ) + + output_config = json.loads((output_dir / "config.json").read_text()) + assert output_config["quantization_config"] == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "ignored_layers": [], + } + output_tensors = load_file(output_dir / "model.safetensors") + assert output_tensors["layers.0.attention.to_q.weight"].dtype == torch.float8_e4m3fn + assert "layers.0.attention.to_q.weight_scale" in output_tensors + assert "layers.0.attention.to_q.input_scale" not in output_tensors + assert stats["added_scale_tensors"] == 1 + + +@pytest.mark.parametrize( + ("module_prefix", "source_modules", "fused_module"), + [ + ("layers.0.attention", ("to_q", "to_k", "to_v"), "to_qkv"), + ("layers.0.feed_forward", ("w1", "w3"), "w13"), + ], +) +def test_llada_conversion_rejects_mismatched_fused_weight_scales( + tmp_path, module_prefix, source_modules, fused_module +): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + { + f"{module_prefix}.{module}.weight": torch.ones(2, 2) + for module in source_modules + }, + { + f"{module_prefix}.{module}.weight_quantizer._amax": torch.tensor( + 224.0 if index == 0 else 112.0 + ) + for index, module in enumerate(source_modules) + }, + ) + + with pytest.raises( + ValueError, + match=( + "LLaDA-Image FP8 conversion currently requires identical weight scales" + f".*{fused_module}" + ), + ): + build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="llada-image", + ) + + +def test_llada_conversion_rejects_static_activation_quantization(tmp_path): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, + { + "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), + "layers.0.attention.to_qkv.input_quantizer._amax": torch.tensor(112.0), + }, + ) + + with pytest.raises( + ValueError, + match=( + "LLaDA-Image FP8 conversion currently supports only dynamic " + "activation quantization" + ), + ): + build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="llada-image", + activation_scheme="static", + ) diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py index a579e2035da8..9c6185b0ece9 100644 --- a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py +++ b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py @@ -2,7 +2,7 @@ The core conversion path is model-agnostic: - read the ModelOpt diffusers transformer export -- rebuild per-layer `weight_scale` / `input_scale` tensors from `backbone.pt` +- rebuild per-layer `weight_scale` tensors (and static `input_scale` tensors when requested) from `backbone.pt` - materialize SGLang-native `float8_e4m3fn` weights - preserve ModelOpt `ignore` layers in their original dtype @@ -297,11 +297,26 @@ def _map_hunyuanvideo_runtime_module_name(module_name: str) -> list[str]: return mapped_names +def _map_llada_image_runtime_module_name(module_name: str) -> list[str]: + replacements = ( + (r"^(.*\.attention)\.to_[qkv]$", r"\1.to_qkv"), + (r"^(.*\.feed_forward)\.w[13]$", r"\1.w13"), + ) + mapped_names: list[str] = [] + for pattern, replacement in replacements: + mapped = re.sub(pattern, replacement, module_name) + if mapped != module_name: + mapped_names.append(mapped) + return mapped_names + + def _get_runtime_module_name_mapper( *, model_type: str, class_name: str | None ) -> Callable[[str], list[str]] | None: if model_type == "hunyuan-video" or class_name == "HunyuanVideoTransformer3DModel": return _map_hunyuanvideo_runtime_module_name + if model_type == "llada-image" or class_name == "LLaDAImageTransformer2DModel": + return _map_llada_image_runtime_module_name return None @@ -346,25 +361,70 @@ def _preferred_module_name( return _module_name_variants(weight_name, runtime_name_mapper)[-1] -def _scale_key_candidates(weight_name: str) -> list[str]: +def _scale_key_candidates( + weight_name: str, + runtime_name_mapper: Callable[[str], list[str]] | None = None, +) -> list[str]: candidates = [weight_name] if weight_name.startswith("model.diffusion_model."): candidates.append( "velocity_model." + weight_name[len("model.diffusion_model.") :] ) + for module_name in _module_name_variants(weight_name, runtime_name_mapper): + candidate = f"{module_name}.weight" + if candidate not in candidates: + candidates.append(candidate) return candidates def _resolve_scale_key( weight_name: str, scale_map: Mapping[str, Mapping[str, torch.Tensor]], + runtime_name_mapper: Callable[[str], list[str]] | None = None, ) -> str | None: - for candidate in _scale_key_candidates(weight_name): + for candidate in _scale_key_candidates(weight_name, runtime_name_mapper): if candidate in scale_map: return candidate return None +def _validate_llada_image_fused_weight_scales( + weight_names: Iterable[str], + scale_map: Mapping[str, Mapping[str, torch.Tensor]], +) -> None: + scales_by_fused_module: dict[str, list[tuple[str, torch.Tensor]]] = defaultdict( + list + ) + for weight_name in weight_names: + mapped_modules = _map_llada_image_runtime_module_name(weight_name[:-7]) + if not mapped_modules: + continue + scale_key = _resolve_scale_key( + weight_name, scale_map, _map_llada_image_runtime_module_name + ) + if scale_key is None: + continue + fused_module = _preferred_module_name( + weight_name, _map_llada_image_runtime_module_name + ) + scales_by_fused_module[fused_module].append( + (scale_key, scale_map[scale_key]["weight_scale"]) + ) + + for fused_module, scale_entries in scales_by_fused_module.items(): + reference_scale = scale_entries[0][1] + if all(torch.equal(scale, reference_scale) for _, scale in scale_entries[1:]): + continue + scale_details = ", ".join( + f"{scale_key}={scale.item():g}" for scale_key, scale in scale_entries + ) + raise ValueError( + "LLaDA-Image FP8 conversion currently requires identical weight scales " + f"for parameters fused into {fused_module!r}; got {scale_details}. " + "Re-export the checkpoint with a fused weight observer for this module." + ) + + def _is_ltx2_x0_export( *, config: Mapping[str, object], @@ -473,6 +533,7 @@ def build_fp8_scale_map( model_state_dict: Mapping[str, torch.Tensor], *, maxbound: float = FP8_E4M3_MAXBOUND, + require_input_scale: bool = True, ) -> dict[str, dict[str, torch.Tensor]]: scale_map: dict[str, dict[str, torch.Tensor]] = {} for key, value in model_state_dict.items(): @@ -487,10 +548,13 @@ def build_fp8_scale_map( value.detach().to(torch.float32).reshape(1).cpu() / maxbound ) + required_scales = ( + {"weight_scale", "input_scale"} if require_input_scale else {"weight_scale"} + ) return { weight_name: scale_tensors for weight_name, scale_tensors in scale_map.items() - if {"weight_scale", "input_scale"} <= set(scale_tensors) + if required_scales <= set(scale_tensors) } @@ -551,6 +615,7 @@ def build_modelopt_fp8_transformer( base_transformer_dir: str | None = None, model_type: str = "auto", keep_bf16_patterns: Sequence[str] | None = None, + activation_scheme: str = "auto", maxbound: float = FP8_E4M3_MAXBOUND, overwrite: bool = False, ) -> dict[str, int]: @@ -580,6 +645,23 @@ def build_modelopt_fp8_transformer( source_weight_map=source_weight_map_all, ) class_name = config.get("_class_name") + is_llada_image = ( + model_type == "llada-image" or class_name == "LLaDAImageTransformer2DModel" + ) + if activation_scheme == "auto": + activation_scheme = "dynamic" if is_llada_image else "static" + if activation_scheme not in {"static", "dynamic"}: + raise ValueError( + "activation_scheme must be one of 'auto', 'static', or 'dynamic', " + f"got {activation_scheme!r}." + ) + if is_llada_image and activation_scheme != "dynamic": + raise ValueError( + "LLaDA-Image FP8 conversion currently supports only dynamic activation " + "quantization; set activation_scheme='dynamic' or leave it as 'auto'." + ) + dynamic_activation = activation_scheme == "dynamic" + runtime_name_mapper = _get_runtime_module_name_mapper( model_type=model_type, class_name=class_name ) @@ -630,7 +712,11 @@ def build_modelopt_fp8_transformer( backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[ "model_state_dict" ] - fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound) + fp8_scale_map = build_fp8_scale_map( + backbone_state, + maxbound=maxbound, + require_input_scale=not dynamic_activation, + ) quant_algo = str(quant_config.get("quant_algo", "")).upper() if quant_algo and "FP8" not in quant_algo: raise ValueError( @@ -655,7 +741,8 @@ def build_modelopt_fp8_transformer( _preferred_module_name(weight_name, runtime_name_mapper) for weight_name in source_weight_map if weight_name.endswith(".weight") - and _resolve_scale_key(weight_name, fp8_scale_map) is None + and _resolve_scale_key(weight_name, fp8_scale_map, runtime_name_mapper) + is None } ) fallback_ignore_modules = sorted( @@ -671,7 +758,27 @@ def build_modelopt_fp8_transformer( *fallback_ignore_modules, } ) - effective_quant_config["ignore"] = ignore_patterns + if is_llada_image: + _validate_llada_image_fused_weight_scales( + ( + weight_name + for weight_name in source_weight_map + if weight_name.endswith(".weight") + and weight_name not in fallback_weight_names_set + and not is_ignored_by_modelopt( + weight_name, ignore_patterns, runtime_name_mapper + ) + ), + fp8_scale_map, + ) + if dynamic_activation: + effective_quant_config = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "ignored_layers": ignore_patterns, + } + else: + effective_quant_config["ignore"] = ignore_patterns serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True) output_config = _build_output_config( source_config=config, @@ -686,11 +793,12 @@ def build_modelopt_fp8_transformer( else {} ) fallback_scale_names = { - scale_name + weight_name[:-7] + scale_suffix for weight_name in fallback_weight_names - for scale_name in ( - weight_name[:-7] + ".weight_scale", - weight_name[:-7] + ".input_scale", + for scale_suffix in ( + (".weight_scale",) + if dynamic_activation + else (".weight_scale", ".input_scale") ) } @@ -726,6 +834,9 @@ def build_modelopt_fp8_transformer( if "_quantizer." in name: del shard_tensors[name] continue + if dynamic_activation and name.endswith(".input_scale"): + del shard_tensors[name] + continue if name in fallback_scale_names: del shard_tensors[name] continue @@ -737,7 +848,7 @@ def build_modelopt_fp8_transformer( ): preserved_ignored_weight_count += 1 continue - scale_key = _resolve_scale_key(name, fp8_scale_map) + scale_key = _resolve_scale_key(name, fp8_scale_map, runtime_name_mapper) if ( name.endswith(".weight") and scale_key is not None @@ -748,11 +859,18 @@ def build_modelopt_fp8_transformer( shard_tensors[name] = quantize_fp8_weight( shard_tensors[name], scale_tensors["weight_scale"] ) + # One fused runtime observer can feed multiple logical source + # weights (for example LLaDA Q/K/V). Safetensors rejects + # multiple dictionary entries that alias the same storage. shard_tensors[name[:-7] + ".weight_scale"] = scale_tensors[ "weight_scale" - ] - shard_tensors[name[:-7] + ".input_scale"] = scale_tensors["input_scale"] - added_scale_count += 2 + ].clone() + added_scale_count += 1 + if not dynamic_activation: + shard_tensors[name[:-7] + ".input_scale"] = scale_tensors[ + "input_scale" + ].clone() + added_scale_count += 1 save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata) @@ -782,7 +900,7 @@ def build_modelopt_fp8_transformer( 1 for name in source_weight_map if name.endswith(".weight") - and _resolve_scale_key(name, fp8_scale_map) is not None + and _resolve_scale_key(name, fp8_scale_map, runtime_name_mapper) is not None and not is_ignored_by_modelopt(name, ignore_patterns, runtime_name_mapper) ), "bf16_fallback_weights": len(fallback_weight_names), @@ -830,6 +948,7 @@ def _parse_args() -> argparse.Namespace: "ltx2", "hunyuan-video", "qwen-image", + "llada-image", "none", ], default="auto", @@ -850,6 +969,16 @@ def _parse_args() -> argparse.Namespace: "staying in FP8." ), ) + parser.add_argument( + "--activation-scheme", + choices=["auto", "static", "dynamic"], + default="auto", + help=( + "Activation quantization scheme in the converted checkpoint. 'auto' " + "uses dynamic activation quantization for LLaDA-Image and preserves " + "the legacy static ModelOpt format for other model families." + ), + ) parser.add_argument( "--maxbound", type=float, @@ -873,6 +1002,7 @@ def main() -> None: base_transformer_dir=args.base_transformer_dir, model_type=args.model_type, keep_bf16_patterns=args.keep_bf16_pattern, + activation_scheme=args.activation_scheme, maxbound=args.maxbound, overwrite=args.overwrite, ) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 2a87a7e0be03..1ef6a7fd546b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1240,6 +1240,56 @@ def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: else: return self.get_num_kv_heads(tensor_parallel_size) + def _parse_llada_fp8_experts_config(self) -> Optional[dict]: + if _hf_arch(self.hf_config) != "LLaDA2MoeModelLM": + return None + + llada_fp8_experts = _hf_attr(self.hf_config, "llada_fp8_experts") + if llada_fp8_experts is None: + return None + if not isinstance(llada_fp8_experts, dict): + llada_fp8_experts = llada_fp8_experts.to_dict() + if not llada_fp8_experts.get("enabled", False): + return None + + fp8_format = llada_fp8_experts.get("format") + weight_granularity = llada_fp8_experts.get("weight_granularity") + activation_granularity = llada_fp8_experts.get("activation_granularity") + activation_method = llada_fp8_experts.get("activation_method", "") + if fp8_format not in {"e4m3", "e4m3fn"}: + raise ValueError( + "llada_fp8_experts.format must be e4m3/e4m3fn, " + f"but got {fp8_format!r}." + ) + if weight_granularity != "per_expert_2d_block": + raise ValueError( + "llada_fp8_experts.weight_granularity must be " + f"'per_expert_2d_block', but got {weight_granularity!r}." + ) + if activation_granularity != "per_token" or not str( + activation_method + ).startswith("dynamic"): + raise ValueError( + "llada_fp8_experts requires dynamic per-token " + "activation quantization." + ) + weight_block_size = llada_fp8_experts.get("weight_block_size") + if ( + not isinstance(weight_block_size, (list, tuple)) + or len(weight_block_size) != 2 + or any(not isinstance(size, int) or size <= 0 for size in weight_block_size) + ): + raise ValueError( + "llada_fp8_experts.weight_block_size must contain " + f"two positive integers, but got {weight_block_size!r}." + ) + return { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": list(weight_block_size), + "llada_experts_only": True, + } + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _parse_quant_hf_config(self): quant_cfg = _quant_config_to_dict( @@ -1292,8 +1342,8 @@ def _parse_quant_hf_config(self): max_delay=5.0, ) if not file_exists: - # File doesn't exist on hub, no need to try downloading - return quant_cfg # None + # File doesn't exist on hub, no need to try downloading. + return self._parse_llada_fp8_experts_config() # Download (online mode) or read from cache (offline mode) if envs.SGLANG_USE_MODELSCOPE.get(): @@ -1336,6 +1386,8 @@ def _parse_quant_hf_config(self): with open(quant_config_file) as f: quant_config_dict = json.load(f) quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) + if quant_cfg is None: + quant_cfg = self._parse_llada_fp8_experts_config() return quant_cfg def _find_quant_modelslim_config(self): diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 1ed1d21c8a30..80d430da94a8 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -244,12 +244,14 @@ def __init__( use_mxfp8: bool = False, is_fp4_experts: bool = False, kv_cache_quant_algo: Optional[str] = None, + llada_experts_only: bool = False, ) -> None: super().__init__() # DSV4 mxfp4-packed (True) vs converted FP8 (False); injected by # model_loader from ModelConfig. Default False off the DSV4 path. self.is_fp4_experts = is_fp4_experts self.dequant_fp4_to_fp8 = False + self.llada_experts_only = llada_experts_only self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized if is_checkpoint_fp8_serialized: log_info_on_rank0(logger, "Detected fp8 checkpoint.") @@ -335,6 +337,7 @@ def from_config(cls, config: Dict[str, Any]) -> Fp8Config: kv_cache_quant_algo = cls.get_from_keys_or( config, ["kv_cache_quant_algo"], None ) + llada_experts_only = cls.get_from_keys_or(config, ["llada_experts_only"], False) if use_mxfp8: # MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only. if weight_block_size is not None and weight_block_size != [1, 32]: @@ -351,6 +354,7 @@ def from_config(cls, config: Dict[str, Any]) -> Fp8Config: packed_modules_mapping=packed_modules_mapping, use_mxfp8=use_mxfp8, kv_cache_quant_algo=kv_cache_quant_algo, + llada_experts_only=llada_experts_only, ) def get_quant_method( @@ -361,6 +365,8 @@ def get_quant_method( from sglang.srt.layers.radix_attention import RadixAttention if isinstance(layer, LinearBase): + if self.llada_experts_only: + return UnquantizedLinearMethod() if is_layer_skipped( prefix, self.ignored_layers, fused_mapping=self.packed_modules_mapping ): @@ -425,6 +431,8 @@ def get_quant_method( return Mxfp4FlashinferTrtllmMoEMethod(fp8_method, prefix=prefix) return fp8_method elif isinstance(layer, RadixAttention): + if self.llada_experts_only: + return None return Fp8KVCacheMethod(self) return None diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index fd659e7c983d..fcf323343eff 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -301,6 +301,12 @@ def get_quant_config( model_config, quant_cls.from_config(hf_quant_config) ) + if model_config.quantization == "fp8": + hf_quant_config = model_config._parse_llada_fp8_experts_config() + if hf_quant_config is not None: + hf_quant_config["packed_modules_mapping"] = packed_modules_mapping + return quant_cls.from_config(hf_quant_config) + # In case of bitsandbytes/QLoRA, get quant config from the adapter model. if model_config.quantization == "bitsandbytes": if ( diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index e0631774ea75..8f9f2091ca18 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -900,8 +900,15 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) params_dict = dict(self.named_parameters()) + validate_llada_fp8_experts = bool( + getattr(self.quant_config, "llada_experts_only", False) + ) + loaded_expert_weight_count = 0 + loaded_expert_scale_count = 0 for name, loaded_weight in prepare_llada2_language_weights( - weights, num_experts=self.config.num_experts + weights, + num_experts=self.config.num_experts, + expand_expert_scales=validate_llada_fp8_experts, ): if ( ("v_head" in name) @@ -958,6 +965,11 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): shard_id=shard_id, expert_id=expert_id, ) + if validate_llada_fp8_experts: + if name.endswith("weight_scale_inv"): + loaded_expert_scale_count += 1 + elif name.endswith("weight"): + loaded_expert_weight_count += 1 break else: # Skip loading extra bias for GPTQ models. @@ -983,6 +995,44 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) } ) + if validate_llada_fp8_experts: + sparse_layers = [ + layer + for layer in self.model.layers + if not isinstance(layer, PPMissingLayer) + and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) + ] + expected_tensor_count = len(sparse_layers) * self.config.num_experts * 3 + if loaded_expert_weight_count != expected_tensor_count: + raise ValueError( + "Incomplete LLaDA FP8 expert weights: loaded " + f"{loaded_expert_weight_count}, expected {expected_tensor_count}." + ) + if loaded_expert_scale_count != expected_tensor_count: + scale_param_names = [ + name + for name in params_dict + if "mlp.experts" in name and "scale" in name + ] + raise ValueError( + "Incomplete LLaDA FP8 expert scales: loaded " + f"{loaded_expert_scale_count}, expected {expected_tensor_count}; " + f"registered scale parameters sample={scale_param_names[:8]}." + ) + + if sparse_layers: + first_experts = sparse_layers[0].mlp.experts + logger.info( + "Loaded LLaDA experts-only FP8: w13=%s/%s, w2=%s/%s, " + "w13_scale=%s, w2_scale=%s", + tuple(first_experts.w13_weight.shape), + first_experts.w13_weight.dtype, + tuple(first_experts.w2_weight.shape), + first_experts.w2_weight.dtype, + tuple(first_experts.w13_weight_scale_inv.shape), + tuple(first_experts.w2_weight_scale_inv.shape), + ) + @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/python/sglang/srt/models/llada2_weight_utils.py b/python/sglang/srt/models/llada2_weight_utils.py index 2e38033ae5d3..a9b554f97731 100644 --- a/python/sglang/srt/models/llada2_weight_utils.py +++ b/python/sglang/srt/models/llada2_weight_utils.py @@ -8,14 +8,22 @@ import torch _EXPERT_PROJECTION_NAMES = frozenset({"gate_proj", "up_proj", "down_proj"}) +_EXPERT_SCALE_NAMES = frozenset( + {f"{projection_name}_scale" for projection_name in _EXPERT_PROJECTION_NAMES} +) def prepare_llada2_language_weights( weights: Iterable[tuple[str, torch.Tensor]], *, num_experts: int, + expand_expert_scales: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Normalize language-model prefixes and expand fused expert tensors.""" + fused_expert_names = _EXPERT_PROJECTION_NAMES + if expand_expert_scales: + fused_expert_names |= _EXPERT_SCALE_NAMES + for name, loaded_weight in weights: if name.startswith("model.language_model."): name = "model." + name.removeprefix("model.language_model.") @@ -26,7 +34,7 @@ def prepare_llada2_language_weights( is_fused_expert = ( separator and expert_prefix.endswith(".mlp.experts") - and projection_name in _EXPERT_PROJECTION_NAMES + and projection_name in fused_expert_names ) if not is_fused_expert: yield name, loaded_weight @@ -37,8 +45,11 @@ def prepare_llada2_language_weights( f"Invalid fused expert weight {name!r}: expected first dimension " f"{num_experts}, got shape={tuple(loaded_weight.shape)}" ) + is_scale = expand_expert_scales and projection_name in _EXPERT_SCALE_NAMES + projection_name = projection_name.removesuffix("_scale") + parameter_name = "weight_scale_inv" if is_scale else "weight" for expert_id in range(num_experts): yield ( - f"{expert_prefix}.{expert_id}.{projection_name}.weight", + f"{expert_prefix}.{expert_id}.{projection_name}.{parameter_name}", loaded_weight[expert_id], ) diff --git a/test/srt/test_llada2_fp8.py b/test/srt/test_llada2_fp8.py new file mode 100644 index 000000000000..22c0b0e3391c --- /dev/null +++ b/test/srt/test_llada2_fp8.py @@ -0,0 +1,152 @@ +import tempfile +import unittest +from types import SimpleNamespace + +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.model_loader.weight_utils import get_quant_config + + +def _experts_only_config(**overrides): + config = { + "enabled": True, + "format": "e4m3fn", + "weight_granularity": "per_expert_2d_block", + "weight_block_size": [128, 128], + "activation_method": "dynamic_absmax_rtn", + "activation_granularity": "per_token", + } + config.update(overrides) + return config + + +class TestLLaDA2Fp8(unittest.TestCase): + def setUp(self): + self.model_dir = tempfile.TemporaryDirectory() + self.addCleanup(self.model_dir.cleanup) + + def _model_config( + self, + custom_config, + *, + architecture="LLaDA2MoeModelLM", + quantization_config=None, + compression_config=None, + ): + model_config = ModelConfig.__new__(ModelConfig) + model_config.model_path = self.model_dir.name + model_config.hf_config = SimpleNamespace( + architectures=[architecture], + quantization_config=quantization_config, + compression_config=compression_config, + llada_fp8_experts=custom_config, + ) + return model_config + + def test_parses_experts_only_block_fp8(self): + model_config = self._model_config(_experts_only_config()) + + parsed = model_config._parse_quant_hf_config() + self.assertEqual( + parsed, + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "llada_experts_only": True, + }, + ) + + fp8_config = Fp8Config.from_config(parsed) + self.assertTrue(fp8_config.is_checkpoint_fp8_serialized) + self.assertTrue(fp8_config.llada_experts_only) + self.assertEqual(fp8_config.weight_block_size, [128, 128]) + self.assertEqual(fp8_config.activation_scheme, "dynamic") + + def test_model_loader_preserves_experts_only_block_config(self): + model_config = self._model_config(_experts_only_config()) + model_config.quantization = "fp8" + + fp8_config = get_quant_config( + model_config, + SimpleNamespace(), + packed_modules_mapping={}, + ) + self.assertIsInstance(fp8_config, Fp8Config) + self.assertTrue(fp8_config.llada_experts_only) + self.assertEqual(fp8_config.weight_block_size, [128, 128]) + + def test_experts_only_mode_requires_llada_specific_opt_in(self): + non_llada_config = self._model_config( + _experts_only_config(), architecture="Qwen2ForCausalLM" + ) + self.assertIsNone(non_llada_config._parse_llada_fp8_experts_config()) + + generic_fp8_config = Fp8Config.from_config( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "moe_only": True, + } + ) + self.assertFalse(generic_fp8_config.llada_experts_only) + + def test_standard_quantization_metadata_takes_precedence(self): + standard_config = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + } + for field in ("quantization_config", "compression_config"): + with self.subTest(field=field): + model_config = self._model_config( + _experts_only_config(), **{field: standard_config.copy()} + ) + self.assertEqual(model_config._parse_quant_hf_config(), standard_config) + + def test_rejects_non_dynamic_per_token_activation(self): + model_config = self._model_config( + _experts_only_config( + activation_method="static_absmax", + activation_granularity="per_tensor", + ) + ) + with self.assertRaisesRegex(ValueError, "dynamic per-token"): + model_config._parse_quant_hf_config() + + def test_experts_only_validation_uses_effective_config_and_local_partition(self): + import torch.nn as nn + + from sglang.srt.models.llada2 import ( + LLaDA2MoeModelLM, + LLaDA2MoeSparseMoeBlock, + ) + + model = LLaDA2MoeModelLM.__new__(LLaDA2MoeModelLM) + nn.Module.__init__(model) + model.config = SimpleNamespace( + num_experts=4, + llada_fp8_experts={"enabled": True}, + ) + model.quant_config = SimpleNamespace(llada_experts_only=True) + model.model = nn.Module() + model.model.layers = nn.ModuleList() + + model.load_weights([]) + self.assertEqual(model.routed_experts_weights_of_layer, {}) + + layer = nn.Module() + layer.mlp = LLaDA2MoeSparseMoeBlock.__new__(LLaDA2MoeSparseMoeBlock) + nn.Module.__init__(layer.mlp) + layer.mlp.get_moe_weights = lambda: () + model.model.layers = nn.ModuleList([layer]) + + model.quant_config = SimpleNamespace(llada_experts_only=False) + model.load_weights([]) + + model.quant_config = SimpleNamespace(llada_experts_only=True) + with self.assertRaisesRegex(ValueError, "Incomplete LLaDA FP8 expert weights"): + model.load_weights([]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/test_llada2_weight_utils.py b/test/srt/test_llada2_weight_utils.py index 349b0c5f043e..7b3e68859c8a 100644 --- a/test/srt/test_llada2_weight_utils.py +++ b/test/srt/test_llada2_weight_utils.py @@ -32,19 +32,81 @@ def test_prepares_clean_language_checkpoint_layout(self): torch.testing.assert_close(expanded[1][1], fused[1]) torch.testing.assert_close(expanded[2][1], lm_head) - def test_rejects_wrong_expert_count(self): - with self.assertRaisesRegex(ValueError, "expected first dimension 2"): - list( - prepare_llada2_language_weights( - [ - ( - "model.language_model.layers.1.mlp.experts.down_proj", - torch.empty(3, 4, 5), - ) - ], - num_experts=2, - ) + def test_expert_scale_expansion_is_opt_in(self): + gate_scale = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + up_scale = gate_scale + 100 + down_scale = torch.arange(16, dtype=torch.float32).reshape(2, 4, 2) + weights = [ + ( + "model.language_model.layers.1.mlp.experts.gate_proj_scale", + gate_scale, + ), + ( + "model.language_model.layers.1.mlp.experts.up_proj_scale", + up_scale, + ), + ( + "model.language_model.layers.1.mlp.experts.down_proj_scale", + down_scale, + ), + ] + + unexpanded = list(prepare_llada2_language_weights(weights, num_experts=2)) + self.assertEqual( + [name for name, _ in unexpanded], + [ + "model.layers.1.mlp.experts.gate_proj_scale", + "model.layers.1.mlp.experts.up_proj_scale", + "model.layers.1.mlp.experts.down_proj_scale", + ], + ) + + expanded = list( + prepare_llada2_language_weights( + weights, + num_experts=2, + expand_expert_scales=True, ) + ) + + self.assertEqual( + [name for name, _ in expanded], + [ + "model.layers.1.mlp.experts.0.gate_proj.weight_scale_inv", + "model.layers.1.mlp.experts.1.gate_proj.weight_scale_inv", + "model.layers.1.mlp.experts.0.up_proj.weight_scale_inv", + "model.layers.1.mlp.experts.1.up_proj.weight_scale_inv", + "model.layers.1.mlp.experts.0.down_proj.weight_scale_inv", + "model.layers.1.mlp.experts.1.down_proj.weight_scale_inv", + ], + ) + torch.testing.assert_close(expanded[0][1], gate_scale[0]) + torch.testing.assert_close(expanded[3][1], up_scale[1]) + torch.testing.assert_close(expanded[5][1], down_scale[1]) + + def test_rejects_invalid_fused_expert_shape(self): + cases = [ + ( + "weight", + "model.language_model.layers.1.mlp.experts.down_proj", + False, + ), + ( + "scale", + "model.language_model.layers.1.mlp.experts.gate_proj_scale", + True, + ), + ] + for label, name, expand_expert_scales in cases: + with self.subTest(label=label): + with self.assertRaisesRegex(ValueError, "expected first dimension 2"): + list( + prepare_llada2_language_weights( + [(name, torch.empty(3, 4, 5))], + num_experts=2, + expand_expert_scales=expand_expert_scales, + ) + ) if __name__ == "__main__": From 2b6c6acfcac5914853c168f8545d4d43b12ce4fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8F=82=E6=B7=BC?= Date: Thu, 3 Sep 2026 17:23:18 +0800 Subject: [PATCH 04/18] PullRequest: 13 fix(llada-image): align RMSNorm dtype semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge branch kemiao/attn-fp32-cast of git@code.alipay.com:LLMInfra/sglang.git into dit-main https://code.alipay.com/LLMInfra/sglang/pull_requests/13 Reviewed-by: 恪问 * fix(llada-image): align RMSNorm dtype semantics * refactor(llada-image): scope RMSNorm alignment --- .../runtime/layers/layernorm.py | 33 ++++++++++++++++--- .../runtime/models/dits/llada_image.py | 2 +- .../multimodal_gen/test/test_llada_image.py | 27 +++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index c22de734a276..0005137117f4 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -21,6 +21,10 @@ can_use_fused_inplace_qknorm, fused_inplace_qknorm, ) +from sglang.kernels.ops.layernorm.rmsnorm_hf import ( + is_supported_rmsnorm_hf_hidden_size, + rmsnorm_hf, +) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -91,11 +95,13 @@ def __init__( eps: float = 1e-6, dtype: torch.dtype = torch.float32, var_hidden_size: Optional[int] = None, + cast_x_before_out_mul: bool = False, ) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.hidden_size = hidden_size + self.cast_x_before_out_mul = cast_x_before_out_mul self.variance_size_override = ( None if var_hidden_size == hidden_size else var_hidden_size ) @@ -115,6 +121,22 @@ def forward_cuda( residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: shape = x.shape + if self.cast_x_before_out_mul: + if residual is not None or self.variance_size_override is not None: + return self.forward_native(x, residual) + x_2d = x.contiguous().reshape(-1, shape[-1]) + if ( + x_2d.dtype in (torch.float16, torch.bfloat16) + and self.weight.dtype == x_2d.dtype + and is_supported_rmsnorm_hf_hidden_size(x_2d.shape[-1]) + ): + return rmsnorm_hf( + x_2d, + self.weight.data, + self.variance_epsilon, + ).view(shape) + return self.forward_native(x) + x = x.reshape(-1, shape[-1]) if residual is not None: residual_shape = residual.shape @@ -189,10 +211,13 @@ def forward_native( variance = x_var.pow(2).mean(dim=-1, keepdim=True) x = x * torch.rsqrt(variance + self.variance_epsilon) - weight = self.weight - if x.device.type == "mps" and weight.dtype != x.dtype: - weight = weight.to(dtype=x.dtype) - x = (x * weight).to(orig_dtype) + if self.cast_x_before_out_mul: + x = self.weight * x.to(orig_dtype) + else: + weight = self.weight + if x.device.type == "mps" and weight.dtype != x.dtype: + weight = weight.to(dtype=x.dtype) + x = (x * weight).to(orig_dtype) if residual is None: return x else: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index 0262e315ff5b..43212856f3ab 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -77,7 +77,7 @@ class LLaDAImageRMSNorm(SGLangRMSNorm): """RMSNorm with a synthesized unit weight for affine-free checkpoints.""" def __init__(self, hidden_size: int, eps: float): - super().__init__(hidden_size, eps=eps) + super().__init__(hidden_size, eps=eps, cast_x_before_out_mul=True) self.weight.missing_param_init = "ones" diff --git a/python/sglang/multimodal_gen/test/test_llada_image.py b/python/sglang/multimodal_gen/test/test_llada_image.py index bbe137c04f2c..eaac46f15a06 100644 --- a/python/sglang/multimodal_gen/test/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/test_llada_image.py @@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.llada_image import ( + LLaDAImageRMSNorm, LLaDAImageTransformerBlock, _LLaDAImageTransformer2DModel, ) @@ -62,6 +63,32 @@ def setUpClass(cls): def tearDownClass(cls): set_global_server_args(cls.previous_server_args) + def test_llada_rmsnorm_uses_diffusers_cast_order(self): + generator = torch.Generator().manual_seed(0) + norm = LLaDAImageRMSNorm(17, eps=1e-5).to(torch.bfloat16) + # Non-unit weights make the cast order observable. + with torch.no_grad(): + norm.weight.copy_( + torch.randn(17, generator=generator, dtype=torch.bfloat16) + ) + hidden_states = torch.randn( + 2, + 17, + generator=generator, + dtype=torch.bfloat16, + ) + + variance = hidden_states.float().pow(2).mean(-1, keepdim=True) + expected = norm.weight * ( + hidden_states.float() * torch.rsqrt(variance + norm.variance_epsilon) + ).to(hidden_states.dtype) + + self.assertEqual(norm.weight.dtype, torch.bfloat16) + torch.testing.assert_close( + norm.forward_native(hidden_states), expected, rtol=0, atol=0 + ) + self.assertTrue(norm.cast_x_before_out_mul) + def test_dit_supports_sglang_flash_attention_and_sdpa(self): backends = ( LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends From 390f794e26aa1144d01bf01e6e11951dc5fbc868 Mon Sep 17 00:00:00 2001 From: "wuchengzhang.wcz" Date: Thu, 3 Sep 2026 20:46:53 +0800 Subject: [PATCH 05/18] test(llada-image): cover FP8 and RMSNorm integration --- .../configs/pipeline_configs/llada_image.py | 9 +- python/sglang/multimodal_gen/registry.py | 5 +- .../multimodal_gen/test/test_llada_image.py | 31 +---- .../unit/test_llada_image_pipeline_config.py | 10 +- .../test/unit/test_llada_image_rmsnorm.py | 39 +++++++ python/sglang/srt/arg_groups/parallel_hook.py | 106 +++++++++--------- python/sglang/srt/configs/model_config.py | 6 +- python/sglang/srt/models/llada2.py | 1 - .../unit/models}/test_llada2_fp8.py | 23 +++- .../unit/models}/test_llada2_weight_utils.py | 6 +- 10 files changed, 129 insertions(+), 107 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py rename test/{srt => registered/unit/models}/test_llada2_fp8.py (87%) rename test/{srt => registered/unit/models}/test_llada2_weight_utils.py (95%) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py index 4fb808faeb3c..a3b4042fe301 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -90,8 +90,7 @@ def validate_server_args(self, server_args) -> None: raise ValueError("LLaDA-Image sequence parallelism requires ring_degree=1") if server_args.ulysses_degree != server_args.sp_degree: raise ValueError( - "LLaDA-Image sequence parallelism requires " - "ulysses_degree == sp_degree" + "LLaDA-Image sequence parallelism requires ulysses_degree == sp_degree" ) if server_args.sp_degree not in (1, 2): raise ValueError("LLaDA-Image currently supports only SP degrees 1 and 2") @@ -127,8 +126,7 @@ def validate_server_args(self, server_args) -> None: or server_args.is_arg_explicitly_set("cpu_offload_components") ): raise ValueError( - "LLaDA-Image requires its embedded text encoder to " - "remain resident" + "LLaDA-Image requires its embedded text encoder to remain resident" ) # Strip the encoder from auto-tuned offload selections the tuner # applies without knowing about the embedded worker. @@ -150,8 +148,7 @@ def validate_server_args(self, server_args) -> None: ) if server_args.residency_mode("text_encoder") != "resident": raise ValueError( - "LLaDA-Image requires its embedded text encoder to " - "remain resident" + "LLaDA-Image requires its embedded text encoder to remain resident" ) area_override = server_args.llada_image_max_pixel_area diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index db3f67cc7e5c..d5e2469f8c2a 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -1073,8 +1073,9 @@ def _register_configs(): pipeline_config_cls=LLaDAImagePipelineConfig, hf_model_paths=["inclusionAI/LLaDA-Image"], model_detectors=[ - lambda hf_id: "lladaimagepipeline" in hf_id.lower() - or "llada-image" in hf_id.lower() + lambda hf_id: ( + "lladaimagepipeline" in hf_id.lower() or "llada-image" in hf_id.lower() + ) ], ) # Qwen-Image diff --git a/python/sglang/multimodal_gen/test/test_llada_image.py b/python/sglang/multimodal_gen/test/test_llada_image.py index eaac46f15a06..902797c0e9f4 100644 --- a/python/sglang/multimodal_gen/test/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/test_llada_image.py @@ -14,7 +14,6 @@ from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.llada_image import ( - LLaDAImageRMSNorm, LLaDAImageTransformerBlock, _LLaDAImageTransformer2DModel, ) @@ -63,36 +62,8 @@ def setUpClass(cls): def tearDownClass(cls): set_global_server_args(cls.previous_server_args) - def test_llada_rmsnorm_uses_diffusers_cast_order(self): - generator = torch.Generator().manual_seed(0) - norm = LLaDAImageRMSNorm(17, eps=1e-5).to(torch.bfloat16) - # Non-unit weights make the cast order observable. - with torch.no_grad(): - norm.weight.copy_( - torch.randn(17, generator=generator, dtype=torch.bfloat16) - ) - hidden_states = torch.randn( - 2, - 17, - generator=generator, - dtype=torch.bfloat16, - ) - - variance = hidden_states.float().pow(2).mean(-1, keepdim=True) - expected = norm.weight * ( - hidden_states.float() * torch.rsqrt(variance + norm.variance_epsilon) - ).to(hidden_states.dtype) - - self.assertEqual(norm.weight.dtype, torch.bfloat16) - torch.testing.assert_close( - norm.forward_native(hidden_states), expected, rtol=0, atol=0 - ) - self.assertTrue(norm.cast_x_before_out_mul) - def test_dit_supports_sglang_flash_attention_and_sdpa(self): - backends = ( - LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends - ) + backends = LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends self.assertEqual( backends, {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py index c7bc45a98092..fb3cd0ff0508 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py @@ -191,8 +191,9 @@ def residency_mode(self, name): ] for overrides, message in invalid_cases: args = defaults | overrides - with self.subTest(overrides=overrides), self.assertRaisesRegex( - ValueError, message + with ( + self.subTest(overrides=overrides), + self.assertRaisesRegex(ValueError, message), ): self.config.validate_server_args(SimpleNamespace(**args)) @@ -314,8 +315,9 @@ def server(**overrides): (params(cfg_gate_step="0.5"), sp1, "must be between 0.0 and 1.0"), ] for sampling_params, server_args, message in invalid_cases: - with self.subTest(message=message), self.assertRaisesRegex( - ValueError, message + with ( + self.subTest(message=message), + self.assertRaisesRegex(ValueError, message), ): self.config.validate_request_sampling_params( sampling_params, server_args diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py new file mode 100644 index 000000000000..16e055303661 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import patch + +import pytest +import torch + +from sglang.multimodal_gen.runtime.layers import layernorm as layernorm_module +from sglang.multimodal_gen.runtime.models.dits.llada_image import LLaDAImageRMSNorm + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_llada_image_rmsnorm_dispatches_to_hf_kernel(): + torch.manual_seed(0) + hidden_size = 128 + norm = LLaDAImageRMSNorm(hidden_size, eps=1e-5).cuda().to(torch.bfloat16) + with torch.no_grad(): + norm.weight.copy_(torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)) + hidden_states = torch.randn( + 2, + 8, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + + with patch.object( + layernorm_module, + "rmsnorm_hf", + wraps=layernorm_module.rmsnorm_hf, + ) as rmsnorm_hf: + output = norm(hidden_states) + + rmsnorm_hf.assert_called_once() + variance = hidden_states.float().pow(2).mean(-1, keepdim=True) + expected = norm.weight * ( + hidden_states.float() * torch.rsqrt(variance + norm.variance_epsilon) + ).to(hidden_states.dtype) + torch.testing.assert_close(output, expected, atol=1e-2, rtol=1e-2) diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 6d8d15b69723..0f4ca42fb2ef 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -60,8 +60,7 @@ def handle_context_parallelism(server_args: Any): and not cfg.language_model_only ): raise ValueError( - "MiMo V2 CP-v2 only supports text inference; add " - "--language-only." + "MiMo V2 CP-v2 only supports text inference; add --language-only." ) if cfg.enable_prefill_cp and cfg.cp_strategy is None: @@ -82,35 +81,35 @@ def handle_context_parallelism(server_args: Any): view = resolved_view(server_args) if view.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % view.attn_cp_size == 0 - ), "tp_size must be divisible by attn_cp_size" - assert ( - cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0 - ), "tp_size must be divisible by dp_size * attn_cp_size" + assert cfg.tp_size % view.attn_cp_size == 0, ( + "tp_size must be divisible by attn_cp_size" + ) + assert cfg.tp_size % (cfg.dp_size * view.attn_cp_size) == 0, ( + "tp_size must be divisible by dp_size * attn_cp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if cfg.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - cfg.tp_size % cfg.moe_dp_size == 0 - ), "tp_size must be divisible by moe_dp_size" - assert ( - view.ep_size * cfg.moe_dp_size <= cfg.tp_size - ), "ep_size * moe_dp_size must be less than or equal to tp_size" + assert cfg.tp_size % cfg.moe_dp_size == 0, ( + "tp_size must be divisible by moe_dp_size" + ) + assert view.ep_size * cfg.moe_dp_size <= cfg.tp_size, ( + "ep_size * moe_dp_size must be less than or equal to tp_size" + ) assert cfg.pp_size == 1, "PP is not supported with context parallelism" if view.ep_size > 1: - assert ( - view.ep_size * cfg.moe_dp_size == cfg.tp_size - ), "ep_size * moe_dp_size must be equal to tp_size" + assert view.ep_size * cfg.moe_dp_size == cfg.tp_size, ( + "ep_size * moe_dp_size must be equal to tp_size" + ) - assert ( - not cfg.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not cfg.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if view.attn_cp_size != cfg.moe_dp_size: assert cfg.moe_dp_size == 1 or ( @@ -258,26 +257,26 @@ def handle_dwdp(server_args: Any): if cfg.dwdp_size <= 1: return - assert ( - cfg.dwdp_size >= 2 - ), f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" - assert ( - cfg.dwdp_size == cfg.tp_size - ), f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + assert cfg.dwdp_size >= 2, ( + f"dwdp_size must be >= 2 when enabled, got {cfg.dwdp_size}" + ) + assert cfg.dwdp_size == cfg.tp_size, ( + f"dwdp_size ({cfg.dwdp_size}) must equal tp_size ({cfg.tp_size})" + ) assert cfg.disaggregation_mode in ( "null", "prefill", ), "DWDP requires --disaggregation-mode null or prefill" - assert ( - not cfg.enable_eplb - ), "EPLB dynamic migration conflicts with static DWDP partitioning" - assert ( - cfg.speculative_algorithm is None - ), "DWDP does not support speculative decoding (MTP/draft workers)" + assert not cfg.enable_eplb, ( + "EPLB dynamic migration conflicts with static DWDP partitioning" + ) + assert cfg.speculative_algorithm is None, ( + "DWDP does not support speculative decoding (MTP/draft workers)" + ) assert cfg.pp_size == 1, "DWDP requires pp_size == 1" - assert ( - not cfg.enable_two_batch_overlap - ), "DWDP's prefetch event protocol does not support two-batch overlap" + assert not cfg.enable_two_batch_overlap, ( + "DWDP's prefetch event protocol does not support two-batch overlap" + ) if cfg.disaggregation_mode == "null": logger.warning( @@ -373,7 +372,9 @@ def handle_elastic_ep(server_args: Any): assert cfg.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", - ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ], ( + "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ) assert cfg.pp_size == 1, "PP size should be set to 1 under elastic EP" @@ -384,9 +385,9 @@ def handle_elastic_ep(server_args: Any): mooncake_ib_device=validate_ib_devices(cfg.mooncake_ib_device), ) if cfg.ep_join_mode is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--elastic-ep-join-mode requires --elastic-ep-backend to be set." + ) if cfg.ep_join_mode == "scale": assert cfg.node_rank == 1, ( "Elastic EP scale-up requires one joining TP group at " @@ -404,9 +405,9 @@ def handle_elastic_ep(server_args: Any): ) assert cfg.ep_join_rank_offset >= 0, "elastic EP join rank offset must be >= 0." if cfg.max_ep_size is not None: - assert ( - cfg.elastic_ep_backend is not None - ), "--max-ep-size requires --elastic-ep-backend to be set." + assert cfg.elastic_ep_backend is not None, ( + "--max-ep-size requires --elastic-ep-backend to be set." + ) assert cfg.max_ep_size > 0, "--max-ep-size must be a positive integer." scaling_active = ( @@ -421,16 +422,15 @@ def handle_elastic_ep(server_args: Any): ) if scaling_active: resolved = resolved_view(server_args) - assert ( - cfg.elastic_ep_scale_timeout > 0 - ), "--elastic-ep-scale-timeout must be greater than zero." + assert cfg.elastic_ep_scale_timeout > 0, ( + "--elastic-ep-scale-timeout must be greater than zero." + ) assert cfg.tokenizer_worker_num == 1, ( - "Elastic EP runtime scale-up currently requires " - "--tokenizer-worker-num 1." + "Elastic EP runtime scale-up currently requires --tokenizer-worker-num 1." + ) + assert not cfg.use_ray, ( + "Elastic EP runtime scale-up does not support --use-ray." ) - assert ( - not cfg.use_ray - ), "Elastic EP runtime scale-up does not support --use-ray." assert not cfg.enable_elastic_expert_backup, ( "Elastic EP runtime scale-up does not support " "--enable-elastic-expert-backup." diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 1ef6a7fd546b..e135abc48fd2 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1258,8 +1258,7 @@ def _parse_llada_fp8_experts_config(self) -> Optional[dict]: activation_method = llada_fp8_experts.get("activation_method", "") if fp8_format not in {"e4m3", "e4m3fn"}: raise ValueError( - "llada_fp8_experts.format must be e4m3/e4m3fn, " - f"but got {fp8_format!r}." + f"llada_fp8_experts.format must be e4m3/e4m3fn, but got {fp8_format!r}." ) if weight_granularity != "per_expert_2d_block": raise ValueError( @@ -1270,8 +1269,7 @@ def _parse_llada_fp8_experts_config(self) -> Optional[dict]: activation_method ).startswith("dynamic"): raise ValueError( - "llada_fp8_experts requires dynamic per-token " - "activation quantization." + "llada_fp8_experts requires dynamic per-token activation quantization." ) weight_block_size = llada_fp8_experts.get("weight_block_size") if ( diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 8f9f2091ca18..bda985ca3c11 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -1033,7 +1033,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): tuple(first_experts.w2_weight_scale_inv.shape), ) - @classmethod def get_model_config_for_expert_location(cls, config): num_groups = getattr(config, "n_group", 0) diff --git a/test/srt/test_llada2_fp8.py b/test/registered/unit/models/test_llada2_fp8.py similarity index 87% rename from test/srt/test_llada2_fp8.py rename to test/registered/unit/models/test_llada2_fp8.py index 22c0b0e3391c..1623bb814bbb 100644 --- a/test/srt/test_llada2_fp8.py +++ b/test/registered/unit/models/test_llada2_fp8.py @@ -2,9 +2,19 @@ import unittest from types import SimpleNamespace -from sglang.srt.configs.model_config import ModelConfig -from sglang.srt.layers.quantization.fp8 import Fp8Config -from sglang.srt.model_loader.weight_utils import get_quant_config +from transformers import PretrainedConfig + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.configs.model_config import ModelConfig # noqa: E402 +from sglang.srt.layers.quantization.fp8 import Fp8Config # noqa: E402 +from sglang.srt.model_loader.weight_utils import get_quant_config # noqa: E402 +from sglang.srt.utils.common import LazyValue # noqa: E402 + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") def _experts_only_config(**overrides): @@ -20,7 +30,7 @@ def _experts_only_config(**overrides): return config -class TestLLaDA2Fp8(unittest.TestCase): +class TestLLaDA2Fp8(CustomTestCase): def setUp(self): self.model_dir = tempfile.TemporaryDirectory() self.addCleanup(self.model_dir.cleanup) @@ -35,7 +45,7 @@ def _model_config( ): model_config = ModelConfig.__new__(ModelConfig) model_config.model_path = self.model_dir.name - model_config.hf_config = SimpleNamespace( + model_config.hf_config = PretrainedConfig( architectures=[architecture], quantization_config=quantization_config, compression_config=compression_config, @@ -132,7 +142,8 @@ def test_experts_only_validation_uses_effective_config_and_local_partition(self) model.model.layers = nn.ModuleList() model.load_weights([]) - self.assertEqual(model.routed_experts_weights_of_layer, {}) + self.assertIsInstance(model.routed_experts_weights_of_layer, LazyValue) + self.assertEqual(model.routed_experts_weights_of_layer.value, {}) layer = nn.Module() layer.mlp = LLaDA2MoeSparseMoeBlock.__new__(LLaDA2MoeSparseMoeBlock) diff --git a/test/srt/test_llada2_weight_utils.py b/test/registered/unit/models/test_llada2_weight_utils.py similarity index 95% rename from test/srt/test_llada2_weight_utils.py rename to test/registered/unit/models/test_llada2_weight_utils.py index 7b3e68859c8a..284bbe4ea816 100644 --- a/test/srt/test_llada2_weight_utils.py +++ b/test/registered/unit/models/test_llada2_weight_utils.py @@ -3,9 +3,13 @@ import torch from sglang.srt.models.llada2_weight_utils import prepare_llada2_language_weights +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase +register_cpu_ci(est_time=2, suite="base-a-test-cpu") -class TestLLaDA2WeightUtils(unittest.TestCase): + +class TestLLaDA2WeightUtils(CustomTestCase): def test_prepares_clean_language_checkpoint_layout(self): fused = torch.arange(24).reshape(2, 3, 4) lm_head = torch.randn(4, 3) From 364d6b1407d902fdf0ae7db06f7092a3302d1b58 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Fri, 4 Sep 2026 01:38:00 +0000 Subject: [PATCH 06/18] fix(llada-image): harden FP8 checkpoint compatibility --- .../stages/llada_image_conditioning.py | 25 +++++++++- .../unit/test_llada_image_conditioning.py | 19 +++++++ python/sglang/srt/configs/model_config.py | 15 +++++- python/sglang/srt/models/llada2.py | 18 +++++++ .../registered/unit/models/test_llada2_fp8.py | 49 ++++++++++++++++++- 5 files changed, 122 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py index 69549bbf60de..af14a8de5289 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/llada_image_conditioning.py @@ -19,6 +19,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + prepare_diffusers_component_path_for_loading, +) + + +def _resolve_llada_image_component_path( + model_root: str, + server_args: ServerArgs, + component_name: str, +) -> str: + override_path = (getattr(server_args, "component_paths", {}) or {}).get( + component_name + ) + if override_path is not None: + return prepare_diffusers_component_path_for_loading(override_path) + return str(Path(model_root) / component_name) def ensure_conditioning_mask_active(attn_backend) -> None: @@ -68,14 +84,19 @@ def __init__( self.queryformer = queryformer self.text_projection = text_projection self.tokenizer = tokenizer - text_encoder_path = str(Path(model_root) / "text_encoder") + text_encoder_path = _resolve_llada_image_component_path( + model_root, server_args, "text_encoder" + ) + tokenizer_path = _resolve_llada_image_component_path( + model_root, server_args, "tokenizer" + ) device = get_local_torch_device() gpu_id = device.index if gpu_id is None: gpu_id = int(os.environ.get("LOCAL_RANK", "0")) srt_args = SRTServerArgs( model_path=text_encoder_path, - tokenizer_path=str(Path(model_root) / "tokenizer"), + tokenizer_path=tokenizer_path, trust_remote_code=server_args.trust_remote_code, revision=server_args.revision, skip_tokenizer_init=True, diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py index 8d0645a09513..c171621dff72 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py @@ -194,6 +194,10 @@ def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.get_sp_parallel_rank", return_value=1, ), + patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.llada_image_conditioning.prepare_diffusers_component_path_for_loading", + side_effect=lambda path: f"resolved:{path}", + ), ): runner = LLaDAImageTextEncoderRunner( model_root="/unused/model", @@ -205,6 +209,10 @@ def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): nccl_port=29500, trust_remote_code=True, revision=None, + component_paths={ + "text_encoder": "/custom/text_encoder", + "tokenizer": "/custom/tokenizer", + }, pipeline_config=SimpleNamespace( text_encoder_mem_fraction_static=0.1 ), @@ -220,6 +228,7 @@ def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): nccl_port=29500, trust_remote_code=False, revision="pinned-rev", + component_paths={}, pipeline_config=SimpleNamespace( text_encoder_mem_fraction_static=0.1 ), @@ -227,6 +236,16 @@ def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): ) self.assertIs(runner.worker, fake_worker) + self.assertEqual(runner.server_args.model_path, "resolved:/custom/text_encoder") + self.assertEqual( + runner.server_args.tokenizer_path, "resolved:/custom/tokenizer" + ) + self.assertEqual( + untrusted_runner.server_args.model_path, "/unused/model/text_encoder" + ) + self.assertEqual( + untrusted_runner.server_args.tokenizer_path, "/unused/model/tokenizer" + ) self.assertIsNone(runner.server_args.page_size) self.assertEqual(runner.page_size, resolved_page_size) self.assertEqual( diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index e135abc48fd2..689981d1032c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1246,7 +1246,20 @@ def _parse_llada_fp8_experts_config(self) -> Optional[dict]: llada_fp8_experts = _hf_attr(self.hf_config, "llada_fp8_experts") if llada_fp8_experts is None: - return None + if _hf_attr(self.hf_config, "use_fp8_experts") is not True: + return None + logger.warning_once( + "The legacy LLaDA2 use_fp8_experts flag is deprecated; " + "publish llada_fp8_experts metadata instead." + ) + llada_fp8_experts = { + "enabled": True, + "format": "e4m3fn", + "weight_granularity": "per_expert_2d_block", + "weight_block_size": [128, 128], + "activation_method": "dynamic_absmax_rtn", + "activation_granularity": "per_token", + } if not isinstance(llada_fp8_experts, dict): llada_fp8_experts = llada_fp8_experts.to_dict() if not llada_fp8_experts.get("enabled", False): diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index bda985ca3c11..5c3d186e1a39 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -94,6 +94,23 @@ _is_cuda = is_cuda() _is_npu = is_npu() +_LEGACY_LLADA2_FP8_CONFIG_DEFAULTS = { + "embedding_dropout": 0.0, + "moe_router_enable_expert_bias": True, + "norm_topk_prob": True, + "router_dtype": "fp32", + "score_function": "sigmoid", +} + + +def _apply_legacy_llada2_fp8_config_defaults(config: PretrainedConfig) -> None: + if not getattr(config, "use_fp8_experts", False): + return + for name, default in _LEGACY_LLADA2_FP8_CONFIG_DEFAULTS.items(): + if not hasattr(config, name): + setattr(config, name, default) + + split_qkv_rmsnorm_rope_pos_cache_half_npu = None if _is_npu: try: @@ -811,6 +828,7 @@ def __init__( prefix: str = "", ): super().__init__() + _apply_legacy_llada2_fp8_config_defaults(config) self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config diff --git a/test/registered/unit/models/test_llada2_fp8.py b/test/registered/unit/models/test_llada2_fp8.py index 1623bb814bbb..3d23c6677357 100644 --- a/test/registered/unit/models/test_llada2_fp8.py +++ b/test/registered/unit/models/test_llada2_fp8.py @@ -42,15 +42,19 @@ def _model_config( architecture="LLaDA2MoeModelLM", quantization_config=None, compression_config=None, + use_fp8_experts=None, ): model_config = ModelConfig.__new__(ModelConfig) model_config.model_path = self.model_dir.name - model_config.hf_config = PretrainedConfig( + config_kwargs = dict( architectures=[architecture], quantization_config=quantization_config, compression_config=compression_config, llada_fp8_experts=custom_config, ) + if use_fp8_experts is not None: + config_kwargs["use_fp8_experts"] = use_fp8_experts + model_config.hf_config = PretrainedConfig(**config_kwargs) return model_config def test_parses_experts_only_block_fp8(self): @@ -86,12 +90,55 @@ def test_model_loader_preserves_experts_only_block_config(self): self.assertTrue(fp8_config.llada_experts_only) self.assertEqual(fp8_config.weight_block_size, [128, 128]) + def test_parses_legacy_experts_only_fp8_flag(self): + model_config = self._model_config(None, use_fp8_experts=True) + + self.assertEqual( + model_config._parse_quant_hf_config(), + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "llada_experts_only": True, + }, + ) + + disabled = self._model_config(None, use_fp8_experts=False) + self.assertIsNone(disabled._parse_llada_fp8_experts_config()) + + def test_legacy_flag_restores_only_missing_runtime_fields(self): + from sglang.srt.models.llada2 import ( + _apply_legacy_llada2_fp8_config_defaults, + ) + + legacy_config = SimpleNamespace( + use_fp8_experts=True, + router_dtype="bfloat16", + ) + _apply_legacy_llada2_fp8_config_defaults(legacy_config) + + self.assertEqual(legacy_config.embedding_dropout, 0.0) + self.assertTrue(legacy_config.moe_router_enable_expert_bias) + self.assertTrue(legacy_config.norm_topk_prob) + self.assertEqual(legacy_config.score_function, "sigmoid") + self.assertEqual(legacy_config.router_dtype, "bfloat16") + + standard_config = SimpleNamespace(use_fp8_experts=False) + _apply_legacy_llada2_fp8_config_defaults(standard_config) + self.assertFalse(hasattr(standard_config, "embedding_dropout")) + self.assertFalse(hasattr(standard_config, "score_function")) + def test_experts_only_mode_requires_llada_specific_opt_in(self): non_llada_config = self._model_config( _experts_only_config(), architecture="Qwen2ForCausalLM" ) self.assertIsNone(non_llada_config._parse_llada_fp8_experts_config()) + legacy_non_llada = self._model_config( + None, architecture="Qwen2ForCausalLM", use_fp8_experts=True + ) + self.assertIsNone(legacy_non_llada._parse_llada_fp8_experts_config()) + generic_fp8_config = Fp8Config.from_config( { "quant_method": "fp8", From 503fecc105ec58f5d7443f8caf6169bbdc8f96a4 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Mon, 7 Sep 2026 07:16:55 +0000 Subject: [PATCH 07/18] [diffusion] Isolate LLaDA-Image encoder runtime context --- .../llada_image/conditioning.py | 120 ++++++------- .../unit/test_llada_image_conditioning.py | 157 +++++++++++++++++- python/sglang/srt/runtime_context.py | 45 +++-- test/registered/unit/test_runtime_context.py | 83 +++++++++ 4 files changed, 334 insertions(+), 71 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index af14a8de5289..78c6c7c81a88 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -72,7 +72,7 @@ def __init__( from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.chunk_cache import ChunkCache - from sglang.srt.runtime_context import publish as srt_publish + from sglang.srt.runtime_context import create_context, use_context from sglang.srt.server_args import ServerArgs as SRTServerArgs # Keep the SRT world aligned with the pure-Ulysses ranks. Within that @@ -123,64 +123,63 @@ def __init__( server_args.pipeline_config.text_encoder_mem_fraction_static ), ) - # TpModelWorker reads the runtime_context config bags the srt - # scheduler process publishes before construction. - srt_publish(srt_args, role="scheduler") - # The diffusion runtime mirrors its TP group into the srt globals. - # Clear it so the worker can install the encoder's real srt groups, - # then restore the mirror and keep the encoder group for forwards. - saved_srt_tp = srt_parallel_state._TP - saved_srt_attn_tp = srt_parallel_state._ATTN_TP - try: - mm_parallel_state._clear_srt_tp_group() - self.worker = TpModelWorker( - server_args=srt_args, - gpu_id=gpu_id, - ps=ParallelState.trivial( - tp_rank=text_tp_rank, - tp_size=text_tp_size, - dp_rank=text_tp_rank if text_dp_attention else 0, - dp_size=text_tp_size if text_dp_attention else 1, - attn_tp_rank=0 if text_dp_attention else text_tp_rank, - attn_tp_size=1 if text_dp_attention else text_tp_size, - attn_cp_rank=0, - attn_cp_size=1, - attn_dcp_rank=0, - attn_dcp_size=1, - attn_dp_rank=text_tp_rank if text_dp_attention else 0, - attn_dp_size=text_tp_size if text_dp_attention else 1, - moe_ep_rank=0, - moe_ep_size=1, - moe_dp_rank=text_tp_rank if text_dp_attention else 0, - moe_dp_size=text_tp_size if text_dp_attention else 1, + self.runtime_context = create_context(srt_args, role="scheduler") + with use_context(self.runtime_context): + # The diffusion runtime mirrors its TP group into the srt globals. + # Clear it so the worker can install the encoder's real srt groups, + # then restore the mirror and keep the encoder group for forwards. + saved_srt_tp = srt_parallel_state._TP + saved_srt_attn_tp = srt_parallel_state._ATTN_TP + try: + mm_parallel_state._clear_srt_tp_group() + self.worker = TpModelWorker( + server_args=srt_args, gpu_id=gpu_id, - ), - nccl_port=server_args.nccl_port or 29500, + ps=ParallelState.trivial( + tp_rank=text_tp_rank, + tp_size=text_tp_size, + dp_rank=text_tp_rank if text_dp_attention else 0, + dp_size=text_tp_size if text_dp_attention else 1, + attn_tp_rank=0 if text_dp_attention else text_tp_rank, + attn_tp_size=1 if text_dp_attention else text_tp_size, + attn_cp_rank=0, + attn_cp_size=1, + attn_dcp_rank=0, + attn_dcp_size=1, + attn_dp_rank=text_tp_rank if text_dp_attention else 0, + attn_dp_size=text_tp_size if text_dp_attention else 1, + moe_ep_rank=0, + moe_ep_size=1, + moe_dp_rank=text_tp_rank if text_dp_attention else 0, + moe_dp_size=text_tp_size if text_dp_attention else 1, + gpu_id=gpu_id, + ), + nccl_port=server_args.nccl_port or 29500, + ) + # Run the post-construction init phases the srt scheduler drives. + self.worker.alloc_memory_pool() + self.worker.init_attention_backends() + self.worker.init_cuda_graphs() + self.encoder_tp_group = srt_parallel_state._TP + self.encoder_attn_tp_group = srt_parallel_state._ATTN_TP + finally: + srt_parallel_state._TP = saved_srt_tp + srt_parallel_state._ATTN_TP = saved_srt_attn_tp + self.server_args = srt_args + self.text_dp_attention = text_dp_attention + self.model_runner = self.worker.model_runner + self.page_size = self.model_runner.page_size + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + self.worker.get_memory_pool() ) - # Run the post-construction init phases the srt scheduler drives. - self.worker.alloc_memory_pool() - self.worker.init_attention_backends() - self.worker.init_cuda_graphs() - self.encoder_tp_group = srt_parallel_state._TP - self.encoder_attn_tp_group = srt_parallel_state._ATTN_TP - finally: - srt_parallel_state._TP = saved_srt_tp - srt_parallel_state._ATTN_TP = saved_srt_attn_tp - self.server_args = srt_args - self.text_dp_attention = text_dp_attention - self.model_runner = self.worker.model_runner - self.page_size = self.model_runner.page_size - self.req_to_token_pool, self.token_to_kv_pool_allocator = ( - self.worker.get_memory_pool() - ) - self.tree_cache = ChunkCache( - CacheInitParams( - disable=True, - req_to_token_pool=self.req_to_token_pool, - token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, - page_size=self.page_size, + self.tree_cache = ChunkCache( + CacheInitParams( + disable=True, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + page_size=self.page_size, + ) ) - ) def _prepare_input_embeds( self, @@ -240,10 +239,13 @@ def encode( ): import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state import sglang.srt.distributed.parallel_state as srt_parallel_state + from sglang.srt.runtime_context import use_context - # srt modules resolve their TP group through the srt globals, which - # hold the diffusion mirror between encodes. - with mm_parallel_state.use_tensor_parallel_group(self.encoder_tp_group): + # The worker serializes stages so encoder state can follow its groups. + with ( + use_context(self.runtime_context), + mm_parallel_state.use_tensor_parallel_group(self.encoder_tp_group), + ): saved_attn_tp = srt_parallel_state._ATTN_TP try: srt_parallel_state._ATTN_TP = self.encoder_attn_tp_group diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py index 869bf464bf80..7220ea684140 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py @@ -14,6 +14,7 @@ LLaDAImageTextConditioningStage, LLaDAImageTextEncoderRunner, ) +from sglang.srt.runtime_context import ParallelContext, RuntimeContext _GLOBAL_ARGS_PATCH = ( "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" @@ -173,7 +174,10 @@ def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): worker_module, return_value=fake_worker, ) as worker_cls, - patch("sglang.srt.runtime_context.publish"), + patch( + "sglang.srt.runtime_context.create_context", + side_effect=lambda *_args, **_kwargs: RuntimeContext(ParallelContext()), + ), patch( "sglang.srt.mem_cache.cache_init_params.CacheInitParams", return_value=object(), @@ -289,6 +293,7 @@ def test_text_runner_scopes_singleton_attention_group_and_restores(self): encoder_group = SimpleNamespace(world_size=2, rank_in_group=1) encoder_attention_group = object() runner = object.__new__(LLaDAImageTextEncoderRunner) + runner.runtime_context = RuntimeContext(ParallelContext()) runner.encoder_tp_group = encoder_group runner.encoder_attn_tp_group = encoder_attention_group observed = {} @@ -439,7 +444,10 @@ def fail_after_installing_encoder_groups(**_kwargs): "sglang.srt.managers.tp_worker.TpModelWorker", side_effect=fail_after_installing_encoder_groups, ), - patch("sglang.srt.runtime_context.publish"), + patch( + "sglang.srt.runtime_context.create_context", + side_effect=lambda *_args, **_kwargs: RuntimeContext(ParallelContext()), + ), patch( "sglang.srt.server_args.ServerArgs", side_effect=lambda **kwargs: SimpleNamespace(page_size=1, **kwargs), @@ -476,6 +484,151 @@ def fail_after_installing_encoder_groups(**_kwargs): self.assertIs(srt_parallel_state._ATTN_TP, diffusion_group) +class TestLLaDAImageRuntimeContext(unittest.TestCase): + def test_encoder_preserves_diffusion_runtime(self): + import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_state + import sglang.srt.distributed.parallel_state as srt_state + from sglang.srt import runtime_context as rc + from sglang.srt.server_args import ServerArgs as SRTServerArgs + + saved = rc.snapshot_context() + self.addCleanup(rc.restore_context, saved) + conditioning = LLaDAImageTextEncoderRunner.__module__ + for fail_init in (False, True): + with self.subTest(fail_init=fail_init): + rc.reset_context() + diffusion_args = SRTServerArgs(model_path="dummy", tp_size=1) + diffusion_context = rc.publish( + diffusion_args, role="diffusion_gpu_worker" + ) + diffusion_context.override("diffusion.setup", grammar_backend="none") + diffusion_bags = diffusion_context._config_bags + diffusion_log = diffusion_context.overrides_log() + diffusion_buffer = rc.get_buffer("conditioning", object) + diffusion_group = SimpleNamespace(world_size=1, rank_in_group=0) + encoder_group = SimpleNamespace(world_size=2, rank_in_group=1) + encoder_attention_group = SimpleNamespace(world_size=1, rank_in_group=0) + observed = {} + + def make_args(**kwargs): + kwargs["model_path"] = "dummy" + return SRTServerArgs(**kwargs) + + def check_encoder_context(): + context = rc.assert_published(observed["args"], role="scheduler") + self.assertEqual(rc.get_parallel().tp_size, 2) + self.assertEqual(rc.get_schedule().max_running_requests, 4) + return context + + def allocate(): + context = check_encoder_context() + context.override("encoder.setup", page_size=16) + rc.get_buffer("conditioning", object) + if fail_init: + raise RuntimeError("allocation failed") + + def make_worker(**kwargs): + observed["args"] = kwargs["server_args"] + observed["context"] = check_encoder_context() + srt_state._TP = encoder_group + srt_state._ATTN_TP = encoder_attention_group + return SimpleNamespace( + model_runner=SimpleNamespace(page_size=16), + get_memory_pool=lambda: (object(), object()), + alloc_memory_pool=allocate, + init_attention_backends=check_encoder_context, + init_cuda_graphs=check_encoder_context, + ) + + def make_cache(_params): + self.assertEqual(rc.get_schedule().page_size, 16) + return object() + + def check_diffusion_context(): + self.assertIs(rc.get_context(), diffusion_context) + self.assertIs(rc.get_server_args(), diffusion_args) + self.assertEqual(rc.publish_role(), "diffusion_gpu_worker") + self.assertIs(rc.get_context()._config_bags, diffusion_bags) + self.assertEqual(rc.get_context().overrides_log(), diffusion_log) + self.assertEqual(rc.get_exec().kernel.grammar_backend, "none") + self.assertEqual(rc.get_parallel().tp_size, 1) + self.assertIs( + rc.get_buffer("conditioning", object), diffusion_buffer + ) + self.assertIs(mm_state._TP, diffusion_group) + self.assertIs(srt_state._TP, diffusion_group) + self.assertIs(srt_state._ATTN_TP, diffusion_group) + + with ( + patch("sglang.srt.server_args.ServerArgs", side_effect=make_args), + patch( + "sglang.srt.managers.tp_worker.TpModelWorker", + side_effect=make_worker, + ), + patch( + "sglang.srt.mem_cache.chunk_cache.ChunkCache", + side_effect=make_cache, + ), + patch( + f"{conditioning}.get_local_torch_device", + return_value=torch.device("cpu"), + ), + patch(f"{conditioning}.get_sp_parallel_rank", return_value=1), + patch.object(mm_state, "_TP", diffusion_group), + patch.object(srt_state, "_TP", diffusion_group), + patch.object(srt_state, "_ATTN_TP", diffusion_group), + ): + kwargs = dict( + model_root="/unused/model", + queryformer=object(), + text_projection=object(), + tokenizer=object(), + server_args=SimpleNamespace( + sp_degree=2, + nccl_port=29500, + trust_remote_code=False, + revision=None, + component_paths={}, + pipeline_config=SimpleNamespace( + text_encoder_mem_fraction_static=0.1 + ), + ), + ) + if fail_init: + with self.assertRaisesRegex(RuntimeError, "allocation failed"): + LLaDAImageTextEncoderRunner(**kwargs) + check_diffusion_context() + continue + runner = LLaDAImageTextEncoderRunner(**kwargs) + check_diffusion_context() + self.assertIsNot(observed["context"], diffusion_context) + + def encode(*_args, **_kwargs): + self.assertIs(check_encoder_context(), observed["context"]) + self.assertEqual(rc.get_schedule().page_size, 16) + self.assertEqual(len(rc.get_context().overrides_log()), 1) + self.assertIsNot( + rc.get_buffer("conditioning", object), diffusion_buffer + ) + self.assertIs(srt_state._ATTN_TP, encoder_attention_group) + if observed["fail_encode"]: + raise RuntimeError("encode failed") + return ["encoded"] + + runner._encode_impl = encode + for fail_encode in (False, True, False): + observed["fail_encode"] = fail_encode + if fail_encode: + with self.assertRaisesRegex(RuntimeError, "encode failed"): + runner.encode(["hello"], max_sequence_length=16) + else: + self.assertEqual( + runner.encode(["hello"], max_sequence_length=16), + ["encoded"], + ) + check_diffusion_context() + + class TestLLaDAImageConditionKwargs(unittest.TestCase): def setUp(self): self.config = LLaDAImagePipelineConfig() diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index c9cd6e120351..3e264d11e534 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -569,6 +569,8 @@ class Resources(_FlagGroupBase): # Accessors with bespoke semantics (grow-only, per-device keys) manage # their entries directly. buffers: dict = dataclasses.field(default_factory=dict) + # The weight data parallel manager belongs to the model using it. + dwdp_manager: Any = None # Persistent reusable CUDA events for non-EP DP TBO, keyed by # (kind, subbatch) — see dp_attention._tbo_event for why reuse matters. tbo_event_pool: dict = dataclasses.field(default_factory=dict) @@ -1478,6 +1480,33 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: (bags re-projected, provenance reset, role overwritten), which is what lets one process rebuild an engine after shutting the previous one down. """ + return _configure_context(_CONTEXT, server_args, role=role) + + +def create_context(server_args, *, role: str) -> RuntimeContext: + """Build a model-owned context without replacing the active process config.""" + context = RuntimeContext(parallel=ParallelContext()) + return _configure_context(context, server_args, role=role) + + +@contextmanager +def use_context(context: RuntimeContext): + """Activate a model context for serialized embedded SRT work. + + Like process group scopes, this changes process globals and requires the + caller to serialize model execution. It does not isolate concurrent threads. + The previous context is restored by identity without re-projecting its bags. + """ + global _CONTEXT, _PARALLEL + previous_context, previous_parallel = _CONTEXT, _PARALLEL + _CONTEXT, _PARALLEL = context, context.parallel + try: + yield context + finally: + _CONTEXT, _PARALLEL = previous_context, previous_parallel + + +def _configure_context(context, server_args, *, role: str) -> RuntimeContext: if _ROLE_NS_MODE == "enforce" and role not in ROLE_NAMESPACE_SETS: # Fail closed at publish time, not at the first stray read. raise ValueError( @@ -1485,8 +1514,8 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: "its namespace set (None for the full tree)." ) server_args.resolve_once() - discarded = _CONTEXT.overrides_log() - _CONTEXT.set_server_args(server_args) + discarded = context.overrides_log() + context.set_server_args(server_args) if discarded: logger.warning( "publish(role=%s) re-projected the config bags and dropped %d " @@ -1497,7 +1526,7 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: f"{source}({', '.join(sorted(fields))})" for source, fields in discarded ), ) - _CONTEXT._publish_role = role + context._publish_role = role if _ROLE_NS_MODE == "record": # The '-' marker distinguishes a zero-read role from a process where # recording never ran (signal teardown skips atexit). @@ -1509,7 +1538,7 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: file=sys.stderr, flush=True, ) - return _CONTEXT + return context def assert_published(server_args, *, role: str) -> RuntimeContext: @@ -1563,16 +1592,12 @@ def get_buffer(name: str, factory: Any) -> Any: return _CONTEXT.get_buffer(name, factory) -_GLOBAL_DWDP_MANAGER: Any = None - - def get_global_dwdp_manager() -> Any: - return _GLOBAL_DWDP_MANAGER + return _CONTEXT.resources.dwdp_manager def set_global_dwdp_manager(manager: Any) -> None: - global _GLOBAL_DWDP_MANAGER - _GLOBAL_DWDP_MANAGER = manager + _CONTEXT.resources.dwdp_manager = manager def _group_leaves(group: _FlagGroupBase) -> dict[str, Any]: diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 5d2a6ba9da45..15e19672f825 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -267,6 +267,89 @@ def test_reset_context_clears_owned_store(self): get_server_args() +class TestEmbeddedRuntimeContext(unittest.TestCase): + def setUp(self): + from sglang.srt import runtime_context as rc + + self.addCleanup(rc.restore_context, rc.snapshot_context()) + rc.reset_context() + + def test_serial_contexts_keep_config_and_resources_across_scopes(self): + from sglang.srt import runtime_context as rc + + outer_args = ServerArgs(model_path="dummy", tp_size=1) + outer = publish(outer_args, role="diffusion_gpu_worker") + outer.override("diffusion.setup", grammar_backend="none") + outer_bags = outer._config_bags + outer_log = outer.overrides_log() + outer_buffer = rc.get_buffer("shared_name", object) + outer_manager = object() + rc.set_global_dwdp_manager(outer_manager) + outer.flags.capture.enable_torch_compile = True + outer.parallel.stamp_derived_widths(attn_tp_size=1) + + inner_args = ServerArgs(model_path="dummy", tp_size=2) + inner = rc.create_context(inner_args, role="scheduler") + self.assertIs(get_context(), outer) + self.assertIs(get_parallel(), outer.parallel) + inner_manager = object() + inner_buffer = object() + for index in range(2): + with self.subTest(index=index), rc.use_context(inner): + self.assertIs(assert_published(inner_args, role="scheduler"), inner) + self.assertIs(get_parallel(), inner.parallel) + self.assertEqual(get_parallel().tp_size, 2) + self.assertFalse(get_flags().capture.enable_torch_compile) + if index == 0: + self.assertIsNone(rc.get_global_dwdp_manager()) + self.assertNotIn("shared_name", inner.resources.buffers) + inner.override("encoder.setup", page_size=16) + inner.resources.buffers["shared_name"] = inner_buffer + inner.parallel.stamp_derived_widths(attn_tp_size=2) + rc.set_global_dwdp_manager(inner_manager) + self.assertEqual(rc.get_schedule().page_size, 16) + self.assertEqual(len(inner.overrides_log()), 1) + self.assertIs(rc.get_buffer("shared_name", object), inner_buffer) + self.assertIs(rc.get_global_dwdp_manager(), inner_manager) + self.assertEqual(get_parallel().attn_tp_size, 2) + with self.assertRaisesRegex(RuntimeError, "nested failure"): + with rc.use_context(outer): + self.assertIs(rc.get_global_dwdp_manager(), outer_manager) + self.assertEqual(get_parallel().attn_tp_size, 1) + raise RuntimeError("nested failure") + self.assertIs(get_context(), inner) + self.assertIs(get_parallel(), inner.parallel) + self.assertIs(get_context(), outer) + self.assertIs(get_server_args(), outer_args) + self.assertIs(get_parallel(), outer.parallel) + self.assertEqual(publish_role(), "diffusion_gpu_worker") + self.assertIs(outer._config_bags, outer_bags) + self.assertEqual(outer.overrides_log(), outer_log) + self.assertEqual(get_exec().kernel.grammar_backend, "none") + self.assertTrue(get_flags().capture.enable_torch_compile) + self.assertIs(rc.get_buffer("shared_name", object), outer_buffer) + self.assertIs(rc.get_global_dwdp_manager(), outer_manager) + + def test_failed_context_creation_preserves_active_context(self): + from sglang.srt import runtime_context as rc + + outer = publish(ServerArgs(model_path="dummy"), role="diffusion_gpu_worker") + outer.override("diffusion.setup", grammar_backend="none") + bags, log = outer._config_bags, outer.overrides_log() + with patch.object(rc, "_ROLE_NS_MODE", "enforce"): + with self.assertRaisesRegex(ValueError, "has no ROLE_NAMESPACE_SETS"): + rc.create_context(ServerArgs(model_path="dummy"), role="invalid_role") + broken_args = ServerArgs(model_path="dummy") + with patch.object( + broken_args, "resolve_once", side_effect=RuntimeError("resolution failed") + ): + with self.assertRaisesRegex(RuntimeError, "resolution failed"): + rc.create_context(broken_args, role="scheduler") + self.assertIs(get_context(), outer) + self.assertIs(outer._config_bags, bags) + self.assertEqual(outer.overrides_log(), log) + + class TestAssertPublished(_IsolatedServerArgs): """Publishing is the process entry's job; the constructors only check. From ac3fc0304fdf73c9ce382fad99902ef27e00c9bd Mon Sep 17 00:00:00 2001 From: Wu Chengzhang Date: Mon, 7 Sep 2026 17:11:17 +0800 Subject: [PATCH 08/18] docs: add LLaDA-Image diffusion cookbook --- .../diffusion/InclusionAI/LLaDA-Image.mdx | 252 +++++++++++++++ docs/cookbook/diffusion/intro.mdx | 6 + docs/docs.json | 7 + .../configs/inclusionAI/llada-image.jsx | 300 ++++++++++++++++++ 4 files changed, 565 insertions(+) create mode 100644 docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx create mode 100644 docs/src/snippets/configs/inclusionAI/llada-image.jsx diff --git a/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx b/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx new file mode 100644 index 000000000000..cc1872be6a63 --- /dev/null +++ b/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx @@ -0,0 +1,252 @@ +--- +title: LLaDA-Image +description: "Deploy LLaDA-Image Base, Turbo, and FP8 checkpoints on NVIDIA H200 for text-to-image generation and single-image editing with SGLang Diffusion." +--- + +import { DiffusionModelTags } from '/src/snippets/diffusion/model-tags.jsx'; +import { Deployment } from '/src/snippets/_deployment.jsx'; +import { config } from '/src/snippets/configs/inclusionAI/llada-image.jsx'; + + + +## 1. Quick start + +Install with `uv pip install "sglang[diffusion]" --prerelease=allow`, then choose +the checkpoint, request mode, and verified H200 topology below. + + + + +LLaDA-Image loads reviewed configuration code from the checkpoint. Keep +`--trust-remote-code` and the generated commit-pinned `--revision` together, +and review remote-code changes before moving to another revision. + + +## 2. Model capabilities + +[LLaDA-Image](https://huggingface.co/inclusionAI/LLaDA-Image) is a 6B-parameter +image generation and editing family. The Base checkpoint targets high-fidelity +text-to-image generation and instruction-guided editing with a 50-step schedule; +the distilled Turbo checkpoint runs the same two SGLang request modes in four +steps. Both variants are designed for English and Chinese text rendering. + +Prompt conditioning is handled by an embedded **LLaDA 2.0 16B-parameter MoE +text encoder**. Its checkpoint config declares the `LLaDA2MoeModelLM` +architecture; this is the LLaDA 2.0 structure, not the earlier LLaDA text +encoder architecture. + +Choose Base when fidelity matters more than latency and Turbo for interactive or +high-throughput serving. SGLang exposes text-to-image generation and exactly one +reference image for editing through one server. The model card's VQ-conditioned +generation mode is not currently exposed by the SGLang Images API. + +The checkpoints are released under Apache-2.0. Review the model card and its +usage terms before production deployment. + +## 3. Deployment details + +### Checkpoint selection + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
CheckpointBest fitSamplingValidation contract
inclusionAI/LLaDA-ImageMaximum fidelity50 steps, CFG 5Full BF16 quality suite
inclusionAI/LLaDA-Image-FP8Lower checkpoint memory50 steps, CFG 5Verified H200 SP1 and SP2 serving
inclusionAI/LLaDA-Image-TurboLowest audited latency4 steps, CFG 1Full BF16 SP1 and SP2 suite
inclusionAI/LLaDA-Image-Turbo-FP8Fast lower-memory serving4 steps, CFG 1Verified H200 SP1 and SP2 serving
+ +### Serving topology + +One `LLaDAImagePipeline` server handles both `/v1/images/generations` and +`/v1/images/edits`; you do not need a second editing deployment. The builder +pins each Hugging Face revision and emits the only two validated H200 resource +shapes. + + + + + + + + + + + + + + + + + + + + + +
ShapeTopologyRequest boundary
1× H200TP1, SP1, Ulysses1, Ring1The API accepts up to 10 outputs; this page validates n=1.
2× H200TP1, SP2, Ulysses2, Ring1Requires n=1; both ranks execute one request.
+ +The embedded **LLaDA 2.0 16B MoE text encoder** (`LLaDA2MoeModelLM`) stays +resident on every SP rank. Generic text encoder CPU offload, memory sleep, and +hot weight updates are unsupported. Use the generated resident recipes rather +than applying placement flags from another model's cookbook. + +## 4. Request examples + +The Request tab in the builder emits the matching cURL command for generation or +editing and automatically switches between the Base and Turbo sampling profiles. +The following Python examples show the same API contracts with Turbo. + + +```python Text-to-image +import base64 +from openai import OpenAI + +client = OpenAI(api_key="EMPTY", base_url="http://127.0.0.1:30000/v1") +response = client.images.generate( + model="inclusionAI/LLaDA-Image-Turbo", + prompt="A bilingual bookstore poster reading '星河书店' and 'Starlight Books'", + size="1024x1024", + n=1, + response_format="b64_json", + extra_body={ + "num_inference_steps": 4, + "guidance_scale": 1.0, + "seed": 42, + }, +) + +with open("llada-image.png", "wb") as output: + output.write(base64.b64decode(response.data[0].b64_json)) +``` + +```python Single-image editing +import base64 +from openai import OpenAI + +client = OpenAI(api_key="EMPTY", base_url="http://127.0.0.1:30000/v1") +with open("input.png", "rb") as source: + response = client.images.edit( + model="inclusionAI/LLaDA-Image-Turbo", + image=source, + prompt="Place the subject on a beach at golden hour.", + size="1024x1024", + n=1, + response_format="b64_json", + extra_body={ + "num_inference_steps": 4, + "guidance_scale": 1.0, + "seed": 42, + }, + ) + +with open("llada-image-edit.png", "wb") as output: + output.write(base64.b64decode(response.data[0].b64_json)) +``` + + +For Base or Base FP8, change the model ID and pass 50 steps with guidance 5.0. +SGLang's class defaults are four steps and guidance 1.0, so omitting these fields +on a Base request selects the wrong sampling profile. + +## 5. Feature details + +### Sequence parallelism + +LLaDA-Image supports SP1 and SP2 with Ulysses only. Output width and height must +be divisible by 16; at SP2, height must also be divisible by 32. A request may +cover at most 2048×2048 pixels and at most 10×1024² total output pixels. Editing +accepts exactly one uploaded image or image URL. + +### H200 performance + +The measurements below use BF16, 1024×1024 output, two warmups, and the median +client-wall latency of ten sequential requests. Turbo uses four steps with CFG 1; +Base uses 50 steps with CFG 5. + + + + + + + + + + + + + + + + + + + + + + + + + + + +
TaskTurbo SP1Turbo SP2Base SP1Base SP2
Text to image0.6216 s0.4629 s12.2036 s7.6757 s
Image editing1.7717 s1.2141 s32.8138 s19.9279 s
+ +### Quality validation + +The BF16 suite generated 6,704 images across GenEval, LongText, GEdit, and +CVTG-2K. Turbo SP1 and SP2 reproduced GenEval 0.8190, LongText English 0.8959, +LongText Chinese 0.9123, and CVTG word accuracy 0.8267. Matched text-to-image +outputs were byte-identical across SP1 and SP2; matched edits reached RGB Pearson +correlation 0.9919. Base BF16 produced GenEval 0.8404, LongText English 0.9120, +LongText Chinese 0.9198, and CVTG word accuracy 0.7792. + + +The published Base FP8 and Turbo FP8 checkpoints are supported on the same H200 +SP1 and SP2 topologies. Select the FP8 checkpoint directly; no server-side +quantization flag is required. + + +### Unsupported controls + +Cache-DiT is rejected at server startup and per request. The current SGLang +pipeline also rejects generic text encoder offload, memory release, and hot +weight updates. The native pipeline accepts `FA` and `TORCH_SDPA` attention +overrides, but this page keeps the platform default because alternative-backend +performance was not measured end to end. + +For implementation status and the full reproduction record, see +[SGLang PR #37907](https://github.com/sgl-project/sglang/pull/37907). diff --git a/docs/cookbook/diffusion/intro.mdx b/docs/cookbook/diffusion/intro.mdx index 37772941be68..f1a600e1258c 100644 --- a/docs/cookbook/diffusion/intro.mdx +++ b/docs/cookbook/diffusion/intro.mdx @@ -31,6 +31,12 @@ Image models generate one image request as a bounded denoising job, usually with href="/cookbook/diffusion/Qwen-Image/Qwen-Image" img="/cards/logos/qwen.png" /> + { +const CHECKPOINTS = { + "base-bf16": { + model: "inclusionAI/LLaDA-Image", + revision: "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", + steps: 50, + guidance: 5.0, + }, + "base-fp8": { + model: "inclusionAI/LLaDA-Image-FP8", + revision: "7678c6071b139e1989565db41122b16d767cc585", + steps: 50, + guidance: 5.0, + }, + "turbo-bf16": { + model: "inclusionAI/LLaDA-Image-Turbo", + revision: "033d068d0702ed81c45ec1bfbe61e318d843d948", + steps: 4, + guidance: 1.0, + }, + "turbo-fp8": { + model: "inclusionAI/LLaDA-Image-Turbo-FP8", + revision: "2b822499be3f33a6f83c988dafab7986ebd1f733", + steps: 4, + guidance: 1.0, + }, +}; + +return { + modelName: "LLaDA-Image", + supportedHardware: ["h200"], + groupHardware: false, + matchDims: [], + + overlayDims: [ + { + id: "weights", + title: "Checkpoint weights", + scope: "base", + description: "Base prioritizes fidelity; Turbo uses a distilled four-step schedule. FP8 changes the checkpoint, not a server-side quantization flag.", + default: "turbo-bf16", + options: [ + { id: "base-bf16", label: "Base BF16", subtitle: "50 steps · CFG 5" }, + { + id: "base-fp8", + label: "Base FP8", + subtitle: "50 steps · CFG 5", + }, + { id: "turbo-bf16", label: "Turbo BF16", subtitle: "4 steps · CFG 1", recommended: true }, + { + id: "turbo-fp8", + label: "Turbo FP8", + subtitle: "4 steps · CFG 1", + }, + ], + }, + { + id: "mode", + title: "Request mode", + scope: "base", + description: "One server handles both endpoints; editing requires exactly one source image.", + default: "generate", + options: [ + { id: "generate", label: "Text to image", recommended: true }, + { id: "edit", label: "Image editing" }, + ], + }, + { + id: "placement", + title: "Placement", + scope: "serve", + description: "The embedded LLaDA 2.0 16B MoE text encoder must remain resident.", + default: "resident", + options: [ + { id: "resident", label: "Resident", recommended: true }, + ], + }, + { + id: "attention", + title: "Attention", + scope: "serve", + description: "Use the platform-selected backend for the verified H200 recipes.", + default: "platform", + options: [ + { id: "platform", label: "Platform default", recommended: true }, + ], + }, + { + id: "precision", + title: "Precision", + scope: "serve", + description: "Precision is encoded by the selected BF16 or FP8 checkpoint.", + default: "native", + options: [ + { id: "native", label: "Checkpoint native", recommended: true }, + ], + }, + { + id: "encoder", + title: "Text encoder", + scope: "serve", + description: "LLaDA 2.0 16B MoE (LLaDA2MoeModelLM) runs as an embedded SRT worker on every SP rank and cannot use generic CPU offload.", + default: "auto", + options: [ + { id: "auto", label: "Resident per rank", recommended: true }, + ], + }, + { + id: "execution", + title: "Execution", + scope: "serve", + default: "eager", + options: [ + { id: "eager", label: "Eager", recommended: true }, + ], + }, + { + id: "outputs", + title: "Outputs", + scope: "request", + description: "The verified request shape uses one output. SP2 rejects n greater than one.", + kind: "number", + min: 1, + max: 1, + unit: "output per prompt", + default: 1, + options: [], + }, + ], + + commandBuilder: { + defaultSelection: { + hw: "h200", + nodes: 1, + gpus_per_node: 1, + topology_mode: "auto", + tp_size: 1, + ulysses_degree: 1, + ring_degree: 1, + }, + resource: { + limits: { + nodes: { min: 1, max: 1 }, + gpus_per_node: { min: 1, max: 2 }, + }, + verifiedRecipes: [ + { + id: "h200-sp1", + hw: "h200", + nodes: 1, + gpus_per_node: 1, + placement: "resident", + tp_size: 1, + ulysses_degree: 1, + ring_degree: 1, + encoder: "auto", + default: true, + }, + { + id: "h200-sp2", + hw: "h200", + nodes: 1, + gpus_per_node: 2, + placement: "resident", + tp_size: 1, + ulysses_degree: 2, + ring_degree: 1, + encoder: "auto", + }, + ], + autoTopology: (s) => ({ + tp_size: 1, + ulysses_degree: Number(s.gpus_per_node), + ring_degree: 1, + }), + validateTopology: (s, topology) => { + const errors = []; + const nodes = Number(s.nodes); + const gpus = Number(s.gpus_per_node); + if (nodes !== 1) errors.push("LLaDA-Image currently supports single-node deployment only."); + if (![1, 2].includes(gpus)) errors.push("LLaDA-Image supports one or two GPUs per server."); + if (Number(topology.tp_size) !== 1) errors.push("LLaDA-Image requires tensor parallel size 1."); + if (Number(topology.ulysses_degree) !== gpus) errors.push("Ulysses degree must equal the GPU count."); + if (Number(topology.ring_degree) !== 1) errors.push("LLaDA-Image requires ring degree 1."); + return errors; + }, + }, + resolveDeployment: (s) => { + const resource = config.commandBuilder.resource; + const topology = s.topology_mode === "manual" + ? { + tp_size: Number(s.tp_size), + ulysses_degree: Number(s.ulysses_degree), + ring_degree: Number(s.ring_degree), + } + : resource.autoTopology(s); + const errors = resource.validateTopology(s, topology); + const checkpoint = CHECKPOINTS[s.weights] || CHECKPOINTS["turbo-bf16"]; + const recipe = resource.verifiedRecipes.find((entry) => entry.hw === s.hw + && entry.nodes === Number(s.nodes) + && entry.gpus_per_node === Number(s.gpus_per_node) + && entry.placement === s.placement + && entry.tp_size === topology.tp_size + && entry.ulysses_degree === topology.ulysses_degree + && entry.ring_degree === topology.ring_degree); + const topologyVerified = !!recipe && errors.length === 0; + const serveVerified = topologyVerified; + const requestVerified = topologyVerified; + + const flags = [ + `--model-path ${checkpoint.model}`, + `--revision ${checkpoint.revision}`, + "--trust-remote-code", + ]; + if (Number(s.gpus_per_node) === 2) { + flags.push( + "--num-gpus 2", + "--sp-degree 2", + "--ulysses-degree 2", + "--ring-degree 1", + ); + } + flags.push("--host {{HOST_IP}}", "--port {{PORT}}"); + + return { + match: { hw: s.hw }, + nnodes: 1, + verified: serveVerified, + verificationStatus: serveVerified ? "verified" : "unverified", + flags, + builder: { + topology, + topologySummary: Number(s.gpus_per_node) === 2 + ? "TP 1 · Ulysses 2 · Ring 1 · SP2 · Single node" + : "TP 1 · Ulysses 1 · Ring 1 · SP1 · Single node", + errors, + warnings: [], + verification: { + serve: errors.length ? "error" : (serveVerified ? "verified" : "unverified"), + request: errors.length ? "error" : (requestVerified ? "verified" : "unverified"), + }, + resolvedSettings: { + placement: "Resident", + attention: "Platform default", + precision: s.weights.endsWith("fp8") ? "Published FP8 checkpoint" : "BF16", + encoder: "LLaDA 2.0 16B MoE, resident per SP rank", + execution: "Eager", + }, + }, + }; + }, + }, + + modelNames: { + default: "inclusionAI/LLaDA-Image-Turbo", + }, + dockerImages: { + h200: "lmsysorg/sglang:dev", + }, + placeholders: { + HOST_IP: { target: "command", label: "Bind host", default: "0.0.0.0" }, + PORT: { target: "command", label: "Bind port", default: "30000" }, + CURL_HOST: { target: "curl", label: "Server host", default: "localhost" }, + CURL_PORT: { target: "curl", label: "Server port", default: "30000" }, + INPUT_IMAGE: { target: "curl", label: "Source image", default: "input.png" }, + }, + curl: (s) => { + const checkpoint = CHECKPOINTS[s.weights] || CHECKPOINTS["turbo-bf16"]; + if (s.mode === "edit") { + return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/images/edits \\ + -F 'image=@{{INPUT_IMAGE}}' \\ + -F 'model=${checkpoint.model}' \\ + -F 'prompt=Place the subject on a beach at golden hour.' \\ + -F 'size=1024x1024' \\ + -F 'n=${Number(s.outputs)}' \\ + -F 'response_format=b64_json' \\ + -F 'num_inference_steps=${checkpoint.steps}' \\ + -F 'guidance_scale=${checkpoint.guidance}' \\ + -F 'seed=42'`; + } + const request = { + model: checkpoint.model, + prompt: "A cinematic photograph of a red fox standing in fresh snow at golden hour, detailed fur, shallow depth of field", + size: "1024x1024", + n: Number(s.outputs), + response_format: "b64_json", + num_inference_steps: checkpoint.steps, + guidance_scale: checkpoint.guidance, + seed: 42, + }; + return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/images/generations \\ + -H 'Content-Type: application/json' \\ + -d '${JSON.stringify(request, null, 2)}'`; + }, + cells: [], + showPlaygroundLink: false, +}; +})(); From b1e28944e0e12885f0f28450db0697b47c7aee2d Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Fri, 11 Sep 2026 21:15:41 +0000 Subject: [PATCH 09/18] Address LLaDA-Image activation metadata and checkpointing feedback Read ModelOpt activation metadata while preserving explicit output conversion overrides and legacy defaults for missing metadata. Remove training-only checkpointing branches and support flags from the LLaDA-Image serving models. --- docs/docs/sglang-diffusion/quantization.mdx | 3 +- .../runtime/models/dits/llada_image.py | 157 +++----------- .../test_build_modelopt_fp8_transformer.py | 194 +++++++++++++++++- .../test/unit/test_llada_image.py | 38 ++++ .../tools/build_modelopt_fp8_transformer.py | 58 +++++- 5 files changed, 308 insertions(+), 142 deletions(-) diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 4ada585a9be4..96d5c67862cf 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -702,8 +702,7 @@ sglang generate \ - On disk, the quantization config stays `quant_method=modelopt` with `quant_algo=FP8`; the `modelopt-fp8` label in this document is a support family name, not a serialized config key. -- To build the converted checkpoint yourself from a ModelOpt diffusers export, - use `python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer`. +- To build the converted checkpoint yourself from a ModelOpt diffusers export, use `python -m sglang.multimodal_gen.tools.build_modelopt_fp8_transformer`. The default `--activation-scheme auto` reads the export's input activation metadata. For a LLaDA-Image export that declares static activations, pass `--activation-scheme dynamic` explicitly. Exports without activation metadata keep the legacy defaults of dynamic for LLaDA-Image and static for other models. ## ModelOpt NVFP4 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index 43212856f3ab..7dfb2f0260d4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -548,7 +548,6 @@ class _LLaDAImageTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): Maximum RoPE positions for sequence, height, and width axes. """ - _supports_gradient_checkpointing = True _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ @@ -597,7 +596,6 @@ def __init__( self.all_patch_size = all_patch_size self.all_f_patch_size = all_f_patch_size self.t_scale = t_scale - self.gradient_checkpointing = False self.all_x_embedder = nn.ModuleDict() self.all_final_layer = nn.ModuleDict() @@ -1258,27 +1256,7 @@ def forward( ) for layer in self.noise_refiner: - if torch.is_grad_enabled() and self.gradient_checkpointing: - if is_editing: - image_features = self._gradient_checkpointing_func( - layer, - image_features, - image_attention_mask, - image_frequencies, - None, - image_noise_mask, - noisy_embedding, - clean_embedding, - ) - else: - image_features = self._gradient_checkpointing_func( - layer, - image_features, - image_attention_mask, - image_frequencies, - adaln_input, - ) - elif is_editing: + if is_editing: image_features = layer( image_features, image_attention_mask, @@ -1319,26 +1297,12 @@ def forward( ) for layer in self.context_refiner: - if torch.is_grad_enabled() and self.gradient_checkpointing: - cap_features = self._gradient_checkpointing_func( - layer, - cap_features, - cap_attention_mask, - cap_frequencies, - None, - None, - None, - None, - 0, - True, - ) - else: - cap_features = layer( - cap_features, - cap_attention_mask, - cap_frequencies, - skip_sequence_parallel_override=True, - ) + cap_features = layer( + cap_features, + cap_attention_mask, + cap_frequencies, + skip_sequence_parallel_override=True, + ) sigvq_lengths = [len(features) for features in sigvq_sequence.features] sigvq_features = self.sigvq_embedder( @@ -1366,26 +1330,12 @@ def forward( if any(sigvq_lengths): for layer in self.sigvq_refiner: - if torch.is_grad_enabled() and self.gradient_checkpointing: - sigvq_features = self._gradient_checkpointing_func( - layer, - sigvq_features, - sigvq_attention_mask, - sigvq_frequencies, - None, - None, - None, - None, - 0, - True, - ) - else: - sigvq_features = layer( - sigvq_features, - sigvq_attention_mask, - sigvq_frequencies, - skip_sequence_parallel_override=True, - ) + sigvq_features = layer( + sigvq_features, + sigvq_attention_mask, + sigvq_frequencies, + skip_sequence_parallel_override=True, + ) if get_sp_world_size() > 1: if batch_size != 1: @@ -1528,26 +1478,12 @@ def forward( ) for layer in self.context_refiner: - if torch.is_grad_enabled() and self.gradient_checkpointing: - condition_features = self._gradient_checkpointing_func( - layer, - condition_features, - condition_attention_mask, - condition_frequencies, - None, - None, - None, - None, - 0, - True, - ) - else: - condition_features = layer( - condition_features, - condition_attention_mask, - condition_frequencies, - skip_sequence_parallel_override=True, - ) + condition_features = layer( + condition_features, + condition_attention_mask, + condition_frequencies, + skip_sequence_parallel_override=True, + ) ( unified_features, @@ -1568,32 +1504,7 @@ def forward( num_replicated_suffix = condition_lengths[0] for layer in self.layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - if is_editing: - unified_features = self._gradient_checkpointing_func( - layer, - unified_features, - unified_attention_mask, - unified_frequencies, - None, - unified_noise_mask, - noisy_embedding, - clean_embedding, - num_replicated_suffix, - ) - else: - unified_features = self._gradient_checkpointing_func( - layer, - unified_features, - unified_attention_mask, - unified_frequencies, - adaln_input, - None, - None, - None, - num_replicated_suffix, - ) - elif is_editing: + if is_editing: unified_features = layer( unified_features, unified_attention_mask, @@ -1784,7 +1695,6 @@ class LLaDAImageQueryFormerModel(ModelMixin, ConfigMixin, AttentionMixin): Epsilon used by parameter-free layer normalization. """ - _supports_gradient_checkpointing = True _no_split_modules: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] _skip_layerwise_casting_patterns: ClassVar[list[str]] = ["norm"] @@ -1820,7 +1730,6 @@ def __init__( for _ in range(num_hidden_layers) ] ) - self.gradient_checkpointing = False def forward( self, @@ -1846,15 +1755,7 @@ def forward( attention_mask = attention_mask.bool() for query_block in self.query_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - query_embeds = self._gradient_checkpointing_func( - query_block, - query_embeds, - inputs_embeds, - attention_mask, - ) - else: - query_embeds = query_block(query_embeds, inputs_embeds, attention_mask) + query_embeds = query_block(query_embeds, inputs_embeds, attention_mask) if not return_dict: return (query_embeds,) @@ -1987,7 +1888,6 @@ class LLaDAImageTextProjectionModel(ModelMixin, ConfigMixin, AttentionMixin): Epsilon used by parameter-free RMS normalization. """ - _supports_gradient_checkpointing = True _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ @@ -2026,7 +1926,6 @@ def __init__( ] ) self.projector = nn.Linear(hidden_size, projection_dim, bias=True) - self.gradient_checkpointing = False def forward( self, @@ -2045,10 +1944,7 @@ def forward( Hidden states projected to the denoiser caption dimension. """ for layer in self.layers: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func(layer, hidden_states) - else: - hidden_states = layer(hidden_states) + hidden_states = layer(hidden_states) hidden_states = self.projector(hidden_states) if not return_dict: @@ -2258,7 +2154,6 @@ class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin): `patch_size`. """ - _supports_gradient_checkpointing = True _no_split_modules: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ @@ -2327,7 +2222,6 @@ def __init__( inner_dim=semantic_embed_dim, activation_fn="linear-silu", ) - self.gradient_checkpointing = False def forward( self, @@ -2373,12 +2267,7 @@ def forward( ) for block in self.visual.blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states = self._gradient_checkpointing_func( - block, hidden_states - ) - else: - hidden_states = block(hidden_states) + hidden_states = block(hidden_states) hidden_states = hidden_states.transpose(1, 2).reshape( pixel_values.shape[0], diff --git a/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py index c440856edcfe..74ac288884f3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py +++ b/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py @@ -4,6 +4,7 @@ import pytest import torch +from safetensors import safe_open from safetensors.torch import load_file, save_file from sglang.multimodal_gen.tools.build_modelopt_fp8_transformer import ( @@ -12,15 +13,23 @@ ) -def _write_llada_modelopt_export(tmp_path, source_tensors, backbone_state): +def _write_llada_modelopt_export( + tmp_path, + source_tensors, + backbone_state, + *, + quantization_config=None, + class_name="LLaDAImageTransformer2DModel", +): source_dir = tmp_path / "source" output_dir = tmp_path / "output" source_dir.mkdir() source_config = { - "_class_name": "LLaDAImageTransformer2DModel", + "_class_name": class_name, "quantization_config": { "quant_method": "modelopt", "quant_algo": "FP8", + **(quantization_config or {}), }, } (source_dir / "config.json").write_text(json.dumps(source_config)) @@ -167,3 +176,184 @@ def test_llada_conversion_rejects_static_activation_quantization(tmp_path): model_type="llada-image", activation_scheme="static", ) + + +@pytest.mark.parametrize( + ("metadata", "requested", "expected"), + [ + ({}, "auto", "static"), + ({"activation_scheme": "dynamic"}, "auto", "dynamic"), + ({"activation_scheme": "static"}, "auto", "static"), + ( + {"config_groups": {"group_0": {"input_activations": {"dynamic": True}}}}, + "auto", + "dynamic", + ), + ( + {"config_groups": {"group_0": {"input_activations": {"dynamic": False}}}}, + "auto", + "static", + ), + ( + {"config_groups": {"group_0": {"weights": {"dynamic": True}}}}, + "auto", + "static", + ), + ( + { + "config_groups": { + "group_0": {"input_activations": {"dynamic": False}}, + "group_1": {"input_activations": {"dynamic": True}}, + } + }, + "dynamic", + "dynamic", + ), + ( + { + "config_groups": { + "group_0": {"input_activations": {"dynamic": False}}, + "group_1": {"input_activations": {"dynamic": True}}, + } + }, + "static", + "static", + ), + ({"activation_scheme": "static"}, "dynamic", "dynamic"), + ({"activation_scheme": "dynamic"}, "static", "static"), + ( + {"config_groups": {"group_0": {"input_activations": {"dynamic": True}}}}, + "static", + "static", + ), + ], +) +def test_conversion_resolves_activation_metadata( + tmp_path, metadata, requested, expected +): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, + { + "layers.0.attention.to_q.weight_quantizer._amax": torch.tensor(224.0), + "layers.0.attention.to_q.input_quantizer._amax": torch.tensor(112.0), + }, + quantization_config=metadata, + class_name="ExampleTransformer2DModel", + ) + source_config = (source_dir / "config.json").read_bytes() + build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="none", + activation_scheme=requested, + ) + quant_config = json.loads((output_dir / "config.json").read_text())[ + "quantization_config" + ] + assert quant_config.get("activation_scheme", "static") == expected + for group in quant_config.get("config_groups", {}).values(): + if group.get("input_activations") is not None: + assert group["input_activations"]["dynamic"] is (expected == "dynamic") + tensors = load_file(output_dir / "model.safetensors") + assert ("layers.0.attention.to_q.input_scale" in tensors) is (expected == "static") + with safe_open(output_dir / "model.safetensors", framework="pt") as checkpoint: + assert json.loads(checkpoint.metadata()["quantization_config"]) == quant_config + assert (source_dir / "config.json").read_bytes() == source_config + + +@pytest.mark.parametrize("requested", ["auto", "dynamic"]) +def test_llada_static_source_requires_explicit_dynamic_output(tmp_path, requested): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, + {"layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0)}, + quantization_config={ + "config_groups": {"group_0": {"input_activations": {"dynamic": False}}} + }, + ) + kwargs = dict( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="llada-image", + activation_scheme=requested, + ) + if requested == "auto": + with pytest.raises(ValueError, match="set activation_scheme='dynamic'"): + build_modelopt_fp8_transformer(**kwargs) + assert not output_dir.exists() + else: + build_modelopt_fp8_transformer(**kwargs) + quant_config = json.loads((output_dir / "config.json").read_text())[ + "quantization_config" + ] + assert quant_config["activation_scheme"] == "dynamic" + assert "layers.0.attention.to_q.input_scale" not in load_file( + output_dir / "model.safetensors" + ) + + +@pytest.mark.parametrize( + "metadata", + [ + {"activation_scheme": "invalid"}, + {"config_groups": {"group_0": {"input_activations": {"dynamic": "false"}}}}, + { + "config_groups": { + "group_0": {"input_activations": {"dynamic": False}}, + "group_1": {"input_activations": {"dynamic": True}}, + } + }, + { + "activation_scheme": "dynamic", + "config_groups": {"group_0": {"input_activations": {"dynamic": False}}}, + }, + ], +) +def test_auto_rejects_invalid_or_conflicting_activation_metadata(tmp_path, metadata): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, + {"layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0)}, + quantization_config=metadata, + ) + with pytest.raises(ValueError, match="activation"): + build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="llada-image", + ) + assert not output_dir.exists() + + +@pytest.mark.parametrize("requested", ["static", "dynamic"]) +@pytest.mark.parametrize( + "metadata", + [ + {"config_groups": None}, + {"config_groups": {"group_0": None}}, + {"config_groups": {"group_0": {"input_activations": "invalid"}}}, + ], +) +def test_explicit_output_rejects_malformed_metadata_before_writing( + tmp_path, requested, metadata +): + source_dir, output_dir, backbone_path = _write_llada_modelopt_export( + tmp_path, + {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, + {"layers.0.attention.to_q.weight_quantizer._amax": torch.tensor(224.0)}, + quantization_config=metadata, + class_name="ExampleTransformer2DModel", + ) + with pytest.raises(ValueError, match="ModelOpt.*dictionary"): + build_modelopt_fp8_transformer( + modelopt_hf_dir=str(source_dir), + modelopt_backbone_ckpt=str(backbone_path), + output_dir=str(output_dir), + model_type="none", + activation_scheme=requested, + ) + assert not output_dir.exists() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image.py b/python/sglang/multimodal_gen/test/unit/test_llada_image.py index 9b443271c8b9..5ac2aae5dafe 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image.py @@ -14,6 +14,9 @@ from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.llada_image import ( + LLaDAImageQueryFormerModel, + LLaDAImageSigVQModel, + LLaDAImageTextProjectionModel, LLaDAImageTransformerBlock, _LLaDAImageTransformer2DModel, ) @@ -62,6 +65,41 @@ def setUpClass(cls): def tearDownClass(cls): set_global_server_args(cls.previous_server_args) + def test_serving_models_reject_gradient_checkpointing(self): + small_config = dict( + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + ) + models = [ + _LLaDAImageTransformer2DModel( + in_channels=4, + dim=128, + n_layers=0, + n_refiner_layers=0, + n_heads=1, + cap_feat_dim=8, + semantic_feat_dim=8, + ), + LLaDAImageQueryFormerModel(num_queries=2, **small_config), + LLaDAImageTextProjectionModel(projection_dim=4, **small_config), + LLaDAImageSigVQModel( + image_size=4, + patch_size=2, + codebook_size=8, + codebook_embed_dim=4, + semantic_embed_dim=4, + **small_config, + ), + ] + for model in models: + with self.subTest(model=type(model).__name__): + with self.assertRaisesRegex( + ValueError, "does not support gradient checkpointing" + ): + model.enable_gradient_checkpointing() + def test_dit_supports_sglang_flash_attention_and_sdpa(self): backends = LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends self.assertEqual( diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py index 9c6185b0ece9..c7a1f4b6b9ed 100644 --- a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py +++ b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py @@ -607,6 +607,44 @@ def _load_selected_tensors( return tensors +def _get_modelopt_activation_scheme( + quant_config: Mapping, *, allow_mixed: bool = False +) -> str | None: + schemes = set() + declared = quant_config.get("activation_scheme") + if declared is not None: + if declared not in ("static", "dynamic"): + raise ValueError(f"Invalid ModelOpt activation_scheme {declared!r}.") + schemes.add(declared) + groups = quant_config.get("config_groups", {}) + if not isinstance(groups, dict): + raise ValueError("ModelOpt activation config_groups must be a dictionary.") + for group in groups.values(): + if not isinstance(group, dict): + raise ValueError( + "Each ModelOpt activation config group must be a dictionary." + ) + activations = group.get("input_activations") + if activations is None: + continue + if not isinstance(activations, dict): + raise ValueError("ModelOpt input_activations must be a dictionary.") + if "dynamic" not in activations: + continue + dynamic = activations["dynamic"] + if not isinstance(dynamic, bool): + raise ValueError("ModelOpt input activation dynamic must be a boolean.") + schemes.add("dynamic" if dynamic else "static") + if len(schemes) > 1: + if allow_mixed: + return None + raise ValueError( + "Conflicting ModelOpt activation schemes. " + "Set activation_scheme explicitly to choose the converted output." + ) + return next(iter(schemes), None) + + def build_modelopt_fp8_transformer( *, modelopt_hf_dir: str, @@ -648,8 +686,13 @@ def build_modelopt_fp8_transformer( is_llada_image = ( model_type == "llada-image" or class_name == "LLaDAImageTransformer2DModel" ) + source_activation_scheme = _get_modelopt_activation_scheme( + quant_config, allow_mixed=activation_scheme != "auto" + ) if activation_scheme == "auto": - activation_scheme = "dynamic" if is_llada_image else "static" + activation_scheme = source_activation_scheme or ( + "dynamic" if is_llada_image else "static" + ) if activation_scheme not in {"static", "dynamic"}: raise ValueError( "activation_scheme must be one of 'auto', 'static', or 'dynamic', " @@ -658,7 +701,7 @@ def build_modelopt_fp8_transformer( if is_llada_image and activation_scheme != "dynamic": raise ValueError( "LLaDA-Image FP8 conversion currently supports only dynamic activation " - "quantization; set activation_scheme='dynamic' or leave it as 'auto'." + "quantization. To convert this checkpoint, set activation_scheme='dynamic'." ) dynamic_activation = activation_scheme == "dynamic" @@ -779,6 +822,12 @@ def build_modelopt_fp8_transformer( } else: effective_quant_config["ignore"] = ignore_patterns + if "activation_scheme" in effective_quant_config: + effective_quant_config["activation_scheme"] = "static" + for group in effective_quant_config.get("config_groups", {}).values(): + activations = group.get("input_activations") + if activations is not None: + activations["dynamic"] = False serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True) output_config = _build_output_config( source_config=config, @@ -975,8 +1024,9 @@ def _parse_args() -> argparse.Namespace: default="auto", help=( "Activation quantization scheme in the converted checkpoint. 'auto' " - "uses dynamic activation quantization for LLaDA-Image and preserves " - "the legacy static ModelOpt format for other model families." + "reads ModelOpt input activation metadata, falling back to dynamic " + "for LLaDA-Image and static for other families when metadata is absent. " + "Use 'dynamic' explicitly to convert a static LLaDA-Image export." ), ) parser.add_argument( From fdc38683e7fc67c1554c03085a67825ad210b092 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Tue, 15 Sep 2026 05:30:20 +0000 Subject: [PATCH 10/18] [diffusion] Limit LLaDA-Image support to core serving --- .../diffusion/InclusionAI/LLaDA-Image.mdx | 247 ---------- docs/cookbook/diffusion/intro.mdx | 6 - docs/docs.json | 7 - docs/docs/sglang-diffusion/quantization.mdx | 3 +- .../configs/inclusionAI/llada-image.jsx | 300 ------------ docs/src/snippets/diffusion/model-catalog.jsx | 5 - .../configs/pipeline_configs/base.py | 29 -- .../configs/pipeline_configs/llada_image.py | 257 ++--------- .../configs/sample/llada_image.py | 11 + .../configs/sample/sampling_params.py | 3 - python/sglang/multimodal_gen/registry.py | 1 - .../runtime/entrypoints/openai/image_api.py | 66 +-- .../runtime/loader/transformer_load_utils.py | 11 +- .../memory_occupation_controller.py | 8 - .../runtime/models/dits/llada_image.py | 86 +--- .../scheduling_flow_match_euler_discrete.py | 70 +-- .../runtime/pipelines/llada_image.py | 28 -- .../llada_image/conditioning.py | 51 +- .../runtime/post_training/weights_updater.py | 18 - .../runtime/server_args/server_args.py | 32 -- .../test/runner/diffusion_suite_runner.py | 16 +- .../multimodal_gen/test/server/gpu_cases.py | 103 ++--- .../test/server/test_server_common.py | 7 +- .../test/server/testcase_configs.py | 10 - .../test_build_modelopt_fp8_transformer.py | 359 --------------- .../test/unit/test_layerwise_offload.py | 23 +- .../test/unit/test_llada_image.py | 435 +----------------- .../test/unit/test_llada_image_api.py | 168 ------- .../unit/test_llada_image_conditioning.py | 257 +---------- .../unit/test_llada_image_pipeline_config.py | 264 ----------- .../test/unit/test_llada_image_validation.py | 234 ---------- ...est_modelopt_fp8_layerwise_offload_load.py | 28 +- .../unit/test_sequential_server_requests.py | 54 --- .../test/unit/test_suite_partitioning.py | 33 +- .../test/unit/test_transformer_quant.py | 253 +++------- .../tools/build_modelopt_fp8_transformer.py | 212 +-------- python/sglang/srt/arg_groups/parallel_hook.py | 18 +- python/sglang/srt/configs/model_config.py | 67 +-- python/sglang/srt/layers/communicator.py | 30 +- python/sglang/srt/layers/quantization/fp8.py | 8 - .../srt/model_loader/llada2_weight_utils.py | 15 +- .../sglang/srt/model_loader/weight_utils.py | 6 - python/sglang/srt/models/llada2.py | 81 +--- .../utils/diffusion/diffusion_case_parser.py | 32 +- .../test_llada2_cfg_flashinfer_backend.py | 163 ------- .../unit/ci/test_diffusion_ci_registration.py | 103 ----- .../test_attention_backend_registry.py | 27 -- .../test_layer_communicator_fusion_gate.py | 117 +---- .../registered/unit/models/test_llada2_fp8.py | 210 --------- .../unit/models/test_llada2_full_replica.py | 125 ----- .../unit/models/test_llada2_weight_utils.py | 84 +--- .../unit/server_args/test_server_args.py | 48 -- 52 files changed, 263 insertions(+), 4566 deletions(-) delete mode 100644 docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx delete mode 100644 docs/src/snippets/configs/inclusionAI/llada-image.jsx delete mode 100644 python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_api.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_validation.py delete mode 100644 test/registered/e2e/attention/test_llada2_cfg_flashinfer_backend.py delete mode 100644 test/registered/unit/ci/test_diffusion_ci_registration.py delete mode 100644 test/registered/unit/layers/attention/test_attention_backend_registry.py delete mode 100644 test/registered/unit/models/test_llada2_fp8.py delete mode 100644 test/registered/unit/models/test_llada2_full_replica.py diff --git a/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx b/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx deleted file mode 100644 index 358e10062181..000000000000 --- a/docs/cookbook/diffusion/InclusionAI/LLaDA-Image.mdx +++ /dev/null @@ -1,247 +0,0 @@ ---- -title: LLaDA-Image -description: "Deploy LLaDA-Image Base, Turbo, and FP8 checkpoints on NVIDIA H200 for text-to-image generation and single-image editing with SGLang Diffusion." ---- - -import { DiffusionModelTags } from '/src/snippets/diffusion/model-tags.jsx'; -import { Deployment } from '/src/snippets/_deployment.jsx'; -import { config } from '/src/snippets/configs/inclusionAI/llada-image.jsx'; - - - -## 1. Quick start - -Install with `uv pip install "sglang[diffusion]" --prerelease=allow`, then choose -the checkpoint, request mode, and verified H200 topology below. - - - - -LLaDA-Image loads reviewed configuration code from the checkpoint. Keep -`--trust-remote-code` and the generated commit-pinned `--revision` together, -and review remote-code changes before moving to another revision. - - -## 2. Model capabilities - -[LLaDA-Image](https://huggingface.co/inclusionAI/LLaDA-Image) is a 6B-parameter -image generation and editing family. The Base checkpoint targets high-fidelity -text-to-image generation and instruction-guided editing with a 50-step schedule; -the distilled Turbo checkpoint runs the same two SGLang request modes in four -steps. Both variants are designed for English and Chinese text rendering. - -Prompt conditioning is handled by an embedded **LLaDA 2.0 16B-parameter MoE -text encoder**. Its checkpoint config declares the `LLaDA2MoeModelLM` -architecture; this is the LLaDA 2.0 structure, not the earlier LLaDA text -encoder architecture. - -Choose Base when fidelity matters more than latency and Turbo for interactive or -high-throughput serving. SGLang exposes text-to-image generation and exactly one -reference image for editing through one server. The model card's VQ-conditioned -generation mode is not currently exposed by the SGLang Images API. - -The checkpoints are released under Apache-2.0. Review the model card and its -usage terms before production deployment. - -## 3. Deployment details - -### Checkpoint selection - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
CheckpointBest fitSamplingValidation contract
inclusionAI/LLaDA-ImageMaximum fidelity50 steps, CFG 5Full BF16 quality suite
inclusionAI/LLaDA-Image-FP8Lower checkpoint memory50 steps, CFG 5Verified H200 SP1 and SP2 serving
inclusionAI/LLaDA-Image-TurboLowest audited latency4 steps, CFG 1Full BF16 SP1 and SP2 suite
inclusionAI/LLaDA-Image-Turbo-FP8Fast lower-memory serving4 steps, CFG 1Verified H200 SP1 and SP2 serving
- -### Serving topology - -One `LLaDAImagePipeline` server handles both `/v1/images/generations` and -`/v1/images/edits`; you do not need a second editing deployment. The builder -pins each Hugging Face revision and emits the only two validated H200 resource -shapes. - - - - - - - - - - - - - - - - - - - - - -
ShapeTopologyRequest boundary
1× H200TP1, SP1, Ulysses1, Ring1The API accepts up to 10 outputs; this page validates n=1.
2× H200TP1, SP2, Ulysses2, Ring1Requires n=1; both ranks execute one request.
- -The embedded **LLaDA 2.0 16B MoE text encoder** (`LLaDA2MoeModelLM`) stays -resident on every SP rank. Generic text encoder CPU offload, memory sleep, and -hot weight updates are unsupported. Use the generated resident recipes rather -than applying placement flags from another model's cookbook. - -## 4. Request examples - -The Request tab in the builder emits the matching cURL command for generation or -editing and automatically switches between the Base and Turbo sampling profiles. -The following Python examples show the same API contracts with Turbo. - - -```python Text-to-image -import base64 -from openai import OpenAI - -client = OpenAI(api_key="EMPTY", base_url="http://127.0.0.1:30000/v1") -response = client.images.generate( - model="inclusionAI/LLaDA-Image-Turbo", - prompt="A bilingual bookstore poster reading '星河书店' and 'Starlight Books'", - size="1024x1024", - n=1, - response_format="b64_json", - extra_body={ - "num_inference_steps": 4, - "guidance_scale": 1.0, - "seed": 42, - }, -) - -with open("llada-image.png", "wb") as output: - output.write(base64.b64decode(response.data[0].b64_json)) -``` - -```python Single-image editing -import base64 -from openai import OpenAI - -client = OpenAI(api_key="EMPTY", base_url="http://127.0.0.1:30000/v1") -with open("input.png", "rb") as source: - response = client.images.edit( - model="inclusionAI/LLaDA-Image-Turbo", - image=source, - prompt="Place the subject on a beach at golden hour.", - size="1024x1024", - n=1, - response_format="b64_json", - extra_body={ - "num_inference_steps": 4, - "guidance_scale": 1.0, - "seed": 42, - }, - ) - -with open("llada-image-edit.png", "wb") as output: - output.write(base64.b64decode(response.data[0].b64_json)) -``` - - -For Base or Base FP8, change the model ID and pass 50 steps with guidance 5.0. -SGLang's class defaults are four steps and guidance 1.0, so omitting these fields -on a Base request selects the wrong sampling profile. - -## 5. Feature details - -### Sequence parallelism - -LLaDA-Image supports SP1 and SP2 with Ulysses only. Output width and height must -be divisible by 16; at SP2, height must also be divisible by 32. A request may -cover at most 2048×2048 pixels and at most 10×1024² total output pixels. Editing -accepts exactly one uploaded image or image URL. - -### Reported L20X performance - -The [upstream run guide](https://github.com/Anmuliar/llada-image-run-guide/blob/bd85618dba129f8bd00e6c654558cb024ce9f54c/README.md#L663-L679) and [model support PR](https://github.com/sgl-project/sglang/pull/37907) report the following BF16 timings on NVIDIA L20X for revision `8ecbcc6`. The reported protocol uses 1024×1024 output, two warmups and the median client-wall latency of ten sequential requests. Turbo uses four steps with CFG 1, and Base uses 50 steps with CFG 5. - -These are historical reported results and have not been independently rerun on L20X. The deployment recipes above are validated on H200. - - - - - - - - - - - - - - - - - - - - - - - - - - - -
TaskTurbo SP1Turbo SP2Base SP1Base SP2
Text to image0.6216 s0.4629 s12.2036 s7.6757 s
Image editing1.7717 s1.2141 s32.8138 s19.9279 s
- -### Quality validation - -Historical BF16 runs generated 6,704 images each for Base SP1, Turbo SP1 and Turbo SP2 across GenEval, LongText, GEdit and CVTG-2K. Turbo SP1 and SP2 reproduced GenEval 0.8190, LongText English 0.8959, LongText Chinese 0.9123 and CVTG word accuracy 0.8267. Matched Turbo text-to-image outputs were byte-identical across SP1 and SP2, and matched Turbo edits reached RGB Pearson correlation 0.9919. Base BF16 produced GenEval 0.8404, LongText English 0.9120, LongText Chinese 0.9198 and CVTG word accuracy 0.7792. These quality scores are historical and have not been rerun for the current revision. - - -The published Base FP8 and Turbo FP8 checkpoints are supported on the same H200 -SP1 and SP2 topologies. Select the FP8 checkpoint directly; no server-side -quantization flag is required. - - -### Unsupported controls - -Cache-DiT is rejected at server startup and per request. The current SGLang -pipeline also rejects generic text encoder offload, memory release, and hot -weight updates. The native pipeline accepts `FA` and `TORCH_SDPA` attention -overrides, but this page keeps the platform default because alternative-backend -performance was not measured end to end. - -For implementation status and the full reproduction record, see -[SGLang PR #37907](https://github.com/sgl-project/sglang/pull/37907). diff --git a/docs/cookbook/diffusion/intro.mdx b/docs/cookbook/diffusion/intro.mdx index 68c41f82f7d0..e5ab876123ab 100644 --- a/docs/cookbook/diffusion/intro.mdx +++ b/docs/cookbook/diffusion/intro.mdx @@ -37,12 +37,6 @@ Image models generate one image request as a bounded denoising job, usually with href="/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT" img="/cards/logos/internvl.png" /> - { -const CHECKPOINTS = { - "base-bf16": { - model: "inclusionAI/LLaDA-Image", - revision: "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", - steps: 50, - guidance: 5.0, - }, - "base-fp8": { - model: "inclusionAI/LLaDA-Image-FP8", - revision: "7678c6071b139e1989565db41122b16d767cc585", - steps: 50, - guidance: 5.0, - }, - "turbo-bf16": { - model: "inclusionAI/LLaDA-Image-Turbo", - revision: "033d068d0702ed81c45ec1bfbe61e318d843d948", - steps: 4, - guidance: 1.0, - }, - "turbo-fp8": { - model: "inclusionAI/LLaDA-Image-Turbo-FP8", - revision: "2b822499be3f33a6f83c988dafab7986ebd1f733", - steps: 4, - guidance: 1.0, - }, -}; - -return { - modelName: "LLaDA-Image", - supportedHardware: ["h200"], - groupHardware: false, - matchDims: [], - - overlayDims: [ - { - id: "weights", - title: "Checkpoint weights", - scope: "base", - description: "Base prioritizes fidelity; Turbo uses a distilled four-step schedule. FP8 changes the checkpoint, not a server-side quantization flag.", - default: "turbo-bf16", - options: [ - { id: "base-bf16", label: "Base BF16", subtitle: "50 steps · CFG 5" }, - { - id: "base-fp8", - label: "Base FP8", - subtitle: "50 steps · CFG 5", - }, - { id: "turbo-bf16", label: "Turbo BF16", subtitle: "4 steps · CFG 1", recommended: true }, - { - id: "turbo-fp8", - label: "Turbo FP8", - subtitle: "4 steps · CFG 1", - }, - ], - }, - { - id: "mode", - title: "Request mode", - scope: "base", - description: "One server handles both endpoints; editing requires exactly one source image.", - default: "generate", - options: [ - { id: "generate", label: "Text to image", recommended: true }, - { id: "edit", label: "Image editing" }, - ], - }, - { - id: "placement", - title: "Placement", - scope: "serve", - description: "The embedded LLaDA 2.0 16B MoE text encoder must remain resident.", - default: "resident", - options: [ - { id: "resident", label: "Resident", recommended: true }, - ], - }, - { - id: "attention", - title: "Attention", - scope: "serve", - description: "Use the platform-selected backend for the verified H200 recipes.", - default: "platform", - options: [ - { id: "platform", label: "Platform default", recommended: true }, - ], - }, - { - id: "precision", - title: "Precision", - scope: "serve", - description: "Precision is encoded by the selected BF16 or FP8 checkpoint.", - default: "native", - options: [ - { id: "native", label: "Checkpoint native", recommended: true }, - ], - }, - { - id: "encoder", - title: "Text encoder", - scope: "serve", - description: "LLaDA 2.0 16B MoE (LLaDA2MoeModelLM) runs as an embedded SRT worker on every SP rank and cannot use generic CPU offload.", - default: "auto", - options: [ - { id: "auto", label: "Resident per rank", recommended: true }, - ], - }, - { - id: "execution", - title: "Execution", - scope: "serve", - default: "eager", - options: [ - { id: "eager", label: "Eager", recommended: true }, - ], - }, - { - id: "outputs", - title: "Outputs", - scope: "request", - description: "The verified request shape uses one output. SP2 rejects n greater than one.", - kind: "number", - min: 1, - max: 1, - unit: "output per prompt", - default: 1, - options: [], - }, - ], - - commandBuilder: { - defaultSelection: { - hw: "h200", - nodes: 1, - gpus_per_node: 1, - topology_mode: "auto", - tp_size: 1, - ulysses_degree: 1, - ring_degree: 1, - }, - resource: { - limits: { - nodes: { min: 1, max: 1 }, - gpus_per_node: { min: 1, max: 2 }, - }, - verifiedRecipes: [ - { - id: "h200-sp1", - hw: "h200", - nodes: 1, - gpus_per_node: 1, - placement: "resident", - tp_size: 1, - ulysses_degree: 1, - ring_degree: 1, - encoder: "auto", - default: true, - }, - { - id: "h200-sp2", - hw: "h200", - nodes: 1, - gpus_per_node: 2, - placement: "resident", - tp_size: 1, - ulysses_degree: 2, - ring_degree: 1, - encoder: "auto", - }, - ], - autoTopology: (s) => ({ - tp_size: 1, - ulysses_degree: Number(s.gpus_per_node), - ring_degree: 1, - }), - validateTopology: (s, topology) => { - const errors = []; - const nodes = Number(s.nodes); - const gpus = Number(s.gpus_per_node); - if (nodes !== 1) errors.push("LLaDA-Image currently supports single-node deployment only."); - if (![1, 2].includes(gpus)) errors.push("LLaDA-Image supports one or two GPUs per server."); - if (Number(topology.tp_size) !== 1) errors.push("LLaDA-Image requires tensor parallel size 1."); - if (Number(topology.ulysses_degree) !== gpus) errors.push("Ulysses degree must equal the GPU count."); - if (Number(topology.ring_degree) !== 1) errors.push("LLaDA-Image requires ring degree 1."); - return errors; - }, - }, - resolveDeployment: (s) => { - const resource = config.commandBuilder.resource; - const topology = s.topology_mode === "manual" - ? { - tp_size: Number(s.tp_size), - ulysses_degree: Number(s.ulysses_degree), - ring_degree: Number(s.ring_degree), - } - : resource.autoTopology(s); - const errors = resource.validateTopology(s, topology); - const checkpoint = CHECKPOINTS[s.weights] || CHECKPOINTS["turbo-bf16"]; - const recipe = resource.verifiedRecipes.find((entry) => entry.hw === s.hw - && entry.nodes === Number(s.nodes) - && entry.gpus_per_node === Number(s.gpus_per_node) - && entry.placement === s.placement - && entry.tp_size === topology.tp_size - && entry.ulysses_degree === topology.ulysses_degree - && entry.ring_degree === topology.ring_degree); - const topologyVerified = !!recipe && errors.length === 0; - const serveVerified = topologyVerified; - const requestVerified = topologyVerified; - - const flags = [ - `--model-path ${checkpoint.model}`, - `--revision ${checkpoint.revision}`, - "--trust-remote-code", - ]; - if (Number(s.gpus_per_node) === 2) { - flags.push( - "--num-gpus 2", - "--sp-degree 2", - "--ulysses-degree 2", - "--ring-degree 1", - ); - } - flags.push("--host {{HOST_IP}}", "--port {{PORT}}"); - - return { - match: { hw: s.hw }, - nnodes: 1, - verified: serveVerified, - verificationStatus: serveVerified ? "verified" : "unverified", - flags, - builder: { - topology, - topologySummary: Number(s.gpus_per_node) === 2 - ? "TP 1 · Ulysses 2 · Ring 1 · SP2 · Single node" - : "TP 1 · Ulysses 1 · Ring 1 · SP1 · Single node", - errors, - warnings: [], - verification: { - serve: errors.length ? "error" : (serveVerified ? "verified" : "unverified"), - request: errors.length ? "error" : (requestVerified ? "verified" : "unverified"), - }, - resolvedSettings: { - placement: "Resident", - attention: "Platform default", - precision: s.weights.endsWith("fp8") ? "Published FP8 checkpoint" : "BF16", - encoder: "LLaDA 2.0 16B MoE, resident per SP rank", - execution: "Eager", - }, - }, - }; - }, - }, - - modelNames: { - default: "inclusionAI/LLaDA-Image-Turbo", - }, - dockerImages: { - h200: "lmsysorg/sglang:dev", - }, - placeholders: { - HOST_IP: { target: "command", label: "Bind host", default: "0.0.0.0" }, - PORT: { target: "command", label: "Bind port", default: "30000" }, - CURL_HOST: { target: "curl", label: "Server host", default: "localhost" }, - CURL_PORT: { target: "curl", label: "Server port", default: "30000" }, - INPUT_IMAGE: { target: "curl", label: "Source image", default: "input.png" }, - }, - curl: (s) => { - const checkpoint = CHECKPOINTS[s.weights] || CHECKPOINTS["turbo-bf16"]; - if (s.mode === "edit") { - return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/images/edits \\ - -F 'image=@{{INPUT_IMAGE}}' \\ - -F 'model=${checkpoint.model}' \\ - -F 'prompt=Place the subject on a beach at golden hour.' \\ - -F 'size=1024x1024' \\ - -F 'n=${Number(s.outputs)}' \\ - -F 'response_format=b64_json' \\ - -F 'num_inference_steps=${checkpoint.steps}' \\ - -F 'guidance_scale=${checkpoint.guidance}' \\ - -F 'seed=42'`; - } - const request = { - model: checkpoint.model, - prompt: "A cinematic photograph of a red fox standing in fresh snow at golden hour, detailed fur, shallow depth of field", - size: "1024x1024", - n: Number(s.outputs), - response_format: "b64_json", - num_inference_steps: checkpoint.steps, - guidance_scale: checkpoint.guidance, - seed: 42, - }; - return `curl -sS -X POST http://{{CURL_HOST}}:{{CURL_PORT}}/v1/images/generations \\ - -H 'Content-Type: application/json' \\ - -d '${JSON.stringify(request, null, 2)}'`; - }, - cells: [], - showPlaygroundLink: false, -}; -})(); diff --git a/docs/src/snippets/diffusion/model-catalog.jsx b/docs/src/snippets/diffusion/model-catalog.jsx index e724cd2d154c..5779f6cfdd94 100644 --- a/docs/src/snippets/diffusion/model-catalog.jsx +++ b/docs/src/snippets/diffusion/model-catalog.jsx @@ -45,11 +45,6 @@ export const DiffusionModelCatalog = ({ category }) => { modelIds: ["krea/Krea-2"], cookbook: "/cookbook/diffusion/Krea/Krea-2", }, - { - name: "LLaDA-Image", - modelIds: ["inclusionAI/LLaDA-Image"], - cookbook: "/cookbook/diffusion/InclusionAI/LLaDA-Image", - }, { name: "LongCat-Image", modelIds: [ diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index e2209a16e132..4c51a5397281 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -443,35 +443,6 @@ def supports_sequential_multi_output_inference(self): """Return whether one request's outputs run through DiT/VAE sequentially.""" return False - def validate_num_outputs_per_prompt( - self, num_outputs_per_prompt: int, server_args - ) -> None: - """Validate a request's per-prompt output count for this pipeline.""" - del num_outputs_per_prompt, server_args - - def supports_memory_release(self) -> bool: - """Return whether release_memory_occupation can offload this pipeline.""" - return True - - def supports_hot_weight_updates(self) -> bool: - """Return whether runtime weight updates cover this pipeline fully.""" - return True - - def validate_request_sampling_params(self, sampling_params, server_args) -> None: - """Validate a request's sampling params for this pipeline. - - May also run before full request materialization with a lightweight - object that only carries width, height, max_sequence_length, - num_outputs_per_prompt, diffusers_kwargs, enable_cache_dit, - cache_dit_params, attention_backend_override, and cfg_gate_step, - so overrides must not require other sampling fields. - """ - del sampling_params, server_args - - def validate_edit_source_count(self, source_count: int, server_args) -> None: - """Validate the number of source images before they are materialized.""" - del source_count, server_args - def estimate_request_cost(self, batch) -> float: """Return the relative cost used for batching admission caps. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py index 3299936a60e1..5e72efa2d8ec 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -12,10 +12,6 @@ ModelTaskType, SpatialImagePipelineConfig, ) -from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) @dataclass @@ -37,11 +33,9 @@ class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): latent_scale_factor: int = 16 text_encoder_mem_fraction_static: float | None = None - max_request_pixel_area: int = 2048 * 2048 # The embedded text worker admits 8192 prefill tokens shared by the two # CFG sequences and each sequence appends 256 query tokens. max_request_text_tokens: int = 3584 - max_request_total_pixel_area: int = 10 * 1024 * 1024 def prepare_sigmas(self, sigmas, num_inference_steps): if sigmas is not None: @@ -72,223 +66,48 @@ def calculate_condition_image_size(self, image, width, height): del image, width, height return None - def shard_latents_for_sp(self, batch, latents): - sp_degree = get_sp_world_size() - if latents.dim() == 4 and latents.shape[2] % sp_degree != 0: - raise ValueError( - f"LLaDA-Image latent height {latents.shape[2]} must be divisible " - f"by SP degree {sp_degree}. Choose a compatible output height" - ) - return super().shard_latents_for_sp(batch, latents) - def validate_server_args(self, server_args) -> None: super().validate_server_args(server_args) - for argument in ( - "component_weights_paths", - "component_precisions", - "component_quantizations", - "component_quantization_ignored_layers", - "component_attention_backends", - "component_direct_gpu_weight_loading", - ): - if "text_encoder" in (getattr(server_args, argument, None) or {}): - raise ValueError( - f"LLaDA-Image embedded text_encoder does not support {argument} " - "overrides. Use --component-paths.text_encoder with a complete " - "model directory containing its config and weights" - ) - if envs.SGLANG_CACHE_DIT_ENABLED: - raise ValueError("LLaDA-Image does not support SGLANG_CACHE_DIT_ENABLED") - # Replicated condition suffixes are currently de-duplicated only by the - # Ulysses attention path. - if server_args.ring_degree != 1: - raise ValueError("LLaDA-Image sequence parallelism requires ring_degree=1") - if server_args.ulysses_degree != server_args.sp_degree: - raise ValueError( - "LLaDA-Image sequence parallelism requires ulysses_degree == sp_degree" + if any( + value != 1 + for value in ( + server_args.num_gpus, + server_args.sp_degree, + server_args.tp_size, + server_args.dp_size, + server_args.cfg_parallel_degree, ) - if server_args.sp_degree not in (1, 2): - raise ValueError("LLaDA-Image currently supports only SP degrees 1 and 2") - if ( - server_args.tp_size != 1 - or server_args.dp_size != 1 - or server_args.cfg_parallel_degree != 1 ): - raise ValueError("LLaDA-Image requires diffusion parallelism TP=DP=CFG=1") - if server_args.num_gpus != server_args.sp_degree: - raise ValueError( - "LLaDA-Image requires num_gpus == sp_degree so every GPU belongs " - "to the sequence-parallel group" - ) + raise ValueError("LLaDA-Image currently supports single-GPU serving") + if server_args.quantization is not None: + raise ValueError("LLaDA-Image currently supports BF16 checkpoints") + if envs.SGLANG_CACHE_DIT_ENABLED: + raise ValueError("LLaDA-Image does not support cache-dit") if server_args.text_encoder_cpu_offload: if server_args.is_arg_explicitly_set("text_encoder_cpu_offload"): - raise ValueError( - "LLaDA-Image runs its text encoder inside an embedded srt " - "worker that the component residency manager cannot " - "offload. Remove --text-encoder-cpu-offload" - ) - # Auto-tuned default, corrected here because the tuner does not - # know about the embedded worker. + raise ValueError("LLaDA-Image requires a resident text encoder") server_args.text_encoder_cpu_offload = False - logger.info( - "LLaDA-Image keeps the embedded text encoder resident and " - "ignores the auto-tuned text_encoder_cpu_offload" - ) - if server_args.residency_mode("text_encoder") != "resident": - if ( - server_args.explicit_residency_mode("text_encoder") is not None - or server_args.is_arg_explicitly_set("layerwise_offload_components") - or server_args.is_arg_explicitly_set("cpu_offload_components") - ): - raise ValueError( - "LLaDA-Image requires its embedded text encoder to remain resident" + for attribute in ("layerwise_offload_components", "cpu_offload_components"): + components = getattr(server_args, attribute) + if components is not None and "text_encoder" in components: + if server_args.is_arg_explicitly_set(attribute): + raise ValueError("LLaDA-Image requires a resident text encoder") + setattr( + server_args, + attribute, + [name for name in components if name != "text_encoder"], ) - # Strip the encoder from auto-tuned offload selections the tuner - # applies without knowing about the embedded worker. - if server_args.layerwise_offload_components is not None: - server_args.layerwise_offload_components = [ - name - for name in server_args.layerwise_offload_components - if name != "text_encoder" - ] - if server_args.cpu_offload_components is not None: - server_args.cpu_offload_components = [ - name - for name in server_args.cpu_offload_components - if name != "text_encoder" - ] - logger.info( - "LLaDA-Image keeps the embedded text encoder resident and " - "removed it from auto-tuned offload selections" - ) - if server_args.residency_mode("text_encoder") != "resident": - raise ValueError( - "LLaDA-Image requires its embedded text encoder to remain resident" - ) - - area_override = server_args.llada_image_max_pixel_area - if area_override is not None and area_override <= 0: - raise ValueError("llada_image_max_pixel_area must be positive") - total_area_override = server_args.llada_image_max_total_pixel_area - if total_area_override is not None and total_area_override <= 0: - raise ValueError("llada_image_max_total_pixel_area must be positive") - tokens_override = server_args.llada_image_max_text_tokens - if tokens_override is not None and not ( - 0 < tokens_override <= self.max_request_text_tokens + if server_args.explicit_residency_mode("text_encoder") not in ( + None, + "resident", ): - raise ValueError( - "llada_image_max_text_tokens must be positive and at most " - f"{self.max_request_text_tokens}, the embedded worker budget" - ) - - # The embedded srt worker and its KV pool are invisible to generic module - # discovery, so partial sleep or update would misreport. - def supports_memory_release(self) -> bool: - return False - - def supports_hot_weight_updates(self) -> bool: - return False + raise ValueError("LLaDA-Image requires a resident text encoder") def supports_disaggregation(self) -> bool: return False - def validate_num_outputs_per_prompt( - self, num_outputs_per_prompt: int, server_args - ) -> None: - if server_args.sp_degree > 1 and num_outputs_per_prompt != 1: - raise ValueError( - "LLaDA-Image sequence parallelism supports only n=1. " - "submit separate requests for multiple images" - ) - - def validate_edit_source_count(self, source_count: int, server_args) -> None: - del server_args - if source_count != 1: - raise ValueError("LLaDA-Image editing requires exactly one source image") - - def validate_request_sampling_params(self, sampling_params, server_args) -> None: - enable_cache_dit = sampling_params.enable_cache_dit - if enable_cache_dit is not None and enable_cache_dit is not False: - raise ValueError("LLaDA-Image does not support enable_cache_dit") - if sampling_params.cache_dit_params is not None: - raise ValueError("LLaDA-Image does not support cache_dit_params") - - attention_backend = sampling_params.attention_backend_override - supported_attention_backends = {"FA", "TORCH_SDPA"} - if attention_backend is not None and ( - not isinstance(attention_backend, str) - or attention_backend.upper() not in supported_attention_backends - ): - raise ValueError( - "LLaDA-Image attention_backend_override must be fa or torch_sdpa" - ) - - cfg_gate_step = sampling_params.cfg_gate_step - if cfg_gate_step is not None and ( - isinstance(cfg_gate_step, bool) - or not isinstance(cfg_gate_step, (int, float)) - or not 0.0 <= cfg_gate_step <= 1.0 - or not np.isfinite(cfg_gate_step) - ): - raise ValueError("LLaDA-Image cfg_gate_step must be between 0.0 and 1.0") - - width = int(sampling_params.width or 0) - height = int(sampling_params.height or 0) - if width <= 0 or height <= 0: - raise ValueError("LLaDA-Image requires positive width and height") - if width % self.latent_scale_factor or height % self.latent_scale_factor: - raise ValueError( - f"LLaDA-Image width and height must be divisible by " - f"{self.latent_scale_factor}" - ) - sp_height_multiple = self.latent_scale_factor * server_args.sp_degree - if height % sp_height_multiple: - raise ValueError( - f"LLaDA-Image output height must be divisible by " - f"{sp_height_multiple} at SP degree {server_args.sp_degree}" - ) - self._validate_spatial_rope_bounds(width, height) - area_cap = ( - server_args.llada_image_max_pixel_area - if server_args.llada_image_max_pixel_area is not None - else self.max_request_pixel_area - ) - if width * height > area_cap: - raise ValueError( - f"LLaDA-Image output area {width}x{height} exceeds the " - f"supported maximum of {area_cap} pixels" - ) - total_area_cap = ( - server_args.llada_image_max_total_pixel_area - if server_args.llada_image_max_total_pixel_area is not None - else self.max_request_total_pixel_area - ) - num_outputs = sampling_params.num_outputs_per_prompt - if num_outputs is not None: - total_area = width * height * max(1, int(num_outputs)) - if total_area > total_area_cap: - raise ValueError( - f"LLaDA-Image total output area {total_area} pixels across " - f"{num_outputs} outputs exceeds the supported maximum of " - f"{total_area_cap} pixels per request" - ) - tokens_cap = ( - server_args.llada_image_max_text_tokens - if server_args.llada_image_max_text_tokens is not None - else self.max_request_text_tokens - ) - self._validate_max_sequence_length( - sampling_params.max_sequence_length, tokens_cap - ) - extra_kwargs = sampling_params.diffusers_kwargs - if isinstance(extra_kwargs, dict) and "max_sequence_length" in extra_kwargs: - # prepare_request later overrides the request value from this channel. - self._validate_max_sequence_length( - extra_kwargs["max_sequence_length"], tokens_cap - ) - def _validate_spatial_rope_bounds(self, width: int, height: int) -> None: - # The served DiT uses patch size 1 and SP preserves global height positions. + # The served DiT uses patch size 1. axes_lens = self.dit_config.arch_config.axes_lens max_height = axes_lens[1] * self.latent_scale_factor max_width = axes_lens[2] * self.latent_scale_factor @@ -298,22 +117,6 @@ def _validate_spatial_rope_bounds(self, width: int, height: int) -> None: f"bounds of {max_width}x{max_height} pixels" ) - @staticmethod - def _validate_max_sequence_length(max_sequence_length, tokens_cap: int) -> None: - if max_sequence_length is None: - return - if ( - isinstance(max_sequence_length, bool) - or not isinstance(max_sequence_length, int) - or max_sequence_length <= 0 - ): - raise ValueError("max_sequence_length must be a positive integer") - if max_sequence_length > tokens_cap: - raise ValueError( - f"max_sequence_length {max_sequence_length} exceeds the " - f"embedded text encoder budget of {tokens_cap} tokens" - ) - @staticmethod def _prepare_condition_list(values, name: str, expected_size: int, device, dtype): if values is None: @@ -333,10 +136,6 @@ def _prepare_source_latents(self, batch, device, dtype): device, dtype, ) - if source_latents and getattr(batch, "did_sp_shard_latents", False): - source_latents = [ - self.shard_latents_for_sp(batch, latent)[0] for latent in source_latents - ] return source_latents def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): diff --git a/python/sglang/multimodal_gen/configs/sample/llada_image.py b/python/sglang/multimodal_gen/configs/sample/llada_image.py index 1c0ac99cb1d3..d6eb300af1d1 100644 --- a/python/sglang/multimodal_gen/configs/sample/llada_image.py +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -18,3 +18,14 @@ class LLaDAImageSamplingParams(SamplingParams): width: int = 1024 guidance_scale: float = 1.0 max_sequence_length: int = 2048 + + def _validate_with_pipeline_config(self, pipeline_config): + super()._validate_with_pipeline_config(pipeline_config) + if self.enable_cache_dit or self.cache_dit_params is not None: + raise ValueError("LLaDA-Image does not support cache-dit") + if self.max_sequence_length is not None and not ( + 0 < self.max_sequence_length <= pipeline_config.max_request_text_tokens + ): + raise ValueError( + "LLaDA-Image max_sequence_length must be between 1 and 3584" + ) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 9b4aaa193744..f3e079bfd649 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -983,9 +983,6 @@ def from_user_sampling_params_args( sampling_params._adjust(server_args) sampling_params._validate_with_pipeline_config(server_args.pipeline_config) - server_args.pipeline_config.validate_request_sampling_params( - sampling_params, server_args - ) return sampling_params diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 4062e99eff96..2c2f2d9a132c 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -1144,7 +1144,6 @@ def _register_configs(): ], ) # Qwen-Image - register_configs( sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImagePipelineConfig, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index fbb967257ca2..59dfd4f83d7e 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -5,7 +5,6 @@ import contextlib import os import time -from types import SimpleNamespace from typing import Any, List, Optional from fastapi import ( @@ -49,7 +48,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args -from sglang.multimodal_gen.runtime.utils.common import parse_size from sglang.srt.observability.trace import extract_trace_headers router = APIRouter(prefix="/v1/images", tags=["images"]) @@ -85,50 +83,6 @@ def _runtime_sampling_quality(quality: str | None) -> str | None: return None if quality in (None, "auto") else quality -def _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt: int) -> None: - try: - server_args.pipeline_config.validate_num_outputs_per_prompt( - num_outputs_per_prompt, server_args - ) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - - -def _build_sampling_params_or_400(request_id: str, **kwargs): - try: - return build_sampling_params(request_id, **kwargs) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - - -def _early_validate_edit_bounds( - server_args, size: Optional[str], num_outputs_per_prompt: int = 1 -) -> None: - """Best-effort size rejection before any source image is handled.""" - if not size: - return - width, height = parse_size(size) - if not width or not height: - return - lightweight = SimpleNamespace( - width=width, - height=height, - max_sequence_length=None, - num_outputs_per_prompt=num_outputs_per_prompt, - diffusers_kwargs=None, - enable_cache_dit=None, - cache_dit_params=None, - attention_backend_override=None, - cfg_gate_step=None, - ) - try: - server_args.pipeline_config.validate_request_sampling_params( - lightweight, server_args - ) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error - - def _read_b64_for_paths(paths: list[str]) -> list[str]: """Read and base64-encode each file. Must be called before cloud upload deletes them.""" result = [] @@ -312,8 +266,6 @@ async def generations( ): request_id = generate_request_id() server_args = get_global_server_args() - num_outputs_per_prompt = max(1, min(int(request.n or 1), 10)) - _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt) sampling_params_cls = resolve_sampling_params_cls(server_args) model_kwargs = _image_request_model_kwargs(request, sampling_params_cls) output_format = ( @@ -324,13 +276,13 @@ async def generations( ext = choose_output_image_ext(output_format, request.background) with temp_dir_if_disabled(server_args.output_path) as output_dir: - sampling = _build_sampling_params_or_400( + sampling = build_sampling_params( request_id, prompt=request.prompt, size=request.size, width=request.width, height=request.height, - num_outputs_per_prompt=num_outputs_per_prompt, + num_outputs_per_prompt=max(1, min(int(request.n or 1), 10)), output_file_name=f"{request_id}.{ext}", output_path=output_dir, num_frames=1, @@ -487,17 +439,7 @@ async def edits( status_code=422, detail="Field 'image' or 'url' is required" ) - num_outputs_per_prompt = max(1, min(int(n or 1), 10)) - _validate_num_outputs_per_prompt(server_args, num_outputs_per_prompt) - _early_validate_edit_bounds(server_args, size, num_outputs_per_prompt) - image_list = merge_image_input_list(images, urls) - try: - server_args.pipeline_config.validate_edit_source_count( - len(image_list), server_args - ) - except ValueError as error: - raise HTTPException(status_code=400, detail=str(error)) from error with contextlib.ExitStack() as stack: uploads_dir = stack.enter_context( @@ -522,11 +464,11 @@ async def edits( ) ext = choose_output_image_ext(output_format, background) - sampling = _build_sampling_params_or_400( + sampling = build_sampling_params( request_id, prompt=prompt, size=size, - num_outputs_per_prompt=num_outputs_per_prompt, + num_outputs_per_prompt=max(1, min(int(n or 1), 10)), output_file_name=f"{request_id}.{ext}", output_path=output_dir, image_path=input_paths, diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index b594ef49e0b6..d6f13135e167 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -382,7 +382,7 @@ def prepare(self) -> None: class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter): - """Disable whole-DiT CPU offload for FP8 modes that require CUDA setup.""" + """Adapter for diffusion ModelOpt FP8 checkpoints.""" def __init__( self, @@ -404,11 +404,7 @@ def _maybe_disable_incompatible_dit_offload_modes( if quant_config is None: return - quant_name = _get_quant_config_name(quant_config) - is_online_fp8 = quant_name == "fp8" and not getattr( - quant_config, "is_checkpoint_fp8_serialized", False - ) - if quant_name != "modelopt_fp8" and not is_online_fp8: + if _get_quant_config_name(quant_config) != "modelopt_fp8": return component_offload = _uses_component_offload( @@ -430,9 +426,8 @@ def _maybe_disable_incompatible_dit_offload_modes( feature_name="ModelOpt FP8 diffusion checkpoints", ) logger.warning( - "%s diffusion quantization keeps the DiT resident instead " + "ModelOpt FP8 diffusion checkpoints keep the DiT resident instead " "of using component offload. Layerwise offload remains supported.", - quant_name, ) def prepare(self) -> None: diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py index 968393855c79..008caaf09081 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py @@ -192,14 +192,6 @@ def release_memory_occupation(self) -> dict[str, bool | str]: success=False, message="pipeline not initialized", ) - if not self.pipeline.server_args.pipeline_config.supports_memory_release(): - return self._memory_occupation_result( - success=False, - message=( - "pipeline holds components outside generic module " - "discovery and does not support memory release" - ), - ) self._sleep_restore_map = self._offload_active_modules_to_cpu() self._sleeping = True diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index 7dfb2f0260d4..2c4d60a2ad6b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -31,8 +31,6 @@ from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig from sglang.multimodal_gen.runtime.distributed import ( - get_sp_parallel_rank, - get_sp_world_size, get_tp_world_size, ) from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul @@ -167,8 +165,6 @@ def __call__( hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, freqs_cis: torch.Tensor | None = None, - num_replicated_suffix: int = 0, - skip_sequence_parallel_override: bool = False, ) -> torch.Tensor: qkv, _ = attn.to_qkv(hidden_states) query, key, value = qkv.split( @@ -224,8 +220,6 @@ def __call__( key, value, attn_mask=attention_mask, - num_replicated_suffix=num_replicated_suffix, - skip_sequence_parallel_override=skip_sequence_parallel_override, ) hidden_states = hidden_states.flatten(2, 3) hidden_states, _ = attn.to_out[0](hidden_states) @@ -297,16 +291,12 @@ def forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor | None, freqs_cis: torch.Tensor, - num_replicated_suffix: int = 0, - skip_sequence_parallel_override: bool = False, ) -> torch.Tensor: return self.processor( self, hidden_states, attention_mask, freqs_cis, - num_replicated_suffix, - skip_sequence_parallel_override, ) @@ -401,8 +391,6 @@ def forward( noise_mask: torch.Tensor | None = None, adaln_noisy: torch.Tensor | None = None, adaln_clean: torch.Tensor | None = None, - num_replicated_suffix: int = 0, - skip_sequence_parallel_override: bool = False, ) -> torch.Tensor: if self.modulation: sequence_length = hidden_states.shape[1] @@ -449,8 +437,6 @@ def forward( self.attention_norm1(hidden_states) * scale_msa, attention_mask, freqs_cis, - num_replicated_suffix, - skip_sequence_parallel_override, ) hidden_states = apply_rmsnorm_tanh_mul_add( attention_output, gate_msa, hidden_states, self.attention_norm2 @@ -466,8 +452,6 @@ def forward( self.attention_norm1(hidden_states), attention_mask, freqs_cis, - num_replicated_suffix, - skip_sequence_parallel_override, ) hidden_states = hidden_states + self.attention_norm2(attention_output) hidden_states = hidden_states + self.ffn_norm2( @@ -833,7 +817,7 @@ def _unpatchify( image_offsets: list[tuple[int, int]] | None = None, ) -> list[torch.Tensor]: outputs = [] - sequence_multiple = SEQUENCE_MULTIPLE // get_sp_world_size() + sequence_multiple = SEQUENCE_MULTIPLE for batch_index, batch_hidden_states in enumerate(hidden_states): if image_offsets is None: batch_sizes = [sizes[batch_index]] @@ -926,13 +910,11 @@ def _prepare_t2i_sequences( latent, patch_size, f_patch_size ) image_height_start = 0 - if get_sp_world_size() > 1: - image_height_start = get_sp_parallel_rank() * token_grid_size[1] padded_features, position_ids, padding_mask, _, _ = self._pad_with_ids( patches, token_grid_size, (position_cursor, image_height_start, 0), - sequence_multiple=SEQUENCE_MULTIPLE // get_sp_world_size(), + sequence_multiple=SEQUENCE_MULTIPLE, ) image_sequence.features.append(padded_features) image_sequence.position_ids.append(position_ids) @@ -1000,15 +982,13 @@ def _prepare_editing_sequences( image, patch_size, f_patch_size ) image_height_start = 0 - if get_sp_world_size() > 1: - image_height_start = get_sp_parallel_rank() * token_grid_size[1] padded_features, position_ids, padding_mask, _, noise_mask = ( self._pad_with_ids( patches, token_grid_size, (position_start, image_height_start, 0), noise_value, - sequence_multiple=SEQUENCE_MULTIPLE // get_sp_world_size(), + sequence_multiple=SEQUENCE_MULTIPLE, ) ) batch_image_features.append(padded_features) @@ -1040,9 +1020,7 @@ def _prepare_editing_sequences( glm_cap_feats[batch_index], (len(glm_cap_feats[batch_index]), 1, 1), ( - len(batch_cap_features) - + len(batch_image_features) * get_sp_world_size() - + 1, + len(batch_cap_features) + len(batch_image_features) + 1, 0, 0, ), @@ -1190,7 +1168,6 @@ def forward( noisy_embedding = None clean_embedding = None image_offsets = None - num_replicated_suffix = 0 if is_editing: if t.shape[0] == 1: @@ -1301,7 +1278,6 @@ def forward( cap_features, cap_attention_mask, cap_frequencies, - skip_sequence_parallel_override=True, ) sigvq_lengths = [len(features) for features in sigvq_sequence.features] @@ -1334,41 +1310,20 @@ def forward( sigvq_features, sigvq_attention_mask, sigvq_frequencies, - skip_sequence_parallel_override=True, ) - if get_sp_world_size() > 1: - if batch_size != 1: - raise NotImplementedError( - "LLaDA-Image sequence parallelism currently supports batch size 1." - ) - feature_groups = (image_features, cap_features, sigvq_features) - frequency_groups = ( - image_frequencies, - cap_frequencies, - sigvq_frequencies, - ) - length_groups = (image_lengths, cap_lengths, sigvq_lengths) - noise_mask_groups = ( - image_noise_mask, - cap_noise_mask, - sigvq_noise_mask, - ) - num_replicated_suffix = cap_lengths[0] + sigvq_lengths[0] - image_offsets = [(0, image_lengths[0])] - else: - feature_groups = (cap_features, image_features, sigvq_features) - frequency_groups = ( - cap_frequencies, - image_frequencies, - sigvq_frequencies, - ) - length_groups = (cap_lengths, image_lengths, sigvq_lengths) - noise_mask_groups = ( - cap_noise_mask, - image_noise_mask, - sigvq_noise_mask, - ) + feature_groups = (cap_features, image_features, sigvq_features) + frequency_groups = ( + cap_frequencies, + image_frequencies, + sigvq_frequencies, + ) + length_groups = (cap_lengths, image_lengths, sigvq_lengths) + noise_mask_groups = ( + cap_noise_mask, + image_noise_mask, + sigvq_noise_mask, + ) ( unified_features, @@ -1482,7 +1437,6 @@ def forward( condition_features, condition_attention_mask, condition_frequencies, - skip_sequence_parallel_override=True, ) ( @@ -1496,12 +1450,6 @@ def forward( (image_frequencies, condition_frequencies), (image_lengths, condition_lengths), ) - if get_sp_world_size() > 1: - if batch_size != 1: - raise NotImplementedError( - "LLaDA-Image sequence parallelism currently supports batch size 1." - ) - num_replicated_suffix = condition_lengths[0] for layer in self.layers: if is_editing: @@ -1512,7 +1460,6 @@ def forward( noise_mask=unified_noise_mask, adaln_noisy=noisy_embedding, adaln_clean=clean_embedding, - num_replicated_suffix=num_replicated_suffix, ) else: unified_features = layer( @@ -1520,7 +1467,6 @@ def forward( unified_attention_mask, unified_frequencies, adaln_input, - num_replicated_suffix=num_replicated_suffix, ) if is_editing: diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 493ecd7ee96a..f0c640081dd0 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -39,7 +39,6 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) -SP_STOCHASTIC_NOISE_KEY = "sp_stochastic_noise" @dataclass @@ -456,68 +455,6 @@ def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: else: self._step_index = self._begin_index - @staticmethod - def _draw_stochastic_noise( - sample: torch.Tensor, - generator: torch.Generator | list[torch.Generator] | None, - batch, - ) -> torch.Tensor: - batch_extra = getattr(batch, "extra", None) - metadata = batch_extra.get(SP_STOCHASTIC_NOISE_KEY) if batch_extra else None - if metadata is None: - return randn_tensor( - sample.shape, - generator=generator, - device=sample.device, - dtype=sample.dtype, - ) - - required_keys = {"full_shape", "dim", "start", "length"} - if not isinstance(metadata, dict) or not required_keys.issubset(metadata): - raise ValueError("Invalid SP stochastic noise metadata fields") - if not getattr(batch, "did_sp_shard_latents", False): - raise ValueError("Invalid SP stochastic noise metadata without sharding") - - full_shape = metadata["full_shape"] - dim = metadata["dim"] - start = metadata["start"] - length = metadata["length"] - if ( - not isinstance(full_shape, tuple) - or len(full_shape) != sample.ndim - or any(type(value) is not int or value <= 0 for value in full_shape) - ): - raise ValueError("Invalid SP stochastic noise metadata full shape") - if type(dim) is not int or dim < 1 or dim >= sample.ndim: - raise ValueError("Invalid SP stochastic noise metadata dimension") - if type(start) is not int or type(length) is not int: - raise ValueError("Invalid SP stochastic noise metadata slice") - if length != sample.shape[dim] or start < 0: - raise ValueError("Invalid SP stochastic noise metadata local slice") - if start + length > full_shape[dim]: - raise ValueError("Invalid SP stochastic noise metadata slice bounds") - raw_latent_shape = getattr(batch, "raw_latent_shape", None) - if raw_latent_shape is None or tuple(raw_latent_shape) != full_shape: - raise ValueError("Invalid SP stochastic noise metadata raw shape") - - expected_local_shape = list(full_shape) - expected_local_shape[dim] = length - if tuple(expected_local_shape) != tuple(sample.shape): - raise ValueError("Invalid SP stochastic noise metadata sample shape") - if isinstance(generator, list) and len(generator) != full_shape[0]: - raise ValueError("Invalid SP stochastic noise metadata generator count") - - full_noise = randn_tensor( - full_shape, - generator=generator, - device=sample.device, - dtype=sample.dtype, - ) - noise = full_noise.narrow(dim, start, length).contiguous() - if noise.shape != sample.shape: - raise ValueError("Invalid SP stochastic noise metadata result shape") - return noise - def step( self, model_output: torch.FloatTensor, @@ -608,7 +545,12 @@ def step( else: if self.config.stochastic_sampling: x0 = sample - current_sigma * model_output - noise = self._draw_stochastic_noise(sample, generator, batch) + noise = randn_tensor( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise else: prev_sample = sample + dt * model_output diff --git a/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py b/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py index 72a610f1fc6d..b9edb8ca4fbd 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py @@ -7,15 +7,10 @@ from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, - get_sp_parallel_rank, - get_sp_world_size, ) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PlainStateDictComponentLoader, ) -from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( - SP_STOCHASTIC_NOISE_KEY, -) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -50,29 +45,6 @@ def forward(self, batch, server_args): dtype=model_dtype, ).float() batch = super().forward(batch, server_args) - batch.extra.pop(SP_STOCHASTIC_NOISE_KEY, None) - - sp_size = get_sp_world_size() - if self.scheduler.config.stochastic_sampling and sp_size > 1: - full_shape = tuple(int(value) for value in batch.raw_latent_shape) - if len(full_shape) != 4: - raise ValueError( - "LLaDA-Image SP noise requires four-dimensional latents" - ) - latent_height = full_shape[2] - if latent_height % sp_size != 0: - raise ValueError( - "LLaDA-Image SP noise requires height divisible by SP size" - ) - local_height = latent_height // sp_size - sp_rank = get_sp_parallel_rank() - batch.extra[SP_STOCHASTIC_NOISE_KEY] = { - "full_shape": full_shape, - "dim": 2, - "start": sp_rank * local_height, - "length": local_height, - } - return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index ab9e256f3fbe..c989ad94207f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -11,7 +11,6 @@ from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, - get_sp_parallel_rank, ) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ComponentUse, @@ -75,12 +74,6 @@ def __init__( from sglang.srt.runtime_context import create_context, use_context from sglang.srt.server_args import ServerArgs as SRTServerArgs - # Keep the SRT world aligned with the pure-Ulysses ranks. Within that - # world, SP ranks become data replicas so conditioning matches SP1. - text_tp_size = int(server_args.sp_degree) - text_tp_rank = get_sp_parallel_rank() - text_dp_attention = text_tp_size > 1 - self.queryformer = queryformer self.text_projection = text_projection self.tokenizer = tokenizer @@ -101,24 +94,14 @@ def __init__( revision=server_args.revision, skip_tokenizer_init=True, dtype="bfloat16", - tp_size=text_tp_size, - dp_size=text_tp_size if text_dp_attention else 1, - enable_dp_attention=text_dp_attention, - enable_dp_lm_head=text_dp_attention, - dcp_size=1, - attn_cp_size=1, - ep_size=1, - moe_dp_size=text_tp_size if text_dp_attention else 1, - moe_dense_tp_size=1 if text_dp_attention else None, - moe_a2a_backend="none", - pp_size=1, + tp_size=1, attention_backend="llada2_cfg_flashinfer", disable_cuda_graph=True, disable_radix_cache=True, chunked_prefill_size=-1, max_prefill_tokens=8192, max_total_tokens=8192, - max_running_requests=4 if text_dp_attention else 2, + max_running_requests=2, mem_fraction_static=( server_args.pipeline_config.text_encoder_mem_fraction_static ), @@ -135,25 +118,7 @@ def __init__( self.worker = TpModelWorker( server_args=srt_args, gpu_id=gpu_id, - ps=ParallelState.trivial( - tp_rank=text_tp_rank, - tp_size=text_tp_size, - dp_rank=text_tp_rank if text_dp_attention else 0, - dp_size=text_tp_size if text_dp_attention else 1, - attn_tp_rank=0 if text_dp_attention else text_tp_rank, - attn_tp_size=1 if text_dp_attention else text_tp_size, - attn_cp_rank=0, - attn_cp_size=1, - attn_dcp_rank=0, - attn_dcp_size=1, - attn_dp_rank=text_tp_rank if text_dp_attention else 0, - attn_dp_size=text_tp_size if text_dp_attention else 1, - moe_ep_rank=0, - moe_ep_size=1, - moe_dp_rank=text_tp_rank if text_dp_attention else 0, - moe_dp_size=text_tp_size if text_dp_attention else 1, - gpu_id=gpu_id, - ), + ps=ParallelState.trivial(gpu_id=gpu_id), nccl_port=server_args.nccl_port or 29500, ) # Run the post-construction init phases the srt scheduler drives. @@ -166,7 +131,6 @@ def __init__( srt_parallel_state._TP = saved_srt_tp srt_parallel_state._ATTN_TP = saved_srt_attn_tp self.server_args = srt_args - self.text_dp_attention = text_dp_attention self.model_runner = self.worker.model_runner self.page_size = self.model_runner.page_size self.req_to_token_pool, self.token_to_kv_pool_allocator = ( @@ -241,8 +205,7 @@ def encode( import sglang.srt.distributed.parallel_state as srt_parallel_state from sglang.srt.runtime_context import use_context - # The TP scope belongs to diffusion. Install the encoder context after it - # so replicated attention and MoE keep the encoder's derived widths. + # Restore the encoder context inside the diffusion TP scope. with ( mm_parallel_state.use_tensor_parallel_group(self.encoder_tp_group), use_context(self.runtime_context), @@ -343,12 +306,6 @@ def _encode_impl( spec_algorithm=SpeculativeAlgorithm.NONE, dllm_config=None, ) - if self.text_dp_attention: - token_count = sum(sequence_lengths) - batch.global_num_tokens = [token_count] * self.server_args.dp_size - batch.global_num_tokens_for_logprob = [len(reqs)] * ( - self.server_args.dp_size - ) try: batch.prepare_for_extend() batch.forward_mode = ForwardMode.EXTEND diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 8b10106a1afd..75359a82cefe 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -371,10 +371,6 @@ def update_weights_from_disk( """Update model weights from disk without restarting the server.""" logger.info(f"Updating weights from disk: {model_path}") - rejected, message = self._reject_if_partial_update() - if rejected: - return False, message - try: modules_to_update = self._collect_modules(target_modules) except ValueError as e: @@ -448,17 +444,6 @@ def update_weights_from_disk( logger.info(message) return success, message - def _reject_if_partial_update(self) -> tuple[bool, str]: - config = self.pipeline.server_args.pipeline_config - if config.supports_hot_weight_updates(): - return False, "" - message = ( - "pipeline holds components outside generic module discovery, a " - "runtime weight update would leave them at the old checkpoint" - ) - logger.error(message) - return True, message - def _collect_modules( self, target_modules: list[str] | None ) -> list[tuple[str, torch.nn.Module]]: @@ -546,9 +531,6 @@ def update_weights_from_tensor( lora_alpha: int | None = None, lora_rank: int | None = None, ) -> tuple[bool, str]: - rejected, message = self._reject_if_partial_update() - if rejected: - return False, message if weight_update_mode == LORA_MERGE_WEIGHT_UPDATE_MODE: return self._update_lora_from_tensor( named_tensors=named_tensors, diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index fa8e7b28ca6f..e1cb364a7c8b 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -431,9 +431,6 @@ class ServerArgs(DisaggServerArgsMixin): use_fsdp_inference: bool | None = None pin_cpu_memory: bool = True ltx2_two_stage_device_mode: str | None = None - llada_image_max_pixel_area: int | None = None - llada_image_max_text_tokens: int | None = None - llada_image_max_total_pixel_area: int | None = None _explicit_arg_names: set[str] = field(default_factory=set, repr=False) _automatic_component_attention_backend_keys: set[str] = field( default_factory=set, init=False, repr=False @@ -2636,35 +2633,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "Default is auto: resident on H200/high-memory CUDA GPUs, otherwise original." ), ) - parser.add_argument( - "--llada-image-max-pixel-area", - type=int, - default=ServerArgs.llada_image_max_pixel_area, - help=( - "LLaDA-Image only. Largest accepted output pixel area per " - "request. Defaults to the pipeline config value of 2048*2048." - ), - ) - parser.add_argument( - "--llada-image-max-text-tokens", - type=int, - default=ServerArgs.llada_image_max_text_tokens, - help=( - "LLaDA-Image only. Largest accepted request " - "max_sequence_length. Defaults to the pipeline config value " - "of 3584 and cannot exceed it." - ), - ) - parser.add_argument( - "--llada-image-max-total-pixel-area", - type=int, - default=ServerArgs.llada_image_max_total_pixel_area, - help=( - "LLaDA-Image only. Largest accepted output pixel area times " - "the per-prompt output count. Defaults to the pipeline " - "config value of ten 1024x1024 outputs." - ), - ) parser.add_argument( "--disable-autocast", action=StoreBoolean, diff --git a/python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py b/python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py index 20f9d5c14c22..63480362927d 100644 --- a/python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py +++ b/python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py @@ -27,7 +27,6 @@ from sglang.multimodal_gen.test.server.testcase_configs import ( BASELINE_CONFIG, DiffusionTestCase, - get_case_full_test_time_estimate, ) # TODO: remove duplicated code @@ -78,21 +77,10 @@ class PartitionAssignment: def get_case_est_time(case_id: str) -> float: scenario = BASELINE_CONFIG.scenarios.get(case_id) - case = next( - ( - case - for groups in PARAMETRIZED_CASE_GROUPS.values() - for _, cases in groups - for case in cases - if case.id == case_id - ), - None, - ) - estimate = get_case_full_test_time_estimate(case, scenario) - if estimate is not None: - return estimate if scenario is None: return DEFAULT_EST_TIME_SECONDS + if scenario.estimated_full_test_time_s is not None: + return scenario.estimated_full_test_time_s return scenario.expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index c0b8f3e5f1b5..91747f9fff25 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -69,42 +69,6 @@ _CACHE_DIT_CONFIG_DIR = Path(__file__).parent / "configs" -def _llada_image_ci_kwargs(sp_degree: int, mode: str) -> dict: - return dict( - server_args=DiffusionServerArgs( - model_path="inclusionAI/LLaDA-Image-Turbo-FP8", - modality="image", - num_gpus=sp_degree, - tp_size=1, - ulysses_degree=sp_degree, - ring_degree=1, - cfg_parallel=False, - extras=[ - "--trust-remote-code", - "--revision", - "2b822499be3f33a6f83c988dafab7986ebd1f733", - ], - ), - sampling_params=DiffusionSamplingParams( - prompt=( - "Convert 2D style to 3D style" - if mode == "edit" - else "A red fox sitting on a snowy hill" - ), - image_path=( - "https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg" - if mode == "edit" - else None - ), - output_size="512x512", - extras={"num_inference_steps": 4, "guidance_scale": 1.0, "seed": 42}, - ), - run_perf_check=False, - run_consistency_check=False, - run_component_accuracy_check=False, - ) - - # All test cases with clean default values # To test different models, simply add more DiffusionCase entries ONE_GPU_CASES: list[DiffusionTestCase] = [ @@ -539,14 +503,55 @@ def _llada_image_ci_kwargs(sp_degree: int, mode: str) -> dict: if not current_platform.is_hip() and not current_platform.is_xpu(): ONE_GPU_CASES += [ DiffusionTestCase( - "llada_image_turbo_fp8_t2i_sp1", - estimated_full_test_time_s=600.0, - **_llada_image_ci_kwargs(1, "t2i"), + "llada_image_bf16_t2i", + DiffusionServerArgs( + model_path="inclusionAI/LLaDA-Image", + modality="image", + num_gpus=1, + tp_size=1, + ulysses_degree=1, + ring_degree=1, + cfg_parallel=False, + extras=[ + "--trust-remote-code", + "--revision", + "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", + ], + ), + DiffusionSamplingParams( + prompt="A red fox sitting on a snowy hill", + output_size="512x512", + extras={"num_inference_steps": 50, "guidance_scale": 5.0, "seed": 42}, + ), + run_perf_check=False, + run_consistency_check=False, + run_component_accuracy_check=False, ), DiffusionTestCase( - "llada_image_turbo_fp8_edit_sp1", - estimated_full_test_time_s=600.0, - **_llada_image_ci_kwargs(1, "edit"), + "llada_image_bf16_edit", + DiffusionServerArgs( + model_path="inclusionAI/LLaDA-Image", + modality="image", + num_gpus=1, + tp_size=1, + ulysses_degree=1, + ring_degree=1, + cfg_parallel=False, + extras=[ + "--trust-remote-code", + "--revision", + "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", + ], + ), + DiffusionSamplingParams( + prompt="Convert 2D style to 3D style", + image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", + output_size="512x512", + extras={"num_inference_steps": 50, "guidance_scale": 5.0, "seed": 42}, + ), + run_perf_check=False, + run_consistency_check=False, + run_component_accuracy_check=False, ), ] @@ -1236,20 +1241,6 @@ def _llada_image_ci_kwargs(sp_degree: int, mode: str) -> dict: ), ] -if not current_platform.is_hip() and not current_platform.is_xpu(): - TWO_GPU_CASES += [ - DiffusionTestCase( - "llada_image_turbo_fp8_t2i_sp2", - estimated_full_test_time_s=600.0, - **_llada_image_ci_kwargs(2, "t2i"), - ), - DiffusionTestCase( - "llada_image_turbo_fp8_edit_sp2", - estimated_full_test_time_s=600.0, - **_llada_image_ci_kwargs(2, "edit"), - ), - ] - if not current_platform.is_hip(): # Flux2 multi-image edit with cache-dit, regression test ONE_GPU_CASES.append( diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index d3780eee870a..504873318914 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -43,7 +43,6 @@ DiffusionTestCase, PerformanceSummary, ScenarioConfig, - get_case_full_test_time_estimate, get_model_task_type_for_server_args, get_perf_baseline_update_path, ) @@ -237,7 +236,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: ) scenario = BASELINE_CONFIG.scenarios.get(case.id) - needs_estimated_time = get_case_full_test_time_estimate(case, scenario) is None + needs_estimated_time = ( + scenario is None or scenario.estimated_full_test_time_s is None + ) if needs_estimated_time and not is_baseline_generation_mode: _MISSING_ESTIMATED_TIME_CASES.add(case.id) @@ -397,7 +398,7 @@ def _validate_and_record( if ( not missing_scenario and not is_baseline_generation_mode - and get_case_full_test_time_estimate(case, scenario) is None + and scenario.estimated_full_test_time_s is None ): _MISSING_ESTIMATED_TIME_CASES.add(case.id) diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 07d948f15ee5..d35dcb0d44ed 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -328,8 +328,6 @@ class DiffusionTestCase: run_lora_dynamic_load_check: bool = False run_lora_dynamic_switch_check: bool = False run_multi_lora_api_check: bool = False - # Scheduling budget for functional cases without a performance baseline - estimated_full_test_time_s: float | None = None def __post_init__(self) -> None: if self.perf_repeat_requests < 1: @@ -366,14 +364,6 @@ def __post_init__(self) -> None: ) -def get_case_full_test_time_estimate( - case: DiffusionTestCase | None, scenario: ScenarioConfig | None -) -> float | None: - if scenario is not None and scenario.estimated_full_test_time_s is not None: - return scenario.estimated_full_test_time_s - return case.estimated_full_test_time_s if case is not None else None - - _REALTIME_MODEL_COMMON_EXTRAS = { "seed": 42, "num_inference_steps": 4, diff --git a/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py deleted file mode 100644 index 74ac288884f3..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_build_modelopt_fp8_transformer.py +++ /dev/null @@ -1,359 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import json - -import pytest -import torch -from safetensors import safe_open -from safetensors.torch import load_file, save_file - -from sglang.multimodal_gen.tools.build_modelopt_fp8_transformer import ( - build_fp8_scale_map, - build_modelopt_fp8_transformer, -) - - -def _write_llada_modelopt_export( - tmp_path, - source_tensors, - backbone_state, - *, - quantization_config=None, - class_name="LLaDAImageTransformer2DModel", -): - source_dir = tmp_path / "source" - output_dir = tmp_path / "output" - source_dir.mkdir() - source_config = { - "_class_name": class_name, - "quantization_config": { - "quant_method": "modelopt", - "quant_algo": "FP8", - **(quantization_config or {}), - }, - } - (source_dir / "config.json").write_text(json.dumps(source_config)) - save_file(source_tensors, source_dir / "model.safetensors") - - backbone_path = tmp_path / "backbone.pt" - torch.save({"model_state_dict": backbone_state}, backbone_path) - return source_dir, output_dir, backbone_path - - -def test_dynamic_activation_scale_map_only_requires_weight_amax(): - state_dict = { - "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), - } - - scale_map = build_fp8_scale_map( - state_dict, - maxbound=448.0, - require_input_scale=False, - ) - - assert set(scale_map) == {"layers.0.attention.to_qkv.weight"} - assert set(scale_map["layers.0.attention.to_qkv.weight"]) == {"weight_scale"} - torch.testing.assert_close( - scale_map["layers.0.attention.to_qkv.weight"]["weight_scale"], - torch.tensor([0.5]), - ) - - -def test_static_activation_scale_map_still_requires_input_amax(): - weight_amax_only = { - "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), - } - assert build_fp8_scale_map(weight_amax_only, maxbound=448.0) == {} - - state_dict = { - **weight_amax_only, - "layers.0.attention.to_qkv.input_quantizer._amax": torch.tensor(112.0), - } - scale_map = build_fp8_scale_map(state_dict, maxbound=448.0) - - assert set(scale_map["layers.0.attention.to_qkv.weight"]) == { - "weight_scale", - "input_scale", - } - torch.testing.assert_close( - scale_map["layers.0.attention.to_qkv.weight"]["input_scale"], - torch.tensor([0.25]), - ) - - -def test_llada_conversion_writes_dynamic_fp8_without_input_scales(tmp_path): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - { - "layers.0.attention.to_q.weight": torch.ones(2, 2), - "layers.0.attention.to_q.input_scale": torch.tensor([0.25]), - }, - { - "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), - }, - ) - - stats = build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="llada-image", - ) - - output_config = json.loads((output_dir / "config.json").read_text()) - assert output_config["quantization_config"] == { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "ignored_layers": [], - } - output_tensors = load_file(output_dir / "model.safetensors") - assert output_tensors["layers.0.attention.to_q.weight"].dtype == torch.float8_e4m3fn - assert "layers.0.attention.to_q.weight_scale" in output_tensors - assert "layers.0.attention.to_q.input_scale" not in output_tensors - assert stats["added_scale_tensors"] == 1 - - -@pytest.mark.parametrize( - ("module_prefix", "source_modules", "fused_module"), - [ - ("layers.0.attention", ("to_q", "to_k", "to_v"), "to_qkv"), - ("layers.0.feed_forward", ("w1", "w3"), "w13"), - ], -) -def test_llada_conversion_rejects_mismatched_fused_weight_scales( - tmp_path, module_prefix, source_modules, fused_module -): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - { - f"{module_prefix}.{module}.weight": torch.ones(2, 2) - for module in source_modules - }, - { - f"{module_prefix}.{module}.weight_quantizer._amax": torch.tensor( - 224.0 if index == 0 else 112.0 - ) - for index, module in enumerate(source_modules) - }, - ) - - with pytest.raises( - ValueError, - match=( - "LLaDA-Image FP8 conversion currently requires identical weight scales" - f".*{fused_module}" - ), - ): - build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="llada-image", - ) - - -def test_llada_conversion_rejects_static_activation_quantization(tmp_path): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, - { - "layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0), - "layers.0.attention.to_qkv.input_quantizer._amax": torch.tensor(112.0), - }, - ) - - with pytest.raises( - ValueError, - match=( - "LLaDA-Image FP8 conversion currently supports only dynamic " - "activation quantization" - ), - ): - build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="llada-image", - activation_scheme="static", - ) - - -@pytest.mark.parametrize( - ("metadata", "requested", "expected"), - [ - ({}, "auto", "static"), - ({"activation_scheme": "dynamic"}, "auto", "dynamic"), - ({"activation_scheme": "static"}, "auto", "static"), - ( - {"config_groups": {"group_0": {"input_activations": {"dynamic": True}}}}, - "auto", - "dynamic", - ), - ( - {"config_groups": {"group_0": {"input_activations": {"dynamic": False}}}}, - "auto", - "static", - ), - ( - {"config_groups": {"group_0": {"weights": {"dynamic": True}}}}, - "auto", - "static", - ), - ( - { - "config_groups": { - "group_0": {"input_activations": {"dynamic": False}}, - "group_1": {"input_activations": {"dynamic": True}}, - } - }, - "dynamic", - "dynamic", - ), - ( - { - "config_groups": { - "group_0": {"input_activations": {"dynamic": False}}, - "group_1": {"input_activations": {"dynamic": True}}, - } - }, - "static", - "static", - ), - ({"activation_scheme": "static"}, "dynamic", "dynamic"), - ({"activation_scheme": "dynamic"}, "static", "static"), - ( - {"config_groups": {"group_0": {"input_activations": {"dynamic": True}}}}, - "static", - "static", - ), - ], -) -def test_conversion_resolves_activation_metadata( - tmp_path, metadata, requested, expected -): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, - { - "layers.0.attention.to_q.weight_quantizer._amax": torch.tensor(224.0), - "layers.0.attention.to_q.input_quantizer._amax": torch.tensor(112.0), - }, - quantization_config=metadata, - class_name="ExampleTransformer2DModel", - ) - source_config = (source_dir / "config.json").read_bytes() - build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="none", - activation_scheme=requested, - ) - quant_config = json.loads((output_dir / "config.json").read_text())[ - "quantization_config" - ] - assert quant_config.get("activation_scheme", "static") == expected - for group in quant_config.get("config_groups", {}).values(): - if group.get("input_activations") is not None: - assert group["input_activations"]["dynamic"] is (expected == "dynamic") - tensors = load_file(output_dir / "model.safetensors") - assert ("layers.0.attention.to_q.input_scale" in tensors) is (expected == "static") - with safe_open(output_dir / "model.safetensors", framework="pt") as checkpoint: - assert json.loads(checkpoint.metadata()["quantization_config"]) == quant_config - assert (source_dir / "config.json").read_bytes() == source_config - - -@pytest.mark.parametrize("requested", ["auto", "dynamic"]) -def test_llada_static_source_requires_explicit_dynamic_output(tmp_path, requested): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, - {"layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0)}, - quantization_config={ - "config_groups": {"group_0": {"input_activations": {"dynamic": False}}} - }, - ) - kwargs = dict( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="llada-image", - activation_scheme=requested, - ) - if requested == "auto": - with pytest.raises(ValueError, match="set activation_scheme='dynamic'"): - build_modelopt_fp8_transformer(**kwargs) - assert not output_dir.exists() - else: - build_modelopt_fp8_transformer(**kwargs) - quant_config = json.loads((output_dir / "config.json").read_text())[ - "quantization_config" - ] - assert quant_config["activation_scheme"] == "dynamic" - assert "layers.0.attention.to_q.input_scale" not in load_file( - output_dir / "model.safetensors" - ) - - -@pytest.mark.parametrize( - "metadata", - [ - {"activation_scheme": "invalid"}, - {"config_groups": {"group_0": {"input_activations": {"dynamic": "false"}}}}, - { - "config_groups": { - "group_0": {"input_activations": {"dynamic": False}}, - "group_1": {"input_activations": {"dynamic": True}}, - } - }, - { - "activation_scheme": "dynamic", - "config_groups": {"group_0": {"input_activations": {"dynamic": False}}}, - }, - ], -) -def test_auto_rejects_invalid_or_conflicting_activation_metadata(tmp_path, metadata): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, - {"layers.0.attention.to_qkv.weight_quantizer._amax": torch.tensor(224.0)}, - quantization_config=metadata, - ) - with pytest.raises(ValueError, match="activation"): - build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="llada-image", - ) - assert not output_dir.exists() - - -@pytest.mark.parametrize("requested", ["static", "dynamic"]) -@pytest.mark.parametrize( - "metadata", - [ - {"config_groups": None}, - {"config_groups": {"group_0": None}}, - {"config_groups": {"group_0": {"input_activations": "invalid"}}}, - ], -) -def test_explicit_output_rejects_malformed_metadata_before_writing( - tmp_path, requested, metadata -): - source_dir, output_dir, backbone_path = _write_llada_modelopt_export( - tmp_path, - {"layers.0.attention.to_q.weight": torch.ones(2, 2)}, - {"layers.0.attention.to_q.weight_quantizer._amax": torch.tensor(224.0)}, - quantization_config=metadata, - class_name="ExampleTransformer2DModel", - ) - with pytest.raises(ValueError, match="ModelOpt.*dictionary"): - build_modelopt_fp8_transformer( - modelopt_hf_dir=str(source_dir), - modelopt_backbone_ckpt=str(backbone_path), - output_dir=str(output_dir), - model_type="none", - activation_scheme=requested, - ) - assert not output_dir.exists() diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index ff3f5f3052e7..f7fb99366f5b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -391,11 +391,10 @@ def test_modelopt_fp8_adapter_keeps_layerwise_offload_enabled(): assert server_args.dit_layerwise_offload is True -@pytest.mark.parametrize("layerwise_offload", [False, True]) -def test_online_fp8_disables_component_offload(layerwise_offload): +def test_modelopt_fp8_adapter_does_not_change_online_fp8_offload(): server_args = SimpleNamespace( dit_cpu_offload=True, - dit_layerwise_offload=layerwise_offload, + dit_layerwise_offload=False, quantization="fp8", ) @@ -404,25 +403,7 @@ def test_online_fp8_disables_component_offload(layerwise_offload): quant_config=Fp8Config(), ) - assert server_args.dit_cpu_offload is False - assert server_args.dit_layerwise_offload is layerwise_offload - - -@pytest.mark.parametrize("layerwise_offload", [False, True]) -def test_serialized_fp8_preserves_component_offload(layerwise_offload): - server_args = SimpleNamespace( - dit_cpu_offload=True, - dit_layerwise_offload=layerwise_offload, - quantization="fp8", - ) - - _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( - server_args=server_args, - quant_config=Fp8Config(is_checkpoint_fp8_serialized=True), - ) - assert server_args.dit_cpu_offload is True - assert server_args.dit_layerwise_offload is layerwise_offload def test_layerwise_capability_selects_layerwise_strategy_for_any_component(): diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image.py b/python/sglang/multimodal_gen/test/unit/test_llada_image.py index 5ac2aae5dafe..f27ececf60b4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image.py @@ -17,16 +17,11 @@ LLaDAImageQueryFormerModel, LLaDAImageSigVQModel, LLaDAImageTextProjectionModel, - LLaDAImageTransformerBlock, _LLaDAImageTransformer2DModel, ) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( - SP_STOCHASTIC_NOISE_KEY, FlowMatchEulerDiscreteScheduler, ) -from sglang.multimodal_gen.runtime.pipelines.llada_image import ( - LLaDAImageLatentPreparationStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( format_llada_image_prompt, ) @@ -236,274 +231,6 @@ def test_stochastic_scheduler_accepts_per_sample_generators(self): self.assertEqual(actual.shape, sample.shape) torch.testing.assert_close(actual, expected, rtol=0, atol=0) - def test_latent_stage_records_sp_stochastic_noise_slice(self): - scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=True) - stage = LLaDAImageLatentPreparationStage( - scheduler=scheduler, - transformer=None, - ) - batch = SimpleNamespace( - latents=torch.zeros(1, 2, 4, 3), - raw_latent_shape=torch.Size((1, 2, 4, 3)), - extra={}, - ) - parent_forward = ( - "sglang.multimodal_gen.runtime.pipelines.llada_image." - "LatentPreparationStage.forward" - ) - module = "sglang.multimodal_gen.runtime.pipelines.llada_image" - - with ( - patch(parent_forward, return_value=batch), - patch(f"{module}.get_sp_world_size", return_value=2), - patch(f"{module}.get_sp_parallel_rank", return_value=1), - ): - actual = stage.forward(batch, server_args=None) - - self.assertIs(actual, batch) - self.assertEqual( - batch.extra[SP_STOCHASTIC_NOISE_KEY], - { - "full_shape": (1, 2, 4, 3), - "dim": 2, - "start": 2, - "length": 2, - }, - ) - - with ( - patch(parent_forward, return_value=batch), - patch(f"{module}.get_sp_world_size", return_value=1), - ): - stage.forward(batch, server_args=None) - - self.assertNotIn(SP_STOCHASTIC_NOISE_KEY, batch.extra) - - batch.extra[SP_STOCHASTIC_NOISE_KEY] = {"stale": True} - stage.scheduler = FlowMatchEulerDiscreteScheduler(stochastic_sampling=False) - with ( - patch(parent_forward, return_value=batch), - patch(f"{module}.get_sp_world_size", return_value=2), - ): - stage.forward(batch, server_args=None) - - self.assertNotIn(SP_STOCHASTIC_NOISE_KEY, batch.extra) - - def test_stochastic_scheduler_matches_spatial_sp_slices(self): - full_shape = (1, 2, 4, 3) - - def make_stream(): - generators = [torch.Generator(device="cpu").manual_seed(71)] - initial = torch.randn( - full_shape, - generator=generators[0], - dtype=torch.float32, - ) - return generators, initial - - full_generators, full_sample = make_stream() - rank0_generators, rank0_full = make_stream() - rank1_generators, rank1_full = make_stream() - rank0_sample = rank0_full[:, :, :2, :].contiguous() - rank1_sample = rank1_full[:, :, 2:, :].contiguous() - - schedulers = [ - FlowMatchEulerDiscreteScheduler( - shift=3.0, - use_uniform_sigmas=True, - stochastic_sampling=True, - ) - for _ in range(3) - ] - for scheduler in schedulers: - scheduler.set_timesteps(4, device="cpu") - full_scheduler, rank0_scheduler, rank1_scheduler = schedulers - - def rank_batch(start): - return SimpleNamespace( - rollout=False, - did_sp_shard_latents=True, - raw_latent_shape=torch.Size(full_shape), - extra={ - SP_STOCHASTIC_NOISE_KEY: { - "full_shape": full_shape, - "dim": 2, - "start": start, - "length": 2, - } - }, - ) - - rank0_batch = rank_batch(0) - rank1_batch = rank_batch(2) - - for step_index in range(4): - full_model_output = full_sample / (step_index + 2) - rank0_model_output = full_model_output[:, :, :2, :].contiguous() - rank1_model_output = full_model_output[:, :, 2:, :].contiguous() - - full_sample = full_scheduler.step( - model_output=full_model_output, - timestep=full_scheduler.timesteps[step_index], - sample=full_sample, - generator=full_generators, - return_dict=False, - )[0] - rank0_sample = rank0_scheduler.step( - model_output=rank0_model_output, - timestep=rank0_scheduler.timesteps[step_index], - sample=rank0_sample, - generator=rank0_generators, - batch=rank0_batch, - return_dict=False, - )[0] - rank1_sample = rank1_scheduler.step( - model_output=rank1_model_output, - timestep=rank1_scheduler.timesteps[step_index], - sample=rank1_sample, - generator=rank1_generators, - batch=rank1_batch, - return_dict=False, - )[0] - - torch.testing.assert_close( - torch.cat([rank0_sample, rank1_sample], dim=2), - full_sample, - rtol=0, - atol=0, - ) - self.assertTrue( - torch.equal( - full_generators[0].get_state(), - rank0_generators[0].get_state(), - ) - ) - self.assertTrue( - torch.equal( - full_generators[0].get_state(), - rank1_generators[0].get_state(), - ) - ) - - def test_stochastic_scheduler_without_opt_in_keeps_local_draw(self): - scheduler = FlowMatchEulerDiscreteScheduler( - shift=3.0, - use_uniform_sigmas=True, - stochastic_sampling=True, - ) - scheduler.set_timesteps(4, device="cpu") - sample = torch.arange(12, dtype=torch.float32).reshape(1, 2, 2, 3) - model_output = sample / 16 - generator = [torch.Generator(device="cpu").manual_seed(23)] - reference_noise = torch.randn( - sample.shape, - generator=torch.Generator(device="cpu").manual_seed(23), - ) - current_sigma, next_sigma = scheduler.sigmas[:2] - expected = (1.0 - next_sigma) * ( - sample - current_sigma * model_output - ) + next_sigma * reference_noise - batch = SimpleNamespace( - rollout=False, - did_sp_shard_latents=True, - raw_latent_shape=torch.Size((1, 2, 4, 3)), - extra={}, - ) - - actual = scheduler.step( - model_output=model_output, - timestep=scheduler.timesteps[0], - sample=sample, - generator=generator, - batch=batch, - return_dict=False, - )[0] - - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - def test_stochastic_scheduler_rejects_invalid_sp_metadata(self): - full_shape = (1, 2, 4, 3) - valid = { - "full_shape": full_shape, - "dim": 2, - "start": 0, - "length": 2, - } - cases = ( - ("metadata type", "bad", True, full_shape, 1), - ("missing field", {"full_shape": full_shape}, True, full_shape, 1), - ( - "full shape type", - {**valid, "full_shape": list(full_shape)}, - True, - full_shape, - 1, - ), - ( - "full shape ndim", - {**valid, "full_shape": (1, 2, 4)}, - True, - full_shape, - 1, - ), - ( - "full shape value", - {**valid, "full_shape": (1, 2, 0, 3)}, - True, - full_shape, - 1, - ), - ("without sharding", valid, False, full_shape, 1), - ("batch dimension", {**valid, "dim": 0}, True, full_shape, 1), - ("dimension type", {**valid, "dim": 2.0}, True, full_shape, 1), - ("dimension range", {**valid, "dim": 4}, True, full_shape, 1), - ("slice type", {**valid, "start": 0.0}, True, full_shape, 1), - ("negative start", {**valid, "start": -1}, True, full_shape, 1), - ("slice bounds", {**valid, "start": 3}, True, full_shape, 1), - ("local length", {**valid, "length": 1}, True, full_shape, 1), - ("raw shape", valid, True, (1, 2, 6, 3), 1), - ( - "sample shape", - {**valid, "full_shape": (1, 3, 4, 3)}, - True, - (1, 3, 4, 3), - 1, - ), - ("generator count", valid, True, full_shape, 2), - ) - - for name, metadata, did_shard, raw_shape, generator_count in cases: - with self.subTest(name=name): - scheduler = FlowMatchEulerDiscreteScheduler( - shift=3.0, - use_uniform_sigmas=True, - stochastic_sampling=True, - ) - scheduler.set_timesteps(4, device="cpu") - sample = torch.zeros(1, 2, 2, 3) - generators = [ - torch.Generator(device="cpu").manual_seed(seed) - for seed in range(generator_count) - ] - batch = SimpleNamespace( - rollout=False, - did_sp_shard_latents=did_shard, - raw_latent_shape=torch.Size(raw_shape), - extra={SP_STOCHASTIC_NOISE_KEY: metadata}, - ) - - with self.assertRaisesRegex( - ValueError, "Invalid SP stochastic noise metadata" - ): - scheduler.step( - model_output=torch.zeros_like(sample), - timestep=scheduler.timesteps[0], - sample=sample, - generator=generators, - batch=batch, - return_dict=False, - ) - def test_pipeline_config_rejects_unaligned_resolution(self): config = LLaDAImagePipelineConfig() with self.assertRaisesRegex(ValueError, "height must be divisible by 16"): @@ -522,103 +249,6 @@ def test_service_sampling_params_use_few_step_defaults_without_vq_mode(self): self.assertNotIn("generation_mode", field_names) self.assertNotIn("vq_token_ids", field_names) - def test_dit_block_forwards_sequence_parallel_controls(self): - class CaptureAttention(torch.nn.Module): - def __init__(self): - super().__init__() - self.received = None - - def forward( - self, - hidden_states, - attention_mask, - freqs_cis, - num_replicated_suffix=0, - skip_sequence_parallel_override=False, - ): - self.received = ( - num_replicated_suffix, - skip_sequence_parallel_override, - ) - return torch.zeros_like(hidden_states) - - block = object.__new__(LLaDAImageTransformerBlock) - torch.nn.Module.__init__(block) - block.modulation = False - block.attention = CaptureAttention() - block.attention_norm1 = torch.nn.Identity() - block.attention_norm2 = torch.nn.Identity() - block.ffn_norm1 = torch.nn.Identity() - block.ffn_norm2 = torch.nn.Identity() - block.feed_forward = torch.nn.Identity() - - block( - torch.ones(1, 4, 8), - attention_mask=None, - freqs_cis=torch.empty(0), - num_replicated_suffix=32, - skip_sequence_parallel_override=True, - ) - - self.assertEqual(block.attention.received, (32, True)) - - def test_edit_sp_shards_images_and_replicates_conditions(self): - model_module = "sglang.multimodal_gen.runtime.models.dits.llada_image" - linear_module = "sglang.multimodal_gen.runtime.layers.linear" - with ( - patch(f"{model_module}.get_tp_world_size", return_value=1), - patch(f"{linear_module}.get_tp_group", return_value=None), - patch(f"{linear_module}.get_group_size", return_value=1), - patch(f"{linear_module}.get_group_rank", return_value=0), - patch( - "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", - return_value=1, - ), - patch( - "sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args", - return_value=SimpleNamespace(attention_backend="torch_sdpa"), - ), - ): - model = _LLaDAImageTransformer2DModel( - in_channels=4, - dim=64, - n_layers=1, - n_refiner_layers=1, - n_heads=2, - cap_feat_dim=8, - semantic_feat_dim=10, - axes_dims=(8, 12, 12), - axes_lens=(256, 32, 32), - ) - - noise_refiner = _CaptureSPBlock() - context_refiner = _CaptureSPBlock() - sigvq_refiner = _CaptureSPBlock() - main_block = _CaptureSPBlock() - model.noise_refiner = torch.nn.ModuleList([noise_refiner]) - model.context_refiner = torch.nn.ModuleList([context_refiner]) - model.sigvq_refiner = torch.nn.ModuleList([sigvq_refiner]) - model.layers = torch.nn.ModuleList([main_block]) - - with ( - patch(f"{model_module}.get_sp_world_size", return_value=2), - patch(f"{model_module}.get_sp_parallel_rank", return_value=0), - torch.no_grad(), - ): - model( - x=[torch.randn(4, 1, 4, 4)], - t=torch.tensor([0.5]), - cap_feats=[torch.randn(3, 8)], - glm_cap_feats=[torch.randn(5, 10)], - source_latents=[torch.randn(4, 1, 4, 4)], - ) - - self.assertEqual(noise_refiner.skip_values, [False]) - self.assertEqual(context_refiner.skip_values, [True]) - self.assertEqual(sigvq_refiner.skip_values, [True]) - self.assertEqual(main_block.skip_values, [False]) - self.assertEqual(main_block.replicated_suffixes, [96]) - def test_edit_skips_sigvq_refiner_for_empty_cfg_condition(self): model_module = "sglang.multimodal_gen.runtime.models.dits.llada_image" linear_module = "sglang.multimodal_gen.runtime.layers.linear" @@ -658,8 +288,6 @@ def test_edit_skips_sigvq_refiner_for_empty_cfg_condition(self): model.layers = torch.nn.ModuleList([main_block]) with ( - patch(f"{model_module}.get_sp_world_size", return_value=1), - patch(f"{model_module}.get_sp_parallel_rank", return_value=0), torch.no_grad(), ): model( @@ -671,71 +299,10 @@ def test_edit_skips_sigvq_refiner_for_empty_cfg_condition(self): ) self.assertEqual(noise_refiner.skip_values, [False]) - self.assertEqual(context_refiner.skip_values, [True]) + self.assertEqual(context_refiner.skip_values, [False]) self.assertEqual(sigvq_refiner.skip_values, []) self.assertEqual(main_block.skip_values, [False]) - def test_sp_uses_global_height_coordinates_for_generation_and_edit(self): - calls = [] - - def pad_with_ids( - features, grid_size, start, noise_value=None, sequence_multiple=32 - ): - del sequence_multiple - length = len(features) - calls.append((grid_size, start)) - return ( - features, - torch.zeros((length, 3), dtype=torch.int32), - torch.zeros(length, dtype=torch.bool), - length, - [noise_value] * length if noise_value is not None else None, - ) - - subject = SimpleNamespace( - _patchify_image=lambda _image, _patch, _f_patch: ( - torch.zeros(8, 1), - (1, 2, 4), - (1, 2, 4), - ), - _pad_with_ids=pad_with_ids, - ) - module = "sglang.multimodal_gen.runtime.models.dits.llada_image" - with ( - patch(f"{module}.get_sp_world_size", return_value=2), - patch(f"{module}.get_sp_parallel_rank", return_value=1), - ): - _LLaDAImageTransformer2DModel._prepare_t2i_sequences( - subject, - [torch.zeros(1, 1, 2, 4)], - cap_feats=None, - glm_features=None, - patch_size=1, - f_patch_size=1, - ) - - self.assertEqual(calls[-1], ((1, 2, 4), (1, 2, 0))) - - calls.clear() - with ( - patch(f"{module}.get_sp_world_size", return_value=2), - patch(f"{module}.get_sp_parallel_rank", return_value=1), - ): - _LLaDAImageTransformer2DModel._prepare_editing_sequences( - subject, - [torch.zeros(1, 1, 2, 4)], - cap_feats=[torch.zeros(1, 1)], - glm_cap_feats=[torch.zeros(1, 1)], - source_latents=[torch.zeros(1, 1, 2, 4)], - patch_size=1, - f_patch_size=1, - ) - - image_calls = [call for call in calls if call[0] == (1, 2, 4)] - self.assertEqual(len(image_calls), 2) - self.assertTrue(all(start[1:] == (2, 0) for _, start in image_calls)) - self.assertEqual(calls[-1], ((1, 1, 1), (35, 0, 0))) - if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py deleted file mode 100644 index ea864a595e4a..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_api.py +++ /dev/null @@ -1,168 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import unittest -from types import SimpleNamespace -from unittest.mock import AsyncMock, patch - -from fastapi import HTTPException - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import ( - _build_sampling_params_or_400, - _early_validate_edit_bounds, - edits, - generations, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( - ImageGenerationsRequest, -) - - -class TestLLaDAImageAPI(unittest.IsolatedAsyncioTestCase): - def setUp(self): - self.server_args = SimpleNamespace( - pipeline_config=LLaDAImagePipelineConfig(), - sp_degree=2, - input_save_path=None, - output_path=None, - llada_image_max_pixel_area=None, - llada_image_max_text_tokens=None, - llada_image_max_total_pixel_area=None, - ) - self.raw_request = SimpleNamespace(headers={}) - - async def test_generation_sp_rejects_multiple_outputs_before_scheduling(self): - request = ImageGenerationsRequest(prompt="a red car", n=2) - fallback = HTTPException(status_code=418, detail="request was not validated") - - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", - return_value=self.server_args, - ), - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", - side_effect=fallback, - ), - self.assertRaises(HTTPException) as context, - ): - await generations(request, self.raw_request) - - self.assertEqual(context.exception.status_code, 400) - self.assertIn("n=1", context.exception.detail) - - async def test_edit_sp_rejects_multiple_outputs_before_saving_input(self): - fallback = HTTPException(status_code=418, detail="request was not validated") - - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", - return_value=self.server_args, - ), - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", - new=AsyncMock(side_effect=fallback), - ), - self.assertRaises(HTTPException) as context, - ): - await edits( - raw_request=self.raw_request, - image=None, - image_array=None, - url=["https://example.com/source.png"], - url_array=None, - prompt="make the car blue", - n=2, - ) - - self.assertEqual(context.exception.status_code, 400) - self.assertIn("n=1", context.exception.detail) - - async def test_edit_rejects_oversized_area_before_source_processing(self): - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", - return_value=self.server_args, - ), - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", - new=AsyncMock(side_effect=AssertionError("source was processed")), - ), - self.assertRaises(HTTPException) as context, - ): - await edits( - raw_request=self.raw_request, - image=None, - image_array=None, - url=["https://example.com/source.png"], - url_array=None, - prompt="make the car blue", - size="4096x4096", - n=1, - ) - - self.assertEqual(context.exception.status_code, 400) - self.assertIn("exceeds the supported maximum", context.exception.detail) - - async def test_edit_rejects_multiple_sources_before_source_processing(self): - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.get_global_server_args", - return_value=self.server_args, - ), - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.save_image_to_path", - new=AsyncMock(side_effect=AssertionError("source was processed")), - ), - self.assertRaises(HTTPException) as context, - ): - await edits( - raw_request=self.raw_request, - image=None, - image_array=None, - url=[ - "https://example.com/source-1.png", - "https://example.com/source-2.png", - ], - url_array=None, - prompt="make the car blue", - n=1, - size=None, - ) - - self.assertEqual(context.exception.status_code, 400) - self.assertIn("exactly one source image", context.exception.detail) - - def test_build_sampling_params_maps_value_error_to_400(self): - with ( - patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", - side_effect=ValueError("bad request field"), - ), - self.assertRaises(HTTPException) as context, - ): - _build_sampling_params_or_400("rid", prompt="x") - - self.assertEqual(context.exception.status_code, 400) - self.assertEqual(context.exception.detail, "bad request field") - - def test_build_sampling_params_passthrough_on_success(self): - marker = object() - with patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.image_api.build_sampling_params", - return_value=marker, - ): - self.assertIs(_build_sampling_params_or_400("rid", prompt="x"), marker) - - def test_early_validate_edit_bounds_skips_absent_or_unparseable_size(self): - _early_validate_edit_bounds(self.server_args, None) - _early_validate_edit_bounds(self.server_args, "banana") - with self.assertRaises(HTTPException) as context: - _early_validate_edit_bounds(self.server_args, "4096x4096") - self.assertEqual(context.exception.status_code, 400) - - -if __name__ == "__main__": - unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py index fbbb54bb4103..a3154c33e45c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py @@ -94,39 +94,6 @@ def test_conditioning_mask_guard_fails_closed(self): SimpleNamespace(conditioning_mask_active=False) ) - def test_runtime_lifecycle_apis_reject_partial_coverage(self): - from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig - from sglang.multimodal_gen.runtime.managers.memory_managers.memory_occupation_controller import ( - MemoryOccupationController, - ) - from sglang.multimodal_gen.runtime.post_training.weights_updater import ( - WeightsUpdater, - ) - - self.assertTrue(PipelineConfig().supports_memory_release()) - self.assertTrue(PipelineConfig().supports_hot_weight_updates()) - config = LLaDAImagePipelineConfig() - self.assertFalse(config.supports_memory_release()) - self.assertFalse(config.supports_hot_weight_updates()) - - pipeline = SimpleNamespace(server_args=SimpleNamespace(pipeline_config=config)) - controller = MemoryOccupationController( - pipeline=pipeline, rank=0, use_fsdp_inference=False - ) - result = controller.release_memory_occupation() - self.assertFalse(result["success"]) - self.assertIn("does not support memory release", result["message"]) - self.assertFalse(controller.is_sleeping()) - - updater = object.__new__(WeightsUpdater) - updater.pipeline = pipeline - ok, message = updater.update_weights_from_disk("/unused/path") - self.assertFalse(ok) - self.assertIn("old checkpoint", message) - ok, message = updater.update_weights_from_tensor(named_tensors=[]) - self.assertFalse(ok) - self.assertIn("old checkpoint", message) - def test_forward_batch_declares_conditioning_text_lens_field(self): import dataclasses @@ -157,147 +124,12 @@ def test_stage_declares_auxiliary_component_residency(self): self.runner.component_names, ["queryformer", "text_projection"] ) - def test_text_runner_uses_sp_group_as_text_encoder_tp_group(self): - resolved_page_size = 16 - fake_worker = SimpleNamespace( - model_runner=SimpleNamespace(page_size=resolved_page_size), - model_config=object(), - get_memory_pool=lambda: (object(), object()), - alloc_memory_pool=lambda: None, - init_attention_backends=lambda: None, - init_cuda_graphs=lambda: None, - ) - srt_args_module = "sglang.srt.server_args.ServerArgs" - worker_module = "sglang.srt.managers.tp_worker.TpModelWorker" - with ( - patch( - worker_module, - return_value=fake_worker, - ) as worker_cls, - patch( - "sglang.srt.runtime_context.create_context", - side_effect=lambda *_args, **_kwargs: RuntimeContext(ParallelContext()), - ), - patch( - "sglang.srt.mem_cache.cache_init_params.CacheInitParams", - return_value=object(), - ) as cache_init_params_cls, - patch( - "sglang.srt.mem_cache.chunk_cache.ChunkCache", - return_value=object(), - ), - patch( - srt_args_module, - side_effect=lambda **kwargs: SimpleNamespace(page_size=None, **kwargs), - ) as srt_args_cls, - patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.get_local_torch_device", - return_value=torch.device("cpu"), - ), - patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.get_sp_parallel_rank", - return_value=1, - ), - patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.prepare_diffusers_component_path_for_loading", - side_effect=lambda path: f"resolved:{path}", - ), - ): - runner = LLaDAImageTextEncoderRunner( - model_root="/unused/model", - queryformer=object(), - text_projection=object(), - tokenizer=object(), - server_args=SimpleNamespace( - sp_degree=2, - nccl_port=29500, - trust_remote_code=True, - revision=None, - component_paths={ - "text_encoder": "/custom/text_encoder", - "tokenizer": "/custom/tokenizer", - }, - pipeline_config=SimpleNamespace( - text_encoder_mem_fraction_static=0.1 - ), - ), - ) - untrusted_runner = LLaDAImageTextEncoderRunner( - model_root="/unused/model", - queryformer=object(), - text_projection=object(), - tokenizer=object(), - server_args=SimpleNamespace( - sp_degree=2, - nccl_port=29500, - trust_remote_code=False, - revision="pinned-rev", - component_paths={}, - pipeline_config=SimpleNamespace( - text_encoder_mem_fraction_static=0.1 - ), - ), - ) - - requested_kwargs, untrusted_kwargs = [ - call.kwargs for call in srt_args_cls.call_args_list - ] - worker_srt_args = worker_cls.call_args_list[0].kwargs["server_args"] - self.assertIs(runner.server_args, worker_srt_args) - self.assertIs( - untrusted_runner.server_args, - worker_cls.call_args_list[1].kwargs["server_args"], - ) - self.assertIs(runner.worker, fake_worker) - self.assertEqual( - requested_kwargs["model_path"], "resolved:/custom/text_encoder" - ) - self.assertEqual( - requested_kwargs["tokenizer_path"], "resolved:/custom/tokenizer" - ) - self.assertEqual(untrusted_kwargs["model_path"], "/unused/model/text_encoder") - self.assertEqual(untrusted_kwargs["tokenizer_path"], "/unused/model/tokenizer") - self.assertIsNone(worker_srt_args.page_size) - self.assertEqual(runner.page_size, resolved_page_size) - self.assertEqual( - [call.kwargs["page_size"] for call in cache_init_params_cls.call_args_list], - [resolved_page_size, resolved_page_size], - ) - self.assertEqual(requested_kwargs["tp_size"], 2) - self.assertEqual(requested_kwargs["dp_size"], 2) - self.assertTrue(requested_kwargs["enable_dp_attention"]) - self.assertTrue(requested_kwargs["enable_dp_lm_head"]) - self.assertEqual(requested_kwargs["attn_cp_size"], 1) - self.assertEqual(requested_kwargs["ep_size"], 1) - self.assertEqual(requested_kwargs["moe_dp_size"], 2) - self.assertEqual(requested_kwargs["moe_dense_tp_size"], 1) - self.assertEqual(requested_kwargs["moe_a2a_backend"], "none") - self.assertEqual(requested_kwargs["max_running_requests"], 4) - parallel_state = worker_cls.call_args_list[0].kwargs["ps"] - self.assertEqual(parallel_state.tp_rank, 1) - self.assertEqual(parallel_state.tp_size, 2) - self.assertEqual(parallel_state.dp_rank, 1) - self.assertEqual(parallel_state.dp_size, 2) - self.assertEqual(parallel_state.attn_tp_rank, 0) - self.assertEqual(parallel_state.attn_tp_size, 1) - self.assertEqual(parallel_state.attn_cp_rank, 0) - self.assertEqual(parallel_state.attn_cp_size, 1) - self.assertEqual(parallel_state.attn_dp_rank, 1) - self.assertEqual(parallel_state.attn_dp_size, 2) - self.assertEqual(parallel_state.moe_ep_rank, 0) - self.assertEqual(parallel_state.moe_ep_size, 1) - self.assertEqual(parallel_state.moe_dp_rank, 1) - self.assertEqual(parallel_state.moe_dp_size, 2) - self.assertIs(requested_kwargs["trust_remote_code"], True) - self.assertIs(untrusted_kwargs["trust_remote_code"], False) - self.assertEqual(untrusted_kwargs["revision"], "pinned-rev") - def test_text_runner_scopes_singleton_attention_group_and_restores(self): import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state import sglang.srt.distributed.parallel_state as srt_parallel_state diffusion_group = object() - encoder_group = SimpleNamespace(world_size=2, rank_in_group=1) + encoder_group = SimpleNamespace(world_size=1, rank_in_group=0) encoder_attention_group = object() runner = object.__new__(LLaDAImageTextEncoderRunner) runner.runtime_context = RuntimeContext(ParallelContext()) @@ -361,76 +193,6 @@ class StopAfterPrefillAdder(Exception): self.assertEqual(prefill_adder_cls.call_args.args[0], runner.page_size) - def test_text_dp_metadata_covers_each_replica(self): - from sglang.srt.managers.schedule_policy import AddReqResult - - class StopAfterMetadata(Exception): - pass - - class FakeReq: - def __init__(self, *_args, **_kwargs): - self.kv = None - self.req_pool_idx = None - - def init_next_round_input(self, _tree_cache): - return None - - class FakePrefillAdder: - def __init__(self, *_args, **_kwargs): - self.can_run_list = [] - - def add_one_req(self, req, **_kwargs): - self.can_run_list.append(req) - return AddReqResult.CONTINUE - - def stop_after_metadata(): - raise StopAfterMetadata - - batch = SimpleNamespace(prepare_for_extend=stop_after_metadata) - runner = object.__new__(LLaDAImageTextEncoderRunner) - runner.tokenizer = lambda *_args, **_kwargs: SimpleNamespace( - input_ids=[[1, 2], [3]] - ) - runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) - runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) - runner.server_args = SimpleNamespace( - chunked_prefill_size=-1, - max_prefill_tokens=8192, - dp_size=2, - ) - runner.text_dp_attention = True - runner.page_size = 1 - runner.tree_cache = object() - runner.token_to_kv_pool_allocator = object() - runner.req_to_token_pool = object() - - with ( - patch( - "sglang.srt.sampling.sampling_params.SamplingParams", - return_value=SimpleNamespace(normalize=lambda _config: None), - ), - patch( - "sglang.srt.managers.schedule_batch.Req", - side_effect=FakeReq, - ), - patch( - "sglang.srt.managers.schedule_policy.PrefillAdder", - side_effect=FakePrefillAdder, - ), - patch( - "sglang.srt.managers.schedule_batch.ScheduleBatch.init_new", - return_value=batch, - ), - ): - with self.assertRaises(StopAfterMetadata): - runner._encode_impl( - ["positive", "negative"], - max_sequence_length=16, - ) - - self.assertEqual(batch.global_num_tokens, [7, 7]) - self.assertEqual(batch.global_num_tokens_for_logprob, [2, 2]) - def test_text_runner_restores_diffusion_groups_when_worker_init_fails(self): import sglang.srt.distributed.parallel_state as srt_parallel_state @@ -463,10 +225,6 @@ def fail_after_installing_encoder_groups(**_kwargs): "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.get_local_torch_device", return_value=torch.device("cpu"), ), - patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.get_sp_parallel_rank", - return_value=0, - ), patch.object(srt_parallel_state, "_TP", diffusion_group), patch.object(srt_parallel_state, "_ATTN_TP", diffusion_group), ): @@ -493,7 +251,7 @@ def fail_after_installing_encoder_groups(**_kwargs): class TestLLaDAImageRuntimeContext(unittest.TestCase): def test_encoder_preserves_diffusion_runtime(self): - """The embedded DP encoder must retain its topology inside the TP scope.""" + """The embedded encoder retains its context inside the diffusion TP scope.""" import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_state import sglang.srt.distributed.parallel_state as srt_state from sglang.srt import runtime_context as rc @@ -514,7 +272,7 @@ def test_encoder_preserves_diffusion_runtime(self): diffusion_log = diffusion_context.overrides_log() diffusion_buffer = rc.get_buffer("conditioning", object) diffusion_group = SimpleNamespace(world_size=1, rank_in_group=0) - encoder_group = SimpleNamespace(world_size=2, rank_in_group=1) + encoder_group = SimpleNamespace(world_size=1, rank_in_group=0) encoder_attention_group = SimpleNamespace(world_size=1, rank_in_group=0) observed = {} @@ -524,11 +282,11 @@ def make_args(**kwargs): def check_encoder_context(): context = rc.assert_published(observed["args"], role="scheduler") - self.assertEqual(rc.get_parallel().tp_size, 2) + self.assertEqual(rc.get_parallel().tp_size, 1) self.assertEqual(rc.get_parallel().attn_tp_size, 1) - self.assertEqual(rc.get_parallel().attn_dp_size, 2) + self.assertEqual(rc.get_parallel().attn_dp_size, 1) self.assertEqual(rc.get_parallel().moe_tp_size, 1) - self.assertEqual(rc.get_schedule().max_running_requests, 4) + self.assertEqual(rc.get_schedule().max_running_requests, 2) return context def allocate(): @@ -587,7 +345,6 @@ def check_diffusion_context(): f"{conditioning}.get_local_torch_device", return_value=torch.device("cpu"), ), - patch(f"{conditioning}.get_sp_parallel_rank", return_value=1), patch.object(mm_state, "_TP", diffusion_group), patch.object(srt_state, "_TP", diffusion_group), patch.object(srt_state, "_ATTN_TP", diffusion_group), @@ -598,7 +355,7 @@ def check_diffusion_context(): text_projection=object(), tokenizer=object(), server_args=SimpleNamespace( - sp_degree=2, + sp_degree=1, nccl_port=29500, trust_remote_code=False, revision=None, diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py index fb3cd0ff0508..b9975da0ba0b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -import os import unittest from types import SimpleNamespace -from unittest.mock import patch import torch from PIL import Image @@ -61,268 +59,6 @@ def __init__(self): self.assertEqual(actual.dtype, torch.bfloat16) torch.testing.assert_close(actual, expected, rtol=0, atol=0) - def test_edit_sp_shards_target_and_source_latents(self): - latents = torch.randn(1, 128, 4, 4) - source_latents = torch.randn(128, 1, 4, 4) - batch = SimpleNamespace( - batch_size=1, - condition_image=Image.new("RGB", (64, 64)), - image_embeds=None, - source_latents=[source_latents], - enable_sequence_shard=False, - ) - base_module = "sglang.multimodal_gen.configs.pipeline_configs.base" - llada_module = "sglang.multimodal_gen.configs.pipeline_configs.llada_image" - with ( - patch(f"{base_module}.get_sp_world_size", return_value=2), - patch(f"{base_module}.get_sp_parallel_rank", return_value=0), - patch(f"{llada_module}.get_sp_world_size", return_value=2), - ): - actual, did_shard = self.config.shard_latents_for_sp(batch, latents) - batch.did_sp_shard_latents = did_shard - source = self.config.prepare_pos_cond_kwargs( - batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float32 - )["source_latents"][0] - - self.assertEqual(actual.shape, (1, 128, 2, 4)) - torch.testing.assert_close(actual, latents[:, :, :2, :]) - self.assertTrue(did_shard) - self.assertEqual(source.shape, (128, 1, 2, 4)) - torch.testing.assert_close(source, source_latents[:, :, :2, :]) - - def test_generation_sp_rejects_latent_height_padding(self): - latents = torch.randn(1, 128, 63, 64) - batch = SimpleNamespace( - condition_image=None, - source_latents=None, - enable_sequence_shard=False, - ) - module = "sglang.multimodal_gen.configs.pipeline_configs.llada_image" - with ( - patch(f"{module}.get_sp_world_size", return_value=2), - self.assertRaisesRegex( - ValueError, "latent height 63 must be divisible by SP degree 2" - ), - ): - self.config.shard_latents_for_sp(batch, latents) - - def test_validates_supported_parallel_topology(self): - defaults = dict( - sp_degree=2, - ulysses_degree=2, - ring_degree=1, - tp_size=1, - dp_size=1, - cfg_parallel_degree=1, - num_gpus=2, - text_encoder_cpu_offload=False, - llada_image_max_pixel_area=None, - llada_image_max_text_tokens=None, - llada_image_max_total_pixel_area=None, - is_arg_explicitly_set=lambda name: False, - residency_mode=lambda name: "resident", - explicit_residency_mode=lambda name: None, - layerwise_offload_components=None, - cpu_offload_components=None, - ) - self.config.validate_server_args(SimpleNamespace(**defaults)) - - with ( - patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}), - self.assertRaisesRegex(ValueError, "does not support"), - ): - self.config.validate_server_args(SimpleNamespace(**defaults)) - - auto_offload = SimpleNamespace( - **(defaults | {"text_encoder_cpu_offload": True}) - ) - self.config.validate_server_args(auto_offload) - self.assertFalse(auto_offload.text_encoder_cpu_offload) - - class _AutoTunedArgs(SimpleNamespace): - def residency_mode(self, name): - components = self.layerwise_offload_components or [] - return "layerwise_offload" if name in components else "resident" - - auto_layerwise = _AutoTunedArgs( - **{key: value for key, value in defaults.items() if key != "residency_mode"} - ) - auto_layerwise.layerwise_offload_components = [ - "text_encoder", - "image_encoder", - ] - self.config.validate_server_args(auto_layerwise) - self.assertEqual(auto_layerwise.layerwise_offload_components, ["image_encoder"]) - - invalid_cases = [ - ({"ulysses_degree": 1, "ring_degree": 2}, "ring_degree=1"), - ( - {"sp_degree": 4, "ulysses_degree": 4, "num_gpus": 4}, - "supports only SP degrees 1 and 2", - ), - ({"cfg_parallel_degree": 2}, "TP=DP=CFG=1"), - ({"sp_degree": 1, "ulysses_degree": 1}, "num_gpus == sp_degree"), - ( - { - "text_encoder_cpu_offload": True, - "is_arg_explicitly_set": lambda name: True, - }, - "residency manager cannot offload", - ), - ( - {"residency_mode": lambda name: "component_offload"}, - "embedded text encoder to remain resident", - ), - ( - { - "residency_mode": lambda name: "layerwise_offload", - "is_arg_explicitly_set": ( - lambda name: name == "layerwise_offload_components" - ), - }, - "embedded text encoder to remain resident", - ), - ({"llada_image_max_pixel_area": 0}, "must be positive"), - ({"llada_image_max_total_pixel_area": 0}, "must be positive"), - ( - {"llada_image_max_text_tokens": 4000}, - "at most 3584", - ), - ] - for overrides, message in invalid_cases: - args = defaults | overrides - with ( - self.subTest(overrides=overrides), - self.assertRaisesRegex(ValueError, message), - ): - self.config.validate_server_args(SimpleNamespace(**args)) - - def test_validates_request_sampling_params_bounds(self): - def params(**overrides): - values = dict( - width=1024, - height=1024, - max_sequence_length=2048, - num_outputs_per_prompt=1, - diffusers_kwargs=None, - enable_cache_dit=None, - cache_dit_params=None, - cfg_gate_step=None, - attention_backend_override=None, - ) - values.update(overrides) - return SimpleNamespace(**values) - - def server(**overrides): - values = dict( - sp_degree=1, - llada_image_max_pixel_area=None, - llada_image_max_text_tokens=None, - llada_image_max_total_pixel_area=None, - ) - values.update(overrides) - return SimpleNamespace(**values) - - sp1 = server() - sp2 = server(sp_degree=2) - - self.config.validate_request_sampling_params(params(), sp1) - self.config.validate_request_sampling_params( - params(max_sequence_length=None), sp1 - ) - self.config.validate_request_sampling_params( - params(width=2048, height=2048, max_sequence_length=3584), sp2 - ) - self.config.validate_request_sampling_params( - params(num_outputs_per_prompt=10), sp1 - ) - self.config.validate_request_sampling_params( - params(enable_cache_dit=False), sp1 - ) - self.config.validate_request_sampling_params(params(cfg_gate_step=0.5), sp1) - self.config.validate_request_sampling_params( - params(attention_backend_override="fa"), sp1 - ) - self.config.validate_request_sampling_params( - params(attention_backend_override="TORCH_SDPA"), sp1 - ) - - invalid_cases = [ - (params(width=0), sp1, "positive width and height"), - (params(width=1000), sp1, "divisible by 16"), - (params(height=1040), sp2, "divisible by 32 at SP degree 2"), - (params(width=2064, height=2064), sp1, "exceeds the supported maximum"), - (params(max_sequence_length=0), sp1, "positive integer"), - (params(max_sequence_length=True), sp1, "positive integer"), - (params(max_sequence_length="2048"), sp1, "positive integer"), - ( - params(max_sequence_length=3585), - sp1, - "exceeds the embedded text encoder budget", - ), - ( - params(), - server(llada_image_max_pixel_area=512 * 512), - "exceeds the supported maximum of 262144", - ), - ( - params(), - server(llada_image_max_text_tokens=1024), - "budget of 1024 tokens", - ), - ( - params(width=2048, height=2048, num_outputs_per_prompt=10), - sp1, - "total output area", - ), - ( - params(num_outputs_per_prompt=2), - server(llada_image_max_total_pixel_area=1024 * 1024), - "total output area", - ), - ( - params(diffusers_kwargs={"max_sequence_length": 4096}), - sp1, - "exceeds the embedded text encoder budget", - ), - ( - params(diffusers_kwargs={"max_sequence_length": "x"}), - sp1, - "positive integer", - ), - (params(enable_cache_dit=True), sp1, "does not support enable_cache_dit"), - ( - params(cache_dit_params={}), - sp1, - "does not support cache_dit_params", - ), - ( - params(attention_backend_override="sage_attn"), - sp1, - "must be fa or torch_sdpa", - ), - ( - params(attention_backend_override=1), - sp1, - "must be fa or torch_sdpa", - ), - (params(cfg_gate_step=-0.1), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step=1.1), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step=True), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step=float("nan")), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step=float("inf")), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step=10**1000), sp1, "must be between 0.0 and 1.0"), - (params(cfg_gate_step="0.5"), sp1, "must be between 0.0 and 1.0"), - ] - for sampling_params, server_args, message in invalid_cases: - with ( - self.subTest(message=message), - self.assertRaisesRegex(ValueError, message), - ): - self.config.validate_request_sampling_params( - sampling_params, server_args - ) - if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_validation.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_validation.py deleted file mode 100644 index 65250f4233ca..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_validation.py +++ /dev/null @@ -1,234 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import json -import tempfile -import unittest -from contextlib import nullcontext -from pathlib import Path -from types import SimpleNamespace -from unittest.mock import patch - -import torch -from fastapi import HTTPException - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import ( - _build_sampling_params_or_400, - _early_validate_edit_bounds, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( - ImageGenerationsRequest, -) -from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request -from sglang.multimodal_gen.runtime.models.dits.llada_image import ( - LLaDAImageRopeEmbedder, - _LLaDAImageTransformer2DModel, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( - _resolve_llada_image_component_path, -) -from sglang.multimodal_gen.runtime.server_args import ServerArgs - - -class TestLLaDAImageValidation(unittest.TestCase): - def setUp(self): - self.directory = tempfile.TemporaryDirectory() - self.addCleanup(self.directory.cleanup) - self.model_root = Path(self.directory.name) - (self.model_root / "model_index.json").write_text( - json.dumps( - {"_class_name": "LLaDAImagePipeline", "_diffusers_version": "0.33.1"} - ) - ) - for component in ("transformer", "vae"): - (self.model_root / component).mkdir() - self.config = LLaDAImagePipelineConfig() - - def server_args(self, **overrides): - defaults = dict( - model_path=str(self.model_root), - pipeline_class_name="LLaDAImagePipeline", - pipeline_config=self.config, - performance_mode="manual", - num_gpus=1, - tp_size=1, - sp_degree=1, - ulysses_degree=1, - ring_degree=1, - enable_cfg_parallel=False, - dit_cpu_offload=False, - text_encoder_cpu_offload=False, - image_encoder_cpu_offload=False, - vae_cpu_offload=False, - attention_backend="torch_sdpa", - log_level="error", - ) - return ServerArgs.from_kwargs(**(defaults | overrides)) - - def sampling_params(self, args, width, height): - request = ImageGenerationsRequest( - prompt="a red cube", size=f"{width}x{height}", seed=42 - ) - with patch( - "sglang.multimodal_gen.runtime.entrypoints.openai.utils.get_global_server_args", - return_value=args, - ): - return _build_sampling_params_or_400( - "llada-image-validation", - prompt=request.prompt, - size=request.size, - seed=request.seed, - num_outputs_per_prompt=1, - generator_device="cpu", - save_output=False, - ) - - def test_rejects_normalized_embedded_text_encoder_overrides(self): - weight_path = str(self.model_root / "encoder.safetensors") - overrides = [ - {"component_paths": {"text_encoder": weight_path}}, - {"component_weights_paths": {"text_encoder": weight_path}}, - {"component_precisions": {"text-encoder": "fp32"}}, - {"component_quantizations": {"text-encoder": "FP8"}}, - {"component_attention_backends": "text-encoder=torch_sdpa"}, - {"component_direct_gpu_weight_loading": {"text-encoder": True}}, - ] - for override in overrides: - cuda_capability = ( - patch( - "sglang.multimodal_gen.runtime.server_args.server_args.current_platform.is_cuda", - return_value=True, - ) - if "component_direct_gpu_weight_loading" in override - else nullcontext() - ) - with ( - self.subTest(override=override), - cuda_capability, - self.assertRaisesRegex( - ValueError, "embedded text_encoder does not support component_" - ), - ): - self.server_args(**override) - - def test_preserves_encoder_directory_and_auxiliary_overrides(self): - encoder_path = self.model_root / "alternate_encoder" - encoder_path.mkdir() - (encoder_path / "config.json").write_text("{}") - auxiliary_paths = { - name: str(self.model_root / f"{name}.safetensors") - for name in ("queryformer", "text_projection", "sigvq") - } - args = self.server_args( - component_paths={"text_encoder": str(encoder_path)} | auxiliary_paths, - component_precisions={"queryformer": "fp32"}, - component_quantizations={"sigvq": "fp8"}, - component_attention_backends={"queryformer": "torch_sdpa"}, - ) - - self.assertEqual(args.component_weights_paths, auxiliary_paths) - self.assertEqual(args.component_precisions, {"queryformer": "fp32"}) - self.assertEqual(args.component_quantizations, {"sigvq": "fp8"}) - self.assertEqual( - args.component_attention_backends, {"queryformer": "torch_sdpa"} - ) - self.assertEqual( - _resolve_llada_image_component_path( - str(self.model_root), args, "text_encoder" - ), - str(encoder_path), - ) - - def test_rejects_disaggregation_during_server_args_construction(self): - args = self.server_args(disagg_role="monolithic") - self.assertEqual(args.disagg_role.value, "monolithic") - for role in ("encoder", "denoiser", "decoder"): - with ( - self.subTest(role=role), - self.assertRaisesRegex( - ValueError, "only supports monolithic deployment" - ), - ): - self.server_args(disagg_role=role) - - def test_rejects_each_spatial_axis_before_scheduling_or_edit_loading(self): - for sp_degree in (1, 2): - args = self.server_args( - num_gpus=sp_degree, - sp_degree=sp_degree, - ulysses_degree=sp_degree, - ) - stride = self.config.latent_scale_factor - max_height = self.config.dit_config.arch_config.axes_lens[1] * stride - max_width = self.config.dit_config.arch_config.axes_lens[2] * stride - for width, height in ( - (max_width + stride, stride * sp_degree), - (stride, max_height + stride * sp_degree), - ): - with self.subTest(sp_degree=sp_degree, width=width, height=height): - with self.assertRaises(HTTPException) as generation: - self.sampling_params(args, width, height) - self.assertEqual(generation.exception.status_code, 400) - self.assertIn("RoPE", generation.exception.detail) - with self.assertRaises(HTTPException) as editing: - _early_validate_edit_bounds(args, f"{width}x{height}") - self.assertEqual(editing.exception.status_code, 400) - self.assertIn("RoPE", editing.exception.detail) - - def test_spatial_limits_follow_checkpoint_config_before_latent_allocation(self): - self.config.dit_config.arch_config.axes_lens = (32768, 4, 8) - args = self.server_args() - for width, height in ((128, 64), (128, 16), (16, 64)): - with self.subTest(width=width, height=height): - self.sampling_params(args, width, height) - shape = self.config.prepare_latent_shape( - SimpleNamespace(width=width, height=height), 1, 1 - ) - self.assertEqual(shape[-2:], (height // 16, width // 16)) - for width, height in ((144, 16), (16, 80)): - with self.subTest(width=width, height=height): - with self.assertRaisesRegex(ValueError, "RoPE"): - self.config.prepare_latent_shape( - SimpleNamespace(width=width, height=height), 1, 1 - ) - with self.assertRaises(HTTPException) as context: - self.sampling_params(args, width, height) - self.assertEqual(context.exception.status_code, 400) - self.assertIn("RoPE", context.exception.detail) - - def test_exact_spatial_boundaries_fit_rope_with_global_sp_positions(self): - model = object.__new__(_LLaDAImageTransformer2DModel) - rope = LLaDAImageRopeEmbedder(256.0, (2, 2, 2), (32768, 1024, 1024)) - module = "sglang.multimodal_gen.runtime.models.dits.llada_image" - for sp_degree in (1, 2): - args = self.server_args( - num_gpus=sp_degree, - sp_degree=sp_degree, - ulysses_degree=sp_degree, - ) - for width, height in ((16384, 16 * sp_degree), (16, 16384)): - with self.subTest(sp_degree=sp_degree, width=width, height=height): - sampling = self.sampling_params(args, width, height) - batch = prepare_request(args, sampling) - shape = self.config.prepare_latent_shape(batch, 1, 1) - latent = torch.zeros(shape[1], 1, shape[2] // sp_degree, shape[3]) - with ( - patch(f"{module}.get_sp_world_size", return_value=sp_degree), - patch( - f"{module}.get_sp_parallel_rank", - return_value=sp_degree - 1, - ), - ): - image_sequence, _, _, _ = model._prepare_t2i_sequences( - [latent], [torch.zeros(1, 1)], None, 1, 1 - ) - position_ids = image_sequence.position_ids[0] - spatial_axis = 2 if width == 16384 else 1 - self.assertEqual(position_ids[:, spatial_axis].max().item(), 1023) - self.assertTrue(torch.isfinite(rope(position_ids)).all()) - - -if __name__ == "__main__": - unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py index 4df2dc88b0e6..62c9f774e5c1 100644 --- a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py @@ -15,7 +15,6 @@ from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp8Config, - is_fp8_fnuz, ) from sglang.multimodal_gen.runtime.loader import fsdp_load from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan @@ -98,15 +97,7 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): state_dict, weight_ref = _make_serialized_fp8_checkpoint() checkpoint_weight = state_dict["qkv.weight"].clone() checkpoint_scales = state_dict["qkv.weight_scale"].clone() - scale_factor = 2 if is_fp8_fnuz() else 1 - expected_dtype = ( - torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn - ) - expected_weight = (checkpoint_weight.float() / scale_factor).to( - expected_dtype - ) - runtime_scales = checkpoint_scales * scale_factor - expected_max_scale = runtime_scales.max() + expected_max_scale = checkpoint_scales.max() with patch( "sglang.multimodal_gen.runtime.layers.quantization." @@ -133,21 +124,16 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): # Both paths rebind the runtime weight transposed. CUTLASS can # consume a channelwise scale, so it preserves the checkpoint's - # FP8 shards in the native format. The fallback uses one max scale. + # FP8 shards; the fallback requantizes them to one max scale. weight = model.qkv.weight - self.assertEqual(weight.dtype, expected_dtype) + self.assertEqual(weight.dtype, torch.float8_e4m3fn) self.assertEqual(tuple(weight.shape), (_IN_FEATURES, 2 * _SHARD_OUT)) weight_scale = model.qkv.weight_scale.flatten() if cutlass_supported: expected_scales = torch.repeat_interleave( - runtime_scales, _SHARD_OUT - ) - torch.testing.assert_close( - weight.t().view(torch.int8), - expected_weight.view(torch.int8), - rtol=0, - atol=0, + checkpoint_scales, _SHARD_OUT ) + self.assertTrue(torch.equal(weight.t(), checkpoint_weight)) else: expected_scales = expected_max_scale.expand(weight_scale.numel()) torch.testing.assert_close( @@ -157,7 +143,7 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): ) torch.testing.assert_close( model.qkv.input_scale.flatten().max(), - torch.tensor(0.5 * scale_factor), + torch.tensor(0.5), check_device=False, ) @@ -171,7 +157,7 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): dequant, weight_ref, rtol=0.5, - atol=float(checkpoint_scales.max()) * 8, + atol=float(expected_max_scale) * 8, ) # Layerwise offload contract: the component lands on CPU diff --git a/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py b/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py index 1ed07b33db8e..aa74cabaa788 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py +++ b/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py @@ -41,7 +41,6 @@ def harness(monkeypatch): monkeypatch.setenv("SGLANG_GEN_BASELINE", "0") monkeypatch.setenv("SGLANG_SKIP_CONSISTENCY", "0") monkeypatch.setattr(test_server_common.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr(test_server_common.current_platform, "is_hip", lambda: False) scenario = ScenarioConfig( stages_ms={"DenoisingStage": 10}, denoise_step_ms={0: 5, 1: 5}, @@ -125,31 +124,6 @@ def test_each_request_failure_fails_case(harness, monkeypatch, bad_request, fail assert [r["request_index"] for r in runner._perf_results] == expected -@pytest.mark.parametrize("metric", ["performance", "load_peak", "runtime_peak"]) -def test_hip_performance_outliers_warn(harness, monkeypatch, caplog, metric): - _, case = harness - monkeypatch.setattr(test_server_common.current_platform, "is_cuda", lambda: False) - monkeypatch.setattr(test_server_common.current_platform, "is_hip", lambda: True) - baseline = test_server_common.BASELINE_CONFIG - scenario = baseline.scenarios[case.id] - validator = test_server_common.PerformanceValidator( - scenario, baseline.tolerances, baseline.step_fractions - ) - record = _perf_record() - if metric == "performance": - record.total_duration_ms = 10000 - validator.validate(record, case.sampling_params.num_frames) - else: - record.memory_snapshots[metric]["peak_reserved_mb"] = 10000 - validator.validate_peak_vram( - validator.collect_metrics(record), - scenario.load_peak_vram_mb, - scenario.runtime_peak_vram_mb, - ) - - assert "AMD PERF WARNING" in caplog.text - - def test_both_requests_pass(harness, monkeypatch): runner, case = harness monkeypatch.setattr( @@ -354,31 +328,3 @@ def test_h3_cases_check_two_short_requests_and_audio(): assert case.sampling_params.expect_audio_output assert case.sampling_params.extras["num_inference_steps"] <= 8 assert case.sampling_params.extras["target"]["duration_seconds"] == 4.0 - - -@pytest.mark.parametrize("estimate", [None, 600.0]) -def test_server_teardown_accepts_explicit_case_estimates( - harness, monkeypatch, estimate -): - _, case = harness - case = replace(case, run_perf_check=False, estimated_full_test_time_s=estimate) - monkeypatch.setattr( - test_server_common, - "BASELINE_CONFIG", - replace(test_server_common.BASELINE_CONFIG, scenarios={}), - ) - missing = set() - monkeypatch.setattr(test_server_common, "_MISSING_ESTIMATED_TIME_CASES", missing) - context = Mock() - manager = Mock() - manager.start.return_value = context - monkeypatch.setattr(test_server_common, "ServerManager", Mock(return_value=manager)) - monkeypatch.setattr(test_server_common, "get_dynamic_server_port", lambda: 19000) - - fixture = test_server_common.diffusion_server.__wrapped__(case) - assert next(fixture) is context - with pytest.raises(StopIteration): - next(fixture) - - context.cleanup.assert_called_once_with() - assert (case.id in missing) == (estimate is None) diff --git a/python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py b/python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py index 5c2c568a4e98..6c4214177215 100644 --- a/python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py @@ -117,42 +117,21 @@ def fake_run_standalone_file(suite, rel, target_dir, extra_filter=None): assert exit_code == 1 -def test_two_gpu_cases_have_full_test_time_estimates(): - """Every two-GPU case needs a baseline or explicit scheduling estimate.""" +def test_two_gpu_cases_have_h100_full_test_time_estimates(): + """Every 2-gpu case must have an h100 estimated_full_test_time_s. + + NVIDIA LPT sharding falls back to 300s when the field is missing, which + unbalances the three 2-gpu partitions. + """ scenarios = json.loads(_H100_BASELINE_PATH.read_text())["scenarios"] missing = [ case.id for case in TWO_GPU_CASES if scenarios.get(case.id, {}).get("estimated_full_test_time_s") is None - and case.estimated_full_test_time_s is None ] assert missing == [] -@pytest.mark.parametrize("scenario_estimate", [None, 420.0]) -def test_explicit_case_estimate_is_used_without_a_performance_baseline( - monkeypatch, scenario_estimate -): - case = SimpleNamespace(id="functional", estimated_full_test_time_s=600.0) - scenarios = ( - {} - if scenario_estimate is None - else {case.id: SimpleNamespace(estimated_full_test_time_s=scenario_estimate)} - ) - monkeypatch.setattr( - run_suite, "BASELINE_CONFIG", SimpleNamespace(scenarios=scenarios) - ) - monkeypatch.setattr( - run_suite, - "PARAMETRIZED_CASE_GROUPS", - {"1-gpu": [("test_functional.py", [case])]}, - ) - - expected = 600.0 if scenario_estimate is None else scenario_estimate - assert run_suite.get_case_est_time(case.id) == expected - assert run_suite.get_case_est_time("unknown") == run_suite.DEFAULT_EST_TIME_SECONDS - - def test_qwen_quality_variants_use_the_same_generation_request(): cases = {case.id: case for case in TWO_GPU_CASES} lossless = cases["qwen_image_t2i_2_gpus"].sampling_params diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index d4fefcf9bf6b..810b7b158536 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -100,7 +100,6 @@ def _loads(input_str, _flags=None): from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( TransformerQuantLoadSpec, _Flux2Nvfp4FallbackAdapter, - _ModelOptFp8OffloadAdapter, _needs_device_weight_postprocess, _resolve_quant_config, _resolve_weight_override_quantization, @@ -116,7 +115,6 @@ def _loads(input_str, _flags=None): from sglang.multimodal_gen.runtime.models.dits.flux_2 import ( Flux2Transformer2DModel, ) -from sglang.multimodal_gen.runtime.models.dits.llada_image import LLaDAImageFeedForward from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import MiniMaxH3DiTModel from sglang.multimodal_gen.runtime.models.dits.qwen_image import ( QwenImageTransformer2DModel, @@ -166,147 +164,81 @@ def _make_quant_config(name: str, **attrs): return quant_config -class _FakeOnlineFp8Config: - is_checkpoint_fp8_serialized = False - - @classmethod - def get_name(cls): - return "fp8" - - class TestTransformerQuantHelpers(unittest.TestCase): def test_modelopt_fp8_packed_cutlass_preserves_checkpoint_shard_scales(self): - for use_fnuz in (False, True): - with ( - self.subTest(use_fnuz=use_fnuz), - patch( - "sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant." - "is_fp8_fnuz", - return_value=use_fnuz, - ), - ): - method = ModelOptFp8LinearMethod( - ModelOptFp8Config(is_checkpoint_fp8_serialized=True) - ) - method.cutlass_fp8_supported = True - layer = torch.nn.Module() - layer.logical_widths = [2, 2, 2] - values = torch.arange(24, dtype=torch.float32).reshape(6, 4) - 12 - values[0, 0] = -0.0 - weight = values.to(torch.float8_e4m3fn) - checkpoint_scales = torch.tensor([0.1, 0.2, 0.3]) - input_scales = torch.tensor([0.5, 1.0, 1.5]) - layer.register_parameter( - "weight", torch.nn.Parameter(weight.clone(), requires_grad=False) - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(checkpoint_scales.clone(), requires_grad=False), - ) - layer.register_parameter( - "input_scale", - torch.nn.Parameter(input_scales.clone(), requires_grad=False), - ) + method = ModelOptFp8LinearMethod( + ModelOptFp8Config(is_checkpoint_fp8_serialized=True) + ) + method.cutlass_fp8_supported = True + layer = torch.nn.Module() + layer.logical_widths = [2, 2, 2] + weight = ( + torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.float8_e4m3fn) + ) + layer.register_parameter( + "weight", torch.nn.Parameter(weight.clone(), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter( + torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32), + requires_grad=False, + ), + ) + layer.register_parameter( + "input_scale", + torch.nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=False), + ) - method.process_weights_after_loading(layer) + method.process_weights_after_loading(layer) - scale_factor = 2 if use_fnuz else 1 - expected_dtype = ( - torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn - ) - expected_weight = (weight.float() / scale_factor).to(expected_dtype) - self.assertEqual(layer.weight.dtype, expected_dtype) - torch.testing.assert_close( - layer.weight.view(torch.int8), - expected_weight.t().view(torch.int8), - rtol=0, - atol=0, - ) - channel_scales = checkpoint_scales.repeat_interleave(2).view(-1, 1) - torch.testing.assert_close( - layer.weight_scale, channel_scales * scale_factor, rtol=0, atol=0 - ) - torch.testing.assert_close( - layer.input_scale, input_scales.max() * scale_factor, rtol=0, atol=0 - ) - torch.testing.assert_close( - layer.weight.t().float() * layer.weight_scale, - weight.float() * channel_scales, - rtol=0, - atol=0, - ) + torch.testing.assert_close(layer.weight, weight.t(), rtol=0, atol=0) + torch.testing.assert_close( + layer.weight_scale, + torch.tensor([[0.1], [0.1], [0.2], [0.2], [0.3], [0.3]]), + ) + torch.testing.assert_close(layer.input_scale, torch.tensor(1.0)) def test_modelopt_fp8_packed_cutlass_requantizes_incomplete_shard_scales(self): - for use_fnuz in (False, True): - with ( - self.subTest(use_fnuz=use_fnuz), - patch( - "sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant." - "is_fp8_fnuz", - return_value=use_fnuz, + method = ModelOptFp8LinearMethod( + ModelOptFp8Config(is_checkpoint_fp8_serialized=True) + ) + method.cutlass_fp8_supported = True + layer = torch.nn.Module() + layer.logical_widths = [2, 2, 2] + weight = ( + torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.float8_e4m3fn) + ) + layer.register_parameter( + "weight", torch.nn.Parameter(weight.clone(), requires_grad=False) + ) + layer.register_parameter( + "weight_scale", + torch.nn.Parameter( + torch.tensor( + [0.1, torch.finfo(torch.float32).min, 0.3], + dtype=torch.float32, ), - ): - method = ModelOptFp8LinearMethod( - ModelOptFp8Config(is_checkpoint_fp8_serialized=True) - ) - method.cutlass_fp8_supported = True - layer = torch.nn.Module() - layer.logical_widths = [2, 2, 2] - weight = ( - torch.arange(24, dtype=torch.float32) - .reshape(6, 4) - .to(torch.float8_e4m3fn) - ) - checkpoint_scales = torch.tensor( - [0.1, torch.finfo(torch.float32).min, 0.3], dtype=torch.float32 - ) - layer.register_parameter( - "weight", torch.nn.Parameter(weight.clone(), requires_grad=False) - ) - layer.register_parameter( - "weight_scale", - torch.nn.Parameter(checkpoint_scales.clone(), requires_grad=False), - ) - layer.register_parameter( - "input_scale", - torch.nn.Parameter(torch.ones(3), requires_grad=False), - ) - scale_factor = 2 if use_fnuz else 1 - expected_dtype = ( - torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn - ) - expected_weight = (weight.float() / scale_factor).to(expected_dtype) - expected_scales = checkpoint_scales * scale_factor - - def requantize_native(runtime_weight, scales, logical_widths): - self.assertEqual(runtime_weight.dtype, expected_dtype) - self.assertEqual(logical_widths, layer.logical_widths) - torch.testing.assert_close( - runtime_weight, expected_weight, rtol=0, atol=0 - ) - torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0) - return scales.max(), runtime_weight.clone() - - with patch( - "sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant." - "requantize_with_max_scale", - side_effect=requantize_native, - ) as requantize: - method.process_weights_after_loading(layer) - - requantize.assert_called_once() - torch.testing.assert_close( - layer.weight, expected_weight.t(), rtol=0, atol=0 - ) - torch.testing.assert_close( - layer.weight_scale, - expected_scales.max().expand(6, 1), - rtol=0, - atol=0, - ) - torch.testing.assert_close( - layer.input_scale, torch.tensor(float(scale_factor)) - ) + requires_grad=False, + ), + ) + layer.register_parameter( + "input_scale", + torch.nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=False), + ) + + with patch( + "sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant." + "requantize_with_max_scale", + return_value=(torch.tensor(0.3), weight.clone()), + ) as requantize: + method.process_weights_after_loading(layer) + + requantize.assert_called_once() + torch.testing.assert_close(layer.weight, weight.t(), rtol=0, atol=0) + torch.testing.assert_close( + layer.weight_scale, torch.full((6, 1), 0.3), rtol=0, atol=0 + ) def test_modelopt_packed_layer_requires_consistent_shard_precision(self): prefix = "blocks.0.attn.to_qkv" @@ -688,28 +620,6 @@ def test_modelopt_fp4_preserves_explicit_attention_backend(self): self.assertIsNone(backend) - def test_explicit_fp8_quantization_uses_online_defaults(self): - config = _resolve_quant_config( - hf_config={}, - server_args=self._make_server_args(quantization="fp8"), - safetensors_list=[], - component_model_path="/unused/component/path", - ) - - self.assertEqual(config.get_name(), "fp8") - self.assertFalse(config.is_checkpoint_fp8_serialized) - self.assertEqual(config.activation_scheme, "dynamic") - - def test_online_fp8_disables_whole_dit_cpu_offload(self): - server_args = self._make_server_args(dit_cpu_offload=True) - - _ModelOptFp8OffloadAdapter._maybe_disable_incompatible_dit_offload_modes( - server_args=server_args, - quant_config=_FakeOnlineFp8Config(), - ) - - self.assertFalse(server_args.dit_cpu_offload) - def test_resolve_transformer_checkpoint_files_uses_single_override_file(self): with tempfile.NamedTemporaryFile(suffix=".safetensors") as f: server_args = self._make_server_args(transformer_weights_path=f.name) @@ -1975,33 +1885,6 @@ def test_explicit_modelopt_method_without_algorithm_is_preserved(self): quant_method, ) - @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) - @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) - @patch( - "sglang.multimodal_gen.runtime.layers.linear.get_tp_group", return_value=None - ) - def test_llada_image_online_fp8_quantizes_fused_linears( - self, - _mock_tp_group, - _mock_group_size, - _mock_group_rank, - ): - fp8_module = "sglang.multimodal_gen.runtime.layers.quantization.fp8" - with patch( - f"{fp8_module}.get_tensor_model_parallel_world_size", - return_value=1, - ): - feed_forward = LLaDAImageFeedForward( - dim=192, - quant_config=Fp8Config(), - prefix="layers.0.feed_forward", - ) - - self.assertIsInstance(feed_forward.w13.quant_method, Fp8LinearMethod) - self.assertIsInstance(feed_forward.w2.quant_method, Fp8LinearMethod) - self.assertEqual(feed_forward.w13.prefix, "layers.0.feed_forward.w13") - self.assertEqual(feed_forward.w2.prefix, "layers.0.feed_forward.w2") - @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) @patch( diff --git a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py index c7a1f4b6b9ed..a579e2035da8 100644 --- a/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py +++ b/python/sglang/multimodal_gen/tools/build_modelopt_fp8_transformer.py @@ -2,7 +2,7 @@ The core conversion path is model-agnostic: - read the ModelOpt diffusers transformer export -- rebuild per-layer `weight_scale` tensors (and static `input_scale` tensors when requested) from `backbone.pt` +- rebuild per-layer `weight_scale` / `input_scale` tensors from `backbone.pt` - materialize SGLang-native `float8_e4m3fn` weights - preserve ModelOpt `ignore` layers in their original dtype @@ -297,26 +297,11 @@ def _map_hunyuanvideo_runtime_module_name(module_name: str) -> list[str]: return mapped_names -def _map_llada_image_runtime_module_name(module_name: str) -> list[str]: - replacements = ( - (r"^(.*\.attention)\.to_[qkv]$", r"\1.to_qkv"), - (r"^(.*\.feed_forward)\.w[13]$", r"\1.w13"), - ) - mapped_names: list[str] = [] - for pattern, replacement in replacements: - mapped = re.sub(pattern, replacement, module_name) - if mapped != module_name: - mapped_names.append(mapped) - return mapped_names - - def _get_runtime_module_name_mapper( *, model_type: str, class_name: str | None ) -> Callable[[str], list[str]] | None: if model_type == "hunyuan-video" or class_name == "HunyuanVideoTransformer3DModel": return _map_hunyuanvideo_runtime_module_name - if model_type == "llada-image" or class_name == "LLaDAImageTransformer2DModel": - return _map_llada_image_runtime_module_name return None @@ -361,70 +346,25 @@ def _preferred_module_name( return _module_name_variants(weight_name, runtime_name_mapper)[-1] -def _scale_key_candidates( - weight_name: str, - runtime_name_mapper: Callable[[str], list[str]] | None = None, -) -> list[str]: +def _scale_key_candidates(weight_name: str) -> list[str]: candidates = [weight_name] if weight_name.startswith("model.diffusion_model."): candidates.append( "velocity_model." + weight_name[len("model.diffusion_model.") :] ) - for module_name in _module_name_variants(weight_name, runtime_name_mapper): - candidate = f"{module_name}.weight" - if candidate not in candidates: - candidates.append(candidate) return candidates def _resolve_scale_key( weight_name: str, scale_map: Mapping[str, Mapping[str, torch.Tensor]], - runtime_name_mapper: Callable[[str], list[str]] | None = None, ) -> str | None: - for candidate in _scale_key_candidates(weight_name, runtime_name_mapper): + for candidate in _scale_key_candidates(weight_name): if candidate in scale_map: return candidate return None -def _validate_llada_image_fused_weight_scales( - weight_names: Iterable[str], - scale_map: Mapping[str, Mapping[str, torch.Tensor]], -) -> None: - scales_by_fused_module: dict[str, list[tuple[str, torch.Tensor]]] = defaultdict( - list - ) - for weight_name in weight_names: - mapped_modules = _map_llada_image_runtime_module_name(weight_name[:-7]) - if not mapped_modules: - continue - scale_key = _resolve_scale_key( - weight_name, scale_map, _map_llada_image_runtime_module_name - ) - if scale_key is None: - continue - fused_module = _preferred_module_name( - weight_name, _map_llada_image_runtime_module_name - ) - scales_by_fused_module[fused_module].append( - (scale_key, scale_map[scale_key]["weight_scale"]) - ) - - for fused_module, scale_entries in scales_by_fused_module.items(): - reference_scale = scale_entries[0][1] - if all(torch.equal(scale, reference_scale) for _, scale in scale_entries[1:]): - continue - scale_details = ", ".join( - f"{scale_key}={scale.item():g}" for scale_key, scale in scale_entries - ) - raise ValueError( - "LLaDA-Image FP8 conversion currently requires identical weight scales " - f"for parameters fused into {fused_module!r}; got {scale_details}. " - "Re-export the checkpoint with a fused weight observer for this module." - ) - - def _is_ltx2_x0_export( *, config: Mapping[str, object], @@ -533,7 +473,6 @@ def build_fp8_scale_map( model_state_dict: Mapping[str, torch.Tensor], *, maxbound: float = FP8_E4M3_MAXBOUND, - require_input_scale: bool = True, ) -> dict[str, dict[str, torch.Tensor]]: scale_map: dict[str, dict[str, torch.Tensor]] = {} for key, value in model_state_dict.items(): @@ -548,13 +487,10 @@ def build_fp8_scale_map( value.detach().to(torch.float32).reshape(1).cpu() / maxbound ) - required_scales = ( - {"weight_scale", "input_scale"} if require_input_scale else {"weight_scale"} - ) return { weight_name: scale_tensors for weight_name, scale_tensors in scale_map.items() - if required_scales <= set(scale_tensors) + if {"weight_scale", "input_scale"} <= set(scale_tensors) } @@ -607,44 +543,6 @@ def _load_selected_tensors( return tensors -def _get_modelopt_activation_scheme( - quant_config: Mapping, *, allow_mixed: bool = False -) -> str | None: - schemes = set() - declared = quant_config.get("activation_scheme") - if declared is not None: - if declared not in ("static", "dynamic"): - raise ValueError(f"Invalid ModelOpt activation_scheme {declared!r}.") - schemes.add(declared) - groups = quant_config.get("config_groups", {}) - if not isinstance(groups, dict): - raise ValueError("ModelOpt activation config_groups must be a dictionary.") - for group in groups.values(): - if not isinstance(group, dict): - raise ValueError( - "Each ModelOpt activation config group must be a dictionary." - ) - activations = group.get("input_activations") - if activations is None: - continue - if not isinstance(activations, dict): - raise ValueError("ModelOpt input_activations must be a dictionary.") - if "dynamic" not in activations: - continue - dynamic = activations["dynamic"] - if not isinstance(dynamic, bool): - raise ValueError("ModelOpt input activation dynamic must be a boolean.") - schemes.add("dynamic" if dynamic else "static") - if len(schemes) > 1: - if allow_mixed: - return None - raise ValueError( - "Conflicting ModelOpt activation schemes. " - "Set activation_scheme explicitly to choose the converted output." - ) - return next(iter(schemes), None) - - def build_modelopt_fp8_transformer( *, modelopt_hf_dir: str, @@ -653,7 +551,6 @@ def build_modelopt_fp8_transformer( base_transformer_dir: str | None = None, model_type: str = "auto", keep_bf16_patterns: Sequence[str] | None = None, - activation_scheme: str = "auto", maxbound: float = FP8_E4M3_MAXBOUND, overwrite: bool = False, ) -> dict[str, int]: @@ -683,28 +580,6 @@ def build_modelopt_fp8_transformer( source_weight_map=source_weight_map_all, ) class_name = config.get("_class_name") - is_llada_image = ( - model_type == "llada-image" or class_name == "LLaDAImageTransformer2DModel" - ) - source_activation_scheme = _get_modelopt_activation_scheme( - quant_config, allow_mixed=activation_scheme != "auto" - ) - if activation_scheme == "auto": - activation_scheme = source_activation_scheme or ( - "dynamic" if is_llada_image else "static" - ) - if activation_scheme not in {"static", "dynamic"}: - raise ValueError( - "activation_scheme must be one of 'auto', 'static', or 'dynamic', " - f"got {activation_scheme!r}." - ) - if is_llada_image and activation_scheme != "dynamic": - raise ValueError( - "LLaDA-Image FP8 conversion currently supports only dynamic activation " - "quantization. To convert this checkpoint, set activation_scheme='dynamic'." - ) - dynamic_activation = activation_scheme == "dynamic" - runtime_name_mapper = _get_runtime_module_name_mapper( model_type=model_type, class_name=class_name ) @@ -755,11 +630,7 @@ def build_modelopt_fp8_transformer( backbone_state = torch.load(backbone_ckpt_path, map_location="cpu")[ "model_state_dict" ] - fp8_scale_map = build_fp8_scale_map( - backbone_state, - maxbound=maxbound, - require_input_scale=not dynamic_activation, - ) + fp8_scale_map = build_fp8_scale_map(backbone_state, maxbound=maxbound) quant_algo = str(quant_config.get("quant_algo", "")).upper() if quant_algo and "FP8" not in quant_algo: raise ValueError( @@ -784,8 +655,7 @@ def build_modelopt_fp8_transformer( _preferred_module_name(weight_name, runtime_name_mapper) for weight_name in source_weight_map if weight_name.endswith(".weight") - and _resolve_scale_key(weight_name, fp8_scale_map, runtime_name_mapper) - is None + and _resolve_scale_key(weight_name, fp8_scale_map) is None } ) fallback_ignore_modules = sorted( @@ -801,33 +671,7 @@ def build_modelopt_fp8_transformer( *fallback_ignore_modules, } ) - if is_llada_image: - _validate_llada_image_fused_weight_scales( - ( - weight_name - for weight_name in source_weight_map - if weight_name.endswith(".weight") - and weight_name not in fallback_weight_names_set - and not is_ignored_by_modelopt( - weight_name, ignore_patterns, runtime_name_mapper - ) - ), - fp8_scale_map, - ) - if dynamic_activation: - effective_quant_config = { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "ignored_layers": ignore_patterns, - } - else: - effective_quant_config["ignore"] = ignore_patterns - if "activation_scheme" in effective_quant_config: - effective_quant_config["activation_scheme"] = "static" - for group in effective_quant_config.get("config_groups", {}).values(): - activations = group.get("input_activations") - if activations is not None: - activations["dynamic"] = False + effective_quant_config["ignore"] = ignore_patterns serialized_quant_config = json.dumps(effective_quant_config, sort_keys=True) output_config = _build_output_config( source_config=config, @@ -842,12 +686,11 @@ def build_modelopt_fp8_transformer( else {} ) fallback_scale_names = { - weight_name[:-7] + scale_suffix + scale_name for weight_name in fallback_weight_names - for scale_suffix in ( - (".weight_scale",) - if dynamic_activation - else (".weight_scale", ".input_scale") + for scale_name in ( + weight_name[:-7] + ".weight_scale", + weight_name[:-7] + ".input_scale", ) } @@ -883,9 +726,6 @@ def build_modelopt_fp8_transformer( if "_quantizer." in name: del shard_tensors[name] continue - if dynamic_activation and name.endswith(".input_scale"): - del shard_tensors[name] - continue if name in fallback_scale_names: del shard_tensors[name] continue @@ -897,7 +737,7 @@ def build_modelopt_fp8_transformer( ): preserved_ignored_weight_count += 1 continue - scale_key = _resolve_scale_key(name, fp8_scale_map, runtime_name_mapper) + scale_key = _resolve_scale_key(name, fp8_scale_map) if ( name.endswith(".weight") and scale_key is not None @@ -908,18 +748,11 @@ def build_modelopt_fp8_transformer( shard_tensors[name] = quantize_fp8_weight( shard_tensors[name], scale_tensors["weight_scale"] ) - # One fused runtime observer can feed multiple logical source - # weights (for example LLaDA Q/K/V). Safetensors rejects - # multiple dictionary entries that alias the same storage. shard_tensors[name[:-7] + ".weight_scale"] = scale_tensors[ "weight_scale" - ].clone() - added_scale_count += 1 - if not dynamic_activation: - shard_tensors[name[:-7] + ".input_scale"] = scale_tensors[ - "input_scale" - ].clone() - added_scale_count += 1 + ] + shard_tensors[name[:-7] + ".input_scale"] = scale_tensors["input_scale"] + added_scale_count += 2 save_file(shard_tensors, os.path.join(output_path, filename), metadata=metadata) @@ -949,7 +782,7 @@ def build_modelopt_fp8_transformer( 1 for name in source_weight_map if name.endswith(".weight") - and _resolve_scale_key(name, fp8_scale_map, runtime_name_mapper) is not None + and _resolve_scale_key(name, fp8_scale_map) is not None and not is_ignored_by_modelopt(name, ignore_patterns, runtime_name_mapper) ), "bf16_fallback_weights": len(fallback_weight_names), @@ -997,7 +830,6 @@ def _parse_args() -> argparse.Namespace: "ltx2", "hunyuan-video", "qwen-image", - "llada-image", "none", ], default="auto", @@ -1018,17 +850,6 @@ def _parse_args() -> argparse.Namespace: "staying in FP8." ), ) - parser.add_argument( - "--activation-scheme", - choices=["auto", "static", "dynamic"], - default="auto", - help=( - "Activation quantization scheme in the converted checkpoint. 'auto' " - "reads ModelOpt input activation metadata, falling back to dynamic " - "for LLaDA-Image and static for other families when metadata is absent. " - "Use 'dynamic' explicitly to convert a static LLaDA-Image export." - ), - ) parser.add_argument( "--maxbound", type=float, @@ -1052,7 +873,6 @@ def main() -> None: base_transformer_dir=args.base_transformer_dir, model_type=args.model_type, keep_bf16_patterns=args.keep_bf16_pattern, - activation_scheme=args.activation_scheme, maxbound=args.maxbound, overwrite=args.overwrite, ) diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index e86eca9b9147..6e2b97b1ba57 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -38,7 +38,6 @@ def handle_context_parallelism(server_args: Any): run_hook(validate_prefill_cp_platform, server_args) cfg = resolving_view(server_args) - model_arch = None if parse_connector_type(cfg.model_path) != ConnectorType.INSTANCE: model_config = model_config_of(server_args) hf_config = model_config.hf_config @@ -109,21 +108,8 @@ def handle_context_parallelism(server_args: Any): ) if view.attn_cp_size != cfg.moe_dp_size: - assert cfg.moe_dp_size == 1 or ( - view.enable_dp_attention - and model_arch == "LLaDA2MoeModelLM" - and view.ep_size == 1 - and view.pp_size == 1 - and cfg.tp_size == cfg.dp_size - and cfg.dp_size == cfg.moe_dp_size - and view.enable_dp_lm_head - and view.attn_cp_size == 1 - and view.dcp_size == 1 - and view.moe_dense_tp_size == 1 - and view.moe_a2a_backend == "none" - ), ( - "attn_cp_size != moe_dp_size requires moe_dp_size == 1 or " - "the LLaDA2 full-replica topology" + assert cfg.moe_dp_size == 1, ( + "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" ) from sglang.srt.layers.cp.base import init_cp_strategy diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 71b18283dc9b..e7ce9277aba2 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1439,67 +1439,6 @@ def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: else: return self.get_num_kv_heads(tensor_parallel_size) - def _parse_llada_fp8_experts_config(self) -> Optional[dict]: - if _hf_arch(self.hf_config) != "LLaDA2MoeModelLM": - return None - - llada_fp8_experts = _hf_attr(self.hf_config, "llada_fp8_experts") - if llada_fp8_experts is None: - if _hf_attr(self.hf_config, "use_fp8_experts") is not True: - return None - logger.warning_once( - "The legacy LLaDA2 use_fp8_experts flag is deprecated; " - "publish llada_fp8_experts metadata instead." - ) - llada_fp8_experts = { - "enabled": True, - "format": "e4m3fn", - "weight_granularity": "per_expert_2d_block", - "weight_block_size": [128, 128], - "activation_method": "dynamic_absmax_rtn", - "activation_granularity": "per_token", - } - if not isinstance(llada_fp8_experts, dict): - llada_fp8_experts = llada_fp8_experts.to_dict() - if not llada_fp8_experts.get("enabled", False): - return None - - fp8_format = llada_fp8_experts.get("format") - weight_granularity = llada_fp8_experts.get("weight_granularity") - activation_granularity = llada_fp8_experts.get("activation_granularity") - activation_method = llada_fp8_experts.get("activation_method", "") - if fp8_format not in {"e4m3", "e4m3fn"}: - raise ValueError( - f"llada_fp8_experts.format must be e4m3/e4m3fn, but got {fp8_format!r}." - ) - if weight_granularity != "per_expert_2d_block": - raise ValueError( - "llada_fp8_experts.weight_granularity must be " - f"'per_expert_2d_block', but got {weight_granularity!r}." - ) - if activation_granularity != "per_token" or not str( - activation_method - ).startswith("dynamic"): - raise ValueError( - "llada_fp8_experts requires dynamic per-token activation quantization." - ) - weight_block_size = llada_fp8_experts.get("weight_block_size") - if ( - not isinstance(weight_block_size, (list, tuple)) - or len(weight_block_size) != 2 - or any(not isinstance(size, int) or size <= 0 for size in weight_block_size) - ): - raise ValueError( - "llada_fp8_experts.weight_block_size must contain " - f"two positive integers, but got {weight_block_size!r}." - ) - return { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "weight_block_size": list(weight_block_size), - "llada_experts_only": True, - } - # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _parse_quant_hf_config(self): quant_cfg = _quant_config_to_dict( @@ -1552,8 +1491,8 @@ def _parse_quant_hf_config(self): max_delay=5.0, ) if not file_exists: - # File doesn't exist on hub, no need to try downloading. - return self._parse_llada_fp8_experts_config() + # File doesn't exist on hub, no need to try downloading + return quant_cfg # None # Download (online mode) or read from cache (offline mode) if envs.SGLANG_USE_MODELSCOPE.get(): @@ -1596,8 +1535,6 @@ def _parse_quant_hf_config(self): with open(quant_config_file) as f: quant_config_dict = json.load(f) quant_cfg = self._parse_modelopt_quant_config(quant_config_dict) - if quant_cfg is None: - quant_cfg = self._parse_llada_fp8_experts_config() return quant_cfg def _find_quant_modelslim_config(self): diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 460b665e406e..323d4c309707 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -430,8 +430,6 @@ def _compute_mlp_mode(cls, context: _LayerModeComputationContext): or enable_dwdp() ): return ScatterMode.SCATTERED - if enable_moe_fully_dp(): - return ScatterMode.SCATTERED # DSA CP and MLA CP both don't support MOE_FULL yet; fall back to FULL. if is_enable_moe_cp_allgather() and not ( is_dsa_enable_prefill_cp() or is_mla_cp_enabled() @@ -487,23 +485,6 @@ def enable_moe_dense_fully_dp(): return get_parallel().moe_dense_tp_size == 1 -def enable_moe_fully_dp(): - parallel = get_parallel() - return ( - is_dp_attention_enabled() - and parallel.pp_size == 1 - and not parallel.dcp_enabled - and parallel.attn_cp_size == 1 - and parallel.attn_dp_size == parallel.tp_size - and parallel.attn_tp_size == 1 - and parallel.moe_ep_size == 1 - and parallel.moe_tp_size == 1 - and parallel.moe_dp_size == parallel.attn_dp_size - and parallel.moe_dense_tp_size == 1 - and get_moe_a2a_backend().is_none() - ) - - def _generic_prefill_cp_shards_tokens() -> bool: """Whether the strategy prefill CP path shards prefill tokens across CP ranks.""" parallel = get_parallel() @@ -1030,15 +1011,6 @@ def init_new(cls): tp_size = get_parallel().tp_size tp_rank = get_parallel().tp_rank moe_cp_size = get_moe_cp_size() - if enable_moe_fully_dp(): - moe_full_size = tp_size // attn_cp_size - else: - if attn_cp_size % moe_cp_size != 0: - raise ValueError( - "Attention CP size must be divisible by the MoE DP group size" - ) - cp_per_moe = attn_cp_size // moe_cp_size - moe_full_size = tp_size // cp_per_moe process_group_sizes = { ScatterMode.SCATTERED: 1, ScatterMode.TP_ATTN_FULL: attn_tp_size, @@ -1046,7 +1018,7 @@ def init_new(cls): # With context parallel enabled, we should exclude # the attn_cp_size from the total tp_size ScatterMode.FULL: tp_size // attn_cp_size, - ScatterMode.MOE_FULL: moe_full_size, + ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size), } return cls( process_group_sizes=process_group_sizes, diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index c0681e689c81..b9ca03d8c25f 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -256,7 +256,6 @@ def __init__( is_fp4_experts: bool = False, kv_cache_quant_algo: Optional[str] = None, scale_fmt: Optional[str] = None, - llada_experts_only: bool = False, ) -> None: super().__init__() # DSV4 mxfp4-packed (True) vs converted FP8 (False); injected by @@ -264,7 +263,6 @@ def __init__( self.is_fp4_experts = is_fp4_experts self.dequant_fp4_to_fp8 = False self.is_dsv4_fp4_experts = False - self.llada_experts_only = llada_experts_only self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized if is_checkpoint_fp8_serialized: log_info_on_rank0(logger, "Detected fp8 checkpoint.") @@ -351,7 +349,6 @@ def from_config(cls, config: Dict[str, Any]) -> Fp8Config: kv_cache_quant_algo = cls.get_from_keys_or( config, ["kv_cache_quant_algo"], None ) - llada_experts_only = cls.get_from_keys_or(config, ["llada_experts_only"], False) scale_fmt = cls.get_from_keys_or(config, ["scale_fmt"], None) if use_mxfp8: # MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only. @@ -369,7 +366,6 @@ def from_config(cls, config: Dict[str, Any]) -> Fp8Config: packed_modules_mapping=packed_modules_mapping, use_mxfp8=use_mxfp8, kv_cache_quant_algo=kv_cache_quant_algo, - llada_experts_only=llada_experts_only, scale_fmt=scale_fmt, ) @@ -381,8 +377,6 @@ def get_quant_method( from sglang.srt.layers.radix_attention import RadixAttention if isinstance(layer, LinearBase): - if self.llada_experts_only: - return UnquantizedLinearMethod() if is_layer_skipped( prefix, self.ignored_layers, fused_mapping=self.packed_modules_mapping ): @@ -461,8 +455,6 @@ def get_quant_method( return Mxfp4FlashinferTrtllmMoEMethod(fp8_method, prefix=prefix) return fp8_method elif isinstance(layer, RadixAttention): - if self.llada_experts_only: - return None return Fp8KVCacheMethod(self) return None diff --git a/python/sglang/srt/model_loader/llada2_weight_utils.py b/python/sglang/srt/model_loader/llada2_weight_utils.py index a9b554f97731..2e38033ae5d3 100644 --- a/python/sglang/srt/model_loader/llada2_weight_utils.py +++ b/python/sglang/srt/model_loader/llada2_weight_utils.py @@ -8,22 +8,14 @@ import torch _EXPERT_PROJECTION_NAMES = frozenset({"gate_proj", "up_proj", "down_proj"}) -_EXPERT_SCALE_NAMES = frozenset( - {f"{projection_name}_scale" for projection_name in _EXPERT_PROJECTION_NAMES} -) def prepare_llada2_language_weights( weights: Iterable[tuple[str, torch.Tensor]], *, num_experts: int, - expand_expert_scales: bool = False, ) -> Iterator[tuple[str, torch.Tensor]]: """Normalize language-model prefixes and expand fused expert tensors.""" - fused_expert_names = _EXPERT_PROJECTION_NAMES - if expand_expert_scales: - fused_expert_names |= _EXPERT_SCALE_NAMES - for name, loaded_weight in weights: if name.startswith("model.language_model."): name = "model." + name.removeprefix("model.language_model.") @@ -34,7 +26,7 @@ def prepare_llada2_language_weights( is_fused_expert = ( separator and expert_prefix.endswith(".mlp.experts") - and projection_name in fused_expert_names + and projection_name in _EXPERT_PROJECTION_NAMES ) if not is_fused_expert: yield name, loaded_weight @@ -45,11 +37,8 @@ def prepare_llada2_language_weights( f"Invalid fused expert weight {name!r}: expected first dimension " f"{num_experts}, got shape={tuple(loaded_weight.shape)}" ) - is_scale = expand_expert_scales and projection_name in _EXPERT_SCALE_NAMES - projection_name = projection_name.removesuffix("_scale") - parameter_name = "weight_scale_inv" if is_scale else "weight" for expert_id in range(num_experts): yield ( - f"{expert_prefix}.{expert_id}.{projection_name}.{parameter_name}", + f"{expert_prefix}.{expert_id}.{projection_name}.weight", loaded_weight[expert_id], ) diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 8908e0b89529..6bc867c35d85 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -358,12 +358,6 @@ def get_quant_config( model_config, quant_cls.from_config(hf_quant_config) ) - if model_config.quantization == "fp8": - hf_quant_config = model_config._parse_llada_fp8_experts_config() - if hf_quant_config is not None: - hf_quant_config["packed_modules_mapping"] = packed_modules_mapping - return quant_cls.from_config(hf_quant_config) - # In case of bitsandbytes/QLoRA, get quant config from the adapter model. if model_config.quantization == "bitsandbytes": if ( diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 04608f65451e..f08e6e39c76e 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -40,7 +40,6 @@ LayerCommunicator, LayerScatterModes, enable_moe_dense_fully_dp, - enable_moe_fully_dp, ) from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, @@ -94,23 +93,6 @@ _is_cuda = is_cuda() _is_npu = is_npu() -_LEGACY_LLADA2_FP8_CONFIG_DEFAULTS = { - "embedding_dropout": 0.0, - "moe_router_enable_expert_bias": True, - "norm_topk_prob": True, - "router_dtype": "fp32", - "score_function": "sigmoid", -} - - -def _apply_legacy_llada2_fp8_config_defaults(config: PretrainedConfig) -> None: - if not getattr(config, "use_fp8_experts", False): - return - for name, default in _LEGACY_LLADA2_FP8_CONFIG_DEFAULTS.items(): - if not hasattr(config, name): - setattr(config, name, default) - - split_qkv_rmsnorm_rope_pos_cache_half_npu = None if _is_npu: try: @@ -214,7 +196,6 @@ def __init__( self.layer_id = layer_id self.alt_stream = alt_stream self.tp_size = get_parallel().tp_size - self.moe_fully_replicated = enable_moe_fully_dp() self.top_k = config.num_experts_per_tok self.norm_topk_prob = config.norm_topk_prob self.hidden_size = config.hidden_size @@ -316,7 +297,7 @@ def __init__( prefix=add_prefix("shared_experts", prefix), **( dict(tp_rank=0, tp_size=1) - if get_moe_a2a_backend().is_deepep() or self.moe_fully_replicated + if get_moe_a2a_backend().is_deepep() else {} ), ) @@ -403,12 +384,8 @@ def forward_normal( if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output - if ( - self.tp_size > 1 - and not self.moe_fully_replicated - and not should_skip_post_experts_all_reduce( - is_tp_path=True, - ) + if self.tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_size) @@ -825,7 +802,6 @@ def __init__( prefix: str = "", ): super().__init__() - _apply_legacy_llada2_fp8_config_defaults(config) self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config @@ -915,15 +891,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) params_dict = dict(self.named_parameters()) - validate_llada_fp8_experts = bool( - getattr(self.quant_config, "llada_experts_only", False) - ) - loaded_expert_weight_count = 0 - loaded_expert_scale_count = 0 for name, loaded_weight in prepare_llada2_language_weights( - weights, - num_experts=self.config.num_experts, - expand_expert_scales=validate_llada_fp8_experts, + weights, num_experts=self.config.num_experts ): if ( ("v_head" in name) @@ -980,11 +949,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): shard_id=shard_id, expert_id=expert_id, ) - if validate_llada_fp8_experts: - if name.endswith("weight_scale_inv"): - loaded_expert_scale_count += 1 - elif name.endswith("weight"): - loaded_expert_weight_count += 1 break else: # Skip loading extra bias for GPTQ models. @@ -1010,43 +974,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) } ) - if validate_llada_fp8_experts: - sparse_layers = [ - layer - for layer in self.model.layers - if not isinstance(layer, PPMissingLayer) - and isinstance(layer.mlp, LLaDA2MoeSparseMoeBlock) - ] - expected_tensor_count = len(sparse_layers) * self.config.num_experts * 3 - if loaded_expert_weight_count != expected_tensor_count: - raise ValueError( - "Incomplete LLaDA FP8 expert weights: loaded " - f"{loaded_expert_weight_count}, expected {expected_tensor_count}." - ) - if loaded_expert_scale_count != expected_tensor_count: - scale_param_names = [ - name - for name in params_dict - if "mlp.experts" in name and "scale" in name - ] - raise ValueError( - "Incomplete LLaDA FP8 expert scales: loaded " - f"{loaded_expert_scale_count}, expected {expected_tensor_count}; " - f"registered scale parameters sample={scale_param_names[:8]}." - ) - - if sparse_layers: - first_experts = sparse_layers[0].mlp.experts - logger.info( - "Loaded LLaDA experts-only FP8: w13=%s/%s, w2=%s/%s, " - "w13_scale=%s, w2_scale=%s", - tuple(first_experts.w13_weight.shape), - first_experts.w13_weight.dtype, - tuple(first_experts.w2_weight.shape), - first_experts.w2_weight.dtype, - tuple(first_experts.w13_weight_scale_inv.shape), - tuple(first_experts.w2_weight_scale_inv.shape), - ) @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/scripts/ci/utils/diffusion/diffusion_case_parser.py b/scripts/ci/utils/diffusion/diffusion_case_parser.py index d6d97cb111ba..fbbc69d8a65c 100755 --- a/scripts/ci/utils/diffusion/diffusion_case_parser.py +++ b/scripts/ci/utils/diffusion/diffusion_case_parser.py @@ -96,7 +96,6 @@ class DiffusionTestCaseVisitor(ast.NodeVisitor): def __init__(self): self.cases: Dict[str, List[str]] = {} # list_name -> [case_id, ...] self.factory_case_ids: Dict[str, str] = {} - self.case_est_times: Dict[str, float] = {} def visit_Module(self, node: ast.Module): for stmt in node.body: @@ -226,15 +225,7 @@ def _extract_case_id_from_call(self, node: ast.AST) -> Optional[str]: "_make_modelopt_ci_case", }: if node.args and isinstance(node.args[0], ast.Constant): - case_id = node.args[0].value - for keyword in node.keywords: - if ( - keyword.arg == "estimated_full_test_time_s" - and isinstance(keyword.value, ast.Constant) - and isinstance(keyword.value.value, (int, float)) - ): - self.case_est_times[case_id] = float(keyword.value.value) - return case_id + return node.args[0].value if isinstance(node.func, ast.Name) and not node.args: return self.factory_case_ids.get(node.func.id) @@ -374,9 +365,7 @@ def _iter_baseline_paths(baseline_path: Path) -> List[Path]: return [path for path in ordered_paths if path.exists()] -def load_baselines( - baseline_path: Path, case_est_times: Optional[Dict[str, float]] = None -) -> Dict[str, float]: +def load_baselines(baseline_path: Path) -> Dict[str, float]: """ Load performance baselines from a JSON file or platform baseline directory. @@ -392,8 +381,6 @@ def load_baselines( for case_id, scenario in scenarios.items(): if scenario.get("estimated_full_test_time_s") is not None: est_time = scenario["estimated_full_test_time_s"] - elif case_est_times and case_id in case_est_times: - est_time = case_est_times[case_id] else: expected_e2e_ms = scenario.get("expected_e2e_ms", 0) est_time = expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS @@ -415,17 +402,14 @@ def parse_testcase_configs(config_path: Path) -> Dict[str, List[str]]: Dictionary mapping list name to case IDs. e.g., {"ONE_GPU_CASES_A": ["qwen_image_t2i", ...], ...} """ - return _parse_testcase_config(config_path).cases - - -def _parse_testcase_config(config_path: Path) -> DiffusionTestCaseVisitor: with open(config_path, "r", encoding="utf-8") as f: content = f.read() tree = ast.parse(content, filename=str(config_path)) visitor = DiffusionTestCaseVisitor() visitor.visit(tree) - return visitor + + return visitor.cases def parse_run_suite_standalone_data( @@ -483,8 +467,7 @@ def collect_diffusion_suites( Dictionary mapping suite name to DiffusionSuiteInfo. """ # Parse case IDs from the single source case config. - parsed_config = _parse_testcase_config(case_config_path) - case_lists = parsed_config.cases + case_lists = parse_testcase_configs(case_config_path) # Parse standalone files from run_suite.py standalone_files, standalone_est_times = parse_run_suite_standalone_data( @@ -495,10 +478,7 @@ def collect_diffusion_suites( ) # Load baselines for time estimation - baselines = { - **parsed_config.case_est_times, - **load_baselines(baseline_path, parsed_config.case_est_times), - } + baselines = load_baselines(baseline_path) # Build suite info suites = {} diff --git a/test/registered/e2e/attention/test_llada2_cfg_flashinfer_backend.py b/test/registered/e2e/attention/test_llada2_cfg_flashinfer_backend.py deleted file mode 100644 index 2b37ccf06524..000000000000 --- a/test/registered/e2e/attention/test_llada2_cfg_flashinfer_backend.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend -from sglang.srt.layers.attention.llada2_attention_utils import ( - build_llada_image_custom_mask, -) -from sglang.srt.layers.attention.llada2_cfg_flashinfer_backend import ( - LLaDA2CFGFlashInferAttnBackend, -) -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.test.ci.ci_register import register_cuda_ci - -register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-small") - - -class _FakeRaggedWrapper: - def __init__(self): - self.args = None - self.kwargs = None - self.forward_args = None - self.forward_kwargs = None - - def begin_forward(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - - def forward(self, *args, **kwargs): - self.forward_args = args - self.forward_kwargs = kwargs - return torch.zeros_like(args[0]) - - -class _FakeTokenToKVPool: - def __init__(self): - self.set_args = None - - def set_kv_buffer(self, *args): - self.set_args = args - - -class _FakePrefillUpdater: - def __init__(self): - self.num_qo_heads = 8 - self.num_kv_heads = 8 - self.head_dim = 64 - self.q_data_type = torch.bfloat16 - self.data_type = torch.bfloat16 - - -def _make_backend(): - backend = object.__new__(LLaDA2CFGFlashInferAttnBackend) - backend._llada_image_conditioning_mask_active = False - backend._cfg_prefill_wrapper_ragged = _FakeRaggedWrapper() - backend.num_wrappers = 1 - backend.prefill_wrappers_paged = [object()] - backend.indices_updater_prefill = _FakePrefillUpdater() - backend.prefill_split_tile_size = None - return backend - - -def _make_forward_batch(prefix_lens=(0, 0)): - seq_lens = torch.tensor([4, 3], dtype=torch.int32) - return SimpleNamespace( - forward_mode=ForwardMode.EXTEND, - seq_lens=seq_lens, - extend_prefix_lens=torch.tensor(prefix_lens, dtype=torch.int32), - llada_image_conditioning_text_lens_cpu=[2, 1], - ) - - -class TestLLaDA2CFGFlashInferBackend(unittest.TestCase): - def test_delegates_batches_without_llada_image_metadata(self): - backend = _make_backend() - backend._cfg_prefill_wrapper_ragged = SimpleNamespace() - forward_batch = _make_forward_batch() - del forward_batch.llada_image_conditioning_text_lens_cpu - - with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: - backend.init_forward_metadata(forward_batch) - - stock_init.assert_called_once_with(forward_batch) - self.assertFalse(backend.conditioning_mask_active) - from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( - ensure_conditioning_mask_active, - ) - - with self.assertRaisesRegex(RuntimeError, "without its block attention mask"): - ensure_conditioning_mask_active(backend) - - def test_delegates_legacy_dllm_padding_metadata(self): - backend = _make_backend() - forward_batch = _make_forward_batch() - del forward_batch.llada_image_conditioning_text_lens_cpu - forward_batch.forward_mode = ForwardMode.DLLM_EXTEND - forward_batch.dllm_left_pad_lens = torch.tensor([1, 1], dtype=torch.int32) - - with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: - backend.init_forward_metadata(forward_batch) - - stock_init.assert_called_once_with(forward_batch) - - def test_consumes_llada_image_text_lengths_as_custom_mask(self): - backend = _make_backend() - forward_batch = _make_forward_batch() - - with patch.object(FlashInferAttnBackend, "init_forward_metadata") as stock_init: - backend.init_forward_metadata(forward_batch) - - stock_init.assert_not_called() - wrapper = backend._cfg_prefill_wrapper_ragged - self.assertEqual(wrapper.args[0].tolist(), [0, 4, 7]) - self.assertEqual(wrapper.args[1].tolist(), [0, 4, 7]) - self.assertFalse(wrapper.kwargs["causal"]) - torch.testing.assert_close( - wrapper.kwargs["custom_mask"], - build_llada_image_custom_mask([2, 1], [4, 3], "cpu"), - ) - self.assertTrue(backend._llada_image_conditioning_mask_active) - self.assertTrue(backend.conditioning_mask_active) - self.assertTrue(backend.forward_metadata.use_ragged) - - def test_masked_forward_uses_ragged_kernel_without_saving_kv(self): - backend = _make_backend() - forward_batch = _make_forward_batch() - forward_batch.token_to_kv_pool = _FakeTokenToKVPool() - backend.init_forward_metadata(forward_batch) - layer = SimpleNamespace( - tp_q_head_num=2, - tp_k_head_num=2, - tp_v_head_num=2, - head_dim=4, - scaling=0.5, - logit_cap=0.0, - is_cross_attention=False, - ) - q = torch.randn(7, 8) - k = torch.randn(7, 8) - v = torch.randn(7, 8) - - output = backend.forward_extend(q, k, v, layer, forward_batch) - - self.assertEqual(output.shape, q.shape) - wrapper = backend._cfg_prefill_wrapper_ragged - torch.testing.assert_close(wrapper.forward_args[0], q.view(7, 2, 4)) - self.assertFalse(wrapper.forward_kwargs["causal"]) - self.assertIsNone(forward_batch.token_to_kv_pool.set_args) - - def test_rejects_cached_prefix_for_one_shot_conditioning(self): - backend = _make_backend() - forward_batch = _make_forward_batch(prefix_lens=(1, 0)) - - with self.assertRaisesRegex(RuntimeError, "does not support cached prefixes"): - backend.init_forward_metadata(forward_batch) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/ci/test_diffusion_ci_registration.py b/test/registered/unit/ci/test_diffusion_ci_registration.py deleted file mode 100644 index 7c0375304d04..000000000000 --- a/test/registered/unit/ci/test_diffusion_ci_registration.py +++ /dev/null @@ -1,103 +0,0 @@ -import importlib -import json -import os -import sys -import tempfile -import unittest -from contextlib import ExitStack -from pathlib import Path -from unittest.mock import patch - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=2, suite="base-a-test-cpu") - -_REPO_ROOT = Path(__file__).resolve().parents[4] - - -class TestDiffusionCIRegistration(unittest.TestCase): - @classmethod - def setUpClass(cls): - stack = ExitStack() - cls.addClassCleanup(stack.close) - stack.enter_context(patch.dict(os.environ, {"USE_NPU_CONFIGS": "0"})) - stack.enter_context(patch.dict(sys.modules)) - parser_dir = _REPO_ROOT / "scripts/ci/utils/diffusion" - stack.enter_context(patch.object(sys, "path", [str(parser_dir), *sys.path])) - for name in ( - "diffusion_case_parser", - "compute_diffusion_partitions", - "verify_diffusion_coverage", - ): - sys.modules.pop(name, None) - cls.parser = importlib.import_module("diffusion_case_parser") - cls.planner = importlib.import_module("compute_diffusion_partitions") - cls.coverage = importlib.import_module("verify_diffusion_coverage") - - def test_llada_cases_reach_plans_and_coverage_with_full_estimates(self): - runner = _REPO_ROOT / self.parser.RUN_SUITE_REL_PATH - config = self.parser.resolve_case_config_path(_REPO_ROOT, runner) - suites = self.parser.collect_diffusion_suites( - config, runner, _REPO_ROOT / self.parser.BASELINE_REL_PATH - ) - coverage = self.coverage.get_expected_cases(_REPO_ROOT) - for degree in (1, 2): - suite = f"{degree}-gpu" - expected = { - f"llada_image_turbo_fp8_{mode}_sp{degree}" for mode in ("t2i", "edit") - } - with self.subTest(suite=suite): - cases = suites[suite].cases - ids = [case.case_id for case in cases] - self.assertEqual(expected.intersection(ids), expected) - self.assertTrue(expected.issubset(coverage[suite])) - for case in cases: - if case.case_id in expected: - self.assertEqual(case.est_time, 600.0) - self.assertEqual(ids.count(case.case_id), 1) - for count in (1, 3, 8): - partitions = self.planner.partition_items_by_lpt( - self.planner.build_partition_items(suites[suite]), count - ) - plan = self.planner.build_partition_plan(suite, partitions) - planned_ids = [ - case_id - for partition in plan["partitions"] - for case_id in partition["case_ids"] - ] - for case_id in expected: - self.assertEqual(planned_ids.count(case_id), 1) - - def test_static_estimate_fallback_and_baseline_precedence(self): - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - config = root / "cases.py" - config.write_text( - 'ONE_GPU_CASES = [DiffusionTestCase("explicit", estimated_full_test_time_s=600.0), ' - 'DiffusionTestCase("fallback")]\n' - ) - runner = root / "runner.py" - runner.write_text("STANDALONE_FILES = {}\n") - baseline = root / "baseline.json" - for scenario, expected in ( - (None, 600.0), - ({"expected_e2e_ms": 60000.0}, 600.0), - ({"estimated_full_test_time_s": 420.0}, 420.0), - ): - with self.subTest(scenario=scenario): - scenarios = {} if scenario is None else {"explicit": scenario} - baseline.write_text(json.dumps({"scenarios": scenarios})) - suites = self.parser.collect_diffusion_suites( - config, runner, baseline - ) - estimates = { - case.case_id: case.est_time for case in suites["1-gpu"].cases - } - self.assertEqual(estimates["explicit"], expected) - self.assertEqual( - estimates["fallback"], self.parser.DEFAULT_EST_TIME_SECONDS - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/layers/attention/test_attention_backend_registry.py b/test/registered/unit/layers/attention/test_attention_backend_registry.py deleted file mode 100644 index 77aa6be50963..000000000000 --- a/test/registered/unit/layers/attention/test_attention_backend_registry.py +++ /dev/null @@ -1,27 +0,0 @@ -import unittest -from types import SimpleNamespace - -from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS -from sglang.srt.server_args import ATTENTION_BACKEND_CHOICES -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=10, suite="base-a-test-cpu") - - -class TestAttentionBackendRegistry(unittest.TestCase): - def test_llada2_cfg_flashinfer_is_registered(self): - backend_name = "llada2_cfg_flashinfer" - - self.assertIn(backend_name, ATTENTION_BACKENDS) - self.assertIn(backend_name, ATTENTION_BACKEND_CHOICES) - - def test_llada2_cfg_flashinfer_rejects_mla(self): - factory = ATTENTION_BACKENDS["llada2_cfg_flashinfer"] - runner = SimpleNamespace(use_mla_backend=True) - - with self.assertRaisesRegex(ValueError, "does not use an MLA backend"): - factory(runner) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py index 269d55a4e00f..aa52f99033e6 100644 --- a/test/registered/unit/layers/test_layer_communicator_fusion_gate.py +++ b/test/registered/unit/layers/test_layer_communicator_fusion_gate.py @@ -3,13 +3,7 @@ from unittest.mock import patch from sglang.srt.layers import communicator as comm -from sglang.srt.layers.communicator import ( - CommunicateContext, - LayerCommunicator, - LayerScatterModes, - ScatterMode, - enable_moe_fully_dp, -) +from sglang.srt.layers.communicator import LayerCommunicator, ScatterMode from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -79,114 +73,5 @@ def test_moe_full_layer_does_not_fuse(self): ) -class TestFullReplicaMoE(CustomTestCase): - def _topology(self, **overrides): - topology = dict( - tp_rank=0, - tp_size=2, - pp_size=1, - dcp_enabled=False, - attn_tp_rank=0, - attn_tp_size=1, - attn_dp_size=2, - attn_cp_rank=0, - attn_cp_size=1, - moe_ep_size=1, - moe_tp_size=1, - moe_dp_size=2, - moe_dense_tp_size=1, - dwdp_size=1, - ) - topology.update(overrides) - return topology - - def test_full_replica_sparse_mlp_stays_local(self): - backend = types.SimpleNamespace(is_none=lambda: True) - mode_args = dict( - num_layers=2, - layer_id=0, - is_layer_sparse=True, - is_previous_layer_sparse=None, - is_next_layer_sparse=True, - ) - with ( - patch.object(comm, "is_dp_attention_enabled", return_value=True), - patch.object(comm, "get_moe_a2a_backend", return_value=backend), - patch.object( - comm, - "should_use_flashinfer_cutlass_moe_fp4_allgather", - return_value=False, - ), - patch.object(comm, "is_enable_moe_cp_allgather", return_value=False), - patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False), - patch.object(comm, "is_mla_cp_enabled", return_value=False), - patch.object( - comm, - "get_parallel", - return_value=types.SimpleNamespace(**self._topology()), - ), - ): - self.assertTrue(enable_moe_fully_dp()) - self.assertIs( - LayerScatterModes.init_new(**mode_args).mlp_mode, - ScatterMode.SCATTERED, - ) - - with ( - patch.object(comm, "is_dp_attention_enabled", return_value=True), - patch.object(comm, "get_moe_a2a_backend", return_value=backend), - patch.object( - comm, - "should_use_flashinfer_cutlass_moe_fp4_allgather", - return_value=False, - ), - patch.object(comm, "is_enable_moe_cp_allgather", return_value=False), - patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False), - patch.object(comm, "is_mla_cp_enabled", return_value=False), - patch.object( - comm, - "get_parallel", - return_value=types.SimpleNamespace(**self._topology(moe_tp_size=2)), - ), - ): - self.assertFalse(enable_moe_fully_dp()) - self.assertIs( - LayerScatterModes.init_new(**mode_args).mlp_mode, - ScatterMode.FULL, - ) - - def test_full_replica_moe_full_width_is_defined(self): - backend = types.SimpleNamespace(is_none=lambda: True) - with ( - patch.object(comm, "is_dp_attention_enabled", return_value=True), - patch.object(comm, "get_moe_a2a_backend", return_value=backend), - patch.object(comm, "get_moe_cp_size", return_value=2), - patch.object( - comm, - "get_parallel", - return_value=types.SimpleNamespace(**self._topology()), - ), - ): - context = CommunicateContext.init_new() - - self.assertEqual(context.process_group_sizes[ScatterMode.FULL], 2) - self.assertEqual(context.process_group_sizes[ScatterMode.MOE_FULL], 2) - - def test_invalid_non_replica_moe_width_fails_early(self): - backend = types.SimpleNamespace(is_none=lambda: True) - with ( - patch.object(comm, "is_dp_attention_enabled", return_value=True), - patch.object(comm, "get_moe_a2a_backend", return_value=backend), - patch.object(comm, "get_moe_cp_size", return_value=2), - patch.object( - comm, - "get_parallel", - return_value=types.SimpleNamespace(**self._topology(moe_tp_size=2)), - ), - ): - with self.assertRaisesRegex(ValueError, "must be divisible"): - CommunicateContext.init_new() - - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_llada2_fp8.py b/test/registered/unit/models/test_llada2_fp8.py deleted file mode 100644 index 3d23c6677357..000000000000 --- a/test/registered/unit/models/test_llada2_fp8.py +++ /dev/null @@ -1,210 +0,0 @@ -import tempfile -import unittest -from types import SimpleNamespace - -from transformers import PretrainedConfig - -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel - -maybe_stub_sgl_kernel() - -from sglang.srt.configs.model_config import ModelConfig # noqa: E402 -from sglang.srt.layers.quantization.fp8 import Fp8Config # noqa: E402 -from sglang.srt.model_loader.weight_utils import get_quant_config # noqa: E402 -from sglang.srt.utils.common import LazyValue # noqa: E402 - -register_cpu_ci(est_time=3, suite="base-a-test-cpu") - - -def _experts_only_config(**overrides): - config = { - "enabled": True, - "format": "e4m3fn", - "weight_granularity": "per_expert_2d_block", - "weight_block_size": [128, 128], - "activation_method": "dynamic_absmax_rtn", - "activation_granularity": "per_token", - } - config.update(overrides) - return config - - -class TestLLaDA2Fp8(CustomTestCase): - def setUp(self): - self.model_dir = tempfile.TemporaryDirectory() - self.addCleanup(self.model_dir.cleanup) - - def _model_config( - self, - custom_config, - *, - architecture="LLaDA2MoeModelLM", - quantization_config=None, - compression_config=None, - use_fp8_experts=None, - ): - model_config = ModelConfig.__new__(ModelConfig) - model_config.model_path = self.model_dir.name - config_kwargs = dict( - architectures=[architecture], - quantization_config=quantization_config, - compression_config=compression_config, - llada_fp8_experts=custom_config, - ) - if use_fp8_experts is not None: - config_kwargs["use_fp8_experts"] = use_fp8_experts - model_config.hf_config = PretrainedConfig(**config_kwargs) - return model_config - - def test_parses_experts_only_block_fp8(self): - model_config = self._model_config(_experts_only_config()) - - parsed = model_config._parse_quant_hf_config() - self.assertEqual( - parsed, - { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "weight_block_size": [128, 128], - "llada_experts_only": True, - }, - ) - - fp8_config = Fp8Config.from_config(parsed) - self.assertTrue(fp8_config.is_checkpoint_fp8_serialized) - self.assertTrue(fp8_config.llada_experts_only) - self.assertEqual(fp8_config.weight_block_size, [128, 128]) - self.assertEqual(fp8_config.activation_scheme, "dynamic") - - def test_model_loader_preserves_experts_only_block_config(self): - model_config = self._model_config(_experts_only_config()) - model_config.quantization = "fp8" - - fp8_config = get_quant_config( - model_config, - SimpleNamespace(), - packed_modules_mapping={}, - ) - self.assertIsInstance(fp8_config, Fp8Config) - self.assertTrue(fp8_config.llada_experts_only) - self.assertEqual(fp8_config.weight_block_size, [128, 128]) - - def test_parses_legacy_experts_only_fp8_flag(self): - model_config = self._model_config(None, use_fp8_experts=True) - - self.assertEqual( - model_config._parse_quant_hf_config(), - { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "weight_block_size": [128, 128], - "llada_experts_only": True, - }, - ) - - disabled = self._model_config(None, use_fp8_experts=False) - self.assertIsNone(disabled._parse_llada_fp8_experts_config()) - - def test_legacy_flag_restores_only_missing_runtime_fields(self): - from sglang.srt.models.llada2 import ( - _apply_legacy_llada2_fp8_config_defaults, - ) - - legacy_config = SimpleNamespace( - use_fp8_experts=True, - router_dtype="bfloat16", - ) - _apply_legacy_llada2_fp8_config_defaults(legacy_config) - - self.assertEqual(legacy_config.embedding_dropout, 0.0) - self.assertTrue(legacy_config.moe_router_enable_expert_bias) - self.assertTrue(legacy_config.norm_topk_prob) - self.assertEqual(legacy_config.score_function, "sigmoid") - self.assertEqual(legacy_config.router_dtype, "bfloat16") - - standard_config = SimpleNamespace(use_fp8_experts=False) - _apply_legacy_llada2_fp8_config_defaults(standard_config) - self.assertFalse(hasattr(standard_config, "embedding_dropout")) - self.assertFalse(hasattr(standard_config, "score_function")) - - def test_experts_only_mode_requires_llada_specific_opt_in(self): - non_llada_config = self._model_config( - _experts_only_config(), architecture="Qwen2ForCausalLM" - ) - self.assertIsNone(non_llada_config._parse_llada_fp8_experts_config()) - - legacy_non_llada = self._model_config( - None, architecture="Qwen2ForCausalLM", use_fp8_experts=True - ) - self.assertIsNone(legacy_non_llada._parse_llada_fp8_experts_config()) - - generic_fp8_config = Fp8Config.from_config( - { - "quant_method": "fp8", - "activation_scheme": "dynamic", - "moe_only": True, - } - ) - self.assertFalse(generic_fp8_config.llada_experts_only) - - def test_standard_quantization_metadata_takes_precedence(self): - standard_config = { - "quant_method": "fp8", - "activation_scheme": "dynamic", - } - for field in ("quantization_config", "compression_config"): - with self.subTest(field=field): - model_config = self._model_config( - _experts_only_config(), **{field: standard_config.copy()} - ) - self.assertEqual(model_config._parse_quant_hf_config(), standard_config) - - def test_rejects_non_dynamic_per_token_activation(self): - model_config = self._model_config( - _experts_only_config( - activation_method="static_absmax", - activation_granularity="per_tensor", - ) - ) - with self.assertRaisesRegex(ValueError, "dynamic per-token"): - model_config._parse_quant_hf_config() - - def test_experts_only_validation_uses_effective_config_and_local_partition(self): - import torch.nn as nn - - from sglang.srt.models.llada2 import ( - LLaDA2MoeModelLM, - LLaDA2MoeSparseMoeBlock, - ) - - model = LLaDA2MoeModelLM.__new__(LLaDA2MoeModelLM) - nn.Module.__init__(model) - model.config = SimpleNamespace( - num_experts=4, - llada_fp8_experts={"enabled": True}, - ) - model.quant_config = SimpleNamespace(llada_experts_only=True) - model.model = nn.Module() - model.model.layers = nn.ModuleList() - - model.load_weights([]) - self.assertIsInstance(model.routed_experts_weights_of_layer, LazyValue) - self.assertEqual(model.routed_experts_weights_of_layer.value, {}) - - layer = nn.Module() - layer.mlp = LLaDA2MoeSparseMoeBlock.__new__(LLaDA2MoeSparseMoeBlock) - nn.Module.__init__(layer.mlp) - layer.mlp.get_moe_weights = lambda: () - model.model.layers = nn.ModuleList([layer]) - - model.quant_config = SimpleNamespace(llada_experts_only=False) - model.load_weights([]) - - model.quant_config = SimpleNamespace(llada_experts_only=True) - with self.assertRaisesRegex(ValueError, "Incomplete LLaDA FP8 expert weights"): - model.load_weights([]) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/models/test_llada2_full_replica.py b/test/registered/unit/models/test_llada2_full_replica.py deleted file mode 100644 index f6562c6ae1ca..000000000000 --- a/test/registered/unit/models/test_llada2_full_replica.py +++ /dev/null @@ -1,125 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel - -maybe_stub_sgl_kernel() - -from sglang.srt.models import llada2 # noqa: E402 - -register_cpu_ci(est_time=3, suite="base-a-test-cpu") - - -class TestLLaDA2FullReplica(CustomTestCase): - @staticmethod - def _config(): - return SimpleNamespace( - num_experts_per_tok=1, - norm_topk_prob=True, - hidden_size=4, - num_shared_experts=1, - hidden_act="silu", - num_experts=2, - moe_intermediate_size=8, - ) - - def _shared_expert_kwargs(self, fully_replicated): - backend = SimpleNamespace(is_deepep=lambda: False) - with ( - patch.object( - llada2, - "get_parallel", - return_value=SimpleNamespace(tp_size=2), - ), - patch.object( - llada2, - "get_exec", - return_value=SimpleNamespace( - moe=SimpleNamespace(ep_num_redundant_experts=0) - ), - ), - patch.object( - llada2, - "enable_moe_fully_dp", - return_value=fully_replicated, - ), - patch.object( - llada2, - "get_moe_a2a_backend", - return_value=backend, - ), - patch.object( - llada2, - "LLaDA2MoeGate", - return_value=SimpleNamespace(expert_bias=None), - ), - patch.object(llada2, "TopK", return_value=object()), - patch.object( - llada2, - "get_moe_impl_class", - return_value=lambda **_kwargs: object(), - ), - patch.object( - llada2, - "LLaDA2MoeMLP", - return_value=object(), - ) as shared_mlp, - ): - llada2.LLaDA2MoeSparseMoeBlock( - layer_id=0, - config=self._config(), - ) - - return shared_mlp.call_args.kwargs - - def test_full_replica_shared_expert_is_tp1(self): - replicated_kwargs = self._shared_expert_kwargs(fully_replicated=True) - self.assertEqual(replicated_kwargs["tp_rank"], 0) - self.assertEqual(replicated_kwargs["tp_size"], 1) - - standard_kwargs = self._shared_expert_kwargs(fully_replicated=False) - self.assertNotIn("tp_rank", standard_kwargs) - self.assertNotIn("tp_size", standard_kwargs) - - def test_full_replica_skips_global_post_expert_allreduce(self): - hidden_states = torch.zeros((2, 3)) - - for fully_replicated, expected_calls in ((True, 0), (False, 1)): - with self.subTest(fully_replicated=fully_replicated): - block = SimpleNamespace( - alt_stream=None, - num_shared_experts=1, - tp_size=2, - moe_fully_replicated=fully_replicated, - _forward_shared_experts=lambda tensor: torch.full_like(tensor, 2), - _forward_router_experts=lambda tensor: torch.full_like(tensor, 3), - ) - with ( - patch.object( - llada2, - "should_skip_post_experts_all_reduce", - return_value=False, - ), - patch.object( - llada2, - "tensor_model_parallel_all_reduce", - side_effect=lambda tensor: tensor, - ) as all_reduce, - ): - output = llada2.LLaDA2MoeSparseMoeBlock.forward_normal( - block, - hidden_states, - ) - - self.assertTrue(torch.equal(output, torch.full_like(output, 5))) - self.assertEqual(all_reduce.call_count, expected_calls) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/models/test_llada2_weight_utils.py b/test/registered/unit/models/test_llada2_weight_utils.py index a988fba5c71a..c6e845251d27 100644 --- a/test/registered/unit/models/test_llada2_weight_utils.py +++ b/test/registered/unit/models/test_llada2_weight_utils.py @@ -36,81 +36,19 @@ def test_prepares_clean_language_checkpoint_layout(self): torch.testing.assert_close(expanded[1][1], fused[1]) torch.testing.assert_close(expanded[2][1], lm_head) - def test_expert_scale_expansion_is_opt_in(self): - gate_scale = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) - up_scale = gate_scale + 100 - down_scale = torch.arange(16, dtype=torch.float32).reshape(2, 4, 2) - weights = [ - ( - "model.language_model.layers.1.mlp.experts.gate_proj_scale", - gate_scale, - ), - ( - "model.language_model.layers.1.mlp.experts.up_proj_scale", - up_scale, - ), - ( - "model.language_model.layers.1.mlp.experts.down_proj_scale", - down_scale, - ), - ] - - unexpanded = list(prepare_llada2_language_weights(weights, num_experts=2)) - self.assertEqual( - [name for name, _ in unexpanded], - [ - "model.layers.1.mlp.experts.gate_proj_scale", - "model.layers.1.mlp.experts.up_proj_scale", - "model.layers.1.mlp.experts.down_proj_scale", - ], - ) - - expanded = list( - prepare_llada2_language_weights( - weights, - num_experts=2, - expand_expert_scales=True, - ) - ) - - self.assertEqual( - [name for name, _ in expanded], - [ - "model.layers.1.mlp.experts.0.gate_proj.weight_scale_inv", - "model.layers.1.mlp.experts.1.gate_proj.weight_scale_inv", - "model.layers.1.mlp.experts.0.up_proj.weight_scale_inv", - "model.layers.1.mlp.experts.1.up_proj.weight_scale_inv", - "model.layers.1.mlp.experts.0.down_proj.weight_scale_inv", - "model.layers.1.mlp.experts.1.down_proj.weight_scale_inv", - ], - ) - torch.testing.assert_close(expanded[0][1], gate_scale[0]) - torch.testing.assert_close(expanded[3][1], up_scale[1]) - torch.testing.assert_close(expanded[5][1], down_scale[1]) - def test_rejects_invalid_fused_expert_shape(self): - cases = [ - ( - "weight", - "model.language_model.layers.1.mlp.experts.down_proj", - False, - ), - ( - "scale", - "model.language_model.layers.1.mlp.experts.gate_proj_scale", - True, - ), - ] - for label, name, expand_expert_scales in cases: - with self.subTest(label=label): - with self.assertRaisesRegex(ValueError, "expected first dimension 2"): - list( - prepare_llada2_language_weights( - [(name, torch.empty(3, 4, 5))], - num_experts=2, - expand_expert_scales=expand_expert_scales, + with self.assertRaisesRegex(ValueError, "expected first dimension 2"): + list( + prepare_llada2_language_weights( + [ + ( + "model.language_model.layers.1.mlp.experts.down_proj", + torch.empty(3, 4, 5), ) - ) + ], + num_experts=2, + ) + ) if __name__ == "__main__": diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index adcf19978bf1..41ba165d9bc9 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1114,12 +1114,8 @@ def _new_cp_args(self, **overrides): tp_size=1, dp_size=1, moe_dp_size=1, - moe_dense_tp_size=None, - moe_a2a_backend="none", ep_size=1, pp_size=1, - enable_dp_attention=False, - enable_dp_lm_head=False, dcp_size=1, enable_aiter_allreduce_fusion=False, ) @@ -1181,50 +1177,6 @@ def test_context_parallel_handler_initializes_cp_strategy(self): self.assertTrue(is_cp_enabled()) self.assertTrue(is_interleave()) - def test_llada2_full_replica_is_the_only_mismatched_cp_exception(self): - cases = [ - ("llada2_exact", "LLaDA2MoeModelLM", {}, True), - ("different_architecture", "Qwen3MoeForCausalLM", {}, False), - ("partial_replication", "LLaDA2MoeModelLM", {"tp_size": 4}, False), - ("dcp_enabled", "LLaDA2MoeModelLM", {"dcp_size": 2}, False), - ( - "sharded_lm_head", - "LLaDA2MoeModelLM", - {"enable_dp_lm_head": False}, - False, - ), - ] - - for name, architecture, overrides, should_pass in cases: - with self.subTest(name=name): - topology = dict( - model_path="dummy", - tp_size=2, - dp_size=2, - enable_dp_attention=True, - enable_dp_lm_head=True, - attn_cp_size=1, - moe_dp_size=2, - moe_dense_tp_size=1, - moe_a2a_backend="none", - ep_size=1, - ) - topology.update(overrides) - server_args = self._new_cp_args(**topology) - server_args._model_config = SimpleNamespace( - hf_config=SimpleNamespace(architectures=[architecture]), - is_multimodal=False, - ) - - if should_pass: - handle_context_parallelism(server_args) - else: - with self.assertRaisesRegex( - AssertionError, - "full-replica topology", - ): - handle_context_parallelism(server_args) - class TestFlashinferA2ADispatchType(CustomTestCase): def setUp(self): From 31463fd44249dedda3984c7f73c4c5389371bad3 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Tue, 15 Sep 2026 16:55:24 +0000 Subject: [PATCH 11/18] [diffusion] Trim LLaDA-Image unit tests --- .../test/unit/test_llada_image.py | 407 ++++-------- .../unit/test_llada_image_conditioning.py | 627 ++++++------------ .../test/unit/test_llada_image_loader.py | 154 ++--- .../unit/test_llada_image_pipeline_config.py | 64 -- .../test/unit/test_llada_image_rmsnorm.py | 78 +-- .../test_llada_image_source_conditioning.py | 251 +++---- 6 files changed, 494 insertions(+), 1087 deletions(-) delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image.py b/python/sglang/multimodal_gen/test/unit/test_llada_image.py index f27ececf60b4..cc004ad79b5f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image.py @@ -1,22 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -import unittest -from dataclasses import fields -from itertools import pairwise from types import SimpleNamespace from unittest.mock import patch +import pytest import torch from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( LLaDAImagePipelineConfig, ) -from sglang.multimodal_gen.configs.sample.llada_image import LLaDAImageSamplingParams from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.llada_image import ( - LLaDAImageQueryFormerModel, - LLaDAImageSigVQModel, - LLaDAImageTextProjectionModel, _LLaDAImageTransformer2DModel, ) from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( @@ -28,281 +22,142 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( TimestepPreparationStage, ) -from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.server_args import ( get_global_server_args, set_global_server_args, ) -class _CaptureSPBlock(torch.nn.Module): - def __init__(self): - super().__init__() - self.skip_values = [] - self.replicated_suffixes = [] - - def forward(self, hidden_states, *args, **kwargs): - self.skip_values.append(kwargs.get("skip_sequence_parallel_override", False)) - self.replicated_suffixes.append(kwargs.get("num_replicated_suffix", 0)) - return hidden_states - - -class TestLLaDAImage(unittest.TestCase): - @classmethod - def setUpClass(cls): - try: - cls.previous_server_args = get_global_server_args() - except ValueError: - cls.previous_server_args = None - set_global_server_args(SimpleNamespace(kv_gather_degree=1, sp_split_auto=False)) - - @classmethod - def tearDownClass(cls): - set_global_server_args(cls.previous_server_args) - - def test_serving_models_reject_gradient_checkpointing(self): - small_config = dict( - hidden_size=8, - intermediate_size=16, - num_hidden_layers=1, - num_attention_heads=2, +@pytest.fixture(autouse=True) +def server_args(): + try: + previous = get_global_server_args() + except ValueError: + previous = None + set_global_server_args(SimpleNamespace(kv_gather_degree=1, sp_split_auto=False)) + yield + set_global_server_args(previous) + + +@pytest.mark.parametrize( + "source,target,shard,total", + [ + ("attention.to_k", "attention.to_qkv", 1, 3), + ("feed_forward.w1", "feed_forward.w13", 0, 2), + ("feed_forward.w3", "feed_forward.w13", 1, 2), + ], +) +def test_dit_weight_mapping(source, target, shard, total): + mapping = get_param_names_mapping( + LLaDAImagePipelineConfig().dit_config.arch_config.param_names_mapping + ) + assert mapping(f"layers.0.{source}.weight") == ( + f"layers.0.{target}.weight", + shard, + total, + ) + + +@pytest.mark.parametrize( + "prompt,instruction", + [("a red car", "Generate an image: a red car"), (None, "Generate an image.")], +) +def test_prompt_format(prompt, instruction): + assert ( + format_llada_image_prompt(prompt) + == f"HUMAN {instruction}\nASSISTANT\n" + ) + + +def test_uniform_scheduler_schedule(): + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0, use_uniform_sigmas=True) + batch = SimpleNamespace( + scheduler=None, + timesteps=None, + sigmas=None, + num_inference_steps=4, + n_tokens=None, + extra={}, + is_warmup=True, + rollout=False, + ) + module = "sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation" + with ( + patch(f"{module}.get_local_torch_device", return_value=torch.device("cpu")), + patch(f"{module}.get_or_create_request_scheduler", return_value=scheduler), + ): + TimestepPreparationStage(scheduler).forward( + batch, SimpleNamespace(pipeline_config=LLaDAImagePipelineConfig()) ) - models = [ - _LLaDAImageTransformer2DModel( - in_channels=4, - dim=128, - n_layers=0, - n_refiner_layers=0, - n_heads=1, - cap_feat_dim=8, - semantic_feat_dim=8, - ), - LLaDAImageQueryFormerModel(num_queries=2, **small_config), - LLaDAImageTextProjectionModel(projection_dim=4, **small_config), - LLaDAImageSigVQModel( - image_size=4, - patch_size=2, - codebook_size=8, - codebook_embed_dim=4, - semantic_embed_dim=4, - **small_config, - ), + assert batch.sigmas is None + torch.testing.assert_close( + scheduler.sigmas, torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0]) + ) + + +@pytest.mark.parametrize("seeds", [(11,), (11, 29)]) +def test_stochastic_scheduler_per_sample_generators(seeds): + scheduler = FlowMatchEulerDiscreteScheduler( + shift=3.0, use_uniform_sigmas=True, stochastic_sampling=True + ) + scheduler.set_timesteps(4, device="cpu") + sample = torch.arange(len(seeds) * 8, dtype=torch.float32).reshape( + len(seeds), 2, 2, 2 + ) + prediction = sample / 16 + generators = [torch.Generator().manual_seed(seed) for seed in seeds] + noise = torch.cat( + [ + torch.randn((1, 2, 2, 2), generator=torch.Generator().manual_seed(seed)) + for seed in seeds ] - for model in models: - with self.subTest(model=type(model).__name__): - with self.assertRaisesRegex( - ValueError, "does not support gradient checkpointing" - ): - model.enable_gradient_checkpointing() - - def test_dit_supports_sglang_flash_attention_and_sdpa(self): - backends = LLaDAImagePipelineConfig().dit_config.arch_config._supported_attention_backends - self.assertEqual( - backends, - {AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, - ) - - def test_dit_weight_mapping_fuses_qkv_and_swiglu_inputs(self): - mapping = get_param_names_mapping( - LLaDAImagePipelineConfig().dit_config.arch_config.param_names_mapping - ) - - self.assertEqual( - mapping("layers.0.attention.to_k.weight"), - ("layers.0.attention.to_qkv.weight", 1, 3), - ) - self.assertEqual( - mapping("layers.0.feed_forward.w1.weight"), - ("layers.0.feed_forward.w13.weight", 0, 2), - ) - self.assertEqual( - mapping("layers.0.feed_forward.w3.weight"), - ("layers.0.feed_forward.w13.weight", 1, 2), - ) - self.assertEqual( - mapping("layers.0.attention.to_q.weight_scale"), - ("layers.0.attention.to_qkv.weight_scale", None, None), - ) - self.assertEqual( - mapping("layers.0.feed_forward.w1.weight_scale"), - ("layers.0.feed_forward.w13.weight_scale", None, None), - ) - - def test_prompt_format_matches_official_pipeline(self): - self.assertEqual( - format_llada_image_prompt("a red car"), - "HUMAN Generate an image: a red car\n" - "ASSISTANT\n", - ) - self.assertEqual( - format_llada_image_prompt(None), - "HUMAN Generate an image.\nASSISTANT\n", - ) - - def test_pipeline_config_defaults_schedule_and_shape(self): - config = LLaDAImagePipelineConfig() - sigmas = config.prepare_sigmas(None, num_inference_steps=8) - - self.assertIsNone(config.text_encoder_mem_fraction_static) - self.assertEqual(len(sigmas), 8) - self.assertTrue(all(left > right for left, right in pairwise(sigmas))) - self.assertEqual( - config.prepare_latent_shape( - SimpleNamespace(height=1024, width=768), - batch_size=1, - num_frames=1, - ), - (1, 128, 64, 48), - ) - - def test_timestep_stage_uses_uniform_scheduler_schedule(self): - scheduler = FlowMatchEulerDiscreteScheduler( - shift=3.0, - use_uniform_sigmas=True, + ) + sigma, next_sigma = scheduler.sigmas[:2] + expected = (1 - next_sigma) * (sample - sigma * prediction) + next_sigma * noise + actual = scheduler.step( + prediction, + scheduler.timesteps[0], + sample, + generator=generators, + return_dict=False, + )[0] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_edit_skips_empty_semantic_refiner(): + class CaptureBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, hidden_states, *args, **kwargs): + self.calls += 1 + return hidden_states + + model = _LLaDAImageTransformer2DModel( + in_channels=4, + dim=64, + n_layers=0, + n_refiner_layers=0, + n_heads=2, + cap_feat_dim=8, + semantic_feat_dim=10, + axes_dims=(8, 12, 12), + axes_lens=(256, 32, 32), + ) + blocks = [CaptureBlock() for _ in range(4)] + for name, block in zip( + ("noise_refiner", "context_refiner", "sigvq_refiner", "layers"), + blocks, + strict=True, + ): + setattr(model, name, torch.nn.ModuleList([block])) + with torch.no_grad(): + model( + x=[torch.randn(4, 1, 4, 4)], + t=torch.tensor([0.5]), + cap_feats=[torch.randn(3, 8)], + glm_cap_feats=[torch.empty(0, 10)], + source_latents=[torch.randn(4, 1, 4, 4)], ) - stage = TimestepPreparationStage(scheduler) - batch = SimpleNamespace( - scheduler=None, - timesteps=None, - sigmas=None, - num_inference_steps=4, - n_tokens=None, - extra={}, - is_warmup=True, - rollout=False, - ) - server_args = SimpleNamespace(pipeline_config=LLaDAImagePipelineConfig()) - module = ( - "sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation" - ) - - with ( - patch(f"{module}.get_local_torch_device", return_value=torch.device("cpu")), - patch(f"{module}.get_or_create_request_scheduler", return_value=scheduler), - ): - stage.forward(batch, server_args) - - self.assertIsNone(batch.sigmas) - torch.testing.assert_close( - scheduler.sigmas, - torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0]), - ) - - def test_stochastic_scheduler_accepts_per_sample_generators(self): - for seeds in ((11,), (11, 29)): - with self.subTest(batch_size=len(seeds)): - scheduler = FlowMatchEulerDiscreteScheduler( - shift=3.0, - use_uniform_sigmas=True, - stochastic_sampling=True, - ) - scheduler.set_timesteps(4, device="cpu") - - sample = torch.arange(len(seeds) * 8, dtype=torch.float32).reshape( - len(seeds), 2, 2, 2 - ) - model_output = sample / 16 - generators = [ - torch.Generator(device="cpu").manual_seed(seed) for seed in seeds - ] - reference_noise = torch.cat( - [ - torch.randn( - (1, *sample.shape[1:]), - generator=torch.Generator(device="cpu").manual_seed(seed), - dtype=sample.dtype, - ) - for seed in seeds - ] - ) - current_sigma, next_sigma = scheduler.sigmas[:2] - expected = (1.0 - next_sigma) * ( - sample - current_sigma * model_output - ) + next_sigma * reference_noise - - actual = scheduler.step( - model_output=model_output, - timestep=scheduler.timesteps[0], - sample=sample, - generator=generators, - return_dict=False, - )[0] - - self.assertEqual(actual.shape, sample.shape) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - def test_pipeline_config_rejects_unaligned_resolution(self): - config = LLaDAImagePipelineConfig() - with self.assertRaisesRegex(ValueError, "height must be divisible by 16"): - config.prepare_latent_shape( - SimpleNamespace(height=1023, width=1024), - batch_size=1, - num_frames=1, - ) - - def test_service_sampling_params_use_few_step_defaults_without_vq_mode(self): - params = LLaDAImageSamplingParams() - field_names = {field.name for field in fields(params)} - - self.assertEqual(params.num_inference_steps, 4) - self.assertEqual(params.guidance_scale, 1.0) - self.assertNotIn("generation_mode", field_names) - self.assertNotIn("vq_token_ids", field_names) - - def test_edit_skips_sigvq_refiner_for_empty_cfg_condition(self): - model_module = "sglang.multimodal_gen.runtime.models.dits.llada_image" - linear_module = "sglang.multimodal_gen.runtime.layers.linear" - with ( - patch(f"{model_module}.get_tp_world_size", return_value=1), - patch(f"{linear_module}.get_tp_group", return_value=None), - patch(f"{linear_module}.get_group_size", return_value=1), - patch(f"{linear_module}.get_group_rank", return_value=0), - patch( - "sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size", - return_value=1, - ), - patch( - "sglang.multimodal_gen.runtime.layers.attention.selector.get_global_server_args", - return_value=SimpleNamespace(attention_backend="torch_sdpa"), - ), - ): - model = _LLaDAImageTransformer2DModel( - in_channels=4, - dim=64, - n_layers=1, - n_refiner_layers=1, - n_heads=2, - cap_feat_dim=8, - semantic_feat_dim=10, - axes_dims=(8, 12, 12), - axes_lens=(256, 32, 32), - ) - - noise_refiner = _CaptureSPBlock() - context_refiner = _CaptureSPBlock() - sigvq_refiner = _CaptureSPBlock() - main_block = _CaptureSPBlock() - model.noise_refiner = torch.nn.ModuleList([noise_refiner]) - model.context_refiner = torch.nn.ModuleList([context_refiner]) - model.sigvq_refiner = torch.nn.ModuleList([sigvq_refiner]) - model.layers = torch.nn.ModuleList([main_block]) - - with ( - torch.no_grad(), - ): - model( - x=[torch.randn(4, 1, 4, 4)], - t=torch.tensor([0.5]), - cap_feats=[torch.randn(3, 8)], - glm_cap_feats=[torch.empty(0, 10)], - source_latents=[torch.randn(4, 1, 4, 4)], - ) - - self.assertEqual(noise_refiner.skip_values, [False]) - self.assertEqual(context_refiner.skip_values, [False]) - self.assertEqual(sigvq_refiner.skip_values, []) - self.assertEqual(main_block.skip_values, [False]) - - -if __name__ == "__main__": - unittest.main() + assert [block.calls for block in blocks] == [1, 1, 0, 1] diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py index a3154c33e45c..31c33e1c689f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -import unittest from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch +import pytest import torch from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( @@ -13,458 +13,225 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( LLaDAImageTextConditioningStage, LLaDAImageTextEncoderRunner, + ensure_conditioning_mask_active, ) -from sglang.srt.runtime_context import ParallelContext, RuntimeContext -_GLOBAL_ARGS_PATCH = ( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" -) - - -class _FakeTextRunner: - def __init__(self): - self.prompts = None - self.component_names = [] - - def encode(self, prompts, max_sequence_length, component_context=None): - self.prompts = (prompts, max_sequence_length) - if component_context is not None: - for component_name in ("queryformer", "text_projection"): - with component_context(component_name=component_name, module=object()): - self.component_names.append(component_name) - return [ - torch.full((index + 2, 4), float(index + 1)) - for index in range(len(prompts)) - ] - - -class TestLLaDAImageTextConditioning(unittest.TestCase): - def setUp(self): - self.runner = _FakeTextRunner() - with patch(_GLOBAL_ARGS_PATCH, return_value=SimpleNamespace()): - self.stage = LLaDAImageTextConditioningStage(self.runner) - - def test_repeats_positive_and_negative_text_for_each_output(self): - batch = SimpleNamespace( - prompt="a red car", - negative_prompt=None, - guidance_scale=5.0, - num_outputs_per_prompt=2, - max_sequence_length=128, - ) - - result = self.stage.forward(batch, server_args=SimpleNamespace()) - - self.assertTrue(result.do_classifier_free_guidance) - self.assertEqual(len(result.prompt_embeds), 2) - self.assertEqual(len(result.negative_prompt_embeds), 2) - self.assertEqual(len(result.prompt_attention_mask), 2) - self.assertEqual(len(result.negative_attention_mask), 2) - self.assertTrue(torch.equal(result.prompt_embeds[0], result.prompt_embeds[1])) - self.assertTrue( - torch.equal( - result.negative_prompt_embeds[0], result.negative_prompt_embeds[1] - ) - ) - - def test_guidance_disabled_has_no_negative_condition(self): - batch = SimpleNamespace( - prompt="a red car", - negative_prompt=None, - guidance_scale=1.0, - num_outputs_per_prompt=2, - max_sequence_length=128, - ) - result = self.stage.forward(batch, server_args=SimpleNamespace()) - - self.assertFalse(result.do_classifier_free_guidance) - self.assertEqual(len(result.prompt_embeds), 2) - self.assertEqual(result.negative_prompt_embeds, []) - self.assertEqual(result.negative_attention_mask, []) - - def test_conditioning_mask_guard_fails_closed(self): - from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( - ensure_conditioning_mask_active, +@pytest.mark.parametrize("guidance", [1.0, 5.0]) +def test_text_conditioning_repeats_outputs_and_masks(guidance): + outputs = [torch.ones(2, 4), torch.full((3, 4), 2.0)] + runner = SimpleNamespace(encode=Mock(return_value=outputs)) + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args", + return_value=SimpleNamespace(), + ): + stage = LLaDAImageTextConditioningStage(runner) + batch = SimpleNamespace( + prompt="a red car", + negative_prompt=None, + guidance_scale=guidance, + num_outputs_per_prompt=2, + max_sequence_length=128, + ) + stage.forward(batch, SimpleNamespace()) + assert batch.do_classifier_free_guidance == (guidance > 1) + assert len(runner.encode.call_args.args[0]) == (2 if guidance > 1 else 1) + for prefix, output in zip(("prompt", "negative"), outputs, strict=True): + embeddings = getattr( + batch, "prompt_embeds" if prefix == "prompt" else "negative_prompt_embeds" ) - - ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=True)) - with self.assertRaisesRegex(RuntimeError, "without its block attention mask"): - ensure_conditioning_mask_active( - SimpleNamespace(conditioning_mask_active=False) + masks = getattr(batch, f"{prefix}_attention_mask") + if prefix == "negative" and guidance == 1: + assert embeddings == masks == [] + continue + assert len(embeddings) == len(masks) == 2 + for embedding, mask in zip(embeddings, masks, strict=True): + torch.testing.assert_close(embedding, output) + torch.testing.assert_close( + mask, torch.ones(output.shape[0], dtype=torch.bool) ) - def test_forward_batch_declares_conditioning_text_lens_field(self): - import dataclasses - - from sglang.srt.model_executor.forward_batch_info import ForwardBatch - - field_names = {field.name for field in dataclasses.fields(ForwardBatch)} - # Batch rebuilds use dataclasses.replace, which drops dynamic attrs. - self.assertIn("llada_image_conditioning_text_lens_cpu", field_names) - - def test_stage_declares_auxiliary_component_residency(self): - uses = self.stage.component_uses(SimpleNamespace(), "conditioning") - - self.assertEqual( - [use.component_name for use in uses], - ["queryformer", "text_projection"], - ) - - batch = SimpleNamespace( - prompt="a red car", - negative_prompt=None, - guidance_scale=1.0, - num_outputs_per_prompt=1, - max_sequence_length=128, - ) - self.stage.forward(batch, server_args=SimpleNamespace()) - - self.assertEqual( - self.runner.component_names, ["queryformer", "text_projection"] - ) - - def test_text_runner_scopes_singleton_attention_group_and_restores(self): - import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state - import sglang.srt.distributed.parallel_state as srt_parallel_state - diffusion_group = object() +def test_conditioning_mask_guard(): + ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=True)) + with pytest.raises(RuntimeError, match="without its block attention mask"): + ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=False)) + + +def test_prefill_uses_resolved_page_size(): + runner = object.__new__(LLaDAImageTextEncoderRunner) + runner.tokenizer = lambda *args, **kwargs: SimpleNamespace(input_ids=[[1, 2]]) + runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) + runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) + runner.server_args = SimpleNamespace( + page_size=None, chunked_prefill_size=-1, max_prefill_tokens=8192 + ) + runner.page_size = 64 + runner.tree_cache = runner.token_to_kv_pool_allocator = object() + with ( + patch( + "sglang.srt.sampling.sampling_params.SamplingParams", + return_value=SimpleNamespace(normalize=lambda _: None), + ), + patch("sglang.srt.managers.schedule_batch.Req", return_value=SimpleNamespace()), + patch( + "sglang.srt.managers.schedule_policy.PrefillAdder", + side_effect=RuntimeError("stop at prefill"), + ) as adder, + pytest.raises(RuntimeError, match="stop at prefill"), + ): + runner._encode_impl(["hello"], max_sequence_length=16) + assert adder.call_args.args[0] == 64 + + +@pytest.mark.parametrize("fail_init", [False, True]) +def test_encoder_context_restores_after_init_and_encode(fail_init): + import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_state + import sglang.srt.distributed.parallel_state as srt_state + from sglang.srt import runtime_context as rc + from sglang.srt.server_args import ServerArgs as SRTServerArgs + + saved = rc.snapshot_context() + rc.reset_context() + try: + outer_args = SRTServerArgs(model_path="dummy", tp_size=1) + outer = rc.publish(outer_args, role="diffusion_gpu_worker") + outer.override("diffusion.setup", grammar_backend="none") + outer_bags, outer_log = outer._config_bags, outer.overrides_log() + outer_buffer = rc.get_buffer("conditioning", object) + diffusion_group = SimpleNamespace(world_size=1, rank_in_group=0) encoder_group = SimpleNamespace(world_size=1, rank_in_group=0) - encoder_attention_group = object() - runner = object.__new__(LLaDAImageTextEncoderRunner) - runner.runtime_context = RuntimeContext(ParallelContext()) - runner.encoder_tp_group = encoder_group - runner.encoder_attn_tp_group = encoder_attention_group + attention_group = object() observed = {} - def fail_encode(*_args, **_kwargs): - observed["tp"] = srt_parallel_state._TP - observed["attn_tp"] = srt_parallel_state._ATTN_TP - raise RuntimeError("encode failed") - - runner._encode_impl = fail_encode - with ( - patch.object(mm_parallel_state, "_TP", diffusion_group), - patch.object(srt_parallel_state, "_TP", diffusion_group), - patch.object(srt_parallel_state, "_ATTN_TP", diffusion_group), - ): - with self.assertRaisesRegex(RuntimeError, "encode failed"): - runner.encode(["hello"], max_sequence_length=16) - - self.assertIs(mm_parallel_state._TP, diffusion_group) - self.assertIs(srt_parallel_state._TP, diffusion_group) - self.assertIs(srt_parallel_state._ATTN_TP, diffusion_group) - - self.assertIs(observed["tp"], encoder_group) - self.assertIs(observed["attn_tp"], encoder_attention_group) - - def test_text_runner_uses_resolved_page_size_for_prefill(self): - class StopAfterPrefillAdder(Exception): - pass - - runner = object.__new__(LLaDAImageTextEncoderRunner) - runner.tokenizer = lambda *args, **kwargs: SimpleNamespace(input_ids=[[1, 2]]) - runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) - runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) - runner.server_args = SimpleNamespace( - page_size=None, - chunked_prefill_size=-1, - max_prefill_tokens=8192, - ) - runner.page_size = 64 - runner.tree_cache = object() - runner.token_to_kv_pool_allocator = object() - with ( - patch( - "sglang.srt.sampling.sampling_params.SamplingParams", - return_value=SimpleNamespace(normalize=lambda _: None), - ), - patch( - "sglang.srt.managers.schedule_batch.Req", - return_value=SimpleNamespace(), - ), - patch( - "sglang.srt.managers.schedule_policy.PrefillAdder", - side_effect=StopAfterPrefillAdder, - ) as prefill_adder_cls, - ): - with self.assertRaises(StopAfterPrefillAdder): - runner._encode_impl(["hello"], max_sequence_length=16) - - self.assertEqual(prefill_adder_cls.call_args.args[0], runner.page_size) - - def test_text_runner_restores_diffusion_groups_when_worker_init_fails(self): - import sglang.srt.distributed.parallel_state as srt_parallel_state + def check_encoder(): + context = rc.assert_published(observed["args"], role="scheduler") + assert rc.get_schedule().max_running_requests == 2 + assert rc.get_parallel().tp_size == 1 + return context + + def allocate(): + check_encoder().override("encoder.setup", page_size=16) + observed["buffer"] = rc.get_buffer("conditioning", object) + if fail_init: + raise RuntimeError("allocation failed") + + def make_worker(server_args, **kwargs): + observed["args"] = server_args + observed["context"] = check_encoder() + srt_state._TP, srt_state._ATTN_TP = encoder_group, attention_group + return SimpleNamespace( + model_runner=SimpleNamespace(page_size=16), + get_memory_pool=lambda: (object(), object()), + alloc_memory_pool=allocate, + init_attention_backends=check_encoder, + init_cuda_graphs=check_encoder, + ) - diffusion_group = object() - encoder_group = object() + def make_cache(params): + assert params.page_size == rc.get_schedule().page_size == 16 + return object() - def fail_after_installing_encoder_groups(**_kwargs): - srt_parallel_state._TP = encoder_group - srt_parallel_state._ATTN_TP = encoder_group - return SimpleNamespace( - alloc_memory_pool=lambda: (_ for _ in ()).throw( - RuntimeError("allocation failed") - ) + def check_diffusion(): + assert rc.get_context() is outer and rc.get_server_args() is outer_args + assert rc.publish_role() == "diffusion_gpu_worker" + assert ( + outer._config_bags is outer_bags and outer.overrides_log() == outer_log + ) + assert rc.get_exec().kernel.grammar_backend == "none" + assert rc.get_buffer("conditioning", object) is outer_buffer + assert all( + group is diffusion_group + for group in (mm_state._TP, srt_state._TP, srt_state._ATTN_TP) ) with ( patch( - "sglang.srt.managers.tp_worker.TpModelWorker", - side_effect=fail_after_installing_encoder_groups, + "sglang.srt.server_args.ServerArgs", + side_effect=lambda **kwargs: SRTServerArgs( + **dict(kwargs, model_path="dummy") + ), ), patch( - "sglang.srt.runtime_context.create_context", - side_effect=lambda *_args, **_kwargs: RuntimeContext(ParallelContext()), + "sglang.srt.managers.tp_worker.TpModelWorker", side_effect=make_worker ), patch( - "sglang.srt.server_args.ServerArgs", - side_effect=lambda **kwargs: SimpleNamespace(page_size=1, **kwargs), + "sglang.srt.mem_cache.chunk_cache.ChunkCache", side_effect=make_cache ), patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning.get_local_torch_device", + f"{LLaDAImageTextEncoderRunner.__module__}.get_local_torch_device", return_value=torch.device("cpu"), ), - patch.object(srt_parallel_state, "_TP", diffusion_group), - patch.object(srt_parallel_state, "_ATTN_TP", diffusion_group), + patch.object(mm_state, "_TP", diffusion_group), + patch.object(srt_state, "_TP", diffusion_group), + patch.object(srt_state, "_ATTN_TP", diffusion_group), ): - with self.assertRaisesRegex(RuntimeError, "allocation failed"): - LLaDAImageTextEncoderRunner( - model_root="/unused/model", - queryformer=object(), - text_projection=object(), - tokenizer=object(), - server_args=SimpleNamespace( - sp_degree=1, - nccl_port=29500, - trust_remote_code=True, - revision=None, - pipeline_config=SimpleNamespace( - text_encoder_mem_fraction_static=0.1 - ), - ), - ) - - self.assertIs(srt_parallel_state._TP, diffusion_group) - self.assertIs(srt_parallel_state._ATTN_TP, diffusion_group) - - -class TestLLaDAImageRuntimeContext(unittest.TestCase): - def test_encoder_preserves_diffusion_runtime(self): - """The embedded encoder retains its context inside the diffusion TP scope.""" - import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_state - import sglang.srt.distributed.parallel_state as srt_state - from sglang.srt import runtime_context as rc - from sglang.srt.server_args import ServerArgs as SRTServerArgs - - saved = rc.snapshot_context() - self.addCleanup(rc.restore_context, saved) - conditioning = LLaDAImageTextEncoderRunner.__module__ - for fail_init in (False, True): - with self.subTest(fail_init=fail_init): - rc.reset_context() - diffusion_args = SRTServerArgs(model_path="dummy", tp_size=1) - diffusion_context = rc.publish( - diffusion_args, role="diffusion_gpu_worker" - ) - diffusion_context.override("diffusion.setup", grammar_backend="none") - diffusion_bags = diffusion_context._config_bags - diffusion_log = diffusion_context.overrides_log() - diffusion_buffer = rc.get_buffer("conditioning", object) - diffusion_group = SimpleNamespace(world_size=1, rank_in_group=0) - encoder_group = SimpleNamespace(world_size=1, rank_in_group=0) - encoder_attention_group = SimpleNamespace(world_size=1, rank_in_group=0) - observed = {} - - def make_args(**kwargs): - kwargs["model_path"] = "dummy" - return SRTServerArgs(**kwargs) - - def check_encoder_context(): - context = rc.assert_published(observed["args"], role="scheduler") - self.assertEqual(rc.get_parallel().tp_size, 1) - self.assertEqual(rc.get_parallel().attn_tp_size, 1) - self.assertEqual(rc.get_parallel().attn_dp_size, 1) - self.assertEqual(rc.get_parallel().moe_tp_size, 1) - self.assertEqual(rc.get_schedule().max_running_requests, 2) - return context - - def allocate(): - context = check_encoder_context() - context.override("encoder.setup", page_size=16) - rc.get_buffer("conditioning", object) - if fail_init: - raise RuntimeError("allocation failed") - - def make_worker(**kwargs): - observed["args"] = kwargs["server_args"] - observed["context"] = check_encoder_context() - srt_state._TP = encoder_group - srt_state._ATTN_TP = encoder_attention_group - return SimpleNamespace( - model_runner=SimpleNamespace(page_size=16), - get_memory_pool=lambda: (object(), object()), - alloc_memory_pool=allocate, - init_attention_backends=check_encoder_context, - init_cuda_graphs=check_encoder_context, - ) - - def make_cache(_params): - self.assertEqual(rc.get_schedule().page_size, 16) - return object() - - def check_diffusion_context(): - self.assertIs(rc.get_context(), diffusion_context) - self.assertIs(rc.get_server_args(), diffusion_args) - self.assertEqual(rc.publish_role(), "diffusion_gpu_worker") - self.assertIs(rc.get_context()._config_bags, diffusion_bags) - self.assertEqual(rc.get_context().overrides_log(), diffusion_log) - self.assertEqual(rc.get_exec().kernel.grammar_backend, "none") - self.assertEqual(rc.get_parallel().tp_size, 1) - self.assertEqual(rc.get_parallel().attn_tp_size, 1) - self.assertEqual(rc.get_parallel().attn_dp_size, 1) - self.assertEqual(rc.get_parallel().moe_tp_size, 1) - self.assertIs( - rc.get_buffer("conditioning", object), diffusion_buffer - ) - self.assertIs(mm_state._TP, diffusion_group) - self.assertIs(srt_state._TP, diffusion_group) - self.assertIs(srt_state._ATTN_TP, diffusion_group) - - with ( - patch("sglang.srt.server_args.ServerArgs", side_effect=make_args), - patch( - "sglang.srt.managers.tp_worker.TpModelWorker", - side_effect=make_worker, - ), - patch( - "sglang.srt.mem_cache.chunk_cache.ChunkCache", - side_effect=make_cache, - ), - patch( - f"{conditioning}.get_local_torch_device", - return_value=torch.device("cpu"), - ), - patch.object(mm_state, "_TP", diffusion_group), - patch.object(srt_state, "_TP", diffusion_group), - patch.object(srt_state, "_ATTN_TP", diffusion_group), - ): - kwargs = dict( - model_root="/unused/model", - queryformer=object(), - text_projection=object(), - tokenizer=object(), - server_args=SimpleNamespace( - sp_degree=1, - nccl_port=29500, - trust_remote_code=False, - revision=None, - component_paths={}, - pipeline_config=SimpleNamespace( - text_encoder_mem_fraction_static=0.1 - ), - ), + args = SimpleNamespace( + nccl_port=29500, + trust_remote_code=False, + revision=None, + pipeline_config=SimpleNamespace(text_encoder_mem_fraction_static=0.1), + ) + if fail_init: + with pytest.raises(RuntimeError, match="allocation failed"): + LLaDAImageTextEncoderRunner( + "/unused", object(), object(), object(), args ) - if fail_init: - with self.assertRaisesRegex(RuntimeError, "allocation failed"): - LLaDAImageTextEncoderRunner(**kwargs) - check_diffusion_context() - continue - runner = LLaDAImageTextEncoderRunner(**kwargs) - check_diffusion_context() - self.assertIsNot(observed["context"], diffusion_context) - - def encode(*_args, **_kwargs): - self.assertIs(check_encoder_context(), observed["context"]) - self.assertEqual(rc.get_schedule().page_size, 16) - self.assertEqual(len(rc.get_context().overrides_log()), 1) - self.assertIsNot( - rc.get_buffer("conditioning", object), diffusion_buffer - ) - self.assertIs(srt_state._ATTN_TP, encoder_attention_group) - if observed["fail_encode"]: - raise RuntimeError("encode failed") - return ["encoded"] - - runner._encode_impl = encode - for fail_encode in (False, True, False): - observed["fail_encode"] = fail_encode - if fail_encode: - with self.assertRaisesRegex(RuntimeError, "encode failed"): - runner.encode(["hello"], max_sequence_length=16) - else: - self.assertEqual( - runner.encode(["hello"], max_sequence_length=16), - ["encoded"], - ) - check_diffusion_context() - - -class TestLLaDAImageConditionKwargs(unittest.TestCase): - def setUp(self): - self.config = LLaDAImagePipelineConfig() - self.semantic = [ - torch.full((3, 5), 1.0), - torch.full((3, 5), 2.0), - ] - self.source = [ - torch.full((8, 1, 2, 3), 3.0), - torch.full((8, 1, 2, 3), 4.0), - ] - self.batch = SimpleNamespace( - batch_size=2, - image_embeds=self.semantic, - source_latents=self.source, - do_classifier_free_guidance=True, - ) - - def test_cfg_uses_semantics_only_on_positive_branch(self): - positive = self.config.prepare_pos_cond_kwargs( - self.batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float64 - ) - negative = self.config.prepare_neg_cond_kwargs( - self.batch, torch.device("cpu"), rotary_emb=None, dtype=torch.float64 - ) - policy = CFGPolicy().build( - self.batch, - {"encoder_hidden_states_image": self.semantic}, - positive, - negative, - ) - - positive_kwargs = policy.branches[0].kwargs - negative_kwargs = policy.branches[1].kwargs - self.assertEqual( - [tuple(x.shape) for x in positive_kwargs["encoder_hidden_states_image"]], - [(3, 5)] * 2, - ) - self.assertEqual( - [tuple(x.shape) for x in negative_kwargs["encoder_hidden_states_image"]], - [(0, 5)] * 2, - ) - self.assertTrue( - all(x.dtype == torch.float64 for x in positive_kwargs["source_latents"]) - ) - for positive_source, negative_source in zip( - positive_kwargs["source_latents"], - negative_kwargs["source_latents"], - strict=True, - ): - torch.testing.assert_close(positive_source, negative_source) - - def test_rejects_condition_batch_length_mismatch(self): - self.batch.image_embeds = self.semantic[:1] - - with self.assertRaisesRegex( - ValueError, "image_embeds has 1 entries, expected 2" - ): - self.config.prepare_pos_cond_kwargs( - self.batch, - torch.device("cpu"), - rotary_emb=None, - dtype=torch.float32, + check_diffusion() + return + runner = LLaDAImageTextEncoderRunner( + "/unused", object(), object(), object(), args ) - - -if __name__ == "__main__": - unittest.main() + check_diffusion() + assert observed["context"] is not outer + + def encode(*args, **kwargs): + assert check_encoder() is observed["context"] + assert rc.get_schedule().page_size == 16 + assert rc.get_buffer("conditioning", object) is observed["buffer"] + assert observed["buffer"] is not outer_buffer + assert ( + srt_state._TP is encoder_group + and srt_state._ATTN_TP is attention_group + ) + if observed["fail_encode"]: + raise RuntimeError("encode failed") + return ["encoded"] + + runner._encode_impl = encode + for fail_encode in (False, True, False): + observed["fail_encode"] = fail_encode + if fail_encode: + with pytest.raises(RuntimeError, match="encode failed"): + runner.encode(["hello"], 16) + else: + assert runner.encode(["hello"], 16) == ["encoded"] + check_diffusion() + finally: + rc.restore_context(saved) + + +def test_edit_cfg_preserves_source_and_removes_negative_semantics(): + config = LLaDAImagePipelineConfig() + batch = SimpleNamespace( + batch_size=2, + image_embeds=[torch.ones(3, 5)] * 2, + source_latents=[torch.ones(8, 1, 2, 3)] * 2, + do_classifier_free_guidance=True, + ) + positive = config.prepare_pos_cond_kwargs(batch, "cpu", None, torch.float64) + negative = config.prepare_neg_cond_kwargs(batch, "cpu", None, torch.float64) + policy = CFGPolicy().build( + batch, {"encoder_hidden_states_image": batch.image_embeds}, positive, negative + ) + for branch, shape in zip(policy.branches, [(3, 5), (0, 5)], strict=True): + assert [ + tuple(value.shape) for value in branch.kwargs["encoder_hidden_states_image"] + ] == [shape] * 2 + for actual, source in zip( + branch.kwargs["source_latents"], batch.source_latents, strict=True + ): + torch.testing.assert_close(actual, source.double()) + batch.image_embeds = batch.image_embeds[:1] + with pytest.raises(ValueError, match="image_embeds has 1 entries, expected 2"): + config.prepare_pos_cond_kwargs(batch, "cpu", None, torch.float32) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py index 19d1ca26bc4b..7b8f74d99668 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 import json -import shutil from unittest.mock import patch import pytest @@ -24,48 +23,49 @@ from sglang.multimodal_gen.runtime.pipelines.llada_image import LLaDAImagePipeline from sglang.multimodal_gen.runtime.server_args import ServerArgs -# ServerArgs uses the physical device to resolve residency policies. pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA is required" + not torch.cuda.is_available(), + reason="ServerArgs residency resolution requires CUDA", ) -_COMPONENTS = { - "queryformer": (LLaDAImageQueryFormerModel, {"num_queries": 2}), - "text_projection": (LLaDAImageTextProjectionModel, {"projection_dim": 4}), - "sigvq": ( - LLaDAImageSigVQModel, - { - "image_size": 4, - "patch_size": 2, - "codebook_size": 8, - "codebook_embed_dim": 4, - "semantic_embed_dim": 4, - }, - ), -} - -@pytest.fixture(params=list(_COMPONENTS)) -def checkpoint(request, tmp_path): - name = request.param - model_cls, extra_config = _COMPONENTS[name] +@pytest.mark.parametrize( + "name,model_cls,extra", + [ + ("queryformer", LLaDAImageQueryFormerModel, {"num_queries": 2}), + ("text_projection", LLaDAImageTextProjectionModel, {"projection_dim": 4}), + ( + "sigvq", + LLaDAImageSigVQModel, + { + "image_size": 4, + "patch_size": 2, + "codebook_size": 8, + "codebook_embed_dim": 4, + "semantic_embed_dim": 4, + }, + ), + ], +) +@pytest.mark.parametrize("missing_weight", [False, True]) +def test_auxiliary_checkpoint_loading(tmp_path, name, model_cls, extra, missing_weight): config = dict( hidden_size=8, intermediate_size=16, num_hidden_layers=1, num_attention_heads=2, - **extra_config, + **extra, ) torch.manual_seed(0) - reference = model_cls(**config) weights = { - key: tensor.contiguous() for key, tensor in reference.state_dict().items() + key: value.contiguous() + for key, value in model_cls(**config).state_dict().items() } component = tmp_path / name component.mkdir() - config.update(_class_name=model_cls.__name__, _diffusers_version="0") - (component / "config.json").write_text(json.dumps(config)) - save_file(weights, component / "model.safetensors") + (component / "config.json").write_text( + json.dumps(dict(config, _class_name=model_cls.__name__, _diffusers_version="0")) + ) (tmp_path / "model_index.json").write_text( json.dumps( { @@ -76,85 +76,29 @@ def checkpoint(request, tmp_path): } ) ) - return name, component, weights - - -def _load(name, root, **overrides): - server_args = ServerArgs( - model_path=str(root), - component_residency={name: COMPONENT_OFFLOAD}, - **overrides, + if missing_weight: + weights.pop(next(iter(weights))) + save_file(weights, component / "model.safetensors") + args = ServerArgs( + model_path=str(tmp_path), component_residency={name: COMPONENT_OFFLOAD} ) pipeline = object.__new__(LLaDAImagePipeline) - pipeline.model_path = str(root) - pipeline.server_args = server_args + pipeline.model_path, pipeline.server_args = str(tmp_path), args pipeline._disagg_role = RoleType.MONOLITHIC pipeline._required_config_modules = [name] pipeline.memory_usages = {} - return pipeline.load_modules(server_args)[name], server_args - - -@pytest.mark.parametrize( - "layout", - ["single", "sharded", "weight_override", "component_override", "architecture"], -) -def test_auxiliary_checkpoint_loading(checkpoint, tmp_path, layout): - name, component, weights = checkpoint - root = component.parent - overrides = {} - expected_dtype = torch.bfloat16 - if layout in ("sharded", "weight_override"): - overrides["component_precisions"] = {name: "fp32"} - expected_dtype = torch.float32 - if layout == "sharded": - names = list(weights) - partitions = (names[::2], names[1::2]) - weight_map = {} - for index, keys in enumerate(partitions): - filename = f"shard-{index}.safetensors" - save_file({key: weights[key] for key in keys}, component / filename) - weight_map.update({key: filename for key in keys}) - (component / "diffusion_pytorch_model.safetensors.index.json").write_text( - json.dumps({"weight_map": weight_map}) - ) - save_file({"unrelated": torch.zeros(1)}, component / "model.safetensors") - elif layout == "weight_override": - weights = {key: value + 1 for key, value in weights.items()} - override = tmp_path / "override.safetensors" - save_file(weights, override) - overrides["component_weights_paths"] = {name: str(override)} - elif layout == "component_override": - override = tmp_path / "alternate-component" - shutil.copytree(component, override) - component = override - weights = {key: value + 1 for key, value in weights.items()} - save_file(weights, component / "model.safetensors") - overrides["component_paths"] = {name: str(component)} - elif layout == "architecture": - config = json.loads((component / "config.json").read_text()) - config.pop("_class_name") - (component / "config.json").write_text(json.dumps(config)) - - model, server_args = _load(name, root, **overrides) - - assert isinstance(model, _COMPONENTS[name][0]) - assert not model.training - assert server_args.model_paths[name] == str(component) - actual = model.state_dict() - assert actual.keys() == weights.keys() - for key, value in actual.items(): - assert value.device.type == "cpu" - assert value.dtype == expected_dtype - torch.testing.assert_close( - value, weights[key].to(expected_dtype), rtol=0, atol=0 - ) - - -def test_incomplete_auxiliary_checkpoint_does_not_fallback(checkpoint): - name, component, weights = checkpoint - weights.pop(next(iter(weights))) - save_file(weights, component / "model.safetensors") - with patch.object(ComponentLoader, "load_native") as native: - with pytest.raises(ComponentCheckpointUnsupportedError, match="Missing:"): - _load(name, component.parent) - native.assert_not_called() + if missing_weight: + with ( + patch.object(ComponentLoader, "load_native") as native, + pytest.raises(ComponentCheckpointUnsupportedError, match="Missing:"), + ): + pipeline.load_modules(args) + native.assert_not_called() + return + model = pipeline.load_modules(args)[name] + assert isinstance(model, model_cls) and not model.training + assert args.model_paths[name] == str(component) + assert model.state_dict().keys() == weights.keys() + for key, value in model.state_dict().items(): + assert value.device.type == "cpu" and value.dtype == torch.bfloat16 + torch.testing.assert_close(value, weights[key].bfloat16(), rtol=0, atol=0) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py deleted file mode 100644 index b9975da0ba0b..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_pipeline_config.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import unittest -from types import SimpleNamespace - -import torch -from PIL import Image - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) - - -class TestLLaDAImagePipelineConfig(unittest.TestCase): - def setUp(self): - self.config = LLaDAImagePipelineConfig() - - def test_edit_keeps_requested_default_output_size(self): - image = Image.new("RGB", (768, 512)) - - self.assertIsNone( - self.config.calculate_condition_image_size(image, image.width, image.height) - ) - self.assertIsNone(self.config.prepare_calculated_size(image)) - - def test_decode_preprocessing_matches_official_vae_dtype_order(self): - class FakeVAE(torch.nn.Module): - def __init__(self): - super().__init__() - self.anchor = torch.nn.Parameter( - torch.zeros((), dtype=torch.bfloat16), requires_grad=False - ) - self.bn = SimpleNamespace( - running_mean=torch.tensor( - [0.01, -0.02, 0.03, -0.04], dtype=torch.float32 - ), - running_var=torch.full((4,), 0.0129973, dtype=torch.float32), - ) - self.config = SimpleNamespace( - arch_config=SimpleNamespace(batch_norm_eps=0.003) - ) - - vae = FakeVAE() - latents = torch.tensor( - [[[[0.501]], [[-0.249]], [[0.126]], [[-0.751]]]], - dtype=torch.float32, - ) - official_latents = latents.to(torch.bfloat16) - latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(official_latents) - latent_std = torch.sqrt( - vae.bn.running_var.view(1, -1, 1, 1) + vae.config.arch_config.batch_norm_eps - ).to(official_latents) - expected = official_latents * latent_std + latent_mean - expected = expected.reshape(1, 1, 2, 2, 1, 1) - expected = expected.permute(0, 1, 4, 2, 5, 3).reshape(1, 1, 2, 2) - - actual = self.config.preprocess_decoding(latents, vae=vae) - - self.assertEqual(actual.dtype, torch.bfloat16) - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - -if __name__ == "__main__": - unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py index c563c51617fe..7b41dd9250ef 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py @@ -5,61 +5,39 @@ import pytest import torch -from sglang.multimodal_gen.runtime.layers import custom_op from sglang.multimodal_gen.runtime.layers import layernorm as layernorm_module from sglang.multimodal_gen.runtime.models.dits.llada_image import LLaDAImageRMSNorm -from sglang.multimodal_gen.runtime.platforms import current_platform -@pytest.mark.skipif( - not current_platform.is_cuda() or not torch.cuda.is_available(), - reason="The HF RMSNorm kernel requires CUDA", +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is required" + ), + ), + ], ) -def test_llada_image_rmsnorm_dispatches_to_hf_kernel(): +def test_rmsnorm_cast_order_and_cuda_dispatch(device): torch.manual_seed(0) - hidden_size = 128 - norm = LLaDAImageRMSNorm(hidden_size, eps=1e-5).cuda().to(torch.bfloat16) + norm = LLaDAImageRMSNorm(128, eps=1e-5).to(device=device, dtype=torch.bfloat16) with torch.no_grad(): - norm.weight.copy_(torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)) - hidden_states = torch.randn( - 2, - 8, - hidden_size, - device="cuda", - dtype=torch.bfloat16, - ) - - with patch.object( - layernorm_module, - "rmsnorm_hf", - wraps=layernorm_module.rmsnorm_hf, - ) as rmsnorm_hf: - output = norm(hidden_states) - - rmsnorm_hf.assert_called_once() - variance = hidden_states.float().pow(2).mean(-1, keepdim=True) - expected = norm.weight * ( - hidden_states.float() * torch.rsqrt(variance + norm.variance_epsilon) - ).to(hidden_states.dtype) - torch.testing.assert_close(output, expected, atol=1e-2, rtol=1e-2) - - -def test_llada_image_rmsnorm_hip_fallback_preserves_cast_order(monkeypatch): - monkeypatch.setattr(custom_op, "_is_cuda", False) - monkeypatch.setattr(current_platform, "is_hip", lambda: True) - monkeypatch.setattr(layernorm_module, "USE_AITER", False) - torch.manual_seed(0) - norm = LLaDAImageRMSNorm(128, eps=1e-5).to(torch.bfloat16) - with torch.no_grad(): - norm.weight.copy_(torch.randn(128, dtype=torch.bfloat16)) - hidden_states = torch.randn(2, 8, 128, dtype=torch.bfloat16) - - with patch.object(norm, "forward_native", wraps=norm.forward_native) as native: - output = norm(hidden_states) - - native.assert_called_once() - variance = hidden_states.float().pow(2).mean(-1, keepdim=True) + norm.weight.copy_(torch.randn(128, device=device, dtype=torch.bfloat16)) + hidden = torch.randn(2, 8, 128, device=device, dtype=torch.bfloat16) + if device == "cuda": + with patch.object( + layernorm_module, "rmsnorm_hf", wraps=layernorm_module.rmsnorm_hf + ) as kernel: + actual = norm(hidden) + kernel.assert_called_once() + else: + actual = norm.forward_native(hidden) + variance = hidden.float().pow(2).mean(-1, keepdim=True) expected = norm.weight * ( - hidden_states.float() * torch.rsqrt(variance + norm.variance_epsilon) - ).to(hidden_states.dtype) - torch.testing.assert_close(output, expected, rtol=0, atol=0) + hidden.float() * torch.rsqrt(variance + norm.variance_epsilon) + ).to(hidden.dtype) + tolerance = 1e-2 if device == "cuda" else 0 + torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py index 720593df7298..47d2899a6d38 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py @@ -1,174 +1,101 @@ # SPDX-License-Identifier: Apache-2.0 -import unittest from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch +import pytest import torch from PIL import Image +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.source import ( LLaDAImageSourceImageConditioningStage, ) -_GLOBAL_ARGS_PATCH = ( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" -) - - -class _FakeImageProcessor: - def __init__(self): - self.calls = [] - self.output = torch.arange(3 * 8 * 12, dtype=torch.float32).reshape(1, 3, 8, 12) - - def preprocess(self, image, height, width, resize_mode): - self.calls.append((image, height, width, resize_mode)) - return self.output.clone() - - -class _FakeSigVQ(torch.nn.Module): - def __init__(self): - super().__init__() - self.anchor = torch.nn.Parameter(torch.zeros(()), requires_grad=False) - self.config = SimpleNamespace(patch_size=16) - self.pixel_values = None - - def forward(self, pixel_values): - self.pixel_values = pixel_values - batch_size = pixel_values.shape[0] - features = torch.arange( - batch_size * 3 * 5, - device=pixel_values.device, - dtype=pixel_values.dtype, - ).reshape(batch_size, 3, 5) - return SimpleNamespace(semantic_features=features) - - -class _FakePosterior: - def __init__(self, value): - self.value = value - - def mode(self): - return self.value - - -class _FakeVAE(torch.nn.Module): - def __init__(self): - super().__init__() - self.anchor = torch.nn.Parameter(torch.zeros(()), requires_grad=False) - self.bn = SimpleNamespace( - running_mean=torch.arange(8, dtype=torch.float32), - running_var=torch.full((8,), 3.75, dtype=torch.float32), - ) - self.config = SimpleNamespace(arch_config=SimpleNamespace(batch_norm_eps=0.25)) - self.encoded_image = None - self.raw_latents = None - - def encode(self, image): - self.encoded_image = image - batch_size = image.shape[0] - self.raw_latents = torch.arange( - batch_size * 2 * 4 * 6, - device=image.device, - dtype=image.dtype, - ).reshape(batch_size, 2, 4, 6) - return _FakePosterior(self.raw_latents) - - -class TestLLaDAImageSourceImageConditioningStage(unittest.TestCase): - def setUp(self): - self.processor = _FakeImageProcessor() - self.sigvq = _FakeSigVQ() - self.vae = _FakeVAE() - with patch(_GLOBAL_ARGS_PATCH, return_value=SimpleNamespace()): - self.stage = LLaDAImageSourceImageConditioningStage( - sigvq=self.sigvq, - vae=self.vae, - image_processor=self.processor, - ) - - def test_text_generation_is_a_noop(self): - batch = SimpleNamespace( - condition_image=None, - image_embeds=[], - source_latents=None, - ) - - result = self.stage.forward(batch, server_args=SimpleNamespace()) - - self.assertIs(result, batch) - self.assertEqual(result.image_embeds, []) - self.assertIsNone(result.source_latents) - self.assertEqual(self.processor.calls, []) - - def test_edit_matches_source_preprocessing_and_repeats_outputs(self): - image = Image.new("RGB", (40, 30), color="red") - batch = SimpleNamespace( - condition_image=image, - height=64, - width=80, - batch_size=2, - image_embeds=[], - source_latents=None, - ) - - result = self.stage.forward(batch, server_args=SimpleNamespace()) - - self.assertEqual(self.processor.calls, [(image, 64, 80, "crop")]) - self.assertEqual(tuple(self.sigvq.pixel_values.shape), (2, 3, 32, 48)) - self.assertTrue( - torch.equal(self.vae.encoded_image[0], self.processor.output[0]) - ) - self.assertTrue( - torch.equal(self.vae.encoded_image[0], self.vae.encoded_image[1]) - ) - self.assertEqual([tuple(x.shape) for x in result.image_embeds], [(3, 5)] * 2) - self.assertEqual( - [tuple(x.shape) for x in result.source_latents], [(8, 1, 2, 3)] * 2 - ) - - raw = self.vae.raw_latents - expected = raw.reshape(2, 2, 2, 2, 3, 2) - expected = expected.permute(0, 1, 3, 5, 2, 4).reshape(2, 8, 2, 3) - expected = (expected - self.vae.bn.running_mean.view(1, -1, 1, 1)) / 2 - expected = [sample.unsqueeze(1) for sample in expected] - for actual, wanted in zip(result.source_latents, expected, strict=True): - torch.testing.assert_close(actual, wanted, rtol=0, atol=0) - - def test_edit_rejects_more_than_one_source_image(self): - batch = SimpleNamespace( - condition_image=[Image.new("RGB", (32, 32))] * 2, - height=64, - width=64, - batch_size=1, - ) - - with self.assertRaisesRegex(ValueError, "exactly one source image"): - self.stage.forward(batch, server_args=SimpleNamespace()) - - def test_source_normalization_matches_official_std_cast_order(self): - latents = torch.tensor([[[[0.5]]]], dtype=torch.bfloat16) - vae = SimpleNamespace( - bn=SimpleNamespace( - running_mean=torch.tensor([0.01], dtype=torch.float32), - running_var=torch.tensor([0.0129973], dtype=torch.float32), - ), - config=SimpleNamespace(arch_config=SimpleNamespace(batch_norm_eps=0.003)), - ) - expected_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) - expected_std = torch.sqrt( - vae.bn.running_var.view(1, -1, 1, 1) + vae.config.arch_config.batch_norm_eps - ).to(latents) - - actual = self.stage._normalize_latents(latents, vae) - - torch.testing.assert_close( - actual, - (latents - expected_mean) / expected_std, - rtol=0, - atol=0, - ) - -if __name__ == "__main__": - unittest.main() +def test_source_image_preprocessing_and_output_replication(): + pixels = torch.arange(3 * 8 * 12, dtype=torch.float32).reshape(1, 3, 8, 12) + raw = torch.arange(2 * 2 * 4 * 6, dtype=torch.float32).reshape(2, 2, 4, 6) + semantics = torch.arange(30, dtype=torch.float32).reshape(2, 3, 5) + processor = SimpleNamespace(preprocess=Mock(return_value=pixels)) + sigvq = Mock( + return_value=SimpleNamespace(semantic_features=semantics), + config=SimpleNamespace(patch_size=16), + parameters=lambda: iter([torch.zeros(1)]), + ) + vae = SimpleNamespace( + parameters=lambda: iter([torch.zeros(1)]), + bn=SimpleNamespace( + running_mean=torch.arange(8), running_var=torch.full((8,), 3.75) + ), + config=SimpleNamespace(batch_norm_eps=0.25), + encode=Mock(return_value=SimpleNamespace(mode=lambda: raw)), + ) + with patch( + "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args", + return_value=SimpleNamespace(), + ): + stage = LLaDAImageSourceImageConditioningStage(sigvq, vae, processor) + batch = SimpleNamespace( + condition_image=None, + image_embeds=[], + source_latents=None, + height=64, + width=80, + batch_size=2, + ) + assert stage.forward(batch, None) is batch + assert batch.image_embeds == [] and batch.source_latents is None + processor.preprocess.assert_not_called() + batch.condition_image = Image.new("RGB", (40, 30)) + stage.forward(batch, None) + processor.preprocess.assert_called_once_with( + batch.condition_image, height=64, width=80, resize_mode="crop" + ) + assert sigvq.call_args.args[0].shape == (2, 3, 32, 48) + torch.testing.assert_close(vae.encode.call_args.args[0], pixels.repeat(2, 1, 1, 1)) + torch.testing.assert_close(torch.stack(batch.image_embeds), semantics) + expected = torch.stack( + [ + raw[:, channel, row::2, column::2] + for channel in range(2) + for row in range(2) + for column in range(2) + ], + dim=1, + ) + expected = (expected - torch.arange(8).reshape(1, 8, 1, 1)) / 2 + torch.testing.assert_close( + torch.stack(batch.source_latents), expected.unsqueeze(2), rtol=0, atol=0 + ) + batch.condition_image = [batch.condition_image] * 2 + with pytest.raises(ValueError, match="exactly one source image"): + stage.forward(batch, None) + + +def test_source_and_decode_preserve_vae_cast_order(): + vae = SimpleNamespace( + parameters=lambda: iter([torch.zeros(1, dtype=torch.bfloat16)]), + bn=SimpleNamespace( + running_mean=torch.tensor([0.01, -0.02, 0.03, -0.04]), + running_var=torch.full((4,), 0.0129973), + ), + config=SimpleNamespace(batch_norm_eps=0.003), + ) + latents = torch.tensor([0.007, -0.249, 0.126, -0.751]).reshape(1, 4, 1, 1) + bf16 = latents.bfloat16() + mean = vae.bn.running_mean.reshape(1, 4, 1, 1).to(bf16) + std = ( + (vae.bn.running_var + vae.config.batch_norm_eps) + .sqrt() + .reshape(1, 4, 1, 1) + .to(bf16) + ) + normalized = LLaDAImageSourceImageConditioningStage._normalize_latents(bf16, vae) + torch.testing.assert_close(normalized, (bf16 - mean) / std, rtol=0, atol=0) + decoded = LLaDAImagePipelineConfig().preprocess_decoding(latents, vae=vae) + torch.testing.assert_close( + decoded, (bf16 * std + mean).reshape(1, 1, 2, 2), rtol=0, atol=0 + ) From 5cc3f71039de08eb8bbc957121b9ae26fd702880 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Tue, 15 Sep 2026 23:55:23 +0000 Subject: [PATCH 12/18] [diffusion] Remove LLaDA-Image unit tests --- .../test/unit/test_llada_image.py | 163 ------------ .../unit/test_llada_image_conditioning.py | 237 ------------------ .../test/unit/test_llada_image_loader.py | 104 -------- .../test/unit/test_llada_image_rmsnorm.py | 43 ---- .../test_llada_image_source_conditioning.py | 101 -------- .../test/unit/test_minimax_h3_adaln_cache.py | 7 +- 6 files changed, 1 insertion(+), 654 deletions(-) delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py delete mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image.py b/python/sglang/multimodal_gen/test/unit/test_llada_image.py deleted file mode 100644 index cc004ad79b5f..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image.py +++ /dev/null @@ -1,163 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace -from unittest.mock import patch - -import pytest -import torch - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) -from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping -from sglang.multimodal_gen.runtime.models.dits.llada_image import ( - _LLaDAImageTransformer2DModel, -) -from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( - format_llada_image_prompt, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( - TimestepPreparationStage, -) -from sglang.multimodal_gen.runtime.server_args import ( - get_global_server_args, - set_global_server_args, -) - - -@pytest.fixture(autouse=True) -def server_args(): - try: - previous = get_global_server_args() - except ValueError: - previous = None - set_global_server_args(SimpleNamespace(kv_gather_degree=1, sp_split_auto=False)) - yield - set_global_server_args(previous) - - -@pytest.mark.parametrize( - "source,target,shard,total", - [ - ("attention.to_k", "attention.to_qkv", 1, 3), - ("feed_forward.w1", "feed_forward.w13", 0, 2), - ("feed_forward.w3", "feed_forward.w13", 1, 2), - ], -) -def test_dit_weight_mapping(source, target, shard, total): - mapping = get_param_names_mapping( - LLaDAImagePipelineConfig().dit_config.arch_config.param_names_mapping - ) - assert mapping(f"layers.0.{source}.weight") == ( - f"layers.0.{target}.weight", - shard, - total, - ) - - -@pytest.mark.parametrize( - "prompt,instruction", - [("a red car", "Generate an image: a red car"), (None, "Generate an image.")], -) -def test_prompt_format(prompt, instruction): - assert ( - format_llada_image_prompt(prompt) - == f"HUMAN {instruction}\nASSISTANT\n" - ) - - -def test_uniform_scheduler_schedule(): - scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0, use_uniform_sigmas=True) - batch = SimpleNamespace( - scheduler=None, - timesteps=None, - sigmas=None, - num_inference_steps=4, - n_tokens=None, - extra={}, - is_warmup=True, - rollout=False, - ) - module = "sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation" - with ( - patch(f"{module}.get_local_torch_device", return_value=torch.device("cpu")), - patch(f"{module}.get_or_create_request_scheduler", return_value=scheduler), - ): - TimestepPreparationStage(scheduler).forward( - batch, SimpleNamespace(pipeline_config=LLaDAImagePipelineConfig()) - ) - assert batch.sigmas is None - torch.testing.assert_close( - scheduler.sigmas, torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0]) - ) - - -@pytest.mark.parametrize("seeds", [(11,), (11, 29)]) -def test_stochastic_scheduler_per_sample_generators(seeds): - scheduler = FlowMatchEulerDiscreteScheduler( - shift=3.0, use_uniform_sigmas=True, stochastic_sampling=True - ) - scheduler.set_timesteps(4, device="cpu") - sample = torch.arange(len(seeds) * 8, dtype=torch.float32).reshape( - len(seeds), 2, 2, 2 - ) - prediction = sample / 16 - generators = [torch.Generator().manual_seed(seed) for seed in seeds] - noise = torch.cat( - [ - torch.randn((1, 2, 2, 2), generator=torch.Generator().manual_seed(seed)) - for seed in seeds - ] - ) - sigma, next_sigma = scheduler.sigmas[:2] - expected = (1 - next_sigma) * (sample - sigma * prediction) + next_sigma * noise - actual = scheduler.step( - prediction, - scheduler.timesteps[0], - sample, - generator=generators, - return_dict=False, - )[0] - torch.testing.assert_close(actual, expected, rtol=0, atol=0) - - -def test_edit_skips_empty_semantic_refiner(): - class CaptureBlock(torch.nn.Module): - def __init__(self): - super().__init__() - self.calls = 0 - - def forward(self, hidden_states, *args, **kwargs): - self.calls += 1 - return hidden_states - - model = _LLaDAImageTransformer2DModel( - in_channels=4, - dim=64, - n_layers=0, - n_refiner_layers=0, - n_heads=2, - cap_feat_dim=8, - semantic_feat_dim=10, - axes_dims=(8, 12, 12), - axes_lens=(256, 32, 32), - ) - blocks = [CaptureBlock() for _ in range(4)] - for name, block in zip( - ("noise_refiner", "context_refiner", "sigvq_refiner", "layers"), - blocks, - strict=True, - ): - setattr(model, name, torch.nn.ModuleList([block])) - with torch.no_grad(): - model( - x=[torch.randn(4, 1, 4, 4)], - t=torch.tensor([0.5]), - cap_feats=[torch.randn(3, 8)], - glm_cap_feats=[torch.empty(0, 10)], - source_latents=[torch.randn(4, 1, 4, 4)], - ) - assert [block.calls for block in blocks] == [1, 1, 0, 1] diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py deleted file mode 100644 index 31c33e1c689f..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_conditioning.py +++ /dev/null @@ -1,237 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest -import torch - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) -from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.conditioning import ( - LLaDAImageTextConditioningStage, - LLaDAImageTextEncoderRunner, - ensure_conditioning_mask_active, -) - - -@pytest.mark.parametrize("guidance", [1.0, 5.0]) -def test_text_conditioning_repeats_outputs_and_masks(guidance): - outputs = [torch.ones(2, 4), torch.full((3, 4), 2.0)] - runner = SimpleNamespace(encode=Mock(return_value=outputs)) - with patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args", - return_value=SimpleNamespace(), - ): - stage = LLaDAImageTextConditioningStage(runner) - batch = SimpleNamespace( - prompt="a red car", - negative_prompt=None, - guidance_scale=guidance, - num_outputs_per_prompt=2, - max_sequence_length=128, - ) - stage.forward(batch, SimpleNamespace()) - assert batch.do_classifier_free_guidance == (guidance > 1) - assert len(runner.encode.call_args.args[0]) == (2 if guidance > 1 else 1) - for prefix, output in zip(("prompt", "negative"), outputs, strict=True): - embeddings = getattr( - batch, "prompt_embeds" if prefix == "prompt" else "negative_prompt_embeds" - ) - masks = getattr(batch, f"{prefix}_attention_mask") - if prefix == "negative" and guidance == 1: - assert embeddings == masks == [] - continue - assert len(embeddings) == len(masks) == 2 - for embedding, mask in zip(embeddings, masks, strict=True): - torch.testing.assert_close(embedding, output) - torch.testing.assert_close( - mask, torch.ones(output.shape[0], dtype=torch.bool) - ) - - -def test_conditioning_mask_guard(): - ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=True)) - with pytest.raises(RuntimeError, match="without its block attention mask"): - ensure_conditioning_mask_active(SimpleNamespace(conditioning_mask_active=False)) - - -def test_prefill_uses_resolved_page_size(): - runner = object.__new__(LLaDAImageTextEncoderRunner) - runner.tokenizer = lambda *args, **kwargs: SimpleNamespace(input_ids=[[1, 2]]) - runner.queryformer = SimpleNamespace(config=SimpleNamespace(num_queries=2)) - runner.worker = SimpleNamespace(model_config=SimpleNamespace(vocab_size=128)) - runner.server_args = SimpleNamespace( - page_size=None, chunked_prefill_size=-1, max_prefill_tokens=8192 - ) - runner.page_size = 64 - runner.tree_cache = runner.token_to_kv_pool_allocator = object() - with ( - patch( - "sglang.srt.sampling.sampling_params.SamplingParams", - return_value=SimpleNamespace(normalize=lambda _: None), - ), - patch("sglang.srt.managers.schedule_batch.Req", return_value=SimpleNamespace()), - patch( - "sglang.srt.managers.schedule_policy.PrefillAdder", - side_effect=RuntimeError("stop at prefill"), - ) as adder, - pytest.raises(RuntimeError, match="stop at prefill"), - ): - runner._encode_impl(["hello"], max_sequence_length=16) - assert adder.call_args.args[0] == 64 - - -@pytest.mark.parametrize("fail_init", [False, True]) -def test_encoder_context_restores_after_init_and_encode(fail_init): - import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_state - import sglang.srt.distributed.parallel_state as srt_state - from sglang.srt import runtime_context as rc - from sglang.srt.server_args import ServerArgs as SRTServerArgs - - saved = rc.snapshot_context() - rc.reset_context() - try: - outer_args = SRTServerArgs(model_path="dummy", tp_size=1) - outer = rc.publish(outer_args, role="diffusion_gpu_worker") - outer.override("diffusion.setup", grammar_backend="none") - outer_bags, outer_log = outer._config_bags, outer.overrides_log() - outer_buffer = rc.get_buffer("conditioning", object) - diffusion_group = SimpleNamespace(world_size=1, rank_in_group=0) - encoder_group = SimpleNamespace(world_size=1, rank_in_group=0) - attention_group = object() - observed = {} - - def check_encoder(): - context = rc.assert_published(observed["args"], role="scheduler") - assert rc.get_schedule().max_running_requests == 2 - assert rc.get_parallel().tp_size == 1 - return context - - def allocate(): - check_encoder().override("encoder.setup", page_size=16) - observed["buffer"] = rc.get_buffer("conditioning", object) - if fail_init: - raise RuntimeError("allocation failed") - - def make_worker(server_args, **kwargs): - observed["args"] = server_args - observed["context"] = check_encoder() - srt_state._TP, srt_state._ATTN_TP = encoder_group, attention_group - return SimpleNamespace( - model_runner=SimpleNamespace(page_size=16), - get_memory_pool=lambda: (object(), object()), - alloc_memory_pool=allocate, - init_attention_backends=check_encoder, - init_cuda_graphs=check_encoder, - ) - - def make_cache(params): - assert params.page_size == rc.get_schedule().page_size == 16 - return object() - - def check_diffusion(): - assert rc.get_context() is outer and rc.get_server_args() is outer_args - assert rc.publish_role() == "diffusion_gpu_worker" - assert ( - outer._config_bags is outer_bags and outer.overrides_log() == outer_log - ) - assert rc.get_exec().kernel.grammar_backend == "none" - assert rc.get_buffer("conditioning", object) is outer_buffer - assert all( - group is diffusion_group - for group in (mm_state._TP, srt_state._TP, srt_state._ATTN_TP) - ) - - with ( - patch( - "sglang.srt.server_args.ServerArgs", - side_effect=lambda **kwargs: SRTServerArgs( - **dict(kwargs, model_path="dummy") - ), - ), - patch( - "sglang.srt.managers.tp_worker.TpModelWorker", side_effect=make_worker - ), - patch( - "sglang.srt.mem_cache.chunk_cache.ChunkCache", side_effect=make_cache - ), - patch( - f"{LLaDAImageTextEncoderRunner.__module__}.get_local_torch_device", - return_value=torch.device("cpu"), - ), - patch.object(mm_state, "_TP", diffusion_group), - patch.object(srt_state, "_TP", diffusion_group), - patch.object(srt_state, "_ATTN_TP", diffusion_group), - ): - args = SimpleNamespace( - nccl_port=29500, - trust_remote_code=False, - revision=None, - pipeline_config=SimpleNamespace(text_encoder_mem_fraction_static=0.1), - ) - if fail_init: - with pytest.raises(RuntimeError, match="allocation failed"): - LLaDAImageTextEncoderRunner( - "/unused", object(), object(), object(), args - ) - check_diffusion() - return - runner = LLaDAImageTextEncoderRunner( - "/unused", object(), object(), object(), args - ) - check_diffusion() - assert observed["context"] is not outer - - def encode(*args, **kwargs): - assert check_encoder() is observed["context"] - assert rc.get_schedule().page_size == 16 - assert rc.get_buffer("conditioning", object) is observed["buffer"] - assert observed["buffer"] is not outer_buffer - assert ( - srt_state._TP is encoder_group - and srt_state._ATTN_TP is attention_group - ) - if observed["fail_encode"]: - raise RuntimeError("encode failed") - return ["encoded"] - - runner._encode_impl = encode - for fail_encode in (False, True, False): - observed["fail_encode"] = fail_encode - if fail_encode: - with pytest.raises(RuntimeError, match="encode failed"): - runner.encode(["hello"], 16) - else: - assert runner.encode(["hello"], 16) == ["encoded"] - check_diffusion() - finally: - rc.restore_context(saved) - - -def test_edit_cfg_preserves_source_and_removes_negative_semantics(): - config = LLaDAImagePipelineConfig() - batch = SimpleNamespace( - batch_size=2, - image_embeds=[torch.ones(3, 5)] * 2, - source_latents=[torch.ones(8, 1, 2, 3)] * 2, - do_classifier_free_guidance=True, - ) - positive = config.prepare_pos_cond_kwargs(batch, "cpu", None, torch.float64) - negative = config.prepare_neg_cond_kwargs(batch, "cpu", None, torch.float64) - policy = CFGPolicy().build( - batch, {"encoder_hidden_states_image": batch.image_embeds}, positive, negative - ) - for branch, shape in zip(policy.branches, [(3, 5), (0, 5)], strict=True): - assert [ - tuple(value.shape) for value in branch.kwargs["encoder_hidden_states_image"] - ] == [shape] * 2 - for actual, source in zip( - branch.kwargs["source_latents"], batch.source_latents, strict=True - ): - torch.testing.assert_close(actual, source.double()) - batch.image_embeds = batch.image_embeds[:1] - with pytest.raises(ValueError, match="image_embeds has 1 entries, expected 2"): - config.prepare_pos_cond_kwargs(batch, "cpu", None, torch.float32) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py deleted file mode 100644 index 7b8f74d99668..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_loader.py +++ /dev/null @@ -1,104 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import json -from unittest.mock import patch - -import pytest -import torch -from safetensors.torch import save_file - -from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType -from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( - ComponentCheckpointUnsupportedError, - ComponentLoader, -) -from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import ( - COMPONENT_OFFLOAD, -) -from sglang.multimodal_gen.runtime.models.dits.llada_image import ( - LLaDAImageQueryFormerModel, - LLaDAImageSigVQModel, - LLaDAImageTextProjectionModel, -) -from sglang.multimodal_gen.runtime.pipelines.llada_image import LLaDAImagePipeline -from sglang.multimodal_gen.runtime.server_args import ServerArgs - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="ServerArgs residency resolution requires CUDA", -) - - -@pytest.mark.parametrize( - "name,model_cls,extra", - [ - ("queryformer", LLaDAImageQueryFormerModel, {"num_queries": 2}), - ("text_projection", LLaDAImageTextProjectionModel, {"projection_dim": 4}), - ( - "sigvq", - LLaDAImageSigVQModel, - { - "image_size": 4, - "patch_size": 2, - "codebook_size": 8, - "codebook_embed_dim": 4, - "semantic_embed_dim": 4, - }, - ), - ], -) -@pytest.mark.parametrize("missing_weight", [False, True]) -def test_auxiliary_checkpoint_loading(tmp_path, name, model_cls, extra, missing_weight): - config = dict( - hidden_size=8, - intermediate_size=16, - num_hidden_layers=1, - num_attention_heads=2, - **extra, - ) - torch.manual_seed(0) - weights = { - key: value.contiguous() - for key, value in model_cls(**config).state_dict().items() - } - component = tmp_path / name - component.mkdir() - (component / "config.json").write_text( - json.dumps(dict(config, _class_name=model_cls.__name__, _diffusers_version="0")) - ) - (tmp_path / "model_index.json").write_text( - json.dumps( - { - "_class_name": "LLaDAImagePipeline", - "_diffusers_version": "0", - name: ["diffusers", model_cls.__name__], - "scheduler": None, - } - ) - ) - if missing_weight: - weights.pop(next(iter(weights))) - save_file(weights, component / "model.safetensors") - args = ServerArgs( - model_path=str(tmp_path), component_residency={name: COMPONENT_OFFLOAD} - ) - pipeline = object.__new__(LLaDAImagePipeline) - pipeline.model_path, pipeline.server_args = str(tmp_path), args - pipeline._disagg_role = RoleType.MONOLITHIC - pipeline._required_config_modules = [name] - pipeline.memory_usages = {} - if missing_weight: - with ( - patch.object(ComponentLoader, "load_native") as native, - pytest.raises(ComponentCheckpointUnsupportedError, match="Missing:"), - ): - pipeline.load_modules(args) - native.assert_not_called() - return - model = pipeline.load_modules(args)[name] - assert isinstance(model, model_cls) and not model.training - assert args.model_paths[name] == str(component) - assert model.state_dict().keys() == weights.keys() - for key, value in model.state_dict().items(): - assert value.device.type == "cpu" and value.dtype == torch.bfloat16 - torch.testing.assert_close(value, weights[key].bfloat16(), rtol=0, atol=0) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py deleted file mode 100644 index 7b41dd9250ef..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_rmsnorm.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from unittest.mock import patch - -import pytest -import torch - -from sglang.multimodal_gen.runtime.layers import layernorm as layernorm_module -from sglang.multimodal_gen.runtime.models.dits.llada_image import LLaDAImageRMSNorm - - -@pytest.mark.parametrize( - "device", - [ - "cpu", - pytest.param( - "cuda", - marks=pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA is required" - ), - ), - ], -) -def test_rmsnorm_cast_order_and_cuda_dispatch(device): - torch.manual_seed(0) - norm = LLaDAImageRMSNorm(128, eps=1e-5).to(device=device, dtype=torch.bfloat16) - with torch.no_grad(): - norm.weight.copy_(torch.randn(128, device=device, dtype=torch.bfloat16)) - hidden = torch.randn(2, 8, 128, device=device, dtype=torch.bfloat16) - if device == "cuda": - with patch.object( - layernorm_module, "rmsnorm_hf", wraps=layernorm_module.rmsnorm_hf - ) as kernel: - actual = norm(hidden) - kernel.assert_called_once() - else: - actual = norm.forward_native(hidden) - variance = hidden.float().pow(2).mean(-1, keepdim=True) - expected = norm.weight * ( - hidden.float() * torch.rsqrt(variance + norm.variance_epsilon) - ).to(hidden.dtype) - tolerance = 1e-2 if device == "cuda" else 0 - torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py deleted file mode 100644 index 47d2899a6d38..000000000000 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_source_conditioning.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest -import torch -from PIL import Image - -from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( - LLaDAImagePipelineConfig, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.llada_image.source import ( - LLaDAImageSourceImageConditioningStage, -) - - -def test_source_image_preprocessing_and_output_replication(): - pixels = torch.arange(3 * 8 * 12, dtype=torch.float32).reshape(1, 3, 8, 12) - raw = torch.arange(2 * 2 * 4 * 6, dtype=torch.float32).reshape(2, 2, 4, 6) - semantics = torch.arange(30, dtype=torch.float32).reshape(2, 3, 5) - processor = SimpleNamespace(preprocess=Mock(return_value=pixels)) - sigvq = Mock( - return_value=SimpleNamespace(semantic_features=semantics), - config=SimpleNamespace(patch_size=16), - parameters=lambda: iter([torch.zeros(1)]), - ) - vae = SimpleNamespace( - parameters=lambda: iter([torch.zeros(1)]), - bn=SimpleNamespace( - running_mean=torch.arange(8), running_var=torch.full((8,), 3.75) - ), - config=SimpleNamespace(batch_norm_eps=0.25), - encode=Mock(return_value=SimpleNamespace(mode=lambda: raw)), - ) - with patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args", - return_value=SimpleNamespace(), - ): - stage = LLaDAImageSourceImageConditioningStage(sigvq, vae, processor) - batch = SimpleNamespace( - condition_image=None, - image_embeds=[], - source_latents=None, - height=64, - width=80, - batch_size=2, - ) - assert stage.forward(batch, None) is batch - assert batch.image_embeds == [] and batch.source_latents is None - processor.preprocess.assert_not_called() - batch.condition_image = Image.new("RGB", (40, 30)) - stage.forward(batch, None) - processor.preprocess.assert_called_once_with( - batch.condition_image, height=64, width=80, resize_mode="crop" - ) - assert sigvq.call_args.args[0].shape == (2, 3, 32, 48) - torch.testing.assert_close(vae.encode.call_args.args[0], pixels.repeat(2, 1, 1, 1)) - torch.testing.assert_close(torch.stack(batch.image_embeds), semantics) - expected = torch.stack( - [ - raw[:, channel, row::2, column::2] - for channel in range(2) - for row in range(2) - for column in range(2) - ], - dim=1, - ) - expected = (expected - torch.arange(8).reshape(1, 8, 1, 1)) / 2 - torch.testing.assert_close( - torch.stack(batch.source_latents), expected.unsqueeze(2), rtol=0, atol=0 - ) - batch.condition_image = [batch.condition_image] * 2 - with pytest.raises(ValueError, match="exactly one source image"): - stage.forward(batch, None) - - -def test_source_and_decode_preserve_vae_cast_order(): - vae = SimpleNamespace( - parameters=lambda: iter([torch.zeros(1, dtype=torch.bfloat16)]), - bn=SimpleNamespace( - running_mean=torch.tensor([0.01, -0.02, 0.03, -0.04]), - running_var=torch.full((4,), 0.0129973), - ), - config=SimpleNamespace(batch_norm_eps=0.003), - ) - latents = torch.tensor([0.007, -0.249, 0.126, -0.751]).reshape(1, 4, 1, 1) - bf16 = latents.bfloat16() - mean = vae.bn.running_mean.reshape(1, 4, 1, 1).to(bf16) - std = ( - (vae.bn.running_var + vae.config.batch_norm_eps) - .sqrt() - .reshape(1, 4, 1, 1) - .to(bf16) - ) - normalized = LLaDAImageSourceImageConditioningStage._normalize_latents(bf16, vae) - torch.testing.assert_close(normalized, (bf16 - mean) / std, rtol=0, atol=0) - decoded = LLaDAImagePipelineConfig().preprocess_decoding(latents, vae=vae) - torch.testing.assert_close( - decoded, (bf16 * std + mean).reshape(1, 1, 2, 2), rtol=0, atol=0 - ) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py index 41bce3ae90d6..d6a27e01efbb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py @@ -528,17 +528,12 @@ def test_lora_ipc_layer_guard_rejects_adaln_in_cache_mode(): def _updater_for(model, model_path: str): from types import SimpleNamespace - from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.runtime.post_training.weights_updater import ( WeightsUpdater, ) model.register_parameter("probe", torch.nn.Parameter(torch.zeros(2))) - pipeline = SimpleNamespace( - modules={"transformer": model}, - model_path=model_path, - server_args=SimpleNamespace(pipeline_config=PipelineConfig()), - ) + pipeline = SimpleNamespace(modules={"transformer": model}, model_path=model_path) return WeightsUpdater(pipeline), pipeline From 80bb1932a76bcc42866357676f947fa8e44051c7 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Sun, 20 Sep 2026 08:29:19 +0000 Subject: [PATCH 13/18] Fix LLaDA-Image initialization and platform CI fixtures --- .../llada_image/conditioning.py | 6 +++--- .../runtime/utils/quantization_utils.py | 9 +++++---- .../test_modelopt_fp8_layerwise_offload_load.py | 16 ++++++++++++---- .../test/unit/test_sequential_server_requests.py | 1 + .../test/unit/test_transformer_quant.py | 11 ++++++++--- .../cpu/test_subblock_sparse_attention.py | 16 ++++++++-------- 6 files changed, 37 insertions(+), 22 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index a3fb8920410a..a52cd4edeeee 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -65,7 +65,6 @@ def __init__( tokenizer, server_args: ServerArgs, ) -> None: - import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state import sglang.srt.distributed.parallel_state as srt_parallel_state from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.tp_worker import TpModelWorker @@ -111,12 +110,13 @@ def __init__( ) with use_context(self.runtime_context): # The diffusion runtime mirrors its TP group into the srt globals. - # Clear it so the worker can install the encoder's real srt groups, + # Detach the groups without clearing the encoder context's ranks, # then restore the mirror and keep the encoder group for forwards. saved_srt_tp = srt_parallel_state._TP saved_srt_attn_tp = srt_parallel_state._ATTN_TP try: - mm_parallel_state._clear_srt_tp_group() + srt_parallel_state._TP = None + srt_parallel_state._ATTN_TP = None self.worker = TpModelWorker( server_args=srt_args, gpu_id=gpu_id, diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 29d79d49b393..ebce62c1b12e 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -10,10 +10,6 @@ from safetensors import safe_open from torch import nn -from sglang.multimodal_gen.runtime.layers.linear import ( - LinearBase, - UnquantizedLinearMethod, -) from sglang.multimodal_gen.runtime.layers.quantization import ( QuantizationConfig, get_quantization_config, @@ -64,6 +60,11 @@ def process_model_weights_after_loading( quantized_only: bool = False, ) -> int: """Process native and SRT layers once, optionally staging one layer at a time.""" + from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, + ) + processed_layers = 0 for module in model.modules(): if not isinstance( diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py index 62c9f774e5c1..612147bc8806 100644 --- a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py @@ -8,6 +8,7 @@ import torch from torch import nn +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.multimodal_gen.runtime.distributed.parallel_state import ( maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, @@ -98,6 +99,13 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): checkpoint_weight = state_dict["qkv.weight"].clone() checkpoint_scales = state_dict["qkv.weight_scale"].clone() expected_max_scale = checkpoint_scales.max() + scale_factor = 2 if is_fp8_fnuz() else 1 + expected_dtype = ( + torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn + ) + expected_weight = (checkpoint_weight.float() / scale_factor).to( + expected_dtype + ) with patch( "sglang.multimodal_gen.runtime.layers.quantization." @@ -126,24 +134,24 @@ def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): # consume a channelwise scale, so it preserves the checkpoint's # FP8 shards; the fallback requantizes them to one max scale. weight = model.qkv.weight - self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(weight.dtype, expected_dtype) self.assertEqual(tuple(weight.shape), (_IN_FEATURES, 2 * _SHARD_OUT)) weight_scale = model.qkv.weight_scale.flatten() if cutlass_supported: expected_scales = torch.repeat_interleave( checkpoint_scales, _SHARD_OUT ) - self.assertTrue(torch.equal(weight.t(), checkpoint_weight)) + self.assertTrue(torch.equal(weight.t(), expected_weight)) else: expected_scales = expected_max_scale.expand(weight_scale.numel()) torch.testing.assert_close( weight_scale, - expected_scales, + expected_scales * scale_factor, check_device=False, ) torch.testing.assert_close( model.qkv.input_scale.flatten().max(), - torch.tensor(0.5), + torch.tensor(0.5 * scale_factor), check_device=False, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py b/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py index aa74cabaa788..1c19efe7f34c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py +++ b/python/sglang/multimodal_gen/test/unit/test_sequential_server_requests.py @@ -41,6 +41,7 @@ def harness(monkeypatch): monkeypatch.setenv("SGLANG_GEN_BASELINE", "0") monkeypatch.setenv("SGLANG_SKIP_CONSISTENCY", "0") monkeypatch.setattr(test_server_common.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(test_server_common.current_platform, "is_hip", lambda: False) scenario = ScenarioConfig( stages_ms={"DenoisingStage": 10}, denoise_step_ms={0: 5, 1: 5}, diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 810b7b158536..bdb3627b42de 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -13,6 +13,7 @@ import torch from safetensors.torch import save_file +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTArchConfig partial_json_parser = types.ModuleType("partial_json_parser") @@ -192,12 +193,16 @@ def test_modelopt_fp8_packed_cutlass_preserves_checkpoint_shard_scales(self): method.process_weights_after_loading(layer) - torch.testing.assert_close(layer.weight, weight.t(), rtol=0, atol=0) + scale_factor = 2 if is_fp8_fnuz() else 1 + expected_weight = weight + if is_fp8_fnuz(): + expected_weight = (weight.float() / 2).to(torch.float8_e4m3fnuz) + torch.testing.assert_close(layer.weight, expected_weight.t(), rtol=0, atol=0) torch.testing.assert_close( layer.weight_scale, - torch.tensor([[0.1], [0.1], [0.2], [0.2], [0.3], [0.3]]), + torch.tensor([[0.1], [0.1], [0.2], [0.2], [0.3], [0.3]]) * scale_factor, ) - torch.testing.assert_close(layer.input_scale, torch.tensor(1.0)) + torch.testing.assert_close(layer.input_scale, torch.tensor(float(scale_factor))) def test_modelopt_fp8_packed_cutlass_requantizes_incomplete_shard_scales(self): method = ModelOptFp8LinearMethod( diff --git a/test/registered/cpu/test_subblock_sparse_attention.py b/test/registered/cpu/test_subblock_sparse_attention.py index 5a8daa8bcdca..fcc911b2b8ea 100644 --- a/test/registered/cpu/test_subblock_sparse_attention.py +++ b/test/registered/cpu/test_subblock_sparse_attention.py @@ -275,7 +275,7 @@ def test_sage_fp8_dependency_is_checked_during_server_validation(self): return_value=DeviceCapability(9, 0), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), @@ -285,7 +285,7 @@ def test_sage_fp8_dependency_is_checked_during_server_validation(self): loader, ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ), ): @@ -304,7 +304,7 @@ def test_missing_sage_fp8_dependency_fails_server_validation(self): return_value=DeviceCapability(9, 0), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), @@ -314,7 +314,7 @@ def test_missing_sage_fp8_dependency_fails_server_validation(self): side_effect=ImportError("Install SpargeAttention"), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ) as get_backend, self.assertRaisesRegex(ImportError, "Install SpargeAttention"), @@ -330,7 +330,7 @@ def test_bf16_does_not_require_sparge_attention(self): with ( patch.object(current_platform, "is_mps", return_value=False), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), @@ -340,7 +340,7 @@ def test_bf16_does_not_require_sparge_attention(self): loader, ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ), ): @@ -359,12 +359,12 @@ def test_sage_fp8_rejects_non_sm90_during_server_validation(self): return_value=DeviceCapability(10, 0), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ) as get_backend, self.assertRaisesRegex(ValueError, "requires SM90.*found 10.0"), From b3d01af42e463f37ab81b703f70feacfbdeb7285 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Tue, 22 Sep 2026 18:42:58 +0000 Subject: [PATCH 14/18] test(diffusion): fix platform-specific CI fixtures --- .../multimodal_gen/test/unit/test_load_inclusive_e2e.py | 2 +- test/registered/cpu/test_subblock_sparse_attention.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py b/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py index 3d2d3b00a7ed..8f4e129a8c0a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py +++ b/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py @@ -24,6 +24,7 @@ @pytest.fixture def validator(monkeypatch): monkeypatch.setenv("SGLANG_GEN_BASELINE", "0") + monkeypatch.setattr(utils.current_platform, "is_hip", lambda: False) scenario = ScenarioConfig.from_dict( { "stages_ms": {}, @@ -135,7 +136,6 @@ def test_server_load_clock_excludes_warmup( clock = [1_000_000_000] output = io.StringIO() monkeypatch.setattr(utils.time, "monotonic_ns", lambda: clock[0]) - monkeypatch.setattr(utils.current_platform, "is_hip", lambda: False) monkeypatch.setattr(utils.tempfile, "gettempdir", lambda: str(tmp_path)) monkeypatch.setattr( utils, "prepare_perf_log", lambda: (tmp_path, tmp_path / "perf.jsonl") diff --git a/test/registered/cpu/test_subblock_sparse_attention.py b/test/registered/cpu/test_subblock_sparse_attention.py index a2529689089c..52d4f0b81e4e 100644 --- a/test/registered/cpu/test_subblock_sparse_attention.py +++ b/test/registered/cpu/test_subblock_sparse_attention.py @@ -371,7 +371,7 @@ def test_sm120_sage_dependency_is_checked_during_server_validation(self): return_value=DeviceCapability(12, 0), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), @@ -381,7 +381,7 @@ def test_sm120_sage_dependency_is_checked_during_server_validation(self): loader, ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ), ): @@ -400,7 +400,7 @@ def test_missing_sm120_sage_dependency_fails_server_validation(self): return_value=DeviceCapability(12, 0), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_global_forced_attn_backend", return_value=None, ), @@ -410,7 +410,7 @@ def test_missing_sm120_sage_dependency_fails_server_validation(self): side_effect=ImportError("FlashInfer SM120 Sage backend is unavailable"), ), patch( - "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3." + "sglang.multimodal_gen.runtime.layers.attention.selector." "get_attn_backend" ) as get_backend, self.assertRaisesRegex( From 3f70123a90b1a966d37b17a31eec1a457a339a4e Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Thu, 8 Oct 2026 22:13:44 +0000 Subject: [PATCH 15/18] Fix LLaDA-Image defaults, edit size limit, and FP8 loading Give the Base checkpoint the model card recipe of 50 steps and guidance 5.0, and register LLaDA-Image-Turbo with its own 4 step defaults without CFG. A default Base request returned a gray blur. Reject an editing request when its sequence RoPE positions would pass the table, before the RoPE gather. A 2048x2048 edit raised a CUDA device-side assert that failed every later request. Refuse quantized transformers and FP8 text encoders. The public FP8 checkpoints loaded FP8 expert codes into BF16 without their scales. Move the conditioning mask helper into the CFG FlashInfer backend and the checkpoint layout helper into llada2.py, their only users. --- .../configs/models/dits/llada_image.py | 18 ++++++ .../configs/pipeline_configs/llada_image.py | 17 +++++- .../configs/sample/llada_image.py | 12 +++- .../runtime/models/dits/llada_image.py | 31 ++++++++-- .../llada_image/conditioning.py | 11 ++++ .../test/unit/test_llada_image_config.py | 56 +++++++++++++++++++ .../attention/llada2_attention_utils.py | 31 ---------- .../llada2_cfg_flashinfer_backend.py | 30 ++++++++-- .../srt/model_loader/llada2_weight_utils.py | 44 --------------- python/sglang/srt/models/llada2.py | 40 ++++++++++++- ... => test_llada2_cfg_flashinfer_backend.py} | 14 ++--- ...ils.py => test_llada2_language_weights.py} | 8 +-- 12 files changed, 212 insertions(+), 100 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_llada_image_config.py delete mode 100644 python/sglang/srt/layers/attention/llada2_attention_utils.py delete mode 100644 python/sglang/srt/model_loader/llada2_weight_utils.py rename test/registered/unit/layers/attention/{test_llada2_attention_utils.py => test_llada2_cfg_flashinfer_backend.py} (68%) rename test/registered/unit/models/{test_llada2_weight_utils.py => test_llada2_language_weights.py} (87%) diff --git a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py index 92984fdc0a80..f52c66d4cbf9 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py @@ -5,6 +5,24 @@ from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +SEQUENCE_MULTIPLE = 32 + + +def editing_rope_rows( + caption_tokens: int, image_tokens: list[int], sigvq_tokens: int +) -> int: + # Must match LLaDAImageTransformer2DModel._prepare_editing_sequences: two + # padded caption copies and the padded image grids precede the SigVQ tokens. + def padded(length: int) -> int: + return length + (-length) % SEQUENCE_MULTIPLE + + return ( + 1 + + 2 * padded(caption_tokens) + + sum(padded(tokens) for tokens in image_tokens) + + sigvq_tokens + ) + @dataclass class LLaDAImageArchConfig(DiTArchConfig): diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py index 74408751a29d..4517beca8494 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -205,16 +205,31 @@ def post_denoising_loop(self, latents, batch): def register(): from sglang.multimodal_gen.configs.sample.llada_image import ( LLaDAImageSamplingParams, + LLaDAImageTurboSamplingParams, ) from sglang.multimodal_gen.registry import register_configs + # Base and Turbo checkpoints share every class name, so Turbo registers + # first and the first matching detector wins. + register_configs( + sampling_param_cls=LLaDAImageTurboSamplingParams, + pipeline_config_cls=LLaDAImagePipelineConfig, + hf_model_paths=["inclusionAI/LLaDA-Image-Turbo"], + model_detectors=[ + lambda hf_id: "llada-image" in hf_id.lower() and "turbo" in hf_id.lower() + ], + ) register_configs( sampling_param_cls=LLaDAImageSamplingParams, pipeline_config_cls=LLaDAImagePipelineConfig, hf_model_paths=["inclusionAI/LLaDA-Image"], model_detectors=[ lambda hf_id: ( - "lladaimagepipeline" in hf_id.lower() or "llada-image" in hf_id.lower() + ( + "lladaimagepipeline" in hf_id.lower() + or "llada-image" in hf_id.lower() + ) + and "turbo" not in hf_id.lower() ) ], ) diff --git a/python/sglang/multimodal_gen/configs/sample/llada_image.py b/python/sglang/multimodal_gen/configs/sample/llada_image.py index d6eb300af1d1..e0b19f6dc23e 100644 --- a/python/sglang/multimodal_gen/configs/sample/llada_image.py +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -12,11 +12,12 @@ class LLaDAImageSamplingParams(SamplingParams): data_type: DataType = DataType.IMAGE negative_prompt: str | None = None - num_inference_steps: int = 4 + # Base checkpoint recipe from the model card. + num_inference_steps: int = 50 num_frames: int = 1 height: int = 1024 width: int = 1024 - guidance_scale: float = 1.0 + guidance_scale: float = 5.0 max_sequence_length: int = 2048 def _validate_with_pipeline_config(self, pipeline_config): @@ -29,3 +30,10 @@ def _validate_with_pipeline_config(self, pipeline_config): raise ValueError( "LLaDA-Image max_sequence_length must be between 1 and 3584" ) + + +@dataclass +class LLaDAImageTurboSamplingParams(LLaDAImageSamplingParams): + # The distilled Turbo checkpoint samples in 4 steps without CFG. + num_inference_steps: int = 4 + guidance_scale: float = 1.0 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index 2c4d60a2ad6b..ebd4b9468a72 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -29,7 +29,11 @@ from torch import nn from torch.nn.utils.rnn import pad_sequence -from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig +from sglang.multimodal_gen.configs.models.dits.llada_image import ( + SEQUENCE_MULTIPLE, + LLaDAImageDitConfig, + editing_rope_rows, +) from sglang.multimodal_gen.runtime.distributed import ( get_tp_world_size, ) @@ -64,7 +68,6 @@ def apply_rmsnorm_tanh_mul_add( return residual + torch.tanh(gate) * norm(x) -SEQUENCE_MULTIPLE = 32 LLADA_IMAGE_ATTENTION_BACKENDS = { AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA, @@ -973,6 +976,7 @@ def _prepare_editing_sequences( batch_image_positions = [] batch_image_padding = [] batch_image_noise = [] + image_token_counts = [] for image, position_start, noise_value in zip( (source_latents[batch_index], latent), cap_end_positions, @@ -981,6 +985,7 @@ def _prepare_editing_sequences( patches, image_size, token_grid_size = self._patchify_image( image, patch_size, f_patch_size ) + image_token_counts.append(len(patches)) image_height_start = 0 padded_features, position_ids, padding_mask, _, noise_mask = ( self._pad_with_ids( @@ -1015,6 +1020,19 @@ def _prepare_editing_sequences( ) ) + rope_rows = editing_rope_rows( + len(cap_feats[batch_index]), + image_token_counts, + len(glm_cap_feats[batch_index]), + ) + if rope_rows > self.rope_embedder.axes_lens[0]: + # An out-of-range RoPE gather would poison the CUDA context. + raise ValueError( + f"LLaDA-Image editing needs {rope_rows} sequence positions " + "for this image size and prompt, above the model limit of " + f"{self.rope_embedder.axes_lens[0]}. Use a smaller image " + "size or a shorter prompt." + ) padded_features, position_ids, padding_mask, _, noise_mask = ( self._pad_with_ids( glm_cap_feats[batch_index], @@ -2252,10 +2270,13 @@ class LLaDAImageTransformer2DModel(_LLaDAImageTransformer2DModel): reverse_param_names_mapping: ClassVar[dict] = {} def __init__(self, config, hf_config: dict, quant_config=None): + if quant_config is not None or "quantization_config" in hf_config: + raise ValueError( + "LLaDA-Image serves only the BF16 checkpoints, such as " + "inclusionAI/LLaDA-Image, and this transformer is quantized" + ) init_kwargs = { - key: value - for key, value in hf_config.items() - if not key.startswith("_") and key != "quantization_config" + key: value for key, value in hf_config.items() if not key.startswith("_") } super().__init__(quant_config=quant_config, **init_kwargs) self.sgl_config = config diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index 6b78e7fbfde1..0f244f08a741 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os from array import array from contextlib import nullcontext @@ -83,6 +84,16 @@ def __init__( text_encoder_path = _resolve_llada_image_component_path( model_root, server_args, "text_encoder" ) + text_encoder_config = json.loads( + (Path(text_encoder_path) / "config.json").read_text() + ) + if "quantization_config" in text_encoder_config or ( + text_encoder_config.get("llada_fp8_experts", {}).get("enabled") + ): + raise ValueError( + "LLaDA-Image serves only the BF16 checkpoints, such as " + "inclusionAI/LLaDA-Image, and this text encoder is quantized" + ) tokenizer_path = _resolve_llada_image_component_path( model_root, server_args, "tokenizer" ) diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py new file mode 100644 index 000000000000..fabbe13536b4 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +import unittest + +from sglang.multimodal_gen.configs.models.dits.llada_image import ( + LLaDAImageArchConfig, + editing_rope_rows, +) +from sglang.multimodal_gen.configs.sample.llada_image import ( + LLaDAImageSamplingParams, + LLaDAImageTurboSamplingParams, +) +from sglang.multimodal_gen.registry import _get_config_info + + +class TestLLaDAImageConfig(unittest.TestCase): + def test_base_and_turbo_checkpoints_get_their_own_defaults(self): + """A default Base request must not run the 4-step Turbo recipe.""" + _get_config_info.cache_clear() + cases = { + "inclusionAI/LLaDA-Image": (LLaDAImageSamplingParams, 50, 5.0), + "inclusionAI/LLaDA-Image-Turbo": (LLaDAImageTurboSamplingParams, 4, 1.0), + "inclusionAI/LLaDA-Image-Turbo-FP8": ( + LLaDAImageTurboSamplingParams, + 4, + 1.0, + ), + } + for model_id, (sampling_cls, steps, guidance) in cases.items(): + with self.subTest(model_id=model_id): + info = _get_config_info(model_id) + self.assertIs(info.sampling_param_cls, sampling_cls) + params = info.sampling_param_cls() + self.assertEqual(params.num_inference_steps, steps) + self.assertEqual(params.guidance_scale, guidance) + + def test_editing_size_limit_follows_sequence_rope_table(self): + """Editing sizes past the sequence RoPE table must be caught on the host.""" + # A short prompt of 20 tokens plus the 256 conditioning query tokens. + caption_tokens = 20 + 256 + limit = LLaDAImageArchConfig().axes_lens[0] + + def rows(side: int) -> int: + image_tokens = (side // 16) ** 2 + sigvq_tokens = math.ceil(side / 32) ** 2 + return editing_rope_rows( + caption_tokens, [image_tokens, image_tokens], sigvq_tokens + ) + + self.assertLessEqual(rows(1904), limit) + self.assertGreater(rows(1920), limit) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/srt/layers/attention/llada2_attention_utils.py b/python/sglang/srt/layers/attention/llada2_attention_utils.py deleted file mode 100644 index 5e4f639227e8..000000000000 --- a/python/sglang/srt/layers/attention/llada2_attention_utils.py +++ /dev/null @@ -1,31 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Mask helpers shared by LLaDA2 attention backends.""" - -from __future__ import annotations - -import torch - - -def build_llada_image_custom_mask( - text_lens: list[int], - sequence_lens: list[int], - device: torch.device | str, -) -> torch.Tensor: - """Build the official one-pass text/query block mask in ragged layout.""" - if len(text_lens) != len(sequence_lens): - raise RuntimeError("LLaDA-Image conditioning metadata batch mismatch") - - flattened_masks: list[torch.Tensor] = [] - for text_len, sequence_len in zip(text_lens, sequence_lens): - text_len = int(text_len) - sequence_len = int(sequence_len) - if text_len <= 0 or text_len >= sequence_len: - raise RuntimeError( - "LLaDA-Image conditioning text length must leave query tokens" - ) - attention_mask = torch.ones( - (sequence_len, sequence_len), dtype=torch.bool, device=device - ) - attention_mask[:text_len, text_len:] = False - flattened_masks.append(attention_mask.flatten()) - return torch.cat(flattened_masks) diff --git a/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py index af4d11254b7b..f7e676c42f1b 100644 --- a/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py @@ -10,9 +10,31 @@ FlashInferAttnBackend, PrefillMetadata, ) -from sglang.srt.layers.attention.llada2_attention_utils import ( - build_llada_image_custom_mask, -) + + +def _build_llada_image_custom_mask( + text_lens: list[int], + sequence_lens: list[int], + device: torch.device | str, +) -> torch.Tensor: + """Build the official one-pass text/query block mask in ragged layout.""" + if len(text_lens) != len(sequence_lens): + raise RuntimeError("LLaDA-Image conditioning metadata batch mismatch") + + flattened_masks: list[torch.Tensor] = [] + for text_len, sequence_len in zip(text_lens, sequence_lens): + text_len = int(text_len) + sequence_len = int(sequence_len) + if text_len <= 0 or text_len >= sequence_len: + raise RuntimeError( + "LLaDA-Image conditioning text length must leave query tokens" + ) + attention_mask = torch.ones( + (sequence_len, sequence_len), dtype=torch.bool, device=device + ) + attention_mask[:text_len, text_len:] = False + flattened_masks.append(attention_mask.flatten()) + return torch.cat(flattened_masks) class LLaDA2CFGFlashInferAttnBackend(FlashInferAttnBackend): @@ -52,7 +74,7 @@ def init_forward_metadata(self, forward_batch): raise RuntimeError( "LLaDA-Image conditioning does not support cached prefixes" ) - custom_mask = build_llada_image_custom_mask( + custom_mask = _build_llada_image_custom_mask( text_lens, seq_lens.tolist(), seq_lens.device, diff --git a/python/sglang/srt/model_loader/llada2_weight_utils.py b/python/sglang/srt/model_loader/llada2_weight_utils.py deleted file mode 100644 index 2e38033ae5d3..000000000000 --- a/python/sglang/srt/model_loader/llada2_weight_utils.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Checkpoint-layout helpers for LLaDA2 models.""" - -from __future__ import annotations - -from collections.abc import Iterable, Iterator - -import torch - -_EXPERT_PROJECTION_NAMES = frozenset({"gate_proj", "up_proj", "down_proj"}) - - -def prepare_llada2_language_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - num_experts: int, -) -> Iterator[tuple[str, torch.Tensor]]: - """Normalize language-model prefixes and expand fused expert tensors.""" - for name, loaded_weight in weights: - if name.startswith("model.language_model."): - name = "model." + name.removeprefix("model.language_model.") - elif name.startswith("model.lm_head."): - name = "lm_head." + name.removeprefix("model.lm_head.") - - expert_prefix, separator, projection_name = name.rpartition(".") - is_fused_expert = ( - separator - and expert_prefix.endswith(".mlp.experts") - and projection_name in _EXPERT_PROJECTION_NAMES - ) - if not is_fused_expert: - yield name, loaded_weight - continue - - if loaded_weight.ndim != 3 or loaded_weight.shape[0] != num_experts: - raise ValueError( - f"Invalid fused expert weight {name!r}: expected first dimension " - f"{num_experts}, got shape={tuple(loaded_weight.shape)}" - ) - for expert_id in range(num_experts): - yield ( - f"{expert_prefix}.{expert_id}.{projection_name}.weight", - loaded_weight[expert_id], - ) diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index c5495a3f49e0..c6bee2df811f 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -67,7 +67,6 @@ ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.runner import get_is_capture_mode -from sglang.srt.model_loader.llada2_weight_utils import prepare_llada2_language_weights from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.utils import ( apply_qk_norm, @@ -740,6 +739,43 @@ def forward( return hidden_states +_EXPERT_PROJECTION_NAMES = frozenset({"gate_proj", "up_proj", "down_proj"}) + + +def _prepare_llada2_language_weights( + weights: Iterable[Tuple[str, torch.Tensor]], + *, + num_experts: int, +) -> Iterable[Tuple[str, torch.Tensor]]: + """Normalize language-model prefixes and expand fused expert tensors.""" + for name, loaded_weight in weights: + if name.startswith("model.language_model."): + name = "model." + name.removeprefix("model.language_model.") + elif name.startswith("model.lm_head."): + name = "lm_head." + name.removeprefix("model.lm_head.") + + expert_prefix, separator, projection_name = name.rpartition(".") + is_fused_expert = ( + separator + and expert_prefix.endswith(".mlp.experts") + and projection_name in _EXPERT_PROJECTION_NAMES + ) + if not is_fused_expert: + yield name, loaded_weight + continue + + if loaded_weight.ndim != 3 or loaded_weight.shape[0] != num_experts: + raise ValueError( + f"Invalid fused expert weight {name!r}: expected first dimension " + f"{num_experts}, got shape={tuple(loaded_weight.shape)}" + ) + for expert_id in range(num_experts): + yield ( + f"{expert_prefix}.{expert_id}.{projection_name}.weight", + loaded_weight[expert_id], + ) + + class LLaDA2MoeModelLM(nn.Module): def __init__( self, @@ -837,7 +873,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) params_dict = dict(self.named_parameters()) - for name, loaded_weight in prepare_llada2_language_weights( + for name, loaded_weight in _prepare_llada2_language_weights( weights, num_experts=self.config.num_experts ): if ( diff --git a/test/registered/unit/layers/attention/test_llada2_attention_utils.py b/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py similarity index 68% rename from test/registered/unit/layers/attention/test_llada2_attention_utils.py rename to test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py index 146bd8262e75..f8c1054ea445 100644 --- a/test/registered/unit/layers/attention/test_llada2_attention_utils.py +++ b/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py @@ -2,17 +2,17 @@ import torch -from sglang.srt.layers.attention.llada2_attention_utils import ( - build_llada_image_custom_mask, +from sglang.srt.layers.attention.llada2_cfg_flashinfer_backend import ( + _build_llada_image_custom_mask, ) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") -class TestLLaDA2AttentionUtils(unittest.TestCase): +class TestLLaDA2CFGFlashInferMask(unittest.TestCase): def test_builds_official_text_query_block_mask(self): - flattened = build_llada_image_custom_mask([2], [4], "cpu") + flattened = _build_llada_image_custom_mask([2], [4], "cpu") actual = flattened.view(4, 4) expected = torch.tensor( [ @@ -26,15 +26,15 @@ def test_builds_official_text_query_block_mask(self): torch.testing.assert_close(actual, expected) def test_builds_ragged_batch_without_cross_request_entries(self): - flattened = build_llada_image_custom_mask([1, 2], [3, 4], "cpu") + flattened = _build_llada_image_custom_mask([1, 2], [3, 4], "cpu") self.assertEqual(flattened.numel(), 3 * 3 + 4 * 4) def test_rejects_invalid_conditioning_spans(self): with self.assertRaisesRegex(RuntimeError, "metadata batch mismatch"): - build_llada_image_custom_mask([1], [2, 3], "cpu") + _build_llada_image_custom_mask([1], [2, 3], "cpu") with self.assertRaisesRegex(RuntimeError, "must leave query tokens"): - build_llada_image_custom_mask([2], [2], "cpu") + _build_llada_image_custom_mask([2], [2], "cpu") if __name__ == "__main__": diff --git a/test/registered/unit/models/test_llada2_weight_utils.py b/test/registered/unit/models/test_llada2_language_weights.py similarity index 87% rename from test/registered/unit/models/test_llada2_weight_utils.py rename to test/registered/unit/models/test_llada2_language_weights.py index c6e845251d27..7d380d3b7b8f 100644 --- a/test/registered/unit/models/test_llada2_weight_utils.py +++ b/test/registered/unit/models/test_llada2_language_weights.py @@ -2,20 +2,20 @@ import torch -from sglang.srt.model_loader.llada2_weight_utils import prepare_llada2_language_weights +from sglang.srt.models.llada2 import _prepare_llada2_language_weights from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=2, suite="base-a-test-cpu") -class TestLLaDA2WeightUtils(CustomTestCase): +class TestLLaDA2LanguageWeights(CustomTestCase): def test_prepares_clean_language_checkpoint_layout(self): fused = torch.arange(24).reshape(2, 3, 4) lm_head = torch.randn(4, 3) expanded = list( - prepare_llada2_language_weights( + _prepare_llada2_language_weights( [ ("model.language_model.layers.1.mlp.experts.gate_proj", fused), ("model.lm_head.weight", lm_head), @@ -39,7 +39,7 @@ def test_prepares_clean_language_checkpoint_layout(self): def test_rejects_invalid_fused_expert_shape(self): with self.assertRaisesRegex(ValueError, "expected first dimension 2"): list( - prepare_llada2_language_weights( + _prepare_llada2_language_weights( [ ( "model.language_model.layers.1.mlp.experts.down_proj", From fa9f33d160366c5f34e3d43b944cc5dbc996e3a4 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Thu, 8 Oct 2026 23:26:58 +0000 Subject: [PATCH 16/18] Validate LLaDA-Image sizes on request and trim the integration Check LLaDA-Image output sizes while the request is validated, and return request validation errors from the image API as 400 instead of 500, as the video API already does. Editing sizes must be multiples of 32 so the half-size source fits the SigVQ patch grid, and an editing size whose image tokens alone exceed the sequence RoPE table is refused before any GPU work. The DiT keeps the exact check with the prompt. Resize the editing source to the full output frame instead of center cropping it, as the reference pipeline does, and drop the SigVQ pad now that validation guarantees an aligned grid. Release text encoder requests by their KV lifecycle state, so a failed KV allocation surfaces its own error and frees the request rows. Remove code that the served checkpoints never reach, without changing any output. This drops the diffusers attention processors, output classes, and return_dict paths, the unused text-to-image SigVQ branch, the FP8 weight scale mappings that the BF16 guard leaves dead, initializers for weights that are always loaded, duplicated Flux2 VAE helpers, and checks for states the pipeline cannot produce. Generated and edited images from both checkpoints are byte identical before and after. Leave the weight data parallel manager and the quantization utils imports as they are on main, since this integration needs neither change. Merge the LLaDA2 router and checkpoint layout tests, cut the mask builder and runtime context tests to their distinct cases, share the GPU case server arguments, and drop an unrelated test edit. --- .../configs/models/dits/llada_image.py | 6 - .../configs/pipeline_configs/llada_image.py | 129 +- .../configs/sample/llada_image.py | 3 + .../runtime/entrypoints/openai/image_api.py | 12 +- .../runtime/models/dits/llada_image.py | 1535 ++++------------- .../llada_image/conditioning.py | 85 +- .../llada_image/source.py | 48 +- .../runtime/utils/quantization_utils.py | 9 +- .../multimodal_gen/test/server/gpu_cases.py | 49 +- .../test/unit/test_llada_image_config.py | 13 + .../test/unit/test_load_inclusive_e2e.py | 1 + .../test/unit/test_openai_image_api.py | 15 + .../llada2_cfg_flashinfer_backend.py | 4 +- python/sglang/srt/runtime_context.py | 10 +- .../test_llada2_cfg_flashinfer_backend.py | 37 +- .../unit/models/test_llada2_language_model.py | 75 + .../models/test_llada2_language_weights.py | 55 - .../unit/models/test_llada2_router.py | 72 - test/registered/unit/test_runtime_context.py | 94 +- 19 files changed, 582 insertions(+), 1670 deletions(-) create mode 100644 test/registered/unit/models/test_llada2_language_model.py delete mode 100644 test/registered/unit/models/test_llada2_language_weights.py delete mode 100644 test/registered/unit/models/test_llada2_router.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py index f52c66d4cbf9..941a89fad7cc 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py @@ -52,14 +52,8 @@ class LLaDAImageArchConfig(DiTArchConfig): r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3), r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3), r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3), - r"(.*)\.attention\.to_[qkv]\.weight_scale$": ( - r"\1.attention.to_qkv.weight_scale" - ), r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2), r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2), - r"(.*)\.feed_forward\.w[13]\.weight_scale$": ( - r"\1.feed_forward.w13.weight_scale" - ), } ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py index 4517beca8494..314748c43ba5 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -6,12 +6,23 @@ import torch from sglang.multimodal_gen import envs -from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig +from sglang.multimodal_gen.configs.models.dits.llada_image import ( + LLaDAImageDitConfig, + editing_rope_rows, +) from sglang.multimodal_gen.configs.models.vaes.flux import Flux2VAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, SpatialImagePipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.flux import _unpatchify_latents + + +def flux2_vae_bn_stats(vae, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = vae.bn.running_mean.view(1, -1, 1, 1).to(like) + variance = vae.bn.running_var.view(1, -1, 1, 1) + std = torch.sqrt(variance + vae.config.arch_config.batch_norm_eps).to(like) + return mean, std @dataclass @@ -20,7 +31,6 @@ class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): should_use_guidance: bool = False dit_config: LLaDAImageDitConfig = field(default_factory=LLaDAImageDitConfig) - dit_precision: str = "bf16" vae_config: Flux2VAEConfig = field(default_factory=Flux2VAEConfig) vae_precision: str = "bf16" vae_tiling: bool = False @@ -32,7 +42,8 @@ class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): postprocess_text_funcs: tuple = () latent_scale_factor: int = 16 - text_encoder_mem_fraction_static: float | None = None + # Editing feeds the source to SigVQ at half size, and its patch is 16 pixels. + editing_size_multiple: int = 32 # The embedded text worker admits 8192 prefill tokens shared by the two # CFG sequences and each sequence appends 256 query tokens. max_request_text_tokens: int = 3584 @@ -46,11 +57,9 @@ def prepare_sigmas(self, sigmas, num_inference_steps): def prepare_latent_shape(self, batch, batch_size, num_frames): del num_frames - if batch.height % self.latent_scale_factor != 0: - raise ValueError("LLaDA-Image height must be divisible by 16") - if batch.width % self.latent_scale_factor != 0: - raise ValueError("LLaDA-Image width must be divisible by 16") - self._validate_spatial_rope_bounds(batch.width, batch.height) + self.validate_output_size( + batch.width, batch.height, editing=batch.condition_image is not None + ) return ( batch_size, self.dit_config.num_channels_latents, @@ -106,6 +115,31 @@ def validate_server_args(self, server_args) -> None: def supports_disaggregation(self) -> bool: return False + def validate_output_size(self, width: int, height: int, editing: bool) -> None: + multiple = self.editing_size_multiple if editing else self.latent_scale_factor + if width % multiple != 0 or height % multiple != 0: + task = "editing" if editing else "generation" + raise ValueError( + f"LLaDA-Image {task} width and height must be divisible by " + f"{multiple}, got {width}x{height}" + ) + self._validate_spatial_rope_bounds(width, height) + if editing: + # Lower bound without the prompt, the DiT checks the exact length. + image_tokens = (width // self.latent_scale_factor) * ( + height // self.latent_scale_factor + ) + sigvq_tokens = (width // self.editing_size_multiple) * ( + height // self.editing_size_multiple + ) + rows = editing_rope_rows(0, [image_tokens, image_tokens], sigvq_tokens) + limit = self.dit_config.arch_config.axes_lens[0] + if rows > limit: + raise ValueError( + f"LLaDA-Image editing at {width}x{height} needs at least {rows} " + f"sequence positions, above the model limit of {limit}" + ) + def _validate_spatial_rope_bounds(self, width: int, height: int) -> None: # The served DiT uses patch size 1. axes_lens = self.dit_config.arch_config.axes_lens @@ -118,84 +152,39 @@ def _validate_spatial_rope_bounds(self, width: int, height: int) -> None: ) @staticmethod - def _prepare_condition_list(values, name: str, expected_size: int, device, dtype): - if values is None: - return None - if len(values) != expected_size: - raise ValueError( - f"LLaDA-Image {name} has {len(values)} entries, " - f"expected {expected_size}" - ) - return [value.to(device=device, dtype=dtype) for value in values] - - def _prepare_source_latents(self, batch, device, dtype): - source_latents = self._prepare_condition_list( - batch.source_latents, - "source_latents", - batch.batch_size, - device, - dtype, - ) - return source_latents + def _cond_kwargs(batch, device, dtype, negative: bool) -> dict: + def to_device(values): + if values is None: + return None + return [value.to(device=device, dtype=dtype) for value in values] - def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): - del rotary_emb image_embeds = batch.image_embeds if image_embeds: - image_embeds = self._prepare_condition_list( - image_embeds, - "image_embeds", - batch.batch_size, - device, - dtype, - ) - source_latents = self._prepare_source_latents(batch, device, dtype) + image_embeds = to_device(image_embeds) + if negative: + # The unconditional edit keeps the source latents without SigVQ. + empty = image_embeds[0].new_zeros((0, image_embeds[0].shape[-1])) + image_embeds = [empty] * batch.batch_size return { "encoder_hidden_states_image": image_embeds, - "source_latents": source_latents, + "source_latents": to_device(batch.source_latents), } + def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): + return self._cond_kwargs(batch, device, dtype, negative=False) + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): - del rotary_emb - image_embeds = batch.image_embeds - if image_embeds: - image_embeds = self._prepare_condition_list( - image_embeds, - "image_embeds", - batch.batch_size, - device, - dtype, - ) - empty_image_embed = image_embeds[0].new_zeros( - (0, image_embeds[0].shape[-1]) - ) - image_embeds = [empty_image_embed] * batch.batch_size - source_latents = self._prepare_source_latents(batch, device, dtype) - return { - "encoder_hidden_states_image": image_embeds, - "source_latents": source_latents, - } + return self._cond_kwargs(batch, device, dtype, negative=True) def get_decode_scale_and_shift(self, device, dtype, vae): del device, dtype, vae return 1.0, None def preprocess_decoding(self, latents, server_args=None, vae=None): - del server_args - if vae is None or not hasattr(vae, "bn"): - raise ValueError("LLaDA-Image decoding requires the Flux2 VAE BN state") vae_parameter = next(vae.parameters()) latents = latents.to(device=vae_parameter.device, dtype=vae_parameter.dtype) - vae_config = getattr(vae.config, "arch_config", vae.config) - latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) - latent_std = torch.sqrt( - vae.bn.running_var.view(1, -1, 1, 1) + vae_config.batch_norm_eps - ).to(latents) - latents = latents * latent_std + latent_mean - batch_size, channels, height, width = latents.shape - latents = latents.reshape(batch_size, channels // 4, 2, 2, height, width) - latents = latents.permute(0, 1, 4, 2, 5, 3) - return latents.reshape(batch_size, channels // 4, height * 2, width * 2) + mean, std = flux2_vae_bn_stats(vae, latents) + return _unpatchify_latents(latents * std + mean) def post_denoising_loop(self, latents, batch): del batch diff --git a/python/sglang/multimodal_gen/configs/sample/llada_image.py b/python/sglang/multimodal_gen/configs/sample/llada_image.py index e0b19f6dc23e..d31133e9dccb 100644 --- a/python/sglang/multimodal_gen/configs/sample/llada_image.py +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -30,6 +30,9 @@ def _validate_with_pipeline_config(self, pipeline_config): raise ValueError( "LLaDA-Image max_sequence_length must be between 1 and 3584" ) + pipeline_config.validate_output_size( + self.width, self.height, editing=bool(self.image_path) + ) @dataclass diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 25f4211e9b01..f3aff7cd2dcb 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -248,6 +248,14 @@ def _get_response_resize( return f"{width}x{height}" +def _build_request_sampling_params(request_id: str, **kwargs): + # Request validation errors are client errors, as in the video API. + try: + return build_sampling_params(request_id, **kwargs) + except (TypeError, ValueError) as e: + raise HTTPException(status_code=400, detail=str(e)) from e + + @router.post("/generations", response_model=ImageResponse) async def generations( request: ImageGenerationsRequest, @@ -266,7 +274,7 @@ async def generations( ) with temp_dir_if_disabled(server_args.output_path) as output_dir: - sampling = build_sampling_params( + sampling = _build_request_sampling_params( request_id, prompt=prompt, task_type=request.task_type, @@ -459,7 +467,7 @@ async def edits( task="image_edit", image_paths=input_paths, ) - sampling = build_sampling_params( + sampling = _build_request_sampling_params( request_id, prompt=prompt, task_type=task_type, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py index ebd4b9468a72..6e0518185af5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -19,13 +19,10 @@ import torch import torch.nn.functional as F from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward +from diffusers.models.attention import FeedForward from diffusers.models.attention_dispatch import dispatch_attention_fn -from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.modeling_utils import ModelMixin from diffusers.models.normalization import RMSNorm -from diffusers.utils import BaseOutput -from diffusers.utils.torch_utils import maybe_allow_in_graph from torch import nn from torch.nn.utils.rnn import pad_sequence @@ -116,10 +113,6 @@ def forward( ) arguments = timestep[:, None].float() * frequencies[None] embedding = torch.cat([torch.cos(arguments), torch.sin(arguments)], dim=-1) - if self.frequency_embedding_dim % 2: - embedding = torch.cat( - [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 - ) return self.mlp(embedding.to(dtype=hidden_dtype)) @@ -158,82 +151,7 @@ def forward(self, position_ids: torch.Tensor) -> torch.Tensor: return torch.cat(frequencies, dim=-1) -class LLaDAImageAttnProcessor: - _attention_backend = None - _parallel_config = None - - def __call__( - self, - attn: "LLaDAImageAttention", - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - freqs_cis: torch.Tensor | None = None, - ) -> torch.Tensor: - qkv, _ = attn.to_qkv(hidden_states) - query, key, value = qkv.split( - [attn.inner_dim, attn.inner_dim, attn.inner_dim], dim=-1 - ) - query = query.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() - key = key.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() - value = value.unflatten(-1, (attn.local_heads, attn.head_dim)).contiguous() - - if freqs_cis is not None: - cos_sin_cache = torch.cat( - [freqs_cis.real.float(), freqs_cis.imag.float()], dim=-1 - ).reshape(-1, attn.head_dim) - positions = torch.arange( - cos_sin_cache.shape[0], device=query.device, dtype=torch.long - ) - if attn.norm_q is not None: - query, key = apply_qk_norm_with_optional_rope( - q=query, - k=key, - q_norm=attn.norm_q, - k_norm=attn.norm_k, - head_dim=attn.head_dim, - cos_sin_cache=cos_sin_cache.contiguous(), - positions=positions, - is_neox=False, - allow_inplace=True, - ) - else: - query, key = apply_flashinfer_rope_qk_inplace( - query, - key, - cos_sin_cache.contiguous(), - head_size=attn.head_dim, - positions=positions, - is_neox=False, - ) - elif attn.norm_q is not None: - query, key = apply_qk_norm_with_optional_rope( - q=query, - k=key, - q_norm=attn.norm_q, - k_norm=attn.norm_k, - head_dim=attn.head_dim, - allow_inplace=True, - ) - - if attention_mask is not None and attention_mask.ndim == 2: - attention_mask = attention_mask[:, None, None, :] - - hidden_states = attn.sgl_attention( - query, - key, - value, - attn_mask=attention_mask, - ) - hidden_states = hidden_states.flatten(2, 3) - hidden_states, _ = attn.to_out[0](hidden_states) - return hidden_states - - -class LLaDAImageAttention(nn.Module, AttentionModuleMixin): - _default_processor_cls = LLaDAImageAttnProcessor - _available_processors: ClassVar[list[type]] = [LLaDAImageAttnProcessor] - _supports_qkv_fusion = False - +class LLaDAImageAttention(nn.Module): def __init__( self, dim: int, @@ -283,11 +201,9 @@ def __init__( input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.to_out.0", - ), - nn.Dropout(0.0), + ) ] ) - self.set_processor(self._default_processor_cls()) def forward( self, @@ -295,12 +211,43 @@ def forward( attention_mask: torch.Tensor | None, freqs_cis: torch.Tensor, ) -> torch.Tensor: - return self.processor( - self, - hidden_states, - attention_mask, - freqs_cis, - ) + qkv, _ = self.to_qkv(hidden_states) + query, key, value = ( + tensor.unflatten(-1, (self.local_heads, self.head_dim)).contiguous() + for tensor in qkv.split([self.inner_dim] * 3, dim=-1) + ) + cos_sin_cache = torch.cat( + [freqs_cis.real.float(), freqs_cis.imag.float()], dim=-1 + ).reshape(-1, self.head_dim) + positions = torch.arange( + cos_sin_cache.shape[0], device=query.device, dtype=torch.long + ) + if self.norm_q is not None: + query, key = apply_qk_norm_with_optional_rope( + q=query, + k=key, + q_norm=self.norm_q, + k_norm=self.norm_k, + head_dim=self.head_dim, + cos_sin_cache=cos_sin_cache.contiguous(), + positions=positions, + is_neox=False, + allow_inplace=True, + ) + else: + query, key = apply_flashinfer_rope_qk_inplace( + query, + key, + cos_sin_cache.contiguous(), + head_size=self.head_dim, + positions=positions, + is_neox=False, + ) + if attention_mask is not None: + attention_mask = attention_mask[:, None, None, :] + hidden_states = self.sgl_attention(query, key, value, attn_mask=attention_mask) + hidden_states, _ = self.to_out[0](hidden_states.flatten(2, 3)) + return hidden_states class LLaDAImageFeedForward(nn.Module): @@ -337,6 +284,20 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states +def _padding_attention_mask( + lengths: list[int], device: torch.device +) -> torch.Tensor | None: + max_length = max(lengths) + if all(length == max_length for length in lengths): + return None + attention_mask = torch.zeros( + (len(lengths), max_length), dtype=torch.bool, device=device + ) + for batch_index, length in enumerate(lengths): + attention_mask[batch_index, :length] = True + return attention_mask + + def _select_per_token( noisy_value: torch.Tensor, clean_value: torch.Tensor, @@ -351,7 +312,6 @@ def _select_per_token( ) -@maybe_allow_in_graph class LLaDAImageTransformerBlock(nn.Module): def __init__( self, @@ -396,48 +356,22 @@ def forward( adaln_clean: torch.Tensor | None = None, ) -> torch.Tensor: if self.modulation: - sequence_length = hidden_states.shape[1] if noise_mask is None: - scale_msa, gate_msa, scale_mlp, gate_mlp = ( + modulation = ( self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) ) - scale_msa = 1.0 + scale_msa else: - noisy_modulation = self.adaLN_modulation(adaln_noisy) - clean_modulation = self.adaLN_modulation(adaln_clean) - noisy_scale_msa, noisy_gate_msa, noisy_scale_mlp, noisy_gate_mlp = ( - noisy_modulation.chunk(4, dim=1) - ) - clean_scale_msa, clean_gate_msa, clean_scale_mlp, clean_gate_mlp = ( - clean_modulation.chunk(4, dim=1) - ) - scale_msa = _select_per_token( - 1.0 + noisy_scale_msa, - 1.0 + clean_scale_msa, - noise_mask, - sequence_length, - ) - scale_mlp = _select_per_token( - noisy_scale_mlp, - clean_scale_mlp, - noise_mask, - sequence_length, - ) - gate_msa = _select_per_token( - noisy_gate_msa, - clean_gate_msa, - noise_mask, - sequence_length, - ) - gate_mlp = _select_per_token( - noisy_gate_mlp, - clean_gate_mlp, - noise_mask, - sequence_length, - ) + modulation = [ + _select_per_token(noisy, clean, noise_mask, hidden_states.shape[1]) + for noisy, clean in zip( + self.adaLN_modulation(adaln_noisy).chunk(4, dim=1), + self.adaLN_modulation(adaln_clean).chunk(4, dim=1), + ) + ] + scale_msa, gate_msa, scale_mlp, gate_mlp = modulation attention_output = self.attention( - self.attention_norm1(hidden_states) * scale_msa, + self.attention_norm1(hidden_states) * (1.0 + scale_msa), attention_mask, freqs_cis, ) @@ -482,67 +416,20 @@ def forward( adaln_clean: torch.Tensor | None = None, ) -> torch.Tensor: if noise_mask is None: - scale = 1.0 + self.adaLN_modulation(adaln_input) - scale = scale.unsqueeze(1) + scale = self.adaLN_modulation(adaln_input).unsqueeze(1) else: - sequence_length = hidden_states.shape[1] - noisy_scale = 1.0 + self.adaLN_modulation(adaln_noisy) - clean_scale = 1.0 + self.adaLN_modulation(adaln_clean) scale = _select_per_token( - noisy_scale, clean_scale, noise_mask, sequence_length + self.adaLN_modulation(adaln_noisy), + self.adaLN_modulation(adaln_clean), + noise_mask, + hidden_states.shape[1], ) - hidden_states = self.norm_final(hidden_states) * scale + hidden_states = self.norm_final(hidden_states) * (1.0 + scale) return self.linear(hidden_states) -class _LLaDAImageTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): - r""" - The denoising transformer used by LLaDAImage for text-to-image generation and single-image editing. - - This component consumes caption features that have already passed through LLaDAImage's QueryFormer, connector, and - projector. For editing, it additionally consumes GLM/SigVQ features and source-image latents. - - Args: - all_patch_size (`tuple[int, ...]`, defaults to `(1,)`): - Supported spatial patch sizes. - all_f_patch_size (`tuple[int, ...]`, defaults to `(1,)`): - Supported temporal patch sizes paired with `all_patch_size`. - in_channels (`int`, defaults to `128`): - Number of channels in the patchified Flux2 VAE latents. - dim (`int`, defaults to `3840`): - Transformer hidden dimension. - n_layers (`int`, defaults to `30`): - Number of main transformer blocks. - n_refiner_layers (`int`, defaults to `2`): - Number of noise, caption, and SigVQ refiner blocks. - n_heads (`int`, defaults to `30`): - Number of attention heads. - norm_eps (`float`, defaults to `1e-5`): - Epsilon used by RMS normalization layers. - qk_norm (`bool`, defaults to `True`): - Whether to apply RMS normalization to query and key tensors. - cap_feat_dim (`int`, defaults to `2560`): - Dimension of projected QueryFormer caption features. - semantic_feat_dim (`int`, defaults to `4096`): - Dimension of GLM/SigVQ semantic features. - rope_theta (`float`, defaults to `256.0`): - RoPE frequency base. - t_scale (`float`, defaults to `1000.0`): - Scale applied to diffusion timesteps. - axes_dims (`tuple[int, ...]`, defaults to `(32, 48, 48)`): - RoPE dimensions for sequence, height, and width axes. - axes_lens (`tuple[int, ...]`, defaults to `(32768, 1024, 1024)`): - Maximum RoPE positions for sequence, height, and width axes. - """ - - _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] - _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTransformerBlock"] - _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ - "t_embedder", - "cap_embedder", - "semantic_embedder", - "sigvq_embedder", - ] +class _LLaDAImageTransformer2DModel(ModelMixin, ConfigMixin): + """Denoiser for text-to-image generation and single-image editing.""" @register_to_config def __init__( @@ -565,23 +452,7 @@ def __init__( quant_config: QuantizationConfig | None = None, ): super().__init__() - if len(all_patch_size) != len(all_f_patch_size): - raise ValueError( - "`all_patch_size` and `all_f_patch_size` must have the same length." - ) - if dim % n_heads != 0: - raise ValueError( - f"`dim` ({dim}) must be divisible by `n_heads` ({n_heads})." - ) - if dim // n_heads != sum(axes_dims): - raise ValueError( - "The attention head dimension must equal the sum of `axes_dims`." - ) - - self.in_channels = in_channels self.out_channels = in_channels - self.all_patch_size = all_patch_size - self.all_f_patch_size = all_f_patch_size self.t_scale = t_scale self.all_x_embedder = nn.ModuleDict() @@ -592,86 +463,40 @@ def __init__( self.all_x_embedder[patch_key] = nn.Linear(patch_dim, dim, bias=True) self.all_final_layer[patch_key] = LLaDAImageFinalLayer(dim, patch_dim) - self.noise_refiner = nn.ModuleList( - [ - LLaDAImageTransformerBlock( - dim, - n_heads, - norm_eps, - qk_norm, - modulation=True, - quant_config=quant_config, - prefix=f"noise_refiner.{layer_id}", - ) - for layer_id in range(n_refiner_layers) - ] - ) - self.context_refiner = nn.ModuleList( - [ - LLaDAImageTransformerBlock( - dim, - n_heads, - norm_eps, - qk_norm, - modulation=False, - quant_config=quant_config, - prefix=f"context_refiner.{layer_id}", - ) - for layer_id in range(n_refiner_layers) - ] - ) - self.sigvq_refiner = nn.ModuleList( - [ - LLaDAImageTransformerBlock( - dim, - n_heads, - norm_eps, - qk_norm, - modulation=False, - quant_config=quant_config, - prefix=f"sigvq_refiner.{layer_id}", - ) - for layer_id in range(n_refiner_layers) - ] - ) - self.layers = nn.ModuleList( - [ + def blocks(name: str, count: int, modulation: bool) -> nn.ModuleList: + return nn.ModuleList( LLaDAImageTransformerBlock( dim, n_heads, norm_eps, qk_norm, - modulation=True, + modulation=modulation, quant_config=quant_config, - prefix=f"layers.{layer_id}", + prefix=f"{name}.{layer_id}", ) - for layer_id in range(n_layers) - ] - ) + for layer_id in range(count) + ) - self.t_embedder = LLaDAImageTimestepEmbedder(min(dim, ADALN_EMBED_DIM)) - self.cap_embedder = nn.Sequential( - RMSNorm(cap_feat_dim, eps=norm_eps, elementwise_affine=False), - nn.Linear(cap_feat_dim, dim, bias=True), - ) - self.semantic_embedder = nn.Sequential( - RMSNorm(semantic_feat_dim, eps=norm_eps, elementwise_affine=False), - nn.Linear(semantic_feat_dim, dim, bias=True), - ) - self.sigvq_embedder = nn.Sequential( - RMSNorm(semantic_feat_dim, eps=norm_eps, elementwise_affine=False), - nn.Linear(semantic_feat_dim, dim, bias=True), - ) + def feature_embedder(feature_dim: int) -> nn.Sequential: + return nn.Sequential( + RMSNorm(feature_dim, eps=norm_eps, elementwise_affine=False), + nn.Linear(feature_dim, dim, bias=True), + ) - nn.init.normal_(self.semantic_embedder[1].weight, mean=0.0, std=0.02) - nn.init.zeros_(self.semantic_embedder[1].bias) - nn.init.normal_(self.sigvq_embedder[1].weight, mean=0.0, std=0.02) - nn.init.zeros_(self.sigvq_embedder[1].bias) + self.noise_refiner = blocks("noise_refiner", n_refiner_layers, True) + self.context_refiner = blocks("context_refiner", n_refiner_layers, False) + self.sigvq_refiner = blocks("sigvq_refiner", n_refiner_layers, False) + self.layers = blocks("layers", n_layers, True) + + self.t_embedder = LLaDAImageTimestepEmbedder(min(dim, ADALN_EMBED_DIM)) + self.cap_embedder = feature_embedder(cap_feat_dim) + # Checkpoint weights that neither generation nor editing reads. + self.semantic_embedder = feature_embedder(semantic_feat_dim) + self.sigvq_embedder = feature_embedder(semantic_feat_dim) self.x_pad_token = nn.Parameter(torch.zeros(1, dim)) self.cap_pad_token = nn.Parameter(torch.zeros(1, dim)) self.sigvq_pad_token = nn.Parameter(torch.zeros(1, dim)) - nn.init.normal_(self.sigvq_pad_token, mean=0.0, std=0.02) self.rope_embedder = LLaDAImageRopeEmbedder(rope_theta, axes_dims, axes_lens) @@ -724,40 +549,25 @@ def _pad_with_ids( position_grid_size: tuple[int, int, int], position_start: tuple[int, int, int], noise_value: int | None = None, - sequence_multiple: int = SEQUENCE_MULTIPLE, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, list[int] | None]: original_length = len(features) - padding_length = (-original_length) % sequence_multiple + padding_length = (-original_length) % SEQUENCE_MULTIPLE padded_length = original_length + padding_length - device = features.device - position_ids = self._create_coordinate_grid( - position_grid_size, - position_start, - device, + position_grid_size, position_start, features.device ).flatten(0, 2) + padding_mask = torch.zeros( + padded_length, dtype=torch.bool, device=features.device + ) if padding_length > 0: - padding_position_ids = ( - self._create_coordinate_grid( - (1, 1, 1), - (0, 0, 0), - device, - ) - .flatten(0, 2) - .repeat(padding_length, 1) + # Padding rows sit at position zero and repeat the last feature. + position_ids = torch.cat( + [position_ids, position_ids.new_zeros((padding_length, 3))], dim=0 ) - position_ids = torch.cat([position_ids, padding_position_ids], dim=0) features = torch.cat( [features, features[-1:].repeat(padding_length, 1)], dim=0 ) - padding_mask = torch.cat( - [ - torch.zeros(original_length, dtype=torch.bool, device=device), - torch.ones(padding_length, dtype=torch.bool, device=device), - ] - ) - else: - padding_mask = torch.zeros(original_length, dtype=torch.bool, device=device) + padding_mask[original_length:] = True noise_mask = [noise_value] * padded_length if noise_value is not None else None return features, position_ids, padding_mask, padded_length, noise_mask @@ -773,7 +583,6 @@ def _batch_sequences( torch.Tensor, torch.Tensor, torch.Tensor | None, list[int], torch.Tensor | None ]: sequence_lengths = [len(item) for item in features] - max_sequence_length = max(sequence_lengths) features = torch.cat(features, dim=0) inner_padding_mask = torch.cat(inner_padding_masks).unsqueeze(-1) features = torch.where( @@ -787,16 +596,7 @@ def _batch_sequences( frequencies = pad_sequence(frequencies, batch_first=True, padding_value=0.0)[ :, : features.shape[1] ] - - attention_mask = None - if not all(length == max_sequence_length for length in sequence_lengths): - attention_mask = torch.zeros( - (len(sequence_lengths), max_sequence_length), - dtype=torch.bool, - device=features.device, - ) - for batch_index, sequence_length in enumerate(sequence_lengths): - attention_mask[batch_index, :sequence_length] = True + attention_mask = _padding_attention_mask(sequence_lengths, features.device) noise_mask = None if noise_masks is not None: @@ -820,7 +620,6 @@ def _unpatchify( image_offsets: list[tuple[int, int]] | None = None, ) -> list[torch.Tensor]: outputs = [] - sequence_multiple = SEQUENCE_MULTIPLE for batch_index, batch_hidden_states in enumerate(hidden_states): if image_offsets is None: batch_sizes = [sizes[batch_index]] @@ -838,7 +637,7 @@ def _unpatchify( * (height // patch_size) * (width // patch_size) ) - padding_length = (-original_length) % sequence_multiple + padding_length = (-original_length) % SEQUENCE_MULTIPLE output = ( image_hidden_states[ current_offset : current_offset + original_length @@ -862,69 +661,36 @@ def _unpatchify( def _prepare_t2i_sequences( self, x: list[torch.Tensor], - cap_feats: list[torch.Tensor] | None, - glm_features: list[torch.Tensor] | None, + cap_feats: list[torch.Tensor], patch_size: int, f_patch_size: int, - ) -> tuple[ - _LLaDAImageSequence, - _LLaDAImageSequence | None, - _LLaDAImageSequence | None, - list[tuple[int, int, int]], - ]: + ) -> tuple[_LLaDAImageSequence, _LLaDAImageSequence, list[tuple[int, int, int]]]: image_sequence = _LLaDAImageSequence([], [], []) - cap_sequence = ( - _LLaDAImageSequence([], [], []) if cap_feats is not None else None - ) - glm_sequence = ( - _LLaDAImageSequence([], [], []) if glm_features is not None else None - ) + cap_sequence = _LLaDAImageSequence([], [], []) image_sizes = [] - for batch_index, latent in enumerate(x): - position_cursor = 1 - if cap_sequence is not None: - padded_features, position_ids, padding_mask, sequence_length, _ = ( - self._pad_with_ids( - cap_feats[batch_index], - (len(cap_feats[batch_index]), 1, 1), - (position_cursor, 0, 0), - ) - ) - cap_sequence.features.append(padded_features) - cap_sequence.position_ids.append(position_ids) - cap_sequence.padding_masks.append(padding_mask) - position_cursor += sequence_length - - if glm_sequence is not None: - padded_features, position_ids, padding_mask, sequence_length, _ = ( - self._pad_with_ids( - glm_features[batch_index], - (len(glm_features[batch_index]), 1, 1), - (position_cursor, 0, 0), - ) + for latent, batch_cap_feats in zip(x, cap_feats): + padded_features, position_ids, padding_mask, cap_length, _ = ( + self._pad_with_ids( + batch_cap_feats, (len(batch_cap_feats), 1, 1), (1, 0, 0) ) - glm_sequence.features.append(padded_features) - glm_sequence.position_ids.append(position_ids) - glm_sequence.padding_masks.append(padding_mask) - position_cursor += sequence_length + ) + cap_sequence.features.append(padded_features) + cap_sequence.position_ids.append(position_ids) + cap_sequence.padding_masks.append(padding_mask) patches, image_size, token_grid_size = self._patchify_image( latent, patch_size, f_patch_size ) - image_height_start = 0 padded_features, position_ids, padding_mask, _, _ = self._pad_with_ids( - patches, - token_grid_size, - (position_cursor, image_height_start, 0), - sequence_multiple=SEQUENCE_MULTIPLE, + patches, token_grid_size, (1 + cap_length, 0, 0) ) image_sequence.features.append(padded_features) image_sequence.position_ids.append(position_ids) image_sequence.padding_masks.append(padding_mask) image_sizes.append(image_size) - return image_sequence, cap_sequence, glm_sequence, image_sizes + return image_sequence, cap_sequence, image_sizes def _prepare_editing_sequences( self, @@ -986,14 +752,9 @@ def _prepare_editing_sequences( image, patch_size, f_patch_size ) image_token_counts.append(len(patches)) - image_height_start = 0 padded_features, position_ids, padding_mask, _, noise_mask = ( self._pad_with_ids( - patches, - token_grid_size, - (position_start, image_height_start, 0), - noise_value, - sequence_multiple=SEQUENCE_MULTIPLE, + patches, token_grid_size, (position_start, 0, 0), noise_value ) ) batch_image_features.append(padded_features) @@ -1058,135 +819,69 @@ def _merge_padded_sequences( frequency_groups: tuple[torch.Tensor, ...], length_groups: tuple[list[int], ...], noise_mask_groups: tuple[torch.Tensor, ...] | None = None, - ) -> tuple[ - torch.Tensor, torch.Tensor, torch.Tensor | None, list[int], torch.Tensor | None - ]: - batch_size = feature_groups[0].shape[0] - merged_features = [] - merged_frequencies = [] - merged_noise_masks = [] if noise_mask_groups is not None else None - - for batch_index in range(batch_size): - device = feature_groups[0].device - merged_features.append( - torch.cat( - [ - features[batch_index, : lengths[batch_index]].to(device) - for features, lengths in zip(feature_groups, length_groups) - ], - dim=0, - ) - ) - merged_frequencies.append( + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + def merge(groups: tuple[torch.Tensor, ...]) -> list[torch.Tensor]: + return [ torch.cat( [ - frequencies[batch_index, : lengths[batch_index]].to(device) - for frequencies, lengths in zip(frequency_groups, length_groups) + group[batch_index, : lengths[batch_index]] + for group, lengths in zip(groups, length_groups) ], dim=0, ) - ) - if merged_noise_masks is not None: - merged_noise_masks.append( - torch.cat( - [ - noise_masks[batch_index, : lengths[batch_index]].to(device) - for noise_masks, lengths in zip( - noise_mask_groups, length_groups - ) - ], - dim=0, - ) - ) + for batch_index in range(feature_groups[0].shape[0]) + ] - merged_lengths = [len(features) for features in merged_features] - merged_features = pad_sequence( - merged_features, batch_first=True, padding_value=0.0 - ) - merged_frequencies = pad_sequence( - merged_frequencies, batch_first=True, padding_value=0.0 + merged_features = merge(feature_groups) + attention_mask = _padding_attention_mask( + [len(features) for features in merged_features], + feature_groups[0].device, ) - - attention_mask = None - max_length = max(merged_lengths) - if not all(length == max_length for length in merged_lengths): - attention_mask = torch.zeros( - (batch_size, max_length), - dtype=torch.bool, - device=merged_features.device, - ) - for batch_index, sequence_length in enumerate(merged_lengths): - attention_mask[batch_index, :sequence_length] = True - noise_mask = None - if merged_noise_masks is not None: + if noise_mask_groups is not None: noise_mask = pad_sequence( - merged_noise_masks, batch_first=True, padding_value=0 - )[:, : merged_features.shape[1]] - + merge(noise_mask_groups), batch_first=True, padding_value=0 + ) return ( - merged_features, - merged_frequencies, + pad_sequence(merged_features, batch_first=True, padding_value=0.0), + pad_sequence(merge(frequency_groups), batch_first=True, padding_value=0.0), attention_mask, - merged_lengths, noise_mask, ) + def _embed_sequence( + self, + sequence: _LLaDAImageSequence, + embedder: nn.Module, + pad_token: torch.Tensor, + ) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor | None, list[int], torch.Tensor | None + ]: + lengths = [len(features) for features in sequence.features] + features = embedder(torch.cat(sequence.features, dim=0)) + frequencies = self.rope_embedder(torch.cat(sequence.position_ids, dim=0)) + return self._batch_sequences( + list(features.split(lengths, dim=0)), + list(frequencies.split(lengths, dim=0)), + sequence.padding_masks, + pad_token, + sequence.noise_masks, + ) + def forward( self, x: list[torch.Tensor], t: torch.Tensor, - cap_feats: list[torch.Tensor] | None, + cap_feats: list[torch.Tensor], glm_cap_feats: list[torch.Tensor] | None = None, source_latents: list[torch.Tensor] | None = None, patch_size: int = 1, f_patch_size: int = 1, - return_dict: bool = True, - ) -> Transformer2DModelOutput | tuple[list[torch.Tensor]]: - r""" - Args: - x (`list[torch.Tensor]`): - Target latents. Each tensor has shape `(channels, frames, height, width)`. - t (`torch.Tensor`): - Denoising timestep for each batch item. - cap_feats (`list[torch.Tensor]`, *optional*): - Projected QueryFormer features, each with shape `(sequence_length, cap_feat_dim)`. - glm_cap_feats (`list[torch.Tensor]`, *optional*): - GLM/SigVQ features, each with shape `(sequence_length, semantic_feat_dim)`. - source_latents (`list[torch.Tensor]`, *optional*): - Source-image latents for editing. When provided, `cap_feats` and `glm_cap_feats` are required. - patch_size (`int`, defaults to `1`): - Spatial patch size. - f_patch_size (`int`, defaults to `1`): - Temporal patch size. - return_dict (`bool`, defaults to `True`): - Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`]. - - Returns: - [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: - The denoised target latents. - """ + ) -> list[torch.Tensor]: patch_key = f"{patch_size}-{f_patch_size}" - if patch_key not in self.all_x_embedder: - raise ValueError( - f"Unsupported patch sizes: patch_size={patch_size}, f_patch_size={f_patch_size}." - ) - if source_latents is None and cap_feats is None and glm_cap_feats is None: - raise ValueError( - "Text-to-image inference requires `cap_feats` or `glm_cap_feats`." - ) - if source_latents is not None and (cap_feats is None or glm_cap_feats is None): - raise ValueError( - "Editing requires `cap_feats`, `glm_cap_feats`, and `source_latents`." - ) - batch_size = len(x) is_editing = source_latents is not None - adaln_input = None - noisy_embedding = None - clean_embedding = None image_offsets = None - if is_editing: if t.shape[0] == 1: t = t.repeat(batch_size) @@ -1194,8 +889,10 @@ def forward( dual_embedding = self.t_embedder( dual_timestep.abs() * self.t_scale, x[0].dtype ) - noisy_embedding = dual_embedding[:batch_size] - clean_embedding = dual_embedding[batch_size:] + adaln = { + "adaln_noisy": dual_embedding[:batch_size], + "adaln_clean": dual_embedding[batch_size:], + } image_sequence, cap_sequence, sigvq_sequence, image_sizes, image_offsets = ( self._prepare_editing_sequences( x, @@ -1207,410 +904,126 @@ def forward( ) ) else: - adaln_input = self.t_embedder(t * self.t_scale, x[0].dtype) - glm_features = ( - [ - self.semantic_embedder(batch_features) - for batch_features in glm_cap_feats - ] - if glm_cap_feats is not None - else None - ) - image_sequence, cap_sequence, glm_sequence, image_sizes = ( - self._prepare_t2i_sequences( - x, - cap_feats, - glm_features, - patch_size, - f_patch_size, - ) + adaln = {"adaln_input": self.t_embedder(t * self.t_scale, x[0].dtype)} + image_sequence, cap_sequence, image_sizes = self._prepare_t2i_sequences( + x, cap_feats, patch_size, f_patch_size ) - image_lengths = [len(features) for features in image_sequence.features] - image_features = self.all_x_embedder[patch_key]( - torch.cat(image_sequence.features, dim=0) - ) - image_frequencies = list( - self.rope_embedder(torch.cat(image_sequence.position_ids, dim=0)).split( - [len(position_ids) for position_ids in image_sequence.position_ids], - dim=0, + image_features, image_frequencies, image_mask, image_lengths, image_noise = ( + self._embed_sequence( + image_sequence, self.all_x_embedder[patch_key], self.x_pad_token ) ) - ( - image_features, - image_frequencies, - image_attention_mask, - image_lengths, - image_noise_mask, - ) = self._batch_sequences( - list(image_features.split(image_lengths, dim=0)), - image_frequencies, - image_sequence.padding_masks, - self.x_pad_token, - image_sequence.noise_masks, - ) - for layer in self.noise_refiner: - if is_editing: - image_features = layer( - image_features, - image_attention_mask, - image_frequencies, - noise_mask=image_noise_mask, - adaln_noisy=noisy_embedding, - adaln_clean=clean_embedding, - ) - else: - image_features = layer( - image_features, - image_attention_mask, - image_frequencies, - adaln_input, - ) - - if is_editing: - cap_lengths = [len(features) for features in cap_sequence.features] - cap_features = self.cap_embedder(torch.cat(cap_sequence.features, dim=0)) - cap_frequencies = list( - self.rope_embedder(torch.cat(cap_sequence.position_ids, dim=0)).split( - [len(position_ids) for position_ids in cap_sequence.position_ids], - dim=0, - ) - ) - ( - cap_features, - cap_frequencies, - cap_attention_mask, - cap_lengths, - cap_noise_mask, - ) = self._batch_sequences( - list(cap_features.split(cap_lengths, dim=0)), - cap_frequencies, - cap_sequence.padding_masks, - self.cap_pad_token, - cap_sequence.noise_masks, + image_features = layer( + image_features, + image_mask, + image_frequencies, + noise_mask=image_noise, + **adaln, ) - for layer in self.context_refiner: - cap_features = layer( - cap_features, - cap_attention_mask, - cap_frequencies, - ) + cap_features, cap_frequencies, cap_mask, cap_lengths, cap_noise = ( + self._embed_sequence(cap_sequence, self.cap_embedder, self.cap_pad_token) + ) + for layer in self.context_refiner: + cap_features = layer(cap_features, cap_mask, cap_frequencies) - sigvq_lengths = [len(features) for features in sigvq_sequence.features] - sigvq_features = self.sigvq_embedder( - torch.cat(sigvq_sequence.features, dim=0) - ) - sigvq_frequencies = list( - self.rope_embedder(torch.cat(sigvq_sequence.position_ids, dim=0)).split( - [len(position_ids) for position_ids in sigvq_sequence.position_ids], - dim=0, - ) - ) + if is_editing: ( sigvq_features, sigvq_frequencies, - sigvq_attention_mask, + sigvq_mask, sigvq_lengths, - sigvq_noise_mask, - ) = self._batch_sequences( - list(sigvq_features.split(sigvq_lengths, dim=0)), - sigvq_frequencies, - sigvq_sequence.padding_masks, - self.sigvq_pad_token, - sigvq_sequence.noise_masks, + sigvq_noise, + ) = self._embed_sequence( + sigvq_sequence, self.sigvq_embedder, self.sigvq_pad_token ) - if any(sigvq_lengths): for layer in self.sigvq_refiner: sigvq_features = layer( - sigvq_features, - sigvq_attention_mask, - sigvq_frequencies, - ) - - feature_groups = (cap_features, image_features, sigvq_features) - frequency_groups = ( - cap_frequencies, - image_frequencies, - sigvq_frequencies, - ) - length_groups = (cap_lengths, image_lengths, sigvq_lengths) - noise_mask_groups = ( - cap_noise_mask, - image_noise_mask, - sigvq_noise_mask, - ) - - ( - unified_features, - unified_frequencies, - unified_attention_mask, - _, - unified_noise_mask, - ) = self._merge_padded_sequences( - feature_groups, - frequency_groups, - length_groups, - noise_mask_groups, - ) - else: - condition_feature_groups = [] - condition_frequency_groups = [] - condition_length_groups = [] - - if cap_sequence is not None: - cap_lengths = [len(features) for features in cap_sequence.features] - cap_features = self.cap_embedder( - torch.cat(cap_sequence.features, dim=0) - ) - cap_padding_mask = ( - torch.cat(cap_sequence.padding_masks) - .unsqueeze(-1) - .to(cap_features.device) - ) - cap_features = torch.where( - cap_padding_mask, - self.cap_pad_token.to( - device=cap_features.device, dtype=cap_features.dtype - ), - cap_features, - ) - cap_features = pad_sequence( - list(cap_features.split(cap_lengths, dim=0)), - batch_first=True, - padding_value=0.0, - ) - cap_frequencies = list( - self.rope_embedder( - torch.cat(cap_sequence.position_ids, dim=0) - ).split( - [ - len(position_ids) - for position_ids in cap_sequence.position_ids - ], - dim=0, - ) - ) - cap_frequencies = pad_sequence( - cap_frequencies, batch_first=True, padding_value=0.0 - ) - condition_feature_groups.append(cap_features) - condition_frequency_groups.append(cap_frequencies) - condition_length_groups.append(cap_lengths) - - if glm_sequence is not None: - glm_lengths = [len(features) for features in glm_sequence.features] - glm_features = torch.cat(glm_sequence.features, dim=0) - glm_padding_mask = ( - torch.cat(glm_sequence.padding_masks) - .unsqueeze(-1) - .to(glm_features.device) - ) - glm_features = torch.where( - glm_padding_mask, - self.cap_pad_token.to( - device=glm_features.device, dtype=glm_features.dtype - ), - glm_features, - ) - glm_features = pad_sequence( - list(glm_features.split(glm_lengths, dim=0)), - batch_first=True, - padding_value=0.0, - ) - glm_frequencies = list( - self.rope_embedder( - torch.cat(glm_sequence.position_ids, dim=0) - ).split( - [ - len(position_ids) - for position_ids in glm_sequence.position_ids - ], - dim=0, + sigvq_features, sigvq_mask, sigvq_frequencies ) + unified_features, unified_frequencies, unified_mask, unified_noise = ( + self._merge_padded_sequences( + (cap_features, image_features, sigvq_features), + (cap_frequencies, image_frequencies, sigvq_frequencies), + (cap_lengths, image_lengths, sigvq_lengths), + (cap_noise, image_noise, sigvq_noise), ) - glm_frequencies = pad_sequence( - glm_frequencies, batch_first=True, padding_value=0.0 - ) - condition_feature_groups.append(glm_features) - condition_frequency_groups.append(glm_frequencies) - condition_length_groups.append(glm_lengths) - - ( - condition_features, - condition_frequencies, - condition_attention_mask, - condition_lengths, - _, - ) = self._merge_padded_sequences( - tuple(condition_feature_groups), - tuple(condition_frequency_groups), - tuple(condition_length_groups), ) - - for layer in self.context_refiner: - condition_features = layer( - condition_features, - condition_attention_mask, - condition_frequencies, + else: + unified_features, unified_frequencies, unified_mask, unified_noise = ( + self._merge_padded_sequences( + (image_features, cap_features), + (image_frequencies, cap_frequencies), + (image_lengths, cap_lengths), ) - - ( - unified_features, - unified_frequencies, - unified_attention_mask, - _, - unified_noise_mask, - ) = self._merge_padded_sequences( - (image_features, condition_features), - (image_frequencies, condition_frequencies), - (image_lengths, condition_lengths), ) for layer in self.layers: - if is_editing: - unified_features = layer( - unified_features, - unified_attention_mask, - unified_frequencies, - noise_mask=unified_noise_mask, - adaln_noisy=noisy_embedding, - adaln_clean=clean_embedding, - ) - else: - unified_features = layer( - unified_features, - unified_attention_mask, - unified_frequencies, - adaln_input, - ) - - if is_editing: - unified_features = self.all_final_layer[patch_key]( - unified_features, - noise_mask=unified_noise_mask, - adaln_noisy=noisy_embedding, - adaln_clean=clean_embedding, - ) - else: - unified_features = self.all_final_layer[patch_key]( + unified_features = layer( unified_features, - adaln_input=adaln_input, + unified_mask, + unified_frequencies, + noise_mask=unified_noise, + **adaln, ) - - output = self._unpatchify( + unified_features = self.all_final_layer[patch_key]( + unified_features, noise_mask=unified_noise, **adaln + ) + return self._unpatchify( list(unified_features.unbind(dim=0)), image_sizes, patch_size, f_patch_size, image_offsets, ) - if not return_dict: - return (output,) - return Transformer2DModelOutput(sample=output) - -@dataclass -class LLaDAImageQueryFormerOutput(BaseOutput): - query_embeds: torch.Tensor - -class LLaDAImageQueryAttnProcessor: - _attention_backend = None - _parallel_config = None - - def __call__( - self, - attn: "LLaDAImageQueryAttention", - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - query = F.linear( - hidden_states, - attn.in_proj_weight[: attn.inner_dim], - attn.in_proj_bias[: attn.inner_dim], - ) - key = F.linear( - encoder_hidden_states, - attn.in_proj_weight[attn.inner_dim : 2 * attn.inner_dim], - attn.in_proj_bias[attn.inner_dim : 2 * attn.inner_dim], - ) - value = F.linear( - encoder_hidden_states, - attn.in_proj_weight[2 * attn.inner_dim :], - attn.in_proj_bias[2 * attn.inner_dim :], - ) - - query = query.unflatten(-1, (attn.heads, attn.head_dim)) - key = key.unflatten(-1, (attn.heads, attn.head_dim)) - value = value.unflatten(-1, (attn.heads, attn.head_dim)) - - if attention_mask is not None: - attention_mask = attention_mask[:, None, None, :] - - hidden_states = dispatch_attention_fn( - query, - key, - value, - attn_mask=attention_mask, - dropout_p=attn.dropout if attn.training else 0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) - hidden_states = hidden_states.flatten(2, 3) - return attn.out_proj(hidden_states) - - -class LLaDAImageQueryAttention(nn.Module, AttentionModuleMixin): - _default_processor_cls = LLaDAImageQueryAttnProcessor - _available_processors: ClassVar[list[type]] = [LLaDAImageQueryAttnProcessor] - _supports_qkv_fusion = False - - def __init__(self, hidden_size: int, num_heads: int, dropout: float): +class LLaDAImageQueryAttention(nn.Module): + def __init__(self, hidden_size: int, num_heads: int): super().__init__() - self.inner_dim = hidden_size self.heads = num_heads self.head_dim = hidden_size // num_heads - self.dropout = dropout - self.in_proj_weight = nn.Parameter(torch.zeros(3 * hidden_size, hidden_size)) self.in_proj_bias = nn.Parameter(torch.zeros(3 * hidden_size)) self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True) - self.set_processor(self._default_processor_cls()) - - nn.init.xavier_uniform_(self.in_proj_weight) - nn.init.zeros_(self.in_proj_bias) def forward( self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor, ) -> torch.Tensor: - return self.processor( - self, hidden_states, encoder_hidden_states, attention_mask + weights = self.in_proj_weight.chunk(3) + biases = self.in_proj_bias.chunk(3) + query, key, value = ( + F.linear(inputs, weight, bias).unflatten(-1, (self.heads, self.head_dim)) + for inputs, weight, bias in zip( + (hidden_states, encoder_hidden_states, encoder_hidden_states), + weights, + biases, + ) ) + hidden_states = dispatch_attention_fn( + query, key, value, attn_mask=attention_mask[:, None, None, :] + ) + return self.out_proj(hidden_states.flatten(2, 3)) -@maybe_allow_in_graph class LLaDAImageQueryFormerBlock(nn.Module): def __init__( self, hidden_size: int, num_heads: int, intermediate_size: int, - dropout: float, norm_eps: float, ): super().__init__() self.norm_q = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) self.norm_k = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) - self.cross_attn = LLaDAImageQueryAttention(hidden_size, num_heads, dropout) - self.dropout = nn.Dropout(dropout) + self.cross_attn = LLaDAImageQueryAttention(hidden_size, num_heads) self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) self.mlp = nn.Module() self.mlp.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) @@ -1620,48 +1033,21 @@ def forward( self, query_embeds: torch.Tensor, encoder_hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor, ) -> torch.Tensor: query_embeds = self.norm_q(query_embeds) encoder_hidden_states = self.norm_k(encoder_hidden_states) - attention_output = self.cross_attn( + query_embeds = query_embeds + self.cross_attn( query_embeds, encoder_hidden_states, attention_mask ) - query_embeds = query_embeds + self.dropout(attention_output) query_embeds = self.norm1(query_embeds) - mlp_output = self.mlp.fc2( + return query_embeds + self.mlp.fc2( F.gelu(self.mlp.fc1(query_embeds), approximate="tanh") ) - return query_embeds + self.dropout(mlp_output) - - -class LLaDAImageQueryFormerModel(ModelMixin, ConfigMixin, AttentionMixin): - r""" - QueryFormer used by LLaDA-Image to derive learnable image-generation queries from LLaDA token embeddings. - - This model is independent from the LLaDA text encoder. It returns refined query embeddings. The pipeline appends - them to the text embeddings and invokes the text encoder backbone. - - Args: - num_queries (`int`, defaults to `256`): - Number of learnable query tokens. - hidden_size (`int`, defaults to `2048`): - Query and LLaDA token embedding dimension. - num_hidden_layers (`int`, defaults to `1`): - Number of QueryFormer blocks. - num_attention_heads (`int`, defaults to `16`): - Number of cross-attention heads. - intermediate_size (`int`, defaults to `8192`): - Hidden dimension of the QueryFormer MLP. - dropout (`float`, defaults to `0.0`): - Dropout probability. - norm_eps (`float`, defaults to `1e-6`): - Epsilon used by parameter-free layer normalization. - """ - - _no_split_modules: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] - _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageQueryFormerBlock"] - _skip_layerwise_casting_patterns: ClassVar[list[str]] = ["norm"] + + +class LLaDAImageQueryFormerModel(ModelMixin, ConfigMixin): + """Derives the image-generation query tokens from LLaDA token embeddings.""" @register_to_config def __init__( @@ -1675,154 +1061,73 @@ def __init__( norm_eps: float = 1e-6, ): super().__init__() - if hidden_size % num_attention_heads != 0: - raise ValueError( - f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." - ) - self.meta_queries = nn.Parameter(torch.zeros(num_queries, hidden_size)) - nn.init.normal_(self.meta_queries, std=1 / math.sqrt(hidden_size)) self.query_blocks = nn.ModuleList( - [ - LLaDAImageQueryFormerBlock( - hidden_size, - num_attention_heads, - intermediate_size, - dropout, - norm_eps, - ) - for _ in range(num_hidden_layers) - ] + LLaDAImageQueryFormerBlock( + hidden_size, num_attention_heads, intermediate_size, norm_eps + ) + for _ in range(num_hidden_layers) ) def forward( - self, - inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor, - return_dict: bool = True, - ) -> LLaDAImageQueryFormerOutput | tuple[torch.Tensor]: - r""" - Args: - inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - LLaDA input token embeddings. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`): - Mask whose nonzero entries identify valid text tokens. - return_dict (`bool`, defaults to `True`): - Whether to return [`LLaDAImageQueryFormerOutput`] instead of a tuple. - - Returns: - [`LLaDAImageQueryFormerOutput`] or `tuple`: - The refined query embeddings. - """ - batch_size = inputs_embeds.shape[0] - query_embeds = self.meta_queries.unsqueeze(0).expand(batch_size, -1, -1) - attention_mask = attention_mask.bool() - - for query_block in self.query_blocks: - query_embeds = query_block(query_embeds, inputs_embeds, attention_mask) - - if not return_dict: - return (query_embeds,) - return LLaDAImageQueryFormerOutput(query_embeds=query_embeds) - - -@dataclass -class LLaDAImageTextProjectionOutput(BaseOutput): - hidden_states: torch.Tensor - - -class LLaDAImageTextProjectionAttnProcessor: - _attention_backend = None - _parallel_config = None - - def __call__( - self, - attn: "LLaDAImageTextProjectionAttention", - hidden_states: torch.Tensor, + self, inputs_embeds: torch.Tensor, attention_mask: torch.Tensor ) -> torch.Tensor: - query = attn.q_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) - key = attn.k_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) - value = attn.v_proj(hidden_states).unflatten(-1, (attn.heads, attn.head_dim)) - - query = attn.q_norm(query) - key = attn.k_norm(key) - - hidden_states = dispatch_attention_fn( - query, - key, - value, - attn_mask=None, - dropout_p=attn.dropout if attn.training else 0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=self._parallel_config, + query_embeds = self.meta_queries.unsqueeze(0).expand( + inputs_embeds.shape[0], -1, -1 ) - hidden_states = hidden_states.flatten(2, 3) - return attn.out_proj(hidden_states) - + for query_block in self.query_blocks: + query_embeds = query_block( + query_embeds, inputs_embeds, attention_mask.bool() + ) + return query_embeds -class LLaDAImageTextProjectionAttention(nn.Module, AttentionModuleMixin): - _default_processor_cls = LLaDAImageTextProjectionAttnProcessor - _available_processors: ClassVar[list[type]] = [ - LLaDAImageTextProjectionAttnProcessor - ] - _supports_qkv_fusion = False - def __init__( - self, - hidden_size: int, - num_attention_heads: int, - attention_dropout: float, - norm_eps: float, - ): +class LLaDAImageTextProjectionAttention(nn.Module): + def __init__(self, hidden_size: int, num_attention_heads: int, norm_eps: float): super().__init__() self.heads = num_attention_heads self.head_dim = hidden_size // num_attention_heads - self.dropout = attention_dropout - self.k_proj = nn.Linear(hidden_size, hidden_size, bias=True) self.v_proj = nn.Linear(hidden_size, hidden_size, bias=True) self.q_proj = nn.Linear(hidden_size, hidden_size, bias=True) self.out_proj = nn.Linear(hidden_size, hidden_size, bias=True) self.q_norm = RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=False) self.k_norm = RMSNorm(self.head_dim, eps=norm_eps, elementwise_affine=False) - self.set_processor(self._default_processor_cls()) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.processor(self, hidden_states) + heads = (self.heads, self.head_dim) + query = self.q_norm(self.q_proj(hidden_states).unflatten(-1, heads)) + key = self.k_norm(self.k_proj(hidden_states).unflatten(-1, heads)) + value = self.v_proj(hidden_states).unflatten(-1, heads) + hidden_states = dispatch_attention_fn(query, key, value) + return self.out_proj(hidden_states.flatten(2, 3)) -class LLaDAImageTextProjectionMLP(nn.Module): - def __init__(self, hidden_size: int, intermediate_size: int): +class LLaDAImageMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, approximate: str): super().__init__() + self.approximate = approximate self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=True) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.fc1(hidden_states) - hidden_states = F.gelu(hidden_states, approximate="tanh") - return self.fc2(hidden_states) + return self.fc2(F.gelu(self.fc1(hidden_states), approximate=self.approximate)) -@maybe_allow_in_graph class LLaDAImageTextProjectionBlock(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, num_attention_heads: int, - attention_dropout: float, norm_eps: float, ): super().__init__() self.self_attn = LLaDAImageTextProjectionAttention( - hidden_size, - num_attention_heads, - attention_dropout, - norm_eps, + hidden_size, num_attention_heads, norm_eps ) self.layer_norm1 = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) - self.mlp = LLaDAImageTextProjectionMLP(hidden_size, intermediate_size) + self.mlp = LLaDAImageMLP(hidden_size, intermediate_size, approximate="tanh") self.layer_norm2 = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -1831,34 +1136,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class LLaDAImageTextProjectionModel(ModelMixin, ConfigMixin, AttentionMixin): - r""" - Connector and output projection used to map LLaDA hidden states to the LLaDA-Image denoiser context dimension. - - Args: - hidden_size (`int`, defaults to `2048`): - Input and connector hidden dimension. - intermediate_size (`int`, defaults to `8960`): - Connector MLP hidden dimension. - num_hidden_layers (`int`, defaults to `6`): - Number of connector layers. - num_attention_heads (`int`, defaults to `32`): - Number of connector self-attention heads. - projection_dim (`int`, defaults to `2560`): - Output dimension expected by the denoising transformer. - attention_dropout (`float`, defaults to `0.0`): - Attention dropout probability. - norm_eps (`float`, defaults to `1e-6`): - Epsilon used by parameter-free RMS normalization. - """ - - _no_split_modules: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] - _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageTextProjectionBlock"] - _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ - "layer_norm", - "q_norm", - "k_norm", - ] +class LLaDAImageTextProjectionModel(ModelMixin, ConfigMixin): + """Maps LLaDA hidden states to the denoiser caption dimension.""" @register_to_config def __init__( @@ -1872,117 +1151,37 @@ def __init__( norm_eps: float = 1e-6, ): super().__init__() - if hidden_size % num_attention_heads != 0: - raise ValueError( - f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." - ) - self.layers = nn.ModuleList( - [ - LLaDAImageTextProjectionBlock( - hidden_size, - intermediate_size, - num_attention_heads, - attention_dropout, - norm_eps, - ) - for _ in range(num_hidden_layers) - ] + LLaDAImageTextProjectionBlock( + hidden_size, intermediate_size, num_attention_heads, norm_eps + ) + for _ in range(num_hidden_layers) ) self.projector = nn.Linear(hidden_size, projection_dim, bias=True) - def forward( - self, - hidden_states: torch.Tensor, - return_dict: bool = True, - ) -> LLaDAImageTextProjectionOutput | tuple[torch.Tensor]: - r""" - Args: - hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): - Hidden states produced by the LLaDA text backbone. - return_dict (`bool`, defaults to `True`): - Whether to return [`LLaDAImageTextProjectionOutput`] instead of a tuple. - - Returns: - [`LLaDAImageTextProjectionOutput`] or `tuple`: - Hidden states projected to the denoiser caption dimension. - """ + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: for layer in self.layers: hidden_states = layer(hidden_states) + return self.projector(hidden_states) - hidden_states = self.projector(hidden_states) - if not return_dict: - return (hidden_states,) - return LLaDAImageTextProjectionOutput(hidden_states=hidden_states) - - -@dataclass -class LLaDAImageSigVQOutput(BaseOutput): - semantic_features: torch.Tensor - token_ids: torch.Tensor - - -class LLaDAImageSigVQAttnProcessor: - _attention_backend = None - _parallel_config = None - - def __call__( - self, attn: "LLaDAImageSigVQAttention", hidden_states: torch.Tensor - ) -> torch.Tensor: - query, key, value = attn.qkv(hidden_states).chunk(3, dim=-1) - query = query.unflatten(-1, (attn.heads, attn.head_dim)) - key = key.unflatten(-1, (attn.heads, attn.head_dim)) - value = value.unflatten(-1, (attn.heads, attn.head_dim)) - - hidden_states = dispatch_attention_fn( - query, - key, - value, - attn_mask=None, - dropout_p=attn.dropout if attn.training else 0.0, - is_causal=False, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) - hidden_states = hidden_states.flatten(2, 3) - return attn.proj(hidden_states) - - -class LLaDAImageSigVQAttention(nn.Module, AttentionModuleMixin): - _default_processor_cls = LLaDAImageSigVQAttnProcessor - _available_processors: ClassVar[list[type]] = [LLaDAImageSigVQAttnProcessor] - _supports_qkv_fusion = False - def __init__( - self, - hidden_size: int, - num_attention_heads: int, - attention_bias: bool, - attention_dropout: float, - ): +class LLaDAImageSigVQAttention(nn.Module): + def __init__(self, hidden_size: int, num_attention_heads: int, bias: bool): super().__init__() self.heads = num_attention_heads self.head_dim = hidden_size // num_attention_heads - self.dropout = attention_dropout - self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=attention_bias) - self.proj = nn.Linear(hidden_size, hidden_size, bias=attention_bias) - self.set_processor(self._default_processor_cls()) + self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=bias) + self.proj = nn.Linear(hidden_size, hidden_size, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.processor(self, hidden_states) - - -class LLaDAImageSigVQMLP(nn.Module): - def __init__(self, hidden_size: int, intermediate_size: int): - super().__init__() - self.fc1 = nn.Linear(hidden_size, intermediate_size, bias=True) - self.fc2 = nn.Linear(intermediate_size, hidden_size, bias=True) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.fc2(F.gelu(self.fc1(hidden_states))) + query, key, value = ( + tensor.unflatten(-1, (self.heads, self.head_dim)) + for tensor in self.qkv(hidden_states).chunk(3, dim=-1) + ) + hidden_states = dispatch_attention_fn(query, key, value) + return self.proj(hidden_states.flatten(2, 3)) -@maybe_allow_in_graph class LLaDAImageSigVQVisionBlock(nn.Module): def __init__( self, @@ -1990,19 +1189,15 @@ def __init__( intermediate_size: int, num_attention_heads: int, attention_bias: bool, - attention_dropout: float, norm_eps: float, ): super().__init__() self.norm1 = nn.LayerNorm(hidden_size, eps=norm_eps) self.norm2 = nn.LayerNorm(hidden_size, eps=norm_eps) self.attn = LLaDAImageSigVQAttention( - hidden_size, - num_attention_heads, - attention_bias, - attention_dropout, + hidden_size, num_attention_heads, attention_bias ) - self.mlp = LLaDAImageSigVQMLP(hidden_size, intermediate_size) + self.mlp = LLaDAImageMLP(hidden_size, intermediate_size, approximate="none") def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) @@ -2013,7 +1208,6 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class LLaDAImageSigVQPatchEmbed(nn.Module): def __init__(self, in_channels: int, hidden_size: int, patch_size: int): super().__init__() - self.in_channels = in_channels self.patch_size = patch_size self.proj = nn.Conv2d( in_channels, hidden_size, kernel_size=patch_size, stride=patch_size @@ -2109,23 +1303,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return torch.argmin(distances, dim=1) -class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin): - r""" - Minimal GLM SigVQ image encoder used by LLaDA-Image editing. - - The model contains only the GLM vision encoder, VQ quantizer, and prior token projection used during inference. - Input images must already be RGB tensors normalized to `[-1, 1]`, have one common size, and be divisible by - `patch_size`. - """ - - _no_split_modules: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] - _repeated_blocks: ClassVar[list[str]] = ["LLaDAImageSigVQVisionBlock"] - _skip_layerwise_casting_patterns: ClassVar[list[str]] = [ - "patch_embed", - "position_embedding", - "norm", - "quantize", - ] +class LLaDAImageSigVQModel(ModelMixin, ConfigMixin): + """GLM SigVQ image encoder, quantizer, and prior projection for editing.""" @register_to_config def __init__( @@ -2145,11 +1324,6 @@ def __init__( semantic_embed_dim: int = 4096, ): super().__init__() - if hidden_size % num_attention_heads != 0: - raise ValueError( - f"`hidden_size` ({hidden_size}) must be divisible by `num_attention_heads` ({num_attention_heads})." - ) - self.visual = nn.Module() self.visual.patch_embed = LLaDAImageSigVQPatchEmbed( in_channels, hidden_size, patch_size @@ -2158,17 +1332,14 @@ def __init__( image_size, patch_size, hidden_size ) self.visual.blocks = nn.ModuleList( - [ - LLaDAImageSigVQVisionBlock( - hidden_size, - intermediate_size, - num_attention_heads, - attention_bias, - attention_dropout, - norm_eps, - ) - for _ in range(num_hidden_layers) - ] + LLaDAImageSigVQVisionBlock( + hidden_size, + intermediate_size, + num_attention_heads, + attention_bias, + norm_eps, + ) + for _ in range(num_hidden_layers) ) self.vqmodel = nn.Module() @@ -2187,74 +1358,22 @@ def __init__( activation_fn="linear-silu", ) - def forward( - self, - pixel_values: torch.Tensor | None = None, - token_ids: torch.Tensor | None = None, - return_dict: bool = True, - ) -> LLaDAImageSigVQOutput | tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - pixel_values (`torch.Tensor` of shape `(batch_size, 3, height, width)`, *optional*): - RGB images normalized to `[-1, 1]`. Mutually exclusive with `token_ids`. - token_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Precomputed VQ codebook IDs. Mutually exclusive with `pixel_values`. - return_dict (`bool`, defaults to `True`): - Whether to return [`LLaDAImageSigVQOutput`] instead of a tuple. - - Returns: - [`LLaDAImageSigVQOutput`] or `tuple`: - The projected semantic features and their discrete token IDs. - """ - if (pixel_values is None) == (token_ids is None): - raise ValueError("Provide exactly one of `pixel_values` or `token_ids`.") - - if pixel_values is not None: - if pixel_values.ndim != 4: - raise ValueError( - f"`pixel_values` must have 4 dimensions, got shape {tuple(pixel_values.shape)}." - ) - height, width = pixel_values.shape[-2:] - if ( - height % self.config.patch_size != 0 - or width % self.config.patch_size != 0 - ): - raise ValueError( - f"Image height and width must be divisible by {self.config.patch_size}, got {height}x{width}." - ) - - grid_height = height // self.config.patch_size - grid_width = width // self.config.patch_size - hidden_states = self.visual.patch_embed(pixel_values) - hidden_states = self.visual.embeddings( - hidden_states, grid_height, grid_width - ) - - for block in self.visual.blocks: - hidden_states = block(hidden_states) - - hidden_states = hidden_states.transpose(1, 2).reshape( - pixel_values.shape[0], - self.config.hidden_size, - grid_height, - grid_width, - ) - hidden_states = self.vqmodel.quant_conv(hidden_states) - token_ids = self.vqmodel.quantize(hidden_states).reshape( - pixel_values.shape[0], -1 - ) - elif token_ids.ndim != 2: - raise ValueError( - f"`token_ids` must have 2 dimensions, got shape {tuple(token_ids.shape)}." - ) - - semantic_features = self.prior_projector(self.prior_token_embedding(token_ids)) + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Map RGB images in [-1, 1] to projected semantic features.""" + batch_size, _, height, width = pixel_values.shape + grid_height = height // self.config.patch_size + grid_width = width // self.config.patch_size + hidden_states = self.visual.patch_embed(pixel_values) + hidden_states = self.visual.embeddings(hidden_states, grid_height, grid_width) + for block in self.visual.blocks: + hidden_states = block(hidden_states) - if not return_dict: - return semantic_features, token_ids - return LLaDAImageSigVQOutput( - semantic_features=semantic_features, token_ids=token_ids + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, self.config.hidden_size, grid_height, grid_width ) + hidden_states = self.vqmodel.quant_conv(hidden_states) + token_ids = self.vqmodel.quantize(hidden_states).reshape(batch_size, -1) + return self.prior_projector(self.prior_token_embedding(token_ids)) class LLaDAImageTransformer2DModel(_LLaDAImageTransformer2DModel): @@ -2288,64 +1407,24 @@ def post_load_weights(self) -> None: """Run model-specific post-load fixups (none are required).""" return - @staticmethod - def _as_feature_list( - features: torch.Tensor | list[torch.Tensor], - attention_mask: torch.Tensor | list[torch.Tensor] | None = None, - ) -> list[torch.Tensor]: - if isinstance(features, list): - return features - if features.ndim == 2: - return [features] - if features.ndim != 3: - raise ValueError( - "LLaDA-Image condition features must have two or three dimensions" - ) - if attention_mask is None: - return list(features.unbind(dim=0)) - if isinstance(attention_mask, list): - attention_mask = torch.stack(attention_mask) - return [sample[mask.bool()] for sample, mask in zip(features, attention_mask)] - def forward( self, hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor | list[torch.Tensor], + encoder_hidden_states: list[torch.Tensor], timestep: torch.Tensor, - encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, - encoder_attention_mask: torch.Tensor | list[torch.Tensor] | None = None, + encoder_hidden_states_image: list[torch.Tensor] | None = None, source_latents: list[torch.Tensor] | None = None, - guidance=None, **kwargs, ) -> torch.Tensor: - del guidance, kwargs - if hidden_states.ndim != 4: - raise ValueError( - "LLaDA-Image latents must have shape [batch, channels, height, width]" - ) - + del kwargs model_dtype = next(self.parameters()).dtype - latent_list = [latent.unsqueeze(1).to(model_dtype) for latent in hidden_states] - cap_feats = self._as_feature_list(encoder_hidden_states, encoder_attention_mask) - glm_cap_feats = ( - None - if encoder_hidden_states_image is None - or ( - isinstance(encoder_hidden_states_image, list) - and not encoder_hidden_states_image - ) - else self._as_feature_list(encoder_hidden_states_image) - ) - output = ( - super() - .forward( - x=latent_list, - t=(timestep / 1000.0).to(model_dtype), - cap_feats=cap_feats, - glm_cap_feats=glm_cap_feats, - source_latents=source_latents, - ) - .sample + output = super().forward( + x=[latent.unsqueeze(1).to(model_dtype) for latent in hidden_states], + t=(timestep / 1000.0).to(model_dtype), + cap_feats=encoder_hidden_states, + # Text-to-image passes no SigVQ features. + glm_cap_feats=encoder_hidden_states_image or None, + source_latents=source_latents, ) return -torch.stack(output, dim=0).squeeze(2).float() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index 0f244f08a741..1d3b3609444b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -3,9 +3,7 @@ from __future__ import annotations import json -import os from array import array -from contextlib import nullcontext from pathlib import Path import torch @@ -29,23 +27,12 @@ def _resolve_llada_image_component_path( server_args: ServerArgs, component_name: str, ) -> str: - override_path = (getattr(server_args, "component_paths", {}) or {}).get( - component_name - ) + override_path = server_args.component_paths.get(component_name) if override_path is not None: return prepare_diffusers_component_path_for_loading(override_path) return str(Path(model_root) / component_name) -def ensure_conditioning_mask_active(attn_backend) -> None: - """Fail closed if the conditioning block mask did not reach the backend.""" - if not attn_backend.conditioning_mask_active: - raise RuntimeError( - "LLaDA-Image conditioning ran without its block attention mask " - "because the forward batch lost the conditioning text lengths" - ) - - def format_llada_image_prompt(prompt: str | None) -> str: if prompt is None: return "HUMAN Generate an image.\nASSISTANT\n" @@ -99,8 +86,6 @@ def __init__( ) device = get_local_torch_device() gpu_id = device.index - if gpu_id is None: - gpu_id = int(os.environ.get("LOCAL_RANK", "0")) nccl_port = server_args.nccl_port or 29500 srt_args = SRTServerArgs( model_path=text_encoder_path, @@ -117,9 +102,6 @@ def __init__( max_prefill_tokens=8192, max_total_tokens=8192, max_running_requests=2, - mem_fraction_static=( - server_args.pipeline_config.text_encoder_mem_fraction_static - ), ) self.runtime_context = create_context( srt_args, @@ -175,20 +157,14 @@ def _prepare_input_embeds( input_ids: torch.Tensor, sequence_lengths: list[int], text_lengths: list[int], - component_context=None, + component_context, ) -> torch.Tensor: - model = self.model_runner.model - embeddings = model.get_input_embeddings() + embeddings = self.model_runner.model.get_input_embeddings() inputs = [] offset = 0 - context = ( - nullcontext(self.queryformer) - if component_context is None - else component_context( - component_name="queryformer", module=self.queryformer - ) - ) - with context as queryformer: + with component_context( + component_name="queryformer", module=self.queryformer + ) as queryformer: assert queryformer is not None queryformer_dtype = queryformer.dtype for sequence_length, text_length in zip( @@ -202,13 +178,7 @@ def _prepare_input_embeds( query_embeds = queryformer( text_embeds.unsqueeze(0).to(dtype=queryformer_dtype), text_mask, - ).query_embeds - query_count = sequence_length - text_length - if query_embeds.shape[1] != query_count: - raise RuntimeError( - f"QueryFormer returned {query_embeds.shape[1]} queries, " - f"expected {query_count}" - ) + ) inputs.append( torch.cat( [text_embeds, query_embeds.squeeze(0).to(text_embeds.dtype)], @@ -216,16 +186,10 @@ def _prepare_input_embeds( ) ) offset += sequence_length - if offset != input_ids.numel(): - raise ValueError( - f"Conditioning spans {offset} tokens, got {input_ids.numel()}" - ) return torch.cat(inputs, dim=0) @torch.no_grad() - def encode( - self, prompts: list[str], max_sequence_length: int, component_context=None - ): + def encode(self, prompts: list[str], max_sequence_length: int, component_context): import sglang.multimodal_gen.runtime.distributed.parallel_state as mm_parallel_state import sglang.srt.distributed.parallel_state as srt_parallel_state from sglang.srt.runtime_context import use_context @@ -245,7 +209,7 @@ def encode( srt_parallel_state._ATTN_TP = saved_attn_tp def _encode_impl( - self, prompts: list[str], max_sequence_length: int, component_context=None + self, prompts: list[str], max_sequence_length: int, component_context ): from sglang.srt.managers.overlap_utils import resolve_forward_inputs from sglang.srt.managers.schedule_batch import Req as SRTReq @@ -349,23 +313,18 @@ def _encode_impl( ) forward_batch.llada_image_conditioning_text_lens_cpu = text_lengths model_output = self.model_runner.forward(forward_batch=forward_batch) - ensure_conditioning_mask_active(self.model_runner.attn_backend) - hidden_states = model_output.logits_output.hidden_states - if not isinstance(hidden_states, torch.Tensor): - raise TypeError( - "SGLang did not return full LLaDA-Image conditioning states" + if not self.model_runner.attn_backend.conditioning_mask_active: + raise RuntimeError( + "LLaDA-Image conditioning ran without its block attention mask " + "because the forward batch lost the conditioning text lengths" ) + hidden_states = model_output.logits_output.hidden_states projected = [] offset = 0 - context = ( - nullcontext(self.text_projection) - if component_context is None - else component_context( - component_name="text_projection", module=self.text_projection - ) - ) - with context as text_projection: + with component_context( + component_name="text_projection", module=self.text_projection + ) as text_projection: assert text_projection is not None projection_dtype = text_projection.dtype for sequence_length in sequence_lengths: @@ -373,17 +332,17 @@ def _encode_impl( hidden_states[offset : offset + sequence_length] .unsqueeze(0) .to(dtype=projection_dtype) - ).hidden_states + ) projected.append(output.squeeze(0)) offset += sequence_length return projected finally: for req in reqs: - if req.kv is not None: + if not req.kv.is_kv_released: release_kv_cache(req, self.tree_cache, checkpoint=False) - elif req.req_pool_idx is not None: - # KV alloc failed after the slot grab. Freeing the bare - # slot avoids the release_kv_cache lifecycle assert. + elif req.kv.holds_kv: + # KV alloc failed after the row grab. Freeing the bare row + # avoids the release_kv_cache lifecycle assert. self.req_to_token_pool.free(req) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py index 138c854ed029..f13ddacb67f2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py @@ -5,6 +5,10 @@ import torch import torch.nn.functional as F +from sglang.multimodal_gen.configs.pipeline_configs.flux import _patchify_latents +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + flux2_vae_bn_stats, +) from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ComponentUse, ) @@ -35,25 +39,6 @@ def component_uses( ComponentUse(stage_name, "vae"), ] - @staticmethod - def _patchify_latents(latents: torch.Tensor) -> torch.Tensor: - batch_size, channels, height, width = latents.shape - if height % 2 != 0 or width % 2 != 0: - raise ValueError("LLaDA-Image source VAE latents must have even dimensions") - latents = latents.reshape(batch_size, channels, height // 2, 2, width // 2, 2) - latents = latents.permute(0, 1, 3, 5, 2, 4) - return latents.reshape(batch_size, channels * 4, height // 2, width // 2) - - @staticmethod - def _normalize_latents(latents: torch.Tensor, vae) -> torch.Tensor: - if not hasattr(vae, "bn"): - raise ValueError("LLaDA-Image editing requires the Flux2 VAE BN state") - vae_config = getattr(vae.config, "arch_config", vae.config) - eps = vae_config.batch_norm_eps - latent_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(latents) - latent_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + eps).to(latents) - return (latents - latent_mean) / latent_std - @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: del server_args @@ -73,15 +58,9 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: source_images[0], height=batch.height, width=batch.width, - resize_mode="crop", ) - if image.shape[0] == 1 and batch.batch_size > 1: + if batch.batch_size > 1: image = image.repeat(batch.batch_size, 1, 1, 1) - if image.shape[0] != batch.batch_size: - raise ValueError( - "LLaDA-Image source image batch must contain one image or match " - f"the effective batch size {batch.batch_size}, got {image.shape[0]}" - ) sigvq_pixels = F.interpolate( image.float(), @@ -93,27 +72,18 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: component_name="sigvq", module=self.sigvq ) as sigvq: assert sigvq is not None - patch_size = sigvq.config.patch_size - pad_height = (-sigvq_pixels.shape[-2]) % patch_size - pad_width = (-sigvq_pixels.shape[-1]) % patch_size - if pad_height or pad_width: - sigvq_pixels = F.pad( - sigvq_pixels, - (0, pad_width, 0, pad_height), - mode="replicate", - ) sigvq_device, sigvq_dtype = _module_device_dtype(sigvq) semantic_features = sigvq( sigvq_pixels.to(device=sigvq_device, dtype=sigvq_dtype) - ).semantic_features + ) with self.use_declared_component(component_name="vae", module=self.vae) as vae: assert vae is not None vae_device, vae_dtype = _module_device_dtype(vae) posterior = vae.encode(image.to(device=vae_device, dtype=vae_dtype)) - latent_dist = getattr(posterior, "latent_dist", posterior) - source_latents = self._patchify_latents(latent_dist.mode()) - source_latents = self._normalize_latents(source_latents, vae) + source_latents = _patchify_latents(posterior.mode()) + mean, std = flux2_vae_bn_stats(vae, source_latents) + source_latents = (source_latents - mean) / std batch.image_embeds = list(semantic_features.unbind(dim=0)) batch.source_latents = [ diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index de9d1a22db98..cd48f0666c8d 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -10,6 +10,10 @@ from safetensors import safe_open from torch import nn +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) from sglang.multimodal_gen.runtime.layers.quantization import ( QuantizationConfig, get_quantization_config, @@ -60,11 +64,6 @@ def process_model_weights_after_loading( quantized_only: bool = False, ) -> int: """Process native and SRT layers once, optionally staging one layer at a time.""" - from sglang.multimodal_gen.runtime.layers.linear import ( - LinearBase, - UnquantizedLinearMethod, - ) - processed_layers = 0 for module in model.modules(): if not isinstance( diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 070c01e54dc5..908485cf091a 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -543,28 +543,30 @@ ), ] +_LLADA_IMAGE_SERVER_ARGS = DiffusionServerArgs( + model_path="inclusionAI/LLaDA-Image", + modality="image", + num_gpus=1, + tp_size=1, + ulysses_degree=1, + ring_degree=1, + cfg_parallel=False, + extras=[ + "--trust-remote-code", + "--revision", + "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", + ], +) +_LLADA_IMAGE_SAMPLING = {"num_inference_steps": 50, "guidance_scale": 5.0, "seed": 42} if not current_platform.is_hip() and not current_platform.is_xpu(): ONE_GPU_CASES += [ DiffusionTestCase( "llada_image_bf16_t2i", - DiffusionServerArgs( - model_path="inclusionAI/LLaDA-Image", - modality="image", - num_gpus=1, - tp_size=1, - ulysses_degree=1, - ring_degree=1, - cfg_parallel=False, - extras=[ - "--trust-remote-code", - "--revision", - "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", - ], - ), + _LLADA_IMAGE_SERVER_ARGS, DiffusionSamplingParams( prompt="A red fox sitting on a snowy hill", output_size="512x512", - extras={"num_inference_steps": 50, "guidance_scale": 5.0, "seed": 42}, + extras=_LLADA_IMAGE_SAMPLING, ), run_perf_check=False, run_consistency_check=False, @@ -572,25 +574,12 @@ ), DiffusionTestCase( "llada_image_bf16_edit", - DiffusionServerArgs( - model_path="inclusionAI/LLaDA-Image", - modality="image", - num_gpus=1, - tp_size=1, - ulysses_degree=1, - ring_degree=1, - cfg_parallel=False, - extras=[ - "--trust-remote-code", - "--revision", - "e4e2703f410f7ddb6ee8d6b09dac6a8ec5093039", - ], - ), + _LLADA_IMAGE_SERVER_ARGS, DiffusionSamplingParams( prompt="Convert 2D style to 3D style", image_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg", output_size="512x512", - extras={"num_inference_steps": 50, "guidance_scale": 5.0, "seed": 42}, + extras=_LLADA_IMAGE_SAMPLING, ), run_perf_check=False, run_consistency_check=False, diff --git a/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py b/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py index fabbe13536b4..33171649111b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py @@ -7,6 +7,9 @@ LLaDAImageArchConfig, editing_rope_rows, ) +from sglang.multimodal_gen.configs.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) from sglang.multimodal_gen.configs.sample.llada_image import ( LLaDAImageSamplingParams, LLaDAImageTurboSamplingParams, @@ -51,6 +54,16 @@ def rows(side: int) -> int: self.assertLessEqual(rows(1904), limit) self.assertGreater(rows(1920), limit) + def test_unservable_editing_sizes_are_rejected_at_validation(self): + """Edits off the SigVQ grid or past the RoPE table fail before any GPU work.""" + config = LLaDAImagePipelineConfig() + config.validate_output_size(1008, 1008, editing=False) + config.validate_output_size(1024, 1024, editing=True) + with self.assertRaisesRegex(ValueError, "divisible by 32"): + config.validate_output_size(1008, 1008, editing=True) + with self.assertRaisesRegex(ValueError, "sequence positions"): + config.validate_output_size(2048, 2048, editing=True) + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py b/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py index a9691342850e..fdadee44a41f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py +++ b/python/sglang/multimodal_gen/test/unit/test_load_inclusive_e2e.py @@ -149,6 +149,7 @@ def test_server_load_clock_excludes_warmup( clock = [1_000_000_000] output = io.StringIO() monkeypatch.setattr(utils.time, "monotonic_ns", lambda: clock[0]) + monkeypatch.setattr(utils.current_platform, "is_hip", lambda: False) monkeypatch.setattr(utils.tempfile, "gettempdir", lambda: str(tmp_path)) monkeypatch.setattr( utils, "prepare_perf_log", lambda: (tmp_path, tmp_path / "perf.jsonl") diff --git a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py b/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py index 865bc1160255..1c28850c2aed 100644 --- a/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_openai_image_api.py @@ -2,6 +2,7 @@ import os from dataclasses import fields +import pytest from fastapi import HTTPException from PIL import Image @@ -15,6 +16,7 @@ LongCatImageSamplingParams, ) from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.entrypoints.openai import image_api from sglang.multimodal_gen.runtime.entrypoints.openai.image_api import ( _build_image_response_kwargs, _fallback_image_urls, @@ -277,3 +279,16 @@ def test_select_image_variant_path_falls_back_to_single_file_path(): assert _select_image_variant_path(item, None) == "only.png" assert _select_image_variant_path(item, "0") == "only.png" + + +def test_request_validation_errors_are_client_errors(monkeypatch): + """A request that fails sampling-parameter validation returns 400, not 500.""" + + def reject(request_id, **kwargs): + raise ValueError("unsupported size") + + monkeypatch.setattr(image_api, "build_sampling_params", reject) + with pytest.raises(HTTPException) as error: + image_api._build_request_sampling_params("request", size="7x7") + assert error.value.status_code == 400 + assert "unsupported size" in error.value.detail diff --git a/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py index f7e676c42f1b..dddba712b1f0 100644 --- a/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py @@ -53,9 +53,7 @@ def conditioning_mask_active(self) -> bool: def init_forward_metadata(self, forward_batch): self._llada_image_conditioning_mask_active = False - text_lens = getattr( - forward_batch, "llada_image_conditioning_text_lens_cpu", None - ) + text_lens = forward_batch.llada_image_conditioning_text_lens_cpu if text_lens is None: return super().init_forward_metadata(forward_batch) diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 7aa4e2c0bb65..7974a02814fe 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -629,8 +629,6 @@ class Resources(_FlagGroupBase): # Accessors with bespoke semantics (grow-only, per-device keys) manage # their entries directly. buffers: dict = msgspec.field(default_factory=dict) - # The weight data parallel manager belongs to the model using it. - dwdp_manager: Any = None # Persistent reusable CUDA events for non-EP DP TBO, keyed by # (kind, subbatch) — see dp_attention._tbo_event for why reuse matters. tbo_event_pool: dict = msgspec.field(default_factory=dict) @@ -1640,12 +1638,16 @@ def get_buffer(name: str, factory: Any) -> Any: return _CONTEXT.get_buffer(name, factory) +_GLOBAL_DWDP_MANAGER: Any = None + + def get_global_dwdp_manager() -> Any: - return _CONTEXT.resources.dwdp_manager + return _GLOBAL_DWDP_MANAGER def set_global_dwdp_manager(manager: Any) -> None: - _CONTEXT.resources.dwdp_manager = manager + global _GLOBAL_DWDP_MANAGER + _GLOBAL_DWDP_MANAGER = manager def _group_leaves(group: _FlagGroupBase) -> dict[str, Any]: diff --git a/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py b/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py index f8c1054ea445..e214bea9a7f1 100644 --- a/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py +++ b/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py @@ -11,24 +11,29 @@ class TestLLaDA2CFGFlashInferMask(unittest.TestCase): - def test_builds_official_text_query_block_mask(self): - flattened = _build_llada_image_custom_mask([2], [4], "cpu") - actual = flattened.view(4, 4) - expected = torch.tensor( - [ - [True, True, False, False], - [True, True, False, False], - [True, True, True, True], - [True, True, True, True], - ] - ) - - torch.testing.assert_close(actual, expected) - - def test_builds_ragged_batch_without_cross_request_entries(self): + def test_builds_per_request_text_query_block_masks(self): + """Text tokens must not attend to query tokens, and requests stay separate.""" flattened = _build_llada_image_custom_mask([1, 2], [3, 4], "cpu") - self.assertEqual(flattened.numel(), 3 * 3 + 4 * 4) + first = flattened[:9].view(3, 3) + second = flattened[9:].view(4, 4) + torch.testing.assert_close( + first, + torch.tensor( + [[True, False, False], [True, True, True], [True, True, True]] + ), + ) + torch.testing.assert_close( + second, + torch.tensor( + [ + [True, True, False, False], + [True, True, False, False], + [True, True, True, True], + [True, True, True, True], + ] + ), + ) def test_rejects_invalid_conditioning_spans(self): with self.assertRaisesRegex(RuntimeError, "metadata batch mismatch"): diff --git a/test/registered/unit/models/test_llada2_language_model.py b/test/registered/unit/models/test_llada2_language_model.py new file mode 100644 index 000000000000..d190876f631d --- /dev/null +++ b/test/registered/unit/models/test_llada2_language_model.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn.functional as F + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.models.llada2 import ( # noqa: E402 + LLaDA2MoeGate, + _prepare_llada2_language_weights, +) + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestLLaDA2LanguageModel(CustomTestCase): + def test_fp32_router_logits_are_not_rounded_to_the_activation_dtype(self): + """An fp32 router must score bf16 activations in fp32, as the reference does.""" + config = SimpleNamespace( + num_experts=4, hidden_size=8, moe_router_enable_expert_bias=True + ) + gate = LLaDA2MoeGate(config, params_dtype=torch.float32) + with torch.no_grad(): + gate.weight.copy_(torch.arange(32).reshape(4, 8) / 37) + hidden = (torch.arange(24).reshape(3, 8) / 19).to(torch.bfloat16) + + logits = gate(hidden) + + self.assertEqual(logits.dtype, torch.float32) + expected = F.linear(hidden.float(), gate.weight) + torch.testing.assert_close(logits, expected, rtol=0, atol=0) + + def test_language_checkpoint_layout_is_normalized(self): + """Prefixed names and fused expert tensors load as the native layout.""" + fused = torch.arange(24).reshape(2, 3, 4) + lm_head = torch.randn(4, 3) + + expanded = list( + _prepare_llada2_language_weights( + [ + ("model.language_model.layers.1.mlp.experts.gate_proj", fused), + ("model.lm_head.weight", lm_head), + ], + num_experts=2, + ) + ) + + self.assertEqual( + [name for name, _ in expanded], + [ + "model.layers.1.mlp.experts.0.gate_proj.weight", + "model.layers.1.mlp.experts.1.gate_proj.weight", + "lm_head.weight", + ], + ) + torch.testing.assert_close(expanded[0][1], fused[0]) + torch.testing.assert_close(expanded[1][1], fused[1]) + torch.testing.assert_close(expanded[2][1], lm_head) + with self.assertRaisesRegex(ValueError, "expected first dimension 2"): + list( + _prepare_llada2_language_weights( + [("model.layers.1.mlp.experts.down_proj", torch.empty(3, 4, 5))], + num_experts=2, + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_llada2_language_weights.py b/test/registered/unit/models/test_llada2_language_weights.py deleted file mode 100644 index 7d380d3b7b8f..000000000000 --- a/test/registered/unit/models/test_llada2_language_weights.py +++ /dev/null @@ -1,55 +0,0 @@ -import unittest - -import torch - -from sglang.srt.models.llada2 import _prepare_llada2_language_weights -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=2, suite="base-a-test-cpu") - - -class TestLLaDA2LanguageWeights(CustomTestCase): - def test_prepares_clean_language_checkpoint_layout(self): - fused = torch.arange(24).reshape(2, 3, 4) - lm_head = torch.randn(4, 3) - - expanded = list( - _prepare_llada2_language_weights( - [ - ("model.language_model.layers.1.mlp.experts.gate_proj", fused), - ("model.lm_head.weight", lm_head), - ], - num_experts=2, - ) - ) - - self.assertEqual( - [name for name, _ in expanded], - [ - "model.layers.1.mlp.experts.0.gate_proj.weight", - "model.layers.1.mlp.experts.1.gate_proj.weight", - "lm_head.weight", - ], - ) - torch.testing.assert_close(expanded[0][1], fused[0]) - torch.testing.assert_close(expanded[1][1], fused[1]) - torch.testing.assert_close(expanded[2][1], lm_head) - - def test_rejects_invalid_fused_expert_shape(self): - with self.assertRaisesRegex(ValueError, "expected first dimension 2"): - list( - _prepare_llada2_language_weights( - [ - ( - "model.language_model.layers.1.mlp.experts.down_proj", - torch.empty(3, 4, 5), - ) - ], - num_experts=2, - ) - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/models/test_llada2_router.py b/test/registered/unit/models/test_llada2_router.py deleted file mode 100644 index 6a0d3ff4fff0..000000000000 --- a/test/registered/unit/models/test_llada2_router.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import unittest -from types import SimpleNamespace - -import torch -import torch.nn.functional as F - -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel - -maybe_stub_sgl_kernel() - -from sglang.srt.layers.moe.topk import biased_grouped_topk_impl # noqa: E402 -from sglang.srt.models.llada2 import LLaDA2MoeGate # noqa: E402 - -register_cpu_ci(est_time=3, suite="base-a-test-cpu") - - -class TestLLaDA2Router(CustomTestCase): - def test_gate_preserves_router_dtype(self): - config = SimpleNamespace( - num_experts=4, hidden_size=8, moe_router_enable_expert_bias=True - ) - for router_dtype in (torch.float32, torch.bfloat16): - for hidden_dtype in (torch.float32, torch.float16, torch.bfloat16): - with self.subTest(router_dtype=router_dtype, hidden_dtype=hidden_dtype): - gate = LLaDA2MoeGate(config, params_dtype=router_dtype) - with torch.no_grad(): - gate.weight.copy_(torch.arange(32).reshape(4, 8) / 37) - hidden = (torch.arange(48).reshape(3, 16) / 19).to(hidden_dtype)[ - :, ::2 - ] - expected = F.linear(hidden.to(router_dtype), gate.weight) - logits = gate(hidden) - self.assertEqual(logits.dtype, router_dtype) - torch.testing.assert_close(logits, expected, rtol=0, atol=0) - - def test_fp32_gate_keeps_near_tie_expert_selection(self): - config = SimpleNamespace( - num_experts=256, hidden_size=4, moe_router_enable_expert_bias=True - ) - gate = LLaDA2MoeGate(config, params_dtype=torch.float32) - with torch.no_grad(): - gate.weight.zero_() - gate.weight[:, 0].copy_(torch.linspace(0.999, 1.001, 256)) - gate.expert_bias.zero_() - hidden = torch.ones((3, 4), dtype=torch.bfloat16) - logits = gate(hidden) - weights, indices = biased_grouped_topk_impl( - hidden, - logits.float(), - gate.expert_bias, - topk=8, - renormalize=True, - num_expert_group=8, - topk_group=4, - routed_scaling_factor=2.5, - apply_routed_scaling_factor_on_output=True, - ) - expected_indices = torch.arange(248, 256, dtype=torch.int32).expand(3, -1) - torch.testing.assert_close( - indices.sort(dim=-1).values, expected_indices, rtol=0, atol=0 - ) - expected_scores = F.linear(hidden.float(), gate.weight).sigmoid() - selected = expected_scores.gather(1, indices.long()) - expected_weights = selected / (selected.sum(-1, keepdim=True) + 1e-20) * 2.5 - torch.testing.assert_close(weights, expected_weights, rtol=0, atol=0) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 9bbecaaadb95..ffdd5a9d09a1 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -809,97 +809,37 @@ def test_reset_context_clears_owned_store(self): class TestEmbeddedRuntimeContext(unittest.TestCase): - def setUp(self): + def test_embedded_context_is_scoped_and_keeps_its_overrides(self): + """An embedded engine context must neither replace nor leak into the process one.""" from sglang.srt import runtime_context as rc self.addCleanup(rc.restore_context, rc.snapshot_context()) rc.reset_context() - - def test_serial_contexts_keep_config_and_resources_across_scopes(self): - from sglang.srt import runtime_context as rc - outer_args = ServerArgs(model_path="dummy", tp_size=1) outer = publish( outer_args, role="diffusion_gpu_worker", ranks=SpawnRanks(world_rank=0) ) - outer.override("diffusion.setup", grammar_backend="none") - outer_bags = outer._config_bags - outer_log = outer.overrides_log() - outer_buffer = rc.get_buffer("shared_name", object) - outer_manager = object() - rc.set_global_dwdp_manager(outer_manager) - outer.flags.capture.enable_torch_compile = True - outer.parallel.override_permanently(attn_tp_size=1) - - inner_args = ServerArgs(model_path="dummy", tp_size=2) inner = rc.create_context( - inner_args, role="scheduler", ranks=SpawnRanks(world_rank=1, dp_rank=2) + ServerArgs(model_path="dummy", tp_size=2), + role="scheduler", + ranks=SpawnRanks(world_rank=1), ) self.assertIs(get_context(), outer) - self.assertIs(get_parallel(), outer.parallel) - self.assertEqual(outer.parallel.tp_rank, 0) - self.assertIsNone(outer.parallel.dp_rank) - inner_manager = object() - inner_buffer = object() - for index in range(2): - with self.subTest(index=index), rc.use_context(inner): - self.assertIs(assert_published(inner_args, role="scheduler"), inner) + + with self.assertRaisesRegex(RuntimeError, "inner failure"): + with rc.use_context(inner): self.assertIs(get_parallel(), inner.parallel) - self.assertEqual(get_parallel().tp_size, 2) self.assertEqual(get_parallel().tp_rank, 1) - self.assertEqual(get_parallel().attn_tp_rank, 1) - self.assertEqual(get_parallel().dp_rank, 2) - self.assertFalse(get_flags().capture.enable_torch_compile) - if index == 0: - self.assertIsNone(rc.get_global_dwdp_manager()) - self.assertNotIn("shared_name", inner.resources.buffers) - inner.override("encoder.setup", page_size=16) - inner.resources.buffers["shared_name"] = inner_buffer - inner.parallel.override_permanently(attn_tp_size=2) - rc.set_global_dwdp_manager(inner_manager) - self.assertEqual(rc.get_schedule().page_size, 16) - self.assertEqual(len(inner.overrides_log()), 1) - self.assertIs(rc.get_buffer("shared_name", object), inner_buffer) - self.assertIs(rc.get_global_dwdp_manager(), inner_manager) - self.assertEqual(get_parallel().attn_tp_size, 2) - with self.assertRaisesRegex(RuntimeError, "nested failure"): - with rc.use_context(outer): - self.assertIs(rc.get_global_dwdp_manager(), outer_manager) - self.assertEqual(get_parallel().attn_tp_size, 1) - raise RuntimeError("nested failure") - self.assertIs(get_context(), inner) - self.assertIs(get_parallel(), inner.parallel) - self.assertIs(get_context(), outer) - self.assertIs(get_server_args(), outer_args) - self.assertIs(get_parallel(), outer.parallel) - self.assertEqual(get_parallel().tp_rank, 0) - self.assertIsNone(get_parallel().dp_rank) - self.assertEqual(publish_role(), "diffusion_gpu_worker") - self.assertIs(outer._config_bags, outer_bags) - self.assertEqual(outer.overrides_log(), outer_log) - self.assertEqual(get_exec().kernel.grammar_backend, "none") - self.assertTrue(get_flags().capture.enable_torch_compile) - self.assertIs(rc.get_buffer("shared_name", object), outer_buffer) - self.assertIs(rc.get_global_dwdp_manager(), outer_manager) - - def test_failed_context_creation_preserves_active_context(self): - from sglang.srt import runtime_context as rc - - outer = publish(ServerArgs(model_path="dummy"), role="diffusion_gpu_worker") - outer.override("diffusion.setup", grammar_backend="none") - bags, log = outer._config_bags, outer.overrides_log() - with patch.object(rc, "_ROLE_NS_MODE", "enforce"): - with self.assertRaisesRegex(ValueError, "has no ROLE_NAMESPACE_SETS"): - rc.create_context(ServerArgs(model_path="dummy"), role="invalid_role") - broken_args = ServerArgs(model_path="dummy") - with patch.object( - broken_args, "resolve_once", side_effect=RuntimeError("resolution failed") - ): - with self.assertRaisesRegex(RuntimeError, "resolution failed"): - rc.create_context(broken_args, role="scheduler") + inner.override("encoder.setup", page_size=16) + raise RuntimeError("inner failure") self.assertIs(get_context(), outer) - self.assertIs(outer._config_bags, bags) - self.assertEqual(outer.overrides_log(), log) + self.assertIs(get_server_args(), outer_args) + self.assertEqual(get_parallel().tp_rank, 0) + self.assertEqual(publish_role(), "diffusion_gpu_worker") + self.assertNotEqual(rc.get_schedule().page_size, 16) + + with rc.use_context(inner): + self.assertEqual(rc.get_schedule().page_size, 16) class TestAssertPublished(_IsolatedServerArgs): From 90fc4104df9b95abbfb64a38a46e91d789c696bb Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Fri, 9 Oct 2026 23:26:39 +0000 Subject: [PATCH 17/18] Add H100 perf baselines for the LLaDA-Image CI cases The llada_image_bf16_t2i and llada_image_bf16_edit server cases fail the H100 performance check only because their baseline entries are missing. Add both entries from the measurements the H100 runner printed for these cases, keeping the slowest value of each metric across seven runs so the cold first run is covered, with the same fields as the most recent entries in the file. --- .../test/server/perf_baselines/h100.json | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json index 8b8040d023f2..aa88bf4e5999 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines/h100.json @@ -3578,6 +3578,152 @@ "runtime_peak_allocated_mb": 25745.0, "estimated_full_test_time_s": 404.0, "warmup_peak_vram_mb": 14598.0 + }, + "llada_image_bf16_t2i": { + "stages_ms": { + "InputValidationStage": 0.07, + "LLaDAImageTextConditioningStage": 36.79, + "LLaDAImageSourceImageConditioningStage": 0.05, + "TimestepPreparationStage": 1.94, + "LLaDAImageLatentPreparationStage": 0.2, + "DenoisingStage": 4052.86, + "DecodingStage": 40.77 + }, + "denoise_step_ms": { + "0": 77.02, + "1": 78.74, + "2": 83.82, + "3": 81.04, + "4": 80.83, + "5": 81.33, + "6": 80.96, + "7": 81.2, + "8": 80.97, + "9": 81.26, + "10": 80.74, + "11": 81.22, + "12": 80.93, + "13": 81.19, + "14": 80.91, + "15": 81.32, + "16": 80.96, + "17": 82.78, + "18": 81.03, + "19": 81.29, + "20": 83.72, + "21": 81.44, + "22": 80.94, + "23": 81.4, + "24": 80.98, + "25": 81.26, + "26": 80.87, + "27": 81.42, + "28": 81.16, + "29": 81.41, + "30": 81.13, + "31": 81.58, + "32": 81.1, + "33": 81.47, + "34": 81.13, + "35": 81.34, + "36": 81.12, + "37": 81.37, + "38": 81.12, + "39": 81.42, + "40": 81.09, + "41": 81.6, + "42": 92.45, + "43": 83.16, + "44": 81.2, + "45": 81.45, + "46": 81.32, + "47": 81.65, + "48": 81.22, + "49": 81.53 + }, + "expected_e2e_ms": 4214.9, + "expected_load_ms": 51491.39, + "expected_avg_denoise_ms": 80.99, + "expected_median_denoise_ms": 81.06, + "load_peak_vram_mb": 48046.0, + "runtime_peak_vram_mb": 50028.0, + "load_peak_allocated_mb": 47799.72, + "runtime_peak_allocated_mb": 48224.27, + "estimated_full_test_time_s": 76.8, + "warmup_peak_vram_mb": 50028.0 + }, + "llada_image_bf16_edit": { + "stages_ms": { + "InputValidationStage": 10.58, + "LLaDAImageTextConditioningStage": 37.47, + "LLaDAImageSourceImageConditioningStage": 70.63, + "TimestepPreparationStage": 1.28, + "LLaDAImageLatentPreparationStage": 0.21, + "DenoisingStage": 8880.41, + "DecodingStage": 55.64 + }, + "denoise_step_ms": { + "0": 156.23, + "1": 177.2, + "2": 177.21, + "3": 177.7, + "4": 177.58, + "5": 179.87, + "6": 178.99, + "7": 180.97, + "8": 177.51, + "9": 177.87, + "10": 177.94, + "11": 178.05, + "12": 184.78, + "13": 178.68, + "14": 178.51, + "15": 178.74, + "16": 177.99, + "17": 179.74, + "18": 179.31, + "19": 179.54, + "20": 178.13, + "21": 178.0, + "22": 179.02, + "23": 178.76, + "24": 180.81, + "25": 179.14, + "26": 178.3, + "27": 178.71, + "28": 178.18, + "29": 178.26, + "30": 179.02, + "31": 179.39, + "32": 180.16, + "33": 179.42, + "34": 179.84, + "35": 179.17, + "36": 179.35, + "37": 178.6, + "38": 179.37, + "39": 179.67, + "40": 178.92, + "41": 179.36, + "42": 179.29, + "43": 179.02, + "44": 179.62, + "45": 178.28, + "46": 179.62, + "47": 178.49, + "48": 178.7, + "49": 180.12 + }, + "expected_e2e_ms": 9060.3, + "expected_load_ms": 48754.68, + "expected_avg_denoise_ms": 177.56, + "expected_median_denoise_ms": 177.89, + "load_peak_vram_mb": 48046.0, + "runtime_peak_vram_mb": 50028.0, + "load_peak_allocated_mb": 47799.72, + "runtime_peak_allocated_mb": 48227.33, + "estimated_full_test_time_s": 78.9, + "warmup_peak_vram_mb": 50028.0 } } } From a135f7df77cb4775efeec31be46472eb44bfff80 Mon Sep 17 00:00:00 2001 From: Richard Wang Date: Sat, 10 Oct 2026 07:15:40 +0000 Subject: [PATCH 18/18] [diffusion] Fix LLaDA-Image encoder startup with mirrored MoE groups The diffusion GPU worker now mirrors its TP group into the srt MoE TP and EP globals as well. The LLaDA-Image text encoder cleared only the TP and attention TP globals before building its own srt groups, so initialize_model_parallel stopped on its expert parallel assertion and the server failed to start. Detach and restore all four mirrored groups around the encoder build, and use the encoder's own attention and MoE groups during its forward, as before the MoE groups were mirrored. --- .../llada_image/conditioning.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py index a1af3455d2c7..993f95e57531 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -21,6 +21,10 @@ prepare_diffusers_component_path_for_loading, ) +# srt group globals the diffusion runtime mirrors from its own TP group. srt's +# initialize_model_parallel refuses to run while any of them is set. +_MIRRORED_SRT_GROUPS = ("_TP", "_ATTN_TP", "_MOE_TP", "_MOE_EP") + def _resolve_llada_image_component_path( model_root: str, @@ -110,14 +114,16 @@ def __init__( ranks=SpawnRanks(world_rank=0, gpu_id=gpu_id), ) with use_context(self.runtime_context): - # The diffusion runtime mirrors its TP group into the srt globals. - # Detach the groups without clearing the encoder context's ranks, - # then restore the mirror and keep the encoder group for forwards. - saved_srt_tp = srt_parallel_state._TP - saved_srt_attn_tp = srt_parallel_state._ATTN_TP + # The diffusion runtime mirrors its TP group into the srt globals, + # the MoE groups included. Detach the groups without clearing the + # encoder context's ranks, then restore the mirror and keep the + # encoder groups for forwards. + saved_groups = { + name: getattr(srt_parallel_state, name) for name in _MIRRORED_SRT_GROUPS + } try: - srt_parallel_state._TP = None - srt_parallel_state._ATTN_TP = None + for name in _MIRRORED_SRT_GROUPS: + setattr(srt_parallel_state, name, None) init_parallel_runtime( server_args=srt_args, device=device.type, @@ -133,11 +139,13 @@ def __init__( self.worker.alloc_memory_pool() self.worker.init_attention_backends() self.worker.init_cuda_graphs() - self.encoder_tp_group = srt_parallel_state._TP - self.encoder_attn_tp_group = srt_parallel_state._ATTN_TP + self.encoder_groups = { + name: getattr(srt_parallel_state, name) + for name in _MIRRORED_SRT_GROUPS + } finally: - srt_parallel_state._TP = saved_srt_tp - srt_parallel_state._ATTN_TP = saved_srt_attn_tp + for name, group in saved_groups.items(): + setattr(srt_parallel_state, name, group) self.server_args = srt_args self.model_runner = self.worker.model_runner self.page_size = self.model_runner.page_size @@ -199,17 +207,22 @@ def encode(self, prompts: list[str], max_sequence_length: int, component_context # Restore the encoder context inside the diffusion TP scope. with ( - mm_parallel_state.use_tensor_parallel_group(self.encoder_tp_group), + mm_parallel_state.use_tensor_parallel_group(self.encoder_groups["_TP"]), use_context(self.runtime_context), ): - saved_attn_tp = srt_parallel_state._ATTN_TP + # The TP scope points _TP and _ATTN_TP at one group. Use every + # encoder group the srt layers read during the forward. + swapped = [name for name in _MIRRORED_SRT_GROUPS if name != "_TP"] + saved_groups = {name: getattr(srt_parallel_state, name) for name in swapped} try: - srt_parallel_state._ATTN_TP = self.encoder_attn_tp_group + for name in swapped: + setattr(srt_parallel_state, name, self.encoder_groups[name]) return self._encode_impl( prompts, max_sequence_length, component_context=component_context ) finally: - srt_parallel_state._ATTN_TP = saved_attn_tp + for name, group in saved_groups.items(): + setattr(srt_parallel_state, name, group) def _encode_impl( self, prompts: list[str], max_sequence_length: int, component_context