[Feature] Add SANA diffusion model (#19234)

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

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Architecture and model configuration for SANA DiT (Diffusion Transformer).
+#
+# SANA uses a linear-attention-based transformer that replaces standard
+# quadratic self-attention with ReLU-based linear attention, enabling
+# efficient high-resolution image synthesis. Cross-attention (standard SDPA)
+# is used for text conditioning via Gemma2 embeddings.
+#
+# Defaults below correspond to the SANA-1.6B / 1024px variant.
+# For 4.8B, override num_layers=36, num_attention_heads=64, etc.
+#
+# Reference: https://arxiv.org/abs/2410.10629
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+@dataclass
+class SanaArchConfig(DiTArchConfig):
+    patch_size: int = 1
+    in_channels: int = 32
+    out_channels: int = 32
+    num_layers: int = 20
+    attention_head_dim: int = 32
+    num_attention_heads: int = 70
+    num_cross_attention_heads: int = 20
+    cross_attention_head_dim: int = 112
+    cross_attention_dim: int = 2240
+    caption_channels: int = 2304
+
+    mlp_ratio: float = 2.5
+    # "rms_norm_across_heads" applies RMSNorm over the full (num_heads * head_dim)
+
+    qk_norm: str = "rms_norm_across_heads"
+    norm_elementwise_affine: bool = False
+    norm_eps: float = 1e-6
+    sample_size: int = 32
+    guidance_embeds: bool = False
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^transformer\.(.*)$": r"\1",
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class SanaConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=SanaArchConfig)
+    prefix: str = "Sana"

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

@@ -10,6 +10,7 @@ from sglang.multimodal_gen.configs.models.encoders.clip import (
     CLIPTextConfig,
     CLIPVisionConfig,
 )
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
 from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
 from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
 from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
@@ -25,5 +26,6 @@ __all__ = [
     "LlamaConfig",
     "Qwen3TextConfig",
     "T5Config",
+    "Gemma2Config",
     "Gemma3Config",
 ]

python/sglang/multimodal_gen/configs/models/encoders/gemma2.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Text encoder configuration for Gemma2 2B, used by SANA for text conditioning.
+#
+# SANA uses the hidden states from Gemma2 (not logits) as the conditioning
+# signal for cross-attention in the DiT. The encoder output dimension (2304)
+# is projected to the DiT's inner_dim via caption_projection.
+#
+# Defaults match google/gemma-2-2b-it (the model used in SANA HF checkpoints).
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m) -> bool:
+    return n.endswith("norm")
+
+
+@dataclass
+class Gemma2ArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 256000
+    hidden_size: int = 2304
+    intermediate_size: int = 9216
+    num_hidden_layers: int = 26
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 4
+    head_dim: int = 256
+    hidden_act: str = "gelu_pytorch_tanh"
+    hidden_activation: str = "gelu_pytorch_tanh"
+    max_position_embeddings: int = 8192
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    eos_token_id: int = 1
+    bos_token_id: int = 2
+    tie_word_embeddings: bool = True
+    rope_theta: float = 10000.0
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+
+    # Gemma2 alternates between global and sliding-window attention
+    # on odd/even layers, respectively.
+    sliding_window: int = 4096
+
+    # query_pre_attn_scalar replaces the standard 1/sqrt(head_dim) scaling.
+    query_pre_attn_scalar: int = 256
+
+    # Softcapping bounds raw attention logits via tanh(logits/cap)*cap.
+    # NOTE: SDPA does not natively support softcapping; the runtime model
+    # currently skips this (see Gemma2Attention.forward). Quality impact
+    # is minimal for short text-encoder sequences but should be revisited
+    # for longer context.
+    attn_logit_softcapping: float = 50.0
+    final_logit_softcapping: float = 30.0
+
+    text_len: int = 300
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", "0"),
+            (".gate_up_proj", ".up_proj", "1"),
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
+    )
+
+
+@dataclass
+class Gemma2Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=Gemma2ArchConfig)
+    prefix: str = "gemma_2"

python/sglang/multimodal_gen/configs/models/vaes/sana.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# VAE configuration for SANA's DC-AE (Deep Compression AutoEncoder).
+#
+# DC-AE achieves a 32x spatial compression ratio (vs. 8x for standard SD VAEs),
+# which means a 1024x1024 image becomes 32x32 latents with 32 channels.
+# This aggressive compression is what allows SANA to run efficiently at
+# high resolutions despite having a relatively small DiT.
+#
+# Reference: https://arxiv.org/abs/2405.17811
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class SanaVAEArchConfig(VAEArchConfig):
+    spatial_compression_ratio: int = 32
+    # DC-AE uses a different scaling factor than standard VAEs;
+    # this value must match the pretrained checkpoint.
+    scaling_factor: float = 0.41407
+    latent_channels: int = 32
+    in_channels: int = 3
+
+
+@dataclass
+class SanaVAEConfig(VAEConfig):
+    arch_config: SanaVAEArchConfig = field(default_factory=SanaVAEArchConfig)
+
+    # DC-AE does not currently support tiling in our wrapper.
+    # Enable these once the diffusers AutoencoderDC adds tiling support.
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    def post_init(self):
+        # Called by VAELoader AFTER update_model_arch() merges the HF config.json
+        # values into arch_config. Must be post_init() (not __post_init__) because
+        # __post_init__ fires at dataclass creation time, before the HF config merge.
+        #
+        # The base VAEConfig.get_vae_scale_factor() derives from block_out_channels,
+        # which DC-AE doesn't have. Set vae_scale_factor directly from the
+        # spatial_compression_ratio (32x for DC-AE).
+        self.arch_config.vae_scale_factor = self.arch_config.spatial_compression_ratio

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

