diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index c87cc5ab8b0a..2b9d769252c9 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 @@ -33,6 +34,7 @@ "Ideogram4DistilledDiTConfig", "LingBotWorldVideoConfig", "LingBotVideoMoEConfig", + "LLaDAImageDitConfig", "LongLive2VideoConfig", "MiniMaxH3DiTConfig", "WanAnimate2Config", 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..941a89fad7cc --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/llada_image.py @@ -0,0 +1,70 @@ +# 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 + +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): + 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 b111249a8801..4b839373f913 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -42,6 +42,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, @@ -83,6 +86,7 @@ "HunyuanImage3PipelineConfig", "Ideogram4PipelineConfig", "Ideogram4DistilledPipelineConfig", + "LLaDAImagePipelineConfig", "FluxPipelineConfig", "Flux2PipelineConfig", "Flux2KleinPipelineConfig", 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..314748c43ba5 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/llada_image.py @@ -0,0 +1,224 @@ +# 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, + 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 +class LLaDAImagePipelineConfig(SpatialImagePipelineConfig): + task_type: ModelTaskType = ModelTaskType.TI2I + should_use_guidance: bool = False + + dit_config: LLaDAImageDitConfig = field(default_factory=LLaDAImageDitConfig) + 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 + # 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 + + 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 + self.validate_output_size( + batch.width, batch.height, editing=batch.condition_image is not None + ) + 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 validate_server_args(self, server_args) -> None: + super().validate_server_args(server_args) + 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, + ) + ): + 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 requires a resident text encoder") + server_args.text_encoder_cpu_offload = False + 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"], + ) + if server_args.explicit_residency_mode("text_encoder") not in ( + None, + "resident", + ): + raise ValueError("LLaDA-Image requires a resident text encoder") + + 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 + max_height = axes_lens[1] * self.latent_scale_factor + max_width = axes_lens[2] * self.latent_scale_factor + if height > max_height or width > max_width: + raise ValueError( + f"LLaDA-Image output size {width}x{height} exceeds spatial RoPE " + f"bounds of {max_width}x{max_height} pixels" + ) + + @staticmethod + 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] + + image_embeds = batch.image_embeds + if image_embeds: + 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": 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): + 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): + vae_parameter = next(vae.parameters()) + latents = latents.to(device=vae_parameter.device, dtype=vae_parameter.dtype) + mean, std = flux2_vae_bn_stats(vae, latents) + return _unpatchify_latents(latents * std + mean) + + def post_denoising_loop(self, latents, batch): + del batch + return latents + + +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() + ) + 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 new file mode 100644 index 000000000000..d31133e9dccb --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/llada_image.py @@ -0,0 +1,42 @@ +# 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 + # 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 = 5.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" + ) + pipeline_config.validate_output_size( + self.width, self.height, editing=bool(self.image_path) + ) + + +@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/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/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 1585d1c52892..0995cc6b4fba 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -26,6 +26,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, @@ -102,11 +106,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 ) @@ -126,6 +132,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 @@ -200,10 +222,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 new file mode 100644 index 000000000000..6e0518185af5 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/llada_image.py @@ -0,0 +1,1437 @@ +# 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 FeedForward +from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import RMSNorm +from torch import nn +from torch.nn.utils.rnn import pad_sequence + +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, +) +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) + + +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, cast_x_before_out_mul=True) + 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) + 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 LLaDAImageAttention(nn.Module): + 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", + ) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + freqs_cis: torch.Tensor, + ) -> torch.Tensor: + 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): + 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 _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, + 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), + ) + + +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, + ) -> torch.Tensor: + if self.modulation: + if noise_mask is None: + modulation = ( + self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2) + ) + else: + 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) * (1.0 + scale_msa), + attention_mask, + freqs_cis, + ) + 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, + ) + 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 = self.adaLN_modulation(adaln_input).unsqueeze(1) + else: + scale = _select_per_token( + self.adaLN_modulation(adaln_noisy), + self.adaLN_modulation(adaln_clean), + noise_mask, + hidden_states.shape[1], + ) + hidden_states = self.norm_final(hidden_states) * (1.0 + scale) + return self.linear(hidden_states) + + +class _LLaDAImageTransformer2DModel(ModelMixin, ConfigMixin): + """Denoiser for text-to-image generation and single-image editing.""" + + @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__() + self.out_channels = in_channels + self.t_scale = t_scale + + 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) + + def blocks(name: str, count: int, modulation: bool) -> nn.ModuleList: + return nn.ModuleList( + LLaDAImageTransformerBlock( + dim, + n_heads, + norm_eps, + qk_norm, + modulation=modulation, + quant_config=quant_config, + prefix=f"{name}.{layer_id}", + ) + for layer_id in range(count) + ) + + 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), + ) + + 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)) + + 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, + ) -> 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 + position_ids = self._create_coordinate_grid( + 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 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 + ) + features = torch.cat( + [features, features[-1:].repeat(padding_length, 1)], dim=0 + ) + 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 + + @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] + 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 = _padding_attention_mask(sequence_lengths, features.device) + + 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 = [] + 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], + patch_size: int, + f_patch_size: int, + ) -> tuple[_LLaDAImageSequence, _LLaDAImageSequence, list[tuple[int, int, int]]]: + image_sequence = _LLaDAImageSequence([], [], []) + cap_sequence = _LLaDAImageSequence([], [], []) + image_sizes = [] + + 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) + ) + ) + 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 + ) + padded_features, position_ids, padding_mask, _, _ = self._pad_with_ids( + 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, 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 = [] + image_token_counts = [] + 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_token_counts.append(len(patches)) + padded_features, position_ids, padding_mask, _, noise_mask = ( + self._pad_with_ids( + patches, token_grid_size, (position_start, 0, 0), noise_value + ) + ) + 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), + ) + ) + + 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], + (len(glm_cap_feats[batch_index]), 1, 1), + ( + len(batch_cap_features) + len(batch_image_features) + 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, torch.Tensor | None]: + def merge(groups: tuple[torch.Tensor, ...]) -> list[torch.Tensor]: + return [ + torch.cat( + [ + group[batch_index, : lengths[batch_index]] + for group, lengths in zip(groups, length_groups) + ], + dim=0, + ) + for batch_index in range(feature_groups[0].shape[0]) + ] + + merged_features = merge(feature_groups) + attention_mask = _padding_attention_mask( + [len(features) for features in merged_features], + feature_groups[0].device, + ) + noise_mask = None + if noise_mask_groups is not None: + noise_mask = pad_sequence( + merge(noise_mask_groups), batch_first=True, padding_value=0 + ) + return ( + 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, + 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], + glm_cap_feats: list[torch.Tensor] | None = None, + source_latents: list[torch.Tensor] | None = None, + patch_size: int = 1, + f_patch_size: int = 1, + ) -> list[torch.Tensor]: + patch_key = f"{patch_size}-{f_patch_size}" + batch_size = len(x) + is_editing = source_latents is not None + image_offsets = None + 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 + ) + 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, + cap_feats, + glm_cap_feats, + source_latents, + patch_size, + f_patch_size, + ) + ) + else: + 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_features, image_frequencies, image_mask, image_lengths, image_noise = ( + self._embed_sequence( + image_sequence, self.all_x_embedder[patch_key], self.x_pad_token + ) + ) + for layer in self.noise_refiner: + image_features = layer( + image_features, + image_mask, + image_frequencies, + noise_mask=image_noise, + **adaln, + ) + + 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) + + if is_editing: + ( + sigvq_features, + sigvq_frequencies, + sigvq_mask, + sigvq_lengths, + 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_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), + ) + ) + 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), + ) + ) + + for layer in self.layers: + unified_features = layer( + unified_features, + unified_mask, + unified_frequencies, + noise_mask=unified_noise, + **adaln, + ) + 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, + ) + + +class LLaDAImageQueryAttention(nn.Module): + def __init__(self, hidden_size: int, num_heads: int): + super().__init__() + self.heads = num_heads + self.head_dim = hidden_size // num_heads + 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) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + 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)) + + +class LLaDAImageQueryFormerBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + intermediate_size: int, + 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) + 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, + ) -> torch.Tensor: + query_embeds = self.norm_q(query_embeds) + encoder_hidden_states = self.norm_k(encoder_hidden_states) + query_embeds = query_embeds + self.cross_attn( + query_embeds, encoder_hidden_states, attention_mask + ) + query_embeds = self.norm1(query_embeds) + return query_embeds + self.mlp.fc2( + F.gelu(self.mlp.fc1(query_embeds), approximate="tanh") + ) + + +class LLaDAImageQueryFormerModel(ModelMixin, ConfigMixin): + """Derives the image-generation query tokens from LLaDA token embeddings.""" + + @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__() + self.meta_queries = nn.Parameter(torch.zeros(num_queries, hidden_size)) + self.query_blocks = nn.ModuleList( + 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 + ) -> torch.Tensor: + query_embeds = self.meta_queries.unsqueeze(0).expand( + inputs_embeds.shape[0], -1, -1 + ) + 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): + 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.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) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + 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 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: + return self.fc2(F.gelu(self.fc1(hidden_states), approximate=self.approximate)) + + +class LLaDAImageTextProjectionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + norm_eps: float, + ): + super().__init__() + self.self_attn = LLaDAImageTextProjectionAttention( + hidden_size, num_attention_heads, norm_eps + ) + self.layer_norm1 = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False) + 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: + 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): + """Maps LLaDA hidden states to the denoiser caption dimension.""" + + @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__() + self.layers = nn.ModuleList( + 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) -> torch.Tensor: + for layer in self.layers: + hidden_states = layer(hidden_states) + return self.projector(hidden_states) + + +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.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: + 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)) + + +class LLaDAImageSigVQVisionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + attention_bias: bool, + 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 + ) + 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)) + 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.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): + """GLM SigVQ image encoder, quantizer, and prior projection for editing.""" + + @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__() + 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, + 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", + ) + + 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) + + 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): + """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): + 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("_") + } + 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 + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: list[torch.Tensor], + timestep: torch.Tensor, + encoder_hidden_states_image: list[torch.Tensor] | None = None, + source_latents: list[torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + del kwargs + model_dtype = next(self.parameters()).dtype + 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() + + +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 b11e1e93d78e..3e95069e59ae 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 @@ -29,6 +29,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.models.schedulers.sigma_schedule import ( @@ -92,6 +93,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): @@ -116,6 +120,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, ): @@ -125,12 +130,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( @@ -341,13 +348,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) @@ -451,7 +463,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, @@ -472,8 +484,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`): @@ -532,7 +544,7 @@ def step( else: if self.config.stochastic_sampling: x0 = sample - current_sigma * model_output - noise = torch.randn( + noise = randn_tensor( sample.shape, generator=generator, device=sample.device, 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..b9edb8ca4fbd --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/llada_image.py @@ -0,0 +1,122 @@ +# 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, +) +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + PlainStateDictComponentLoader, +) +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.model_specific_stages.llada_image.conditioning import ( + LLaDAImageTextConditioningStage, + LLaDAImageTextEncoderRunner, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_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) + return batch + + +class LLaDAImagePipeline(ComposedPipelineBase): + pipeline_name = "LLaDAImagePipeline" + component_loaders = { + "queryformer": PlainStateDictComponentLoader, + "text_projection": PlainStateDictComponentLoader, + "sigvq": PlainStateDictComponentLoader, + } + + _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 3567a6d2d10a..de10a1666e60 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -135,6 +135,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/model_specific_stages/llada_image/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/__init__.py new file mode 100644 index 000000000000..9881313609aa --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 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 new file mode 100644 index 000000000000..993f95e57531 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/conditioning.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from array import array +from pathlib import Path + +import torch + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, +) +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 +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + 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, + server_args: ServerArgs, + component_name: str, +) -> str: + 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 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.srt.distributed.parallel_state as srt_parallel_state + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.distributed.bootstrap import ( + init_layer_runtime, + init_parallel_runtime, + ) + from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.mem_cache.unified_cache.components import ComponentType + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + from sglang.srt.runtime_context import SpawnRanks, create_context, use_context + from sglang.srt.server_args import ServerArgs as SRTServerArgs + + self.queryformer = queryformer + self.text_projection = text_projection + self.tokenizer = tokenizer + 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" + ) + device = get_local_torch_device() + gpu_id = device.index + nccl_port = server_args.nccl_port or 29500 + srt_args = SRTServerArgs( + model_path=text_encoder_path, + tokenizer_path=tokenizer_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + skip_tokenizer_init=True, + dtype="bfloat16", + 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=2, + ) + self.runtime_context = create_context( + srt_args, + role="scheduler", + 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, + # 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: + for name in _MIRRORED_SRT_GROUPS: + setattr(srt_parallel_state, name, None) + init_parallel_runtime( + server_args=srt_args, + device=device.type, + dist_port=nccl_port, + ) + init_layer_runtime(model_config=ModelConfig.from_server_args(srt_args)) + self.worker = TpModelWorker( + server_args=srt_args, + gpu_id=gpu_id, + nccl_port=nccl_port, + ) + # 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_groups = { + name: getattr(srt_parallel_state, name) + for name in _MIRRORED_SRT_GROUPS + } + finally: + 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 + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + self.worker.get_memory_pool() + ) + # The radix-disabled cache srt builds for a full-attention model. + self.tree_cache = UnifiedRadixCache( + 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, + tree_components=(ComponentType.FULL,), + ) + ) + + def _prepare_input_embeds( + self, + input_ids: torch.Tensor, + sequence_lengths: list[int], + text_lengths: list[int], + component_context, + ) -> torch.Tensor: + embeddings = self.model_runner.model.get_input_embeddings() + inputs = [] + offset = 0 + 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( + 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, + ) + inputs.append( + torch.cat( + [text_embeds, query_embeds.squeeze(0).to(text_embeds.dtype)], + dim=0, + ) + ) + offset += sequence_length + return torch.cat(inputs, dim=0) + + @torch.no_grad() + 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 + + # Restore the encoder context inside the diffusion TP scope. + with ( + mm_parallel_state.use_tensor_parallel_group(self.encoder_groups["_TP"]), + use_context(self.runtime_context), + ): + # 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: + 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: + 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 + ): + 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, + ) + 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) + 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 + 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: + output = text_projection( + hidden_states[offset : offset + sequence_length] + .unsqueeze(0) + .to(dtype=projection_dtype) + ) + projected.append(output.squeeze(0)) + offset += sequence_length + return projected + finally: + for req in reqs: + if not req.kv.is_kv_released: + release_kv_cache(req, self.tree_cache, checkpoint=False) + 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) + + +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/model_specific_stages/llada_image/source.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py new file mode 100644 index 000000000000..f13ddacb67f2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/llada_image/source.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +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, +) +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"), + ] + + @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, + ) + if batch.batch_size > 1: + image = image.repeat(batch.batch_size, 1, 1, 1) + + 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 + sigvq_device, sigvq_dtype = _module_device_dtype(sigvq) + semantic_features = sigvq( + sigvq_pixels.to(device=sigvq_device, dtype=sigvq_dtype) + ) + + 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)) + 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 = [ + 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 ec2993208eee..fbb319e6303f 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 @@ -98,7 +98,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/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 9a25008ef18d..66e4afa2e8ed 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -562,6 +562,50 @@ ), ] +_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", + _LLADA_IMAGE_SERVER_ARGS, + DiffusionSamplingParams( + prompt="A red fox sitting on a snowy hill", + output_size="512x512", + extras=_LLADA_IMAGE_SAMPLING, + ), + run_perf_check=False, + run_consistency_check=False, + run_component_accuracy_check=False, + ), + DiffusionTestCase( + "llada_image_bf16_edit", + _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=_LLADA_IMAGE_SAMPLING, + ), + run_perf_check=False, + run_consistency_check=False, + run_component_accuracy_check=False, + ), + ] + # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm. if not current_platform.is_hip(): ONE_GPU_CASES.append( 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 } } } 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..33171649111b --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_llada_image_config.py @@ -0,0 +1,69 @@ +# 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.pipeline_configs.llada_image import ( + LLaDAImagePipelineConfig, +) +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) + + 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_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/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index ffb363904a04..b17ae4dd6ca3 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -80,6 +80,7 @@ "fa3", "fa4", "flashinfer", + "llada2_cfg_flashinfer", "flashmla", "trtllm_mla", "cutedsl_mla", diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 0e070228761c..e79d5b014ebf 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -77,6 +77,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_cfg_flashinfer_backend.py b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py new file mode 100644 index 000000000000..dddba712b1f0 --- /dev/null +++ b/python/sglang/srt/layers/attention/llada2_cfg_flashinfer_backend.py @@ -0,0 +1,135 @@ +# 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, +) + + +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): + """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 = forward_batch.llada_image_conditioning_text_lens_cpu + 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/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0532121a0fa9..4ba4a6347db7 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -728,6 +728,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 e57db9909603..aa4d4512c710 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -1295,6 +1295,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, @@ -1336,6 +1373,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 @@ -1389,7 +1429,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/runtime_context.py b/python/sglang/srt/runtime_context.py index f6792cf9e756..bcea02a32df9 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -1505,6 +1505,37 @@ def publish( (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, ranks=ranks) + + +def create_context( + server_args, *, role: str, ranks: SpawnRanks | None = None +) -> RuntimeContext: + """Build a model-owned context without replacing the active process config.""" + context = RuntimeContext(parallel=ParallelContext()) + return _configure_context(context, server_args, role=role, ranks=ranks) + + +@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, ranks: SpawnRanks | None = None +) -> 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( @@ -1512,8 +1543,8 @@ def publish( "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 " @@ -1524,16 +1555,16 @@ def publish( f"{source}({', '.join(sorted(fields))})" for source, fields in discarded ), ) - _CONTEXT._publish_role = role + context._publish_role = role # Disabled DCP has rank zero even in processes without a rank bundle. - if not _CONTEXT.parallel.dcp_enabled: - _CONTEXT.parallel.override_permanently(attn_dcp_rank=0) + if not context.parallel.dcp_enabled: + context.parallel.override_permanently(attn_dcp_rank=0) # The device is assigned by the launcher; it is not a config field. - _CONTEXT.config_bag("device")._set( + context.config_bag("device")._set( "gpu_id", ranks.gpu_id if ranks is not None else None ) if ranks is not None: - parallel = _CONTEXT.parallel + parallel = context.parallel placement = derive_spawn_ranks( world_rank=ranks.world_rank, tp_size=parallel.tp_size, @@ -1567,7 +1598,7 @@ def publish( file=sys.stderr, flush=True, ) - return _CONTEXT + return context def _attention_ranks(parallel, tp_rank: int) -> dict: 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 new file mode 100644 index 000000000000..e214bea9a7f1 --- /dev/null +++ b/test/registered/unit/layers/attention/test_llada2_cfg_flashinfer_backend.py @@ -0,0 +1,46 @@ +import unittest + +import torch + +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 TestLLaDA2CFGFlashInferMask(unittest.TestCase): + 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") + + 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"): + _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/registered/unit/models/test_llada2.py b/test/registered/unit/models/test_llada2.py index ff92303fb8db..84ff53d4c2cc 100644 --- a/test/registered/unit/models/test_llada2.py +++ b/test/registered/unit/models/test_llada2.py @@ -6,6 +6,7 @@ from unittest.mock import patch import torch +import torch.nn.functional as F from torch import nn from sglang.srt.dllm.config import DllmConfig @@ -16,9 +17,11 @@ ) from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.models.llada2 import ( + LLaDA2MoeGate, LLaDA2MoeSparseMoeBlock, _get_effective_moe_runner_backend, _make_block_routing_triton_output, + _prepare_llada2_language_weights, _require_block_routing_ep1, _require_block_routing_runner_compatibility, ) @@ -344,5 +347,57 @@ def test_block_routing_rejects_mismatched_dllm_block_size(self, from_server_args DllmConfig.from_server_args(self._server_args(config_path)) +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/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 8dc0a8ecf50a..ba6b60dfdd3b 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -808,6 +808,40 @@ def test_reset_context_clears_owned_store(self): get_server_args() +class TestEmbeddedRuntimeContext(unittest.TestCase): + 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() + outer_args = ServerArgs(model_path="dummy", tp_size=1) + outer = publish( + outer_args, role="diffusion_gpu_worker", ranks=SpawnRanks(world_rank=0) + ) + inner = rc.create_context( + ServerArgs(model_path="dummy", tp_size=2), + role="scheduler", + ranks=SpawnRanks(world_rank=1), + ) + self.assertIs(get_context(), outer) + + with self.assertRaisesRegex(RuntimeError, "inner failure"): + with rc.use_context(inner): + self.assertIs(get_parallel(), inner.parallel) + self.assertEqual(get_parallel().tp_rank, 1) + inner.override("encoder.setup", page_size=16) + raise RuntimeError("inner failure") + self.assertIs(get_context(), outer) + 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): """Publishing is the process entry's job; the constructors only check.