[diffusion] chore: use native ernie prompt enhancer (#34951)

docs/cookbook/diffusion/Ernie-Image/Ernie-Image.mdx

@@ -77,5 +77,5 @@ with open("ernie_image.png", "wb") as f:
 
 - ERNIE-Image is a text-to-image pipeline; do not pass `--image-path`.
 - `--performance-mode auto` keeps conservative defaults while preserving explicit user flags.
-- If the checkpoint includes a PE component, SGLang loads it automatically from `model_index.json`.
+- If the checkpoint includes a PE component, SGLang loads it automatically with the native Ministral3 runtime. Use `--layerwise-offload-components pe` when the local PE decoder needs to trade latency for lower GPU memory usage.
 - Treat FSDP, SP/Ulysses/Ring, and TP as explicit benchmark knobs. Measure the target resolution, step count, and GPU type before making them production defaults.

docs/docs/sglang-diffusion/models_with_pe.mdx

@@ -5,9 +5,9 @@ description: "Run ERNIE-Image with built-in prompt enhancement or a separate SGL
 
 ## Quick Start
 
-By default, the PE model is loaded by the diffusion model server, which may not provide optimal performance. For higher performance, the PE model can be deployed as a separate SGLang server. This document uses `baidu/ERNIE-Image` as an example.
+By default, the diffusion server loads the PE model in-process with SGLang's native Ministral3 implementation. Deploy the PE model as a separate SGLang server when it needs independent resources or scaling. This document uses `baidu/ERNIE-Image` as an example.
 
-Run the model with the built-in Transformers PE implementation (default):
+Run the model with the built-in native PE implementation (default):
 
 ```bash
 # Terminal 1: launch server
@@ -27,7 +27,7 @@ curl -X POST http://${HOST}:${PORT}/v1/images/generations \
   }'
 ```
 
-Run the model with an SGLang-served PE model (high performance):
+Run the model with an SGLang-served PE model (high performance):
 
 ```bash
 # Terminal 1: launch SGLang PE model server
@@ -51,9 +51,21 @@ curl -X POST http://${HOST}:${PORT}/v1/images/generations \
     "guidance_scale": 4.0,
     "use_pe": true
   }'
-```
-
-## Support matrix
+```
+
+For a memory-constrained in-process deployment, the native PE decoder supports
+layerwise offload:
+
+```bash
+sglang serve --model-path baidu/ERNIE-Image \
+  --layerwise-offload-components pe \
+  --port ${PORT}
+```
+
+This option streams PE decoder layers from CPU and can increase prompt-enhancement
+latency. It does not apply when `--pe-server-url` selects an external PE server.
+
+## Support matrix
 
 | Model | Built-in PE | SGLang PE Server |
 |-------|-------------|------------------|

python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py

@@ -79,10 +79,24 @@ class SDPAImpl(AttentionImpl):
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)
         value = value.transpose(1, 2)
+        attn_mask = None
+        is_causal = self.causal
+        if self.causal and query.shape[-2] != key.shape[-2]:
+            is_causal = False
+            if query.shape[-2] > 1:
+                query_length = query.shape[-2]
+                key_length = key.shape[-2]
+                attn_mask = torch.ones(
+                    query_length,
+                    key_length,
+                    dtype=torch.bool,
+                    device=query.device,
+                ).tril(diagonal=key_length - query_length)
+
         attn_kwargs = {
-            "attn_mask": None,
+            "attn_mask": attn_mask,
             "dropout_p": self.dropout,
-            "is_causal": self.causal,
+            "is_causal": is_causal,
             "scale": self.softmax_scale,
         }
         if query.shape[1] != key.shape[1]:

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

@@ -4,12 +4,20 @@ import os
 
 import requests
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch import nn
+from transformers import AutoTokenizer
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     ComponentLoader,
 )
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
+    LayerwiseOffloadableModuleMixin,
+)
+from sglang.multimodal_gen.runtime.models.encoders.mistral_3 import (
+    Ministral3ForCausalLM,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -33,9 +41,12 @@ def _read_model_max_length(model_path: str) -> int | None:
     return None
 
 
-class PEModelWrapper:
+class PEModelWrapper(nn.Module, LayerwiseOffloadableModuleMixin):
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["model.model.layers"]
 
     def __init__(self, model, tokenizer, device, model_max_length: int):
+        super().__init__()
         self.model = model
         self.pe_tokenizer = tokenizer
         self.device = device
@@ -64,7 +75,8 @@ class PEModelWrapper:
             generate_kwargs["top_p"] = top_p
 
         with torch.no_grad():
-            output_ids = self.model.generate(**generate_kwargs)
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                output_ids = self.model.generate(**generate_kwargs)
 
         new_tokens = output_ids[0, input_len:]
         text = self.pe_tokenizer.decode(new_tokens, skip_special_tokens=True)
@@ -72,11 +84,10 @@ class PEModelWrapper:
 
     def to(self, *args, **kwargs):
         """Move underlying model to device."""
-        self.model = self.model.to(*args, **kwargs)
-        if args:
-            device = args[0]
-            if isinstance(device, (str, torch.device)):
-                self.device = torch.device(device)
+        super().to(*args, **kwargs)
+        device = args[0] if args else kwargs.get("device")
+        if isinstance(device, (str, torch.device)):
+            self.device = torch.device(device)
         return self
 
 
@@ -157,33 +168,19 @@ class PELoader(ComponentLoader):
         if tokenizer.pad_token_id is None:
             tokenizer.pad_token_id = tokenizer.eos_token_id
 
-        attn_impl = "flash_attention_2"
-        try:
-            model = AutoModelForCausalLM.from_pretrained(
-                component_model_path,
-                torch_dtype=torch.bfloat16,
-                trust_remote_code=server_args.trust_remote_code,
-                attn_implementation=attn_impl,
-            )
-            logger.info("PE model: using Flash Attention 2")
-        except (ValueError, ImportError):
-            logger.warning("Flash Attention 2 not available, falling back to SDPA")
-            attn_impl = "sdpa"
-            model = AutoModelForCausalLM.from_pretrained(
-                component_model_path,
-                torch_dtype=torch.bfloat16,
-                trust_remote_code=server_args.trust_remote_code,
-                attn_implementation=attn_impl,
-            )
+        model = Ministral3ForCausalLM.from_pretrained(
+            component_model_path,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=server_args.trust_remote_code,
+        )
 
         device = get_local_torch_device()
         model = model.to(device).eval()
 
         logger.info(
-            "PE model loaded on %s: %s (attn=%s)",
+            "PE model loaded on %s: %s",
             device,
             model.__class__.__name__,
-            attn_impl,
         )
 
         return PEModelWrapper(

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

@@ -20,13 +20,27 @@ from typing import Iterable, Optional, Union
 import torch
 from torch import nn
 from torch.nn.attention import SDPBackend, sdpa_kernel
-from transformers import Cache, DynamicCache, LlavaConfig, Mistral3Config, MistralConfig
+from transformers import (
+    Cache,
+    DynamicCache,
+    GenerationMixin,
+    LlavaConfig,
+    Ministral3Config,
+    Mistral3Config,
+    MistralConfig,
+)
 from transformers.activations import ACT2FN
 from transformers.masking_utils import (
     create_causal_mask,
     create_sliding_window_causal_mask,
 )
-from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.models.ministral3.modeling_ministral3 import (
+    Ministral3PreTrainedModel,
+)
 from transformers.models.mistral3.modeling_mistral3 import (
     Mistral3CausalLMOutputWithPast,
     Mistral3ModelOutputWithPast,
@@ -126,10 +140,29 @@ def _can_use_unmasked_causal_attention(
     return bool(torch.all(attention_mask > 0).item())
 
 
+def _get_llama_4_attn_scale(
+    position_ids: torch.Tensor,
+    beta: float | None,
+    original_max_position_embeddings: int | None,
+) -> torch.Tensor | None:
+    if beta is None or original_max_position_embeddings is None:
+        return None
+    scale = 1 + beta * torch.log(
+        1 + torch.floor(position_ids / original_max_position_embeddings)
+    )
+    return scale[:, None, :, None]
+
+
 class MistralAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: MistralConfig, layer_idx: int):
+    def __init__(
+        self,
+        config: MistralConfig,
+        layer_idx: int,
+        *,
+        allow_cudnn_sdp: bool = True,
+    ):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -139,6 +172,11 @@ class MistralAttention(nn.Module):
         )
         self.scaling = self.head_dim**-0.5
         self.attention_dropout = config.attention_dropout
+        rope_parameters = getattr(config, "rope_parameters", {}) or {}
+        self.llama_4_scaling_beta = rope_parameters.get("llama_4_scaling_beta")
+        self.original_max_position_embeddings = rope_parameters.get(
+            "original_max_position_embeddings"
+        )
         self.total_num_heads = config.num_attention_heads
         self.total_num_key_value_heads = config.num_key_value_heads
         tp_size = _tp_world_size()
@@ -197,7 +235,7 @@ class MistralAttention(nn.Module):
                 AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
             },
-            allow_cudnn_sdp=True,
+            allow_cudnn_sdp=allow_cudnn_sdp,
         )
 
     def forward(
@@ -205,6 +243,7 @@ class MistralAttention(nn.Module):
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: Optional[torch.Tensor],
+        position_ids: Optional[torch.LongTensor] = None,
         past_key_values: Optional[Cache] = None,
         cache_position: Optional[torch.LongTensor] = None,
         **kwargs,
@@ -232,6 +271,14 @@ class MistralAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(
             query_states, key_states, cos, sin
         )
+        if position_ids is not None:
+            query_scale = _get_llama_4_attn_scale(
+                position_ids,
+                self.llama_4_scaling_beta,
+                self.original_max_position_embeddings,
+            )
+            if query_scale is not None:
+                query_states = query_states * query_scale.to(query_states.dtype)
 
         if past_key_values is not None:
             # sin and cos are specific to RoPE models; cache_position needed for the static cache
@@ -284,10 +331,20 @@ class MistralTPMLP(nn.Module):
 
 
 class MistralDecoderLayer(nn.Module):
-    def __init__(self, config: MistralConfig, layer_idx: int):
+    def __init__(
+        self,
+        config: MistralConfig,
+        layer_idx: int,
+        *,
+        allow_cudnn_sdp: bool = True,
+    ):
         super().__init__()
         self.hidden_size = config.hidden_size
-        self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)
+        self.self_attn = MistralAttention(
+            config=config,
+            layer_idx=layer_idx,
+            allow_cudnn_sdp=allow_cudnn_sdp,
+        )
         self.mlp = MistralTPMLP(config)
         self.input_layernorm = MistralRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
@@ -333,7 +390,7 @@ class MistralDecoderLayer(nn.Module):
 
 
 class MistralModel(MistralPreTrainedModel):
-    def __init__(self, config: MistralConfig):
+    def __init__(self, config: MistralConfig, *, allow_cudnn_sdp: bool = True):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -343,7 +400,11 @@ class MistralModel(MistralPreTrainedModel):
         )
         self.layers = nn.ModuleList(
             [
-                MistralDecoderLayer(config, layer_idx)
+                MistralDecoderLayer(
+                    config,
+                    layer_idx,
+                    allow_cudnn_sdp=allow_cudnn_sdp,
+                )
                 for layer_idx in range(config.num_hidden_layers)
             ]
         )
@@ -435,6 +496,78 @@ class MistralModel(MistralPreTrainedModel):
         )
 
 
+class Ministral3ForCausalLM(
+    Ministral3PreTrainedModel, GenerationMixin, LayerwiseOffloadableModuleMixin
+):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    layerwise_offload_dit_group_enabled = False
+    layer_names = ["model.layers"]
+
+    def __init__(self, config: Ministral3Config):
+        super().__init__(config)
+        self.model = MistralModel(config, allow_cudnn_sdp=False)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(outputs.last_hidden_state[:, slice_indices, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
 class Mistral3Model(nn.Module):
     _checkpoint_conversion_mapping = {"language_model.model": "language_model"}
 

python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py

@@ -0,0 +1,135 @@
+import pytest
+import torch
+import torch.nn.functional as F
+from transformers import Ministral3Config
+from transformers.models.ministral3.modeling_ministral3 import (
+    Ministral3ForCausalLM as HFMinistral3ForCausalLM,
+)
+
+from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPAImpl
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    global_force_attn_backend_context_manager,
+)
+from sglang.multimodal_gen.runtime.loader.component_loaders.pe_loader import (
+    PEModelWrapper,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.models.encoders.mistral_3 import (
+    Ministral3ForCausalLM,
+    _get_llama_4_attn_scale,
+)
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+def _config() -> Ministral3Config:
+    return Ministral3Config(
+        vocab_size=64,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        max_position_embeddings=128,
+        rope_parameters={
+            "rope_type": "yarn",
+            "rope_theta": 1_000_000.0,
+            "factor": 4.0,
+            "beta_fast": 32.0,
+            "beta_slow": 1.0,
+            "mscale": 1.0,
+            "mscale_all_dim": 1.0,
+            "llama_4_scaling_beta": 0.1,
+            "original_max_position_embeddings": 16,
+        },
+        tie_word_embeddings=True,
+        use_cache=True,
+    )
+
+
+def test_ministral3_query_scale_matches_llama4_rule():
+    position_ids = torch.tensor([[0, 15, 16, 32]])
+
+    scale = _get_llama_4_attn_scale(position_ids, 0.1, 16)
+
+    expected = 1 + 0.1 * torch.log(1 + torch.floor(position_ids.float() / 16))
+    torch.testing.assert_close(scale[:, 0, :, 0], expected)
+
+
+@pytest.mark.parametrize("query_length", [1, 2])
+def test_causal_sdpa_uses_lower_right_alignment_for_cached_keys(query_length):
+    torch.manual_seed(11)
+    key_length = 5
+    query = torch.randn(1, query_length, 2, 8)
+    key = torch.randn(1, key_length, 2, 8)
+    value = torch.randn(1, key_length, 2, 8)
+    attention = SDPAImpl(
+        num_heads=2,
+        head_size=8,
+        causal=True,
+        softmax_scale=8**-0.5,
+    )
+
+    actual = attention.forward(query, key, value, attn_metadata=None)
+
+    mask = None
+    if query_length > 1:
+        mask = torch.ones(query_length, key_length, dtype=torch.bool).tril(
+            diagonal=key_length - query_length
+        )
+    expected = F.scaled_dot_product_attention(
+        query.transpose(1, 2),
+        key.transpose(1, 2),
+        value.transpose(1, 2),
+        attn_mask=mask,
+        is_causal=False,
+        scale=8**-0.5,
+    ).transpose(1, 2)
+    torch.testing.assert_close(actual, expected)
+
+
+def test_native_ministral3_matches_hf_prefill_and_generation():
+    torch.manual_seed(7)
+    with global_force_attn_backend_context_manager(AttentionBackendEnum.TORCH_SDPA):
+        config = _config()
+        reference = HFMinistral3ForCausalLM(config).eval()
+        native = Ministral3ForCausalLM(config).eval()
+        native.load_state_dict(reference.state_dict(), strict=True)
+        input_ids = (torch.arange(20).unsqueeze(0) + 1) % config.vocab_size
+
+        context = set_forward_context(current_timestep=0, attn_metadata=None)
+        with torch.no_grad(), context:
+            reference_output = reference(input_ids=input_ids, use_cache=True)
+            native_output = native(input_ids=input_ids, use_cache=True)
+
+        torch.testing.assert_close(native_output.logits, reference_output.logits)
+        assert len(native_output.past_key_values.layers) == config.num_hidden_layers
+
+        with torch.no_grad(), set_forward_context(
+            current_timestep=0, attn_metadata=None
+        ):
+            native_ids = native.generate(input_ids, max_new_tokens=2, do_sample=False)
+        with torch.no_grad():
+            reference_ids = reference.generate(
+                input_ids, max_new_tokens=2, do_sample=False
+            )
+        torch.testing.assert_close(native_ids, reference_ids)
+
+
+def test_native_ministral3_exposes_decoder_layers_for_offload():
+    assert Ministral3ForCausalLM.layer_names == ["model.layers"]
+
+
+def test_pe_wrapper_exposes_native_decoder_layers_for_offload():
+    with global_force_attn_backend_context_manager(AttentionBackendEnum.TORCH_SDPA):
+        model = Ministral3ForCausalLM(_config())
+    wrapper = PEModelWrapper(
+        model=model,
+        tokenizer=None,
+        device=torch.device("cpu"),
+        model_max_length=128,
+    )
+
+    assert PEModelWrapper.layer_names == ["model.model.layers"]
+    decoder_layers = dict(wrapper.named_modules())["model.model.layers"]
+    assert isinstance(decoder_layers, torch.nn.ModuleList)