@@ -29,6 +29,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
 )
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     SelfForcingWanT2V480PConfig,
     WanI2V480PConfig,
@@ -51,6 +52,7 @@ __all__ = [
     "Flux2KleinPipelineConfig",
     "Flux2FinetunedPipelineConfig",
     "PipelineConfig",
+    "SanaPipelineConfig",
     "SlidingTileAttnConfig",
     "MOVAPipelineConfig",
     "WanT2V480PConfig",

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

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Pipeline configuration for SANA text-to-image generation.
+#
+# SANA produces 4D spatial latents (B, C, H', W') directly — unlike Flux/QwenImage
+# which use packed token-style latents (B, S, D). This means:
+#   - We inherit SpatialImagePipelineConfig (not ImagePipelineConfig)
+#   - prepare_latent_shape returns 4D, not 5D
+#   - post_denoising_loop is a no-op (no un-packing needed)
+#   - shard_latents_for_sp shards along the H' dimension
+#
+# SANA does NOT use rotary position embeddings, so prepare_pos/neg_cond_kwargs
+# return empty dicts (the DiT only needs hidden_states + encoder_hidden_states + timestep).
+#
+# CFG is handled by the denoising stage via guidance_scale in sampling params.
+# should_use_guidance=False means no embedded guidance (no extra guidance token in forward),
+# but negative_prompt + guidance_scale > 1.0 still enables standard classifier-free guidance.
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.base import EncoderConfig
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
+from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    SpatialImagePipelineConfig,
+    preprocess_text,
+)
+
+
+def sana_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
+    # SANA uses the final hidden state from Gemma2 directly as text conditioning.
+    # No intermediate-layer extraction or masking needed (unlike QwenImage/ZImage).
+    return outputs.last_hidden_state
+
+
+@dataclass
+class SanaPipelineConfig(SpatialImagePipelineConfig):
+
+    task_type: ModelTaskType = ModelTaskType.T2I
+
+    # should_use_guidance=False disables *embedded* guidance (timestep-conditioned
+    # guidance token). Standard CFG via guidance_scale is still active.
+    should_use_guidance: bool = False
+    enable_autocast: bool = False
+
+    # DC-AE does not support tiling or SP VAE decode yet.
+    vae_tiling: bool = False
+    vae_sp: bool = False
+    vae_precision: str = "bf16"
+
+    dit_config: DiTConfig = field(default_factory=SanaConfig)
+    vae_config: VAEConfig = field(default_factory=SanaVAEConfig)
+
+    # Single text encoder: Gemma2 (unlike Flux which uses CLIP + T5)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Gemma2Config(),)
+    )
+
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text,),
+    )
+
+    postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (sana_postprocess_text,)
+    )
+
+    def prepare_latent_shape(self, batch, batch_size, num_frames):
+        # 4D latent shape: (B, C, H', W') — no temporal dim for T2I.
+        # DC-AE compresses 1024x1024 -> 32x32 with 32 channels.
+        compression = self.vae_config.arch_config.spatial_compression_ratio
+        height = batch.height // compression
+        width = batch.width // compression
+        num_channels = self.dit_config.arch_config.num_channels_latents
+        shape = (batch_size, num_channels, height, width)
+        return shape
+
+    def get_pos_prompt_embeds(self, batch):
+        # Single encoder -> index [0] (Flux uses [1] because T5 is encoder #2)
+        return batch.prompt_embeds[0]
+
+    def get_neg_prompt_embeds(self, batch):
+        return batch.negative_prompt_embeds[0]
+
+    def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        # encoder_attention_mask: batch stores list-of-tensors; diffusers' SanaTransformer
+        # expects a single tensor (sglang's has list handling). Override with [0].
+        out = {}
+        m = batch.prompt_attention_mask
+        if isinstance(m, (list, tuple)):
+            out["encoder_attention_mask"] = m[0] if m else None
+        elif m is not None:
+            out["encoder_attention_mask"] = m
+        return out
+
+    def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype):
+        out = {}
+        m = batch.negative_attention_mask
+        if isinstance(m, (list, tuple)):
+            out["encoder_attention_mask"] = m[0] if m else None
+        elif m is not None:
+            out["encoder_attention_mask"] = m
+        return out
+
+    def post_denoising_loop(self, latents, batch):
+        return latents

