From d106e8b23af9245bb40fe15c433384803d42ac82 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 16 Aug 2026 09:59:51 +0800 Subject: [PATCH] [diffusion] chore: use native ernie prompt enhancer (#34951) --- .../diffusion/Ernie-Image/Ernie-Image.mdx | 2 +- docs/docs/sglang-diffusion/models_with_pe.mdx | 24 ++- .../runtime/layers/attention/backends/sdpa.py | 18 ++- .../loader/component_loaders/pe_loader.py | 53 +++---- .../runtime/models/encoders/mistral_3.py | 149 +++++++++++++++++- .../test/unit/test_ministral3_generation.py | 135 ++++++++++++++++ 6 files changed, 336 insertions(+), 45 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py diff --git a/docs/cookbook/diffusion/Ernie-Image/Ernie-Image.mdx b/docs/cookbook/diffusion/Ernie-Image/Ernie-Image.mdx index 34d41285f..0d6df4e51 100644 --- a/docs/cookbook/diffusion/Ernie-Image/Ernie-Image.mdx +++ b/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. diff --git a/docs/docs/sglang-diffusion/models_with_pe.mdx b/docs/docs/sglang-diffusion/models_with_pe.mdx index 9de6c99ca..777228f83 100644 --- a/docs/docs/sglang-diffusion/models_with_pe.mdx +++ b/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 | |-------|-------------|------------------| diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 77807ad68..693dadc62 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/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]: diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py index 699436774..54fc62f82 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py index 9e2f74fb0..4fa0b304a 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py +++ b/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"} diff --git a/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py b/python/sglang/multimodal_gen/test/unit/test_ministral3_generation.py new file mode 100644 index 000000000..9922743a8 --- /dev/null +++ b/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)