[diffusion] Support LLaDA-Image BF16 generation and editing

python/sglang/multimodal_gen/configs/models/dits/__init__.py

@@ -14,6 +14,7 @@
 from sglang.multimodal_gen.configs.models.dits.lingbot_world import (
     LingBotWorldVideoConfig,
 )
+from sglang.multimodal_gen.configs.models.dits.llada_image import LLaDAImageDitConfig
 from sglang.multimodal_gen.configs.models.dits.longlive2 import LongLive2VideoConfig
 from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
 from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
@@ -32,6 +33,7 @@
     "Ideogram4DistilledDiTConfig",
     "LingBotWorldVideoConfig",
     "LingBotVideoMoEConfig",
+    "LLaDAImageDitConfig",
     "LongLive2VideoConfig",
     "MiniMaxH3DiTConfig",
     "WanVideoConfig",

python/sglang/multimodal_gen/configs/models/dits/llada_image.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class LLaDAImageArchConfig(DiTArchConfig):
+    all_patch_size: tuple[int, ...] = (1,)
+    all_f_patch_size: tuple[int, ...] = (1,)
+    axes_dims: tuple[int, ...] = (32, 48, 48)
+    axes_lens: tuple[int, ...] = (32768, 1024, 1024)
+    cap_feat_dim: int = 2560
+    dim: int = 3840
+    in_channels: int = 128
+    n_heads: int = 30
+    n_layers: int = 30
+    n_refiner_layers: int = 2
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    rope_theta: float = 256.0
+    semantic_feat_dim: int = 4096
+    t_scale: float = 1000.0
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.TORCH_SDPA,
+        }
+    )
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"(.*)\.attention\.to_q\.weight$": (r"\1.attention.to_qkv.weight", 0, 3),
+            r"(.*)\.attention\.to_k\.weight$": (r"\1.attention.to_qkv.weight", 1, 3),
+            r"(.*)\.attention\.to_v\.weight$": (r"\1.attention.to_qkv.weight", 2, 3),
+            r"(.*)\.attention\.to_[qkv]\.weight_scale$": (
+                r"\1.attention.to_qkv.weight_scale"
+            ),
+            r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
+            r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
+            r"(.*)\.feed_forward\.w[13]\.weight_scale$": (
+                r"\1.feed_forward.w13.weight_scale"
+            ),
+        }
+    )
+
+    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"

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -39,6 +39,9 @@
     LingBotWorldCausalDMDConfig,
     LingBotWorldV2CausalDMDConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.llada_image import (
+    LLaDAImagePipelineConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
     LTX2PipelineConfig,
     LTX23PipelineConfig,
@@ -76,6 +79,7 @@
     "Hunyuan3D2PipelineConfig",
     "Ideogram4PipelineConfig",
     "Ideogram4DistilledPipelineConfig",
+    "LLaDAImagePipelineConfig",
     "FluxPipelineConfig",
     "Flux2PipelineConfig",
     "Flux2KleinPipelineConfig",

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -439,6 +439,35 @@ def supports_sequential_multi_output_inference(self):
         """Return whether one request's outputs run through DiT/VAE sequentially."""
         return False
 
+    def validate_num_outputs_per_prompt(
+        self, num_outputs_per_prompt: int, server_args
+    ) -> None:
+        """Validate a request's per-prompt output count for this pipeline."""
+        del num_outputs_per_prompt, server_args
+
+    def supports_memory_release(self) -> bool:
+        """Return whether release_memory_occupation can offload this pipeline."""
+        return True
+
+    def supports_hot_weight_updates(self) -> bool:
+        """Return whether runtime weight updates cover this pipeline fully."""
+        return True
+
+    def validate_request_sampling_params(self, sampling_params, server_args) -> None:
+        """Validate a request's sampling params for this pipeline.
+
+        May also run before full request materialization with a lightweight
+        object that only carries width, height, max_sequence_length,
+        num_outputs_per_prompt, diffusers_kwargs, enable_cache_dit,
+        cache_dit_params, attention_backend_override, and cfg_gate_step,
+        so overrides must not require other sampling fields.
+        """
+        del sampling_params, server_args
+
+    def validate_edit_source_count(self, source_count: int, server_args) -> None:
+        """Validate the number of source images before they are materialized."""
+        del source_count, server_args
+
     def estimate_request_cost(self, batch) -> float:
         """Return the relative cost used for batching admission caps.