python/sglang/multimodal_gen/configs/sample/sana.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sampling parameters for SANA image generation (T2I)."""
+
+from dataclasses import dataclass
+
+from sglang.multimodal_gen.configs.sample.sampling_params import (
+    DataType,
+    SamplingParams,
+)
+
+
+@dataclass
+class SanaSamplingParams(SamplingParams):
+    """Defaults for SANA 1.5 1024px variant.
+
+    guidance_scale=4.5 enables standard classifier-free guidance.
+    """
+
+    data_type: DataType = DataType.IMAGE
+    num_frames: int = 1
+    guidance_scale: float = 4.5
+    num_inference_steps: int = 20
+    height: int = 1024
+    width: int = 1024
+    negative_prompt: str = (
+        "low quality, low resolution, blurry, overexposed, underexposed, "
+        "distorted, deformed, disfigured, bad anatomy, extra limbs, "
+        "watermark, text, signature, ugly, noisy, artifacts"
+    )

python/sglang/multimodal_gen/registry.py

@@ -63,6 +63,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageLayeredPipelineConfig,
     QwenImagePipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     FastWan2_1_T2V_480P_Config,
     FastWan2_2_TI2V_5B_Config,
@@ -98,6 +99,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
     QwenImageLayeredSamplingParams,
     QwenImageSamplingParams,
 )
+from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams
 from sglang.multimodal_gen.configs.sample.wan import (
     FastWanT2V480PConfig,
     Turbo_Wan2_2_I2V_A14B_SamplingParam,
@@ -796,6 +798,21 @@ def _register_configs():
         ],
     )
 
+    # SANA
+    register_configs(
+        sampling_param_cls=SanaSamplingParams,
+        pipeline_config_cls=SanaPipelineConfig,
+        hf_model_paths=[
+            "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
+            "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers",
+            "Efficient-Large-Model/Sana_1600M_1024px_diffusers",
+            "Efficient-Large-Model/Sana_600M_1024px_diffusers",
+            "Efficient-Large-Model/Sana_1600M_512px_diffusers",
+            "Efficient-Large-Model/Sana_600M_512px_diffusers",
+        ],
+        model_detectors=[lambda hf_id: "sana" in hf_id.lower()],
+    )
+
 
 _register_configs()
 

python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py

@@ -128,9 +128,11 @@ class VAELoader(ComponentLoader):
 
         safetensors_list = _list_safetensors_files(component_model_path)
         assert (
-            len(safetensors_list) == 1
-        ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
-        loaded = safetensors_load_file(safetensors_list[0])
+            len(safetensors_list) >= 1
+        ), f"Found no safetensors files in {component_model_path}"
+        loaded = {}
+        for sf_path in safetensors_list:
+            loaded.update(safetensors_load_file(sf_path))
         vae.load_state_dict(loaded, strict=False)
 
         state_keys = set(vae.state_dict().keys())

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -369,15 +369,16 @@ def load_model_from_full_model_state_dict(
     if unused_keys:
         logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
 
-    # for nunchaku
+    # for nunchaku; norm_q/norm_k for SANA QK normalization layers
     ALLOWED_NEW_PARAM_PATTERNS = [
         "gate_compress",
         "wcscales",
         "wtscale",
         "bias",
+        "norm_q",
+        "norm_k",
     ]
     for new_param_name in unused_keys:
-        # check unallowed missing params
         if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
             logger.error(
                 "Unsupported new parameter: %s. Allowed patterns: %s",
@@ -392,7 +393,9 @@ def load_model_from_full_model_state_dict(
         meta_sharded_param = meta_sd.get(new_param_name)
         meta_sharded_param_dtype = meta_sharded_param.dtype
 
-        if "wcscales" in new_param_name or "wtscale" in new_param_name:
+        if any(
+            p in new_param_name for p in ("wcscales", "wtscale", "norm_q", "norm_k")
+        ):
             init_like = torch.ones_like
         else:
             init_like = torch.zeros_like

python/sglang/multimodal_gen/runtime/models/dits/sana.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding
+
+from sglang.multimodal_gen.configs.models.dits.sana import SanaConfig
+from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
+from sglang.multimodal_gen.runtime.layers.visual_embedding import Timesteps
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SanaCombinedTimestepSizeEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.time_proj = Timesteps(
+            num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
+        )
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim
+        )
+
+    def forward(self, timestep, hidden_dtype=None):
+        timesteps_proj = self.time_proj(timestep)
+        if hidden_dtype is not None:
+            timesteps_proj = timesteps_proj.to(dtype=hidden_dtype)
+        timesteps_emb = self.timestep_embedder(timesteps_proj)
+        return timesteps_emb
+
+
+class SanaAdaLayerNormSingle(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+        self.emb = SanaCombinedTimestepSizeEmbeddings(embedding_dim)
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
+
+    def forward(self, timestep, hidden_dtype=None):
+        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+        out = self.linear(self.silu(embedded_timestep))
+        return out, embedded_timestep
+
+
+class SanaModulatedNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+
+    def forward(self, x, temb, scale_shift_table):
+        x = self.norm(x)
+        shift, scale = (scale_shift_table[None] + temb[:, None]).chunk(2, dim=1)
+        x = x * (1 + scale) + shift
+        return x
+
+
+class GLUMBConv(nn.Module):
+    """Gated Linear Unit with Multi-Branch Convolution."""
+
+    def __init__(self, in_channels, out_channels, expand_ratio=2.5):
+        super().__init__()
+        hidden_channels = int(expand_ratio * in_channels)
+        self.nonlinearity = nn.SiLU()
+        self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0)
+        self.conv_depth = nn.Conv2d(
+            hidden_channels * 2,
+            hidden_channels * 2,
+            3,
+            1,
+            1,
+            groups=hidden_channels * 2,
+        )
+        self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False)
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv_inverted(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv_depth(hidden_states)
+        hidden_states, gate = torch.chunk(hidden_states, 2, dim=1)
+        hidden_states = hidden_states * self.nonlinearity(gate)
+        hidden_states = self.conv_point(hidden_states)
+        return hidden_states
+
+
+class SanaLinearAttention(nn.Module):
+    """Linear attention with O(N*D^2) complexity instead of O(N^2*D)."""
+
+    def __init__(self, query_dim, num_heads, head_dim, qk_norm_dim, bias=False):
+        super().__init__()
+        inner_dim = num_heads * head_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
+        self.to_k = nn.Linear(query_dim, inner_dim, bias=bias)
+        self.to_v = nn.Linear(query_dim, inner_dim, bias=bias)
+        self.to_out = nn.ModuleList(
+            [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()]
+        )
+        self.norm_q = RMSNorm(qk_norm_dim)
+        self.norm_k = RMSNorm(qk_norm_dim)
+
+    def forward(self, hidden_states):
+        B, S, _ = hidden_states.shape
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        query = F.relu(query)
+        key = F.relu(key)
+
+        kv = torch.matmul(key.transpose(-2, -1), value)  # (B, H, D, D)
+        qkv = torch.matmul(query, kv)  # (B, H, S, D)
+
+        key_sum = key.sum(dim=-2, keepdim=True)  # (B, H, 1, D)
+        normalizer = torch.matmul(query, key_sum.transpose(-2, -1)).clamp(min=1e-6)
+        hidden_states = qkv / normalizer
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1)
+        hidden_states = self.to_out[0](hidden_states)
+        return hidden_states
+
+
+class SanaCrossAttention(nn.Module):
+    def __init__(self, query_dim, cross_attention_dim, num_heads, head_dim, bias=False):
+        super().__init__()
+        inner_dim = num_heads * head_dim
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
+        self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
+        self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias)
+        self.to_out = nn.ModuleList(
+            [nn.Linear(inner_dim, query_dim, bias=True), nn.Identity()]
+        )
+
+        self.norm_q = RMSNorm(inner_dim)
+        self.norm_k = RMSNorm(inner_dim)
+
+    def forward(
+        self, hidden_states, encoder_hidden_states, encoder_attention_mask=None
+    ):
+        B, S, _ = hidden_states.shape
+        T = encoder_hidden_states.shape[1]
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        query = query.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn_mask = None
+        if encoder_attention_mask is not None:
+            attn_mask = encoder_attention_mask.bool()
+            attn_mask = attn_mask[:, None, None, :].expand(B, self.num_heads, S, T)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attn_mask
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(B, S, -1)
+        hidden_states = self.to_out[0](hidden_states)
+        return hidden_states
+
+
+class SanaTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_attention_heads,
+        attention_head_dim,
+        num_cross_attention_heads,
+        cross_attention_head_dim,
+        cross_attention_dim,
+        mlp_ratio,
+        norm_eps,
+        attention_bias=False,
+    ):
+        super().__init__()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps)
+        self.attn1 = SanaLinearAttention(
+            query_dim=dim,
+            num_heads=num_attention_heads,
+            head_dim=attention_head_dim,
+            qk_norm_dim=num_attention_heads * attention_head_dim,
+            bias=attention_bias,
+        )
+
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps)
+        self.attn2 = SanaCrossAttention(
+            query_dim=dim,
+            cross_attention_dim=cross_attention_dim,
+            num_heads=num_cross_attention_heads,
+            head_dim=cross_attention_head_dim,
+            bias=True,
+        )
+
+        self.ff = GLUMBConv(in_channels=dim, out_channels=dim, expand_ratio=mlp_ratio)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states,
+        timestep,
+        height,
+        width,
+        encoder_attention_mask=None,
+    ):
+        batch_size = hidden_states.shape[0]
+
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+        ).chunk(6, dim=1)
+
+        norm_hidden = self.norm1(hidden_states)
+        norm_hidden = norm_hidden * (1 + scale_msa) + shift_msa
+        attn_output = self.attn1(norm_hidden)
+        hidden_states = hidden_states + gate_msa * attn_output
+
+        attn_output = self.attn2(
+            hidden_states, encoder_hidden_states, encoder_attention_mask
+        )
+        hidden_states = hidden_states + attn_output
+
+        norm_hidden = self.norm2(hidden_states)
+        norm_hidden = norm_hidden * (1 + scale_mlp) + shift_mlp
+        norm_hidden = norm_hidden.unflatten(1, (height, width)).permute(0, 3, 1, 2)
+        ff_output = self.ff(norm_hidden)
+        ff_output = ff_output.flatten(2, 3).permute(0, 2, 1)
+        hidden_states = hidden_states + gate_mlp * ff_output
+
+        return hidden_states
+
+
+class SanaTransformer2DModel(CachableDiT, OffloadableDiTMixin):
+
+    _fsdp_shard_conditions = [
+        lambda n, m: isinstance(m, SanaTransformerBlock),
+    ]
+    _compile_conditions = [
+        lambda n, m: isinstance(m, SanaTransformerBlock),
+    ]
+    param_names_mapping = SanaConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = {}
+
+    def __init__(self, config: SanaConfig, hf_config=None, **kwargs):
+        super().__init__(config, hf_config=hf_config or {}, **kwargs)
+
+        arch = config.arch_config
+        self.out_channels = arch.out_channels
+        self.patch_size = arch.patch_size
+        self.inner_dim = arch.num_attention_heads * arch.attention_head_dim
+
+        self.hidden_size = self.inner_dim
+        self.num_attention_heads = arch.num_attention_heads
+        self.num_channels_latents = arch.num_channels_latents
+
+        self.patch_embed = nn.ModuleDict(
+            {
+                "proj": nn.Conv2d(
+                    arch.in_channels,
+                    self.inner_dim,
+                    kernel_size=arch.patch_size,
+                    stride=arch.patch_size,
+                    bias=True,
+                ),
+            }
+        )
+        self.time_embed = SanaAdaLayerNormSingle(self.inner_dim)
+        self.caption_projection = PixArtAlphaTextProjection(
+            in_features=arch.caption_channels,
+            hidden_size=self.inner_dim,
+        )
+
+        self.caption_norm = RMSNorm(self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                SanaTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=arch.num_attention_heads,
+                    attention_head_dim=arch.attention_head_dim,
+                    num_cross_attention_heads=arch.num_cross_attention_heads,
+                    cross_attention_head_dim=arch.cross_attention_head_dim,
+                    cross_attention_dim=arch.cross_attention_dim,
+                    mlp_ratio=arch.mlp_ratio,
+                    norm_eps=arch.norm_eps,
+                    attention_bias=False,
+                )
+                for _ in range(arch.num_layers)
+            ]
+        )
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, self.inner_dim) / self.inner_dim**0.5
+        )
+
+        self.norm_out = SanaModulatedNorm(self.inner_dim, eps=arch.norm_eps)
+
+        self.proj_out = nn.Linear(
+            self.inner_dim,
+            arch.patch_size * arch.patch_size * self.out_channels,
+            bias=True,
+        )
+
+        self.layer_names = ["transformer_blocks"]
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        guidance: torch.Tensor = None,
+        encoder_attention_mask: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+
+        # Input validation - fail fast
+        if encoder_hidden_states is None:
+            raise ValueError("SANA forward pass requires encoder_hidden_states")
+
+        batch_size, channels, height, width = hidden_states.shape
+        p = self.patch_size
+        post_patch_height = height // p
+        post_patch_width = width // p
+
+        hidden_states = self.patch_embed["proj"](hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+
+        timestep_emb, embedded_timestep = self.time_embed(
+            timestep, hidden_dtype=hidden_states.dtype
+        )
+
+        if isinstance(encoder_attention_mask, (list, tuple)):
+            encoder_attention_mask = encoder_attention_mask[0]
+
+        encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+        if encoder_hidden_states.shape[0] != batch_size:
+            encoder_hidden_states = encoder_hidden_states.expand(
+                batch_size, -1, -1
+            ).contiguous()
+        encoder_hidden_states = encoder_hidden_states.view(
+            batch_size, -1, hidden_states.shape[-1]
+        )
+        encoder_hidden_states = self.caption_norm(encoder_hidden_states)
+
+        if (
+            encoder_attention_mask is not None
+            and encoder_attention_mask.shape[0] != batch_size
+        ):
+            encoder_attention_mask = encoder_attention_mask.expand(
+                batch_size, -1
+            ).contiguous()
+
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                timestep_emb,
+                post_patch_height,
+                post_patch_width,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+        hidden_states = self.norm_out(
+            hidden_states, embedded_timestep, self.scale_shift_table
+        )
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states.reshape(
+            batch_size, post_patch_height, post_patch_width, p, p, self.out_channels
+        )
+        hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
+        hidden_states = hidden_states.reshape(
+            batch_size, self.out_channels, height, width
+        )
+
+        return hidden_states
+
+
+EntryClass = SanaTransformer2DModel

python/sglang/multimodal_gen/runtime/models/encoders/gemma2.py

@@ -0,0 +1,451 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Gemma2 2B text encoder for SANA.
+#
+# This is a decoder-only language model used as a text encoder: we feed
+# in tokenized text and extract the final hidden states (not logits) as
+# the conditioning signal for SANA's cross-attention layers.
+#
+# Architecture follows google/gemma-2-2b-it:
+#   - 26 layers, alternating global / sliding-window attention
+#   - GQA with 8 query heads, 4 KV heads, head_dim=256
+#   - Pre/post attention + pre/post feedforward LayerNorm (Gemma2-style)
+#   - GeGLU activation (gelu_pytorch_tanh)
+#
+# Adapted from the Gemma3 text model implementation in this codebase.
+
+import logging
+from typing import Any, Iterable
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.gemma2 import Gemma2Config
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul
+from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+
+logger = logging.getLogger(__name__)
+
+
+class Gemma2RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class Gemma2MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "gelu_pytorch_tanh":
+            raise ValueError(
+                "Gemma2 uses `gelu_pytorch_tanh` as the hidden activation. "
+                f"Got: {hidden_act}"
+            )
+        self.act_fn = GeluAndMul(approximate="tanh")
+
+    def forward(self, x):
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Gemma2Attention(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        config: Gemma2Config,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.layer_id = layer_id
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        arch = config.arch_config
+        self.head_dim = arch.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = arch.query_pre_attn_scalar**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=arch.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=arch.attention_bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # Gemma2 interleaves global (even layers) and sliding-window (odd layers)
+        # attention. This pattern reduces memory for long sequences while
+        # maintaining global context every other layer.
+        self.is_sliding = (layer_id % 2) == 1
+        if self.is_sliding:
+            self.sliding_window = arch.sliding_window
+        else:
+            self.sliding_window = None
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.head_dim,
+            max_position=arch.max_position_embeddings,
+            base=arch.rope_theta,
+            is_neox_style=True,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        batch_size, seq_len, _ = q.shape
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        query = q.transpose(1, 2)
+        key = k.transpose(1, 2)
+        value = v.transpose(1, 2)
+
+        attn_mask = torch.zeros(
+            (seq_len, seq_len), device=hidden_states.device, dtype=torch.float32
+        )
+        causal = torch.triu(
+            torch.ones(
+                (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool
+            ),
+            diagonal=1,
+        )
+        attn_mask = attn_mask.masked_fill(causal, float("-inf"))
+        if self.is_sliding and self.sliding_window is not None:
+            idx = torch.arange(seq_len, device=hidden_states.device)
+            dist = idx[None, :] - idx[:, None]
+            too_far = dist > self.sliding_window
+            attn_mask = attn_mask.masked_fill(too_far, float("-inf"))
+
+        if attention_mask is not None:
+            key_pad = ~attention_mask.to(torch.bool)
+            attn_mask = attn_mask[None, None, :, :].expand(
+                batch_size, 1, seq_len, seq_len
+            )
+            attn_mask = attn_mask.masked_fill(
+                key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len),
+                float("-inf"),
+            )
+
+        attn_kwargs = {
+            "attn_mask": attn_mask,
+            "dropout_p": 0.0,
+            "is_causal": False,
+            "scale": self.scaling,
+        }
+        if query.shape[1] != key.shape[1]:
+            attn_kwargs["enable_gqa"] = True
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query, key, value, **attn_kwargs
+        )
+
+        # NOTE: Gemma2 specifies attn_logit_softcapping (tanh(logits/cap)*cap) but
+        # PyTorch's scaled_dot_product_attention does not support it natively.
+        # For short text-encoder sequences (~300 tokens), the quality impact is
+        # negligible. A custom attention kernel would be needed for full fidelity.
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(
+            batch_size, seq_len, self.num_heads * self.head_dim
+        )
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Gemma2DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        config: Gemma2Config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        arch = config.arch_config
+        self.hidden_size = arch.hidden_size
+        self.self_attn = Gemma2Attention(
+            layer_id=layer_id,
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=arch.num_attention_heads,
+            num_kv_heads=arch.num_key_value_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = Gemma2MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=arch.intermediate_size,
+            hidden_act=arch.hidden_activation,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = Gemma2RMSNorm(self.hidden_size, eps=arch.rms_norm_eps)
+        self.post_attention_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+        self.pre_feedforward_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = Gemma2RMSNorm(
+            self.hidden_size, eps=arch.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(positions, hidden_states, attention_mask)
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class Gemma2Model(nn.Module):
+    """Gemma2 text encoder model for SANA pipeline."""
+
+    _fsdp_shard_conditions = []
+
+    def __init__(self, config: Gemma2Config, **kwargs):
+        super().__init__()
+        self.config = config
+        arch = config.arch_config
+        self.quant_config = None
+
+        self.vocab_size = arch.vocab_size
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            arch.hidden_size,
+            org_num_embeddings=arch.vocab_size,
+            quant_config=self.quant_config,
+        )
+        self.embed_scale = arch.hidden_size**0.5
+
+        self.layers = nn.ModuleList(
+            [
+                Gemma2DecoderLayer(
+                    layer_id=i,
+                    config=config,
+                    quant_config=self.quant_config,
+                    prefix=f"model.layers.{i}",
+                )
+                for i in range(arch.num_hidden_layers)
+            ]
+        )
+
+        self.norm = Gemma2RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids) * self.embed_scale
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs,
+    ) -> BaseEncoderOutput:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else getattr(self.config.arch_config, "output_hidden_states", False)
+        )
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            hidden_states = self.get_input_embeddings(input_ids)
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1], device=hidden_states.device
+            ).unsqueeze(0)
+
+        all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (hidden_states,)
+
+            hidden_states = layer(position_ids, hidden_states, attention_mask)
+
+        hidden_states = self.norm(hidden_states)
+
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        stacked_params_mapping = getattr(
+            self.config.arch_config, "stacked_params_mapping", None
+        )
+        if stacked_params_mapping is None:
+            stacked_params_mapping = [
+                (".qkv_proj", ".q_proj", "q"),
+                (".qkv_proj", ".k_proj", "k"),
+                (".qkv_proj", ".v_proj", "v"),
+                (".gate_up_proj", ".gate_proj", "0"),
+                (".gate_up_proj", ".up_proj", "1"),
+            ]
+
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            # HF Gemma2Model stores weights as model.layers.X... / model.embed_tokens...
+            # Strip "model." prefix if present to match our naming
+            if name.startswith("model."):
+                name = name[len("model.") :]
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                self._load_with_shard_id(weight_loader, param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+        return loaded_params
+
+    @staticmethod
+    def _load_with_shard_id(weight_loader, param, loaded_weight, shard_id):
+        try:
+            weight_loader(param, loaded_weight, shard_id)
+            return
+        except (AssertionError, TypeError):
+            pass
+
+        if isinstance(shard_id, str):
+            mapping = {"q": 0, "k": 1, "v": 2}
+            if shard_id in mapping:
+                weight_loader(param, loaded_weight, mapping[shard_id])
+                return
+            if shard_id.isdigit():
+                weight_loader(param, loaded_weight, int(shard_id))
+                return
+        elif isinstance(shard_id, int):
+            mapping = {0: "q", 1: "k", 2: "v"}
+            if shard_id in mapping:
+                weight_loader(param, loaded_weight, mapping[shard_id])
+                return
+
+        raise TypeError(
+            f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader}"
+        )
+
+
+EntryClass = Gemma2Model

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_dpm_solver_multistep.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# DPM-Solver++ multistep scheduler wrapper for SANA.
+#
+# SANA uses DPM-Solver++ (Lu et al., 2022) as its noise scheduler, which
+# is a high-order ODE solver that converges in fewer steps than DDIM.
+# With solver_order=2 and 20 steps, SANA achieves high-quality results.
+#
+# This wrapper delegates all numerical work to diffusers' implementation
+# and only adapts the interface for sglang's denoising stage.
+
+import torch
+from diffusers import (
+    DPMSolverMultistepScheduler as DiffusersDPMSolverMultistepScheduler,
+)
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+
+
+class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
+    """DPM-Solver++ multistep scheduler wrapper for sglang's BaseScheduler interface."""
+
+    order = 1
+    num_train_timesteps = 1000
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "scaled_linear",
+        trained_betas=None,
+        solver_order: int = 2,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        use_karras_sigmas: bool = False,
+        use_lu_lambdas: bool = False,
+        use_exponential_sigmas: bool = False,
+        use_beta_sigmas: bool = False,
+        use_flow_sigmas: bool = False,
+        final_sigmas_type: str = "zero",
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: str | None = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        rescale_betas_zero_snr: bool = False,
+        flow_shift: float | None = None,
+        **kwargs,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self._inner = DiffusersDPMSolverMultistepScheduler(
+            num_train_timesteps=num_train_timesteps,
+            beta_start=beta_start,
+            beta_end=beta_end,
+            beta_schedule=beta_schedule,
+            trained_betas=trained_betas,
+            solver_order=solver_order,
+            prediction_type=prediction_type,
+            thresholding=thresholding,
+            dynamic_thresholding_ratio=dynamic_thresholding_ratio,
+            sample_max_value=sample_max_value,
+            algorithm_type=algorithm_type,
+            solver_type=solver_type,
+            lower_order_final=lower_order_final,
+            euler_at_final=euler_at_final,
+            use_karras_sigmas=use_karras_sigmas,
+            use_lu_lambdas=use_lu_lambdas,
+            use_exponential_sigmas=use_exponential_sigmas,
+            use_beta_sigmas=use_beta_sigmas,
+            use_flow_sigmas=use_flow_sigmas,
+            flow_shift=flow_shift,
+            final_sigmas_type=final_sigmas_type,
+            lambda_min_clipped=lambda_min_clipped,
+            variance_type=variance_type,
+            timestep_spacing=timestep_spacing,
+            steps_offset=steps_offset,
+            rescale_betas_zero_snr=rescale_betas_zero_snr,
+        )
+        self.timesteps = self._inner.timesteps
+        self.order = solver_order
+        self._flow_shift = flow_shift
+        self._begin_index: int | None = None
+        BaseScheduler.__init__(self)
+
+    def set_shift(self, shift: float) -> None:
+        self._flow_shift = shift
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = begin_index
+
+    @property
+    def begin_index(self) -> int | None:
+        return self._begin_index
+
+    def set_timesteps(self, num_inference_steps: int, device=None, **kwargs):
+        self._inner.set_timesteps(num_inference_steps, device=device, **kwargs)
+        self.timesteps = self._inner.timesteps
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        return self._inner.scale_model_input(sample, timestep)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        **kwargs,
+    ):
+        return self._inner.step(model_output, timestep, sample, **kwargs)
+
+    @property
+    def sigmas(self):
+        return getattr(self._inner, "sigmas", None)
+
+    @property
+    def init_noise_sigma(self):
+        return self._inner.init_noise_sigma
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        return self._inner.add_noise(original_samples, noise, timesteps)
+
+
+EntryClass = DPMSolverMultistepScheduler

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_dc.py

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.sana import SanaVAEConfig
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class AutoencoderDC(nn.Module):
+    """Deep Compression Autoencoder wrapper with 32x spatial compression."""
+
+    def __init__(self, config: SanaVAEConfig = None, **kwargs):
+        super().__init__()
+        self._config = config
+        self._inner_model = None
+        self._loaded_state_dict: dict[str, torch.Tensor] = {}
+
+    def _ensure_inner_model(self, state_dict: dict[str, torch.Tensor] | None = None):
+        if self._inner_model is not None:
+            return
+
+        from diffusers import AutoencoderDC as DiffusersAutoencoderDC
+
+        device = "cpu"
+        state_to_load = (
+            state_dict if state_dict is not None else self._loaded_state_dict
+        )
+        if state_to_load:
+            first_tensor = next(iter(state_to_load.values()))
+            device = first_tensor.device
+        hf_config = {}
+        if self._config is not None:
+            arch = self._config.arch_config
+            for key, value in vars(arch).items():
+                if key == "extra_attrs" and isinstance(value, dict):
+                    for ek, ev in value.items():
+                        hf_config[ek] = ev
+                elif not key.startswith("_") and not callable(value):
+                    hf_config[key] = value
+
+        self._inner_model = DiffusersAutoencoderDC.from_config(hf_config)
+
+        if state_to_load:
+            missing, unexpected = self._inner_model.load_state_dict(
+                state_to_load, strict=False
+            )
+            if missing:
+                logger.warning(
+                    "AutoencoderDC missing keys when loading: %d keys", len(missing)
+                )
+                if len(missing) > 10:
+                    logger.debug("First 10 missing keys: %s", list(missing)[:10])
+                else:
+                    logger.debug("Missing keys: %s", list(missing))
+            if unexpected:
+                logger.debug(
+                    "AutoencoderDC unexpected keys when loading: %d keys",
+                    len(unexpected),
+                )
+            if state_dict is None:
+                self._loaded_state_dict.clear()
+
+        self._inner_model = self._inner_model.to(device)
+
+    @property
+    def config(self):
+        if self._inner_model is not None:
+            return self._inner_model.config
+        return self._config
+
+    @property
+    def dtype(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).dtype
+        return torch.float32
+
+    @property
+    def device(self):
+        if self._inner_model is not None:
+            return next(self._inner_model.parameters()).device
+        return torch.device("cpu")
+
+    def encode(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model.encode(x, **kwargs)
+
+    def decode(self, z: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        z = z.to(dtype=self.dtype)
+        return self._inner_model.decode(z, **kwargs)
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        self._ensure_inner_model()
+        return self._inner_model(x, **kwargs)
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        strict: bool = True,
+        assign: bool = False,
+    ):
+        """Intercept load_state_dict to route weights into the inner diffusers model."""
+        self._ensure_inner_model(state_dict=state_dict)
+
+    def state_dict(self, *args, **kwargs) -> dict[str, torch.Tensor]:
+        self._ensure_inner_model()
+        return self._inner_model.state_dict(*args, **kwargs)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Buffer weights for deferred loading. The inner model is built lazily."""
+        loaded_params: set[str] = set()
+        for name, weight in weights:
+            self._loaded_state_dict[name] = weight
+            loaded_params.add(name)
+        return loaded_params
+
+    def to(self, *args, **kwargs):
+        if self._inner_model is not None:
+            self._inner_model = self._inner_model.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+
+EntryClass = AutoencoderDC

python/sglang/multimodal_gen/runtime/pipelines/sana.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# SANA text-to-image pipeline.
+#
+# Stage order matches Flux (InputValidation -> TextEncoding -> TimestepPrep ->
+# LatentPrep -> Denoising -> Decoding) rather than the add_standard_t2i_stages
+# helper (which puts LatentPrep before TimestepPrep). Both orderings are
+# functionally equivalent since these stages are independent.
+#
+# SANA uses a single text encoder (Gemma2), so only one text_encoder + tokenizer
+# pair is registered — unlike Flux which has text_encoder + text_encoder_2.
+# The pipeline_name must match the _class_name in HF model_index.json.
+
+from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    InputValidationStage,
+    TextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class SanaPipeline(LoRAPipeline, ComposedPipelineBase):
+    pipeline_name = "SanaPipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        self.add_stage(InputValidationStage())
+
+        self.add_stage(
+            TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+            "prompt_encoding_stage_primary",
+        )
+
+        self.add_standard_timestep_preparation_stage()
+        self.add_standard_latent_preparation_stage()
+        self.add_standard_denoising_stage()
+        self.add_standard_decoding_stage()
+
+
+EntryClass = SanaPipeline

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py

@@ -118,10 +118,9 @@ class DecodingStage(PipelineStage):
             Decoded video tensor with shape (batch, channels, frames, height, width),
             normalized to [0, 1] range and moved to CPU as float32
         """
-        self.vae = self.vae.to(get_local_torch_device())
-        latents = latents.to(get_local_torch_device())
-        # Setup VAE precision
         vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        self.vae = self.vae.to(device=get_local_torch_device(), dtype=vae_dtype)
+        latents = latents.to(get_local_torch_device())
         vae_autocast_enabled = (
             vae_dtype != torch.float32
         ) and not server_args.disable_autocast

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -459,6 +459,15 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
         ),
         T2I_sampling_params,
     ),
+    DiffusionTestCase(
+        "sana_image_t2i",
+        DiffusionServerArgs(
+            model_path="Efficient-Large-Model/Sana_600M_1024px_diffusers",
+            modality="image",
+        ),
+        T2I_sampling_params,
+        run_perf_check=False,
+    ),
     # === Text and Image to Image (TI2I) ===
     DiffusionTestCase(
         "qwen_image_edit_ti2i",