sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/encoders/mistral_3.py

# coding=utf-8
# Copyright 2025 HuggingFace Inc. team. All rights reserved.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
from contextlib import nullcontext
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.activations import ACT2FN
from transformers.masking_utils import (
    create_causal_mask,
    create_sliding_window_causal_mask,
)
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.mistral3.modeling_mistral3 import (
    Mistral3CausalLMOutputWithPast,
    Mistral3ModelOutputWithPast,
)
from transformers.models.mistral.modeling_mistral import (
    MistralPreTrainedModel,
    MistralRMSNorm,
    MistralRotaryEmbedding,
    apply_rotary_pos_emb,
)

from sglang.multimodal_gen.runtime.distributed import (
    get_tp_world_size,
    model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.platforms import (
    AttentionBackendEnum,
    current_platform,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)
_CREATE_CAUSAL_MASK_ARG = (
    "inputs_embeds"
    if "inputs_embeds" in inspect.signature(create_causal_mask).parameters
    else "input_embeds"
)


def _tp_world_size() -> int:
    if not model_parallel_is_initialized():
        return 1
    return get_tp_world_size()


def _linear_output(linear: nn.Module, x: torch.Tensor) -> torch.Tensor:
    output = linear(x)
    return output[0] if isinstance(output, tuple) else output


def _make_column_linear(
    in_features: int,
    out_features: int,
    *,
    bias: bool,
    use_tensor_parallel: bool,
):
    if use_tensor_parallel:
        return ColumnParallelLinear(
            in_features,
            out_features,
            bias=bias,
            gather_output=False,
        )
    return nn.Linear(in_features, out_features, bias=bias)


def _make_row_linear(
    in_features: int,
    out_features: int,
    *,
    bias: bool,
    use_tensor_parallel: bool,
):
    if use_tensor_parallel:
        return RowParallelLinear(
            in_features,
            out_features,
            bias=bias,
        )
    return nn.Linear(in_features, out_features, bias=bias)


def _can_use_unmasked_causal_attention(
    attention_mask: Optional[torch.Tensor],
    config: MistralConfig,
    past_key_values: Optional[Cache],
) -> bool:
    if (
        getattr(config, "sliding_window", None) is not None
        or past_key_values is not None
    ):
        return False
    if attention_mask is None:
        return True
    if attention_mask.dim() != 2:
        return False
    return bool(torch.all(attention_mask > 0).item())


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
    The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to
    (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_key_value_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


class MistralAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: MistralConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = (
            getattr(config, "head_dim", None)
            or config.hidden_size // config.num_attention_heads
        )
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.total_num_heads = config.num_attention_heads
        self.total_num_key_value_heads = config.num_key_value_heads
        tp_size = _tp_world_size()
        q_size = self.total_num_heads * self.head_dim
        kv_size = self.total_num_key_value_heads * self.head_dim
        self.use_tensor_parallel = (
            tp_size > 1
            and self.total_num_heads % tp_size == 0
            and self.total_num_key_value_heads % tp_size == 0
            and q_size % tp_size == 0
            and kv_size % tp_size == 0
        )
        self.num_heads = (
            self.total_num_heads // tp_size
            if self.use_tensor_parallel
            else self.total_num_heads
        )
        self.num_key_value_heads = (
            self.total_num_key_value_heads // tp_size
            if self.use_tensor_parallel
            else self.total_num_key_value_heads
        )
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.q_proj = _make_column_linear(
            config.hidden_size,
            q_size,
            bias=False,
            use_tensor_parallel=self.use_tensor_parallel,
        )
        self.k_proj = _make_column_linear(
            config.hidden_size,
            kv_size,
            bias=False,
            use_tensor_parallel=self.use_tensor_parallel,
        )
        self.v_proj = _make_column_linear(
            config.hidden_size,
            kv_size,
            bias=False,
            use_tensor_parallel=self.use_tensor_parallel,
        )
        self.o_proj = _make_row_linear(
            q_size,
            config.hidden_size,
            bias=False,
            use_tensor_parallel=self.use_tensor_parallel,
        )
        self.is_causal = True
        self.attn = LocalAttention(
            self.num_heads,
            self.head_dim,
            self.num_key_value_heads,
            softmax_scale=self.scaling,
            causal=True,
            supported_attention_backends={
                AttentionBackendEnum.FA,
                AttentionBackendEnum.TORCH_SDPA,
            },
            allow_cudnn_sdp=True,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = (
            _linear_output(self.q_proj, hidden_states)
            .view(hidden_shape)
            .transpose(1, 2)
        )
        key_states = (
            _linear_output(self.k_proj, hidden_states)
            .view(hidden_shape)
            .transpose(1, 2)
        )
        value_states = (
            _linear_output(self.v_proj, hidden_states)
            .view(hidden_shape)
            .transpose(1, 2)
        )

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin
        )

        if past_key_values is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(
                key_states, value_states, self.layer_idx, cache_kwargs
            )

        attn_output = self.attn(
            query_states.transpose(1, 2),
            key_states.transpose(1, 2),
            value_states.transpose(1, 2),
            attn_mask=attention_mask,
        )
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = _linear_output(self.o_proj, attn_output)
        return attn_output, None


class MistralTPMLP(nn.Module):
    def __init__(self, config: MistralConfig):
        super().__init__()
        tp_size = _tp_world_size()
        use_tensor_parallel = tp_size > 1 and config.intermediate_size % tp_size == 0
        self.gate_proj = _make_column_linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
            use_tensor_parallel=use_tensor_parallel,
        )
        self.up_proj = _make_column_linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
            use_tensor_parallel=use_tensor_parallel,
        )
        self.down_proj = _make_row_linear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
            use_tensor_parallel=use_tensor_parallel,
        )
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act_fn(_linear_output(self.gate_proj, x)) * _linear_output(
            self.up_proj, x
        )
        return _linear_output(self.down_proj, x)


class MistralDecoderLayer(nn.Module):
    def __init__(self, config: MistralConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = MistralAttention(config=config, layer_idx=layer_idx)
        self.mlp = MistralTPMLP(config)
        self.input_layernorm = MistralRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = MistralRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[
            tuple[torch.Tensor, torch.Tensor]
        ] = None,  # necessary, but kept here for BC
        **kwargs,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Self Attention
        hidden_states, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states


class MistralModel(MistralPreTrainedModel):
    def __init__(self, config: MistralConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                MistralDecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = MistralRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.post_init()

    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,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You must specify exactly one of input_ids or inputs_embeds"
            )

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if cache_position is None:
            past_seen_tokens = (
                past_key_values.get_seq_length() if past_key_values is not None else 0
            )
            cache_position = torch.arange(
                past_seen_tokens,
                past_seen_tokens + inputs_embeds.shape[1],
                device=inputs_embeds.device,
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)
        if _can_use_unmasked_causal_attention(
            attention_mask, self.config, past_key_values
        ):
            causal_mask = None
        else:
            mask_function = (
                create_causal_mask
                if getattr(self.config, "sliding_window", None) is None
                else create_sliding_window_causal_mask
            )
            mask_kwargs = {
                "config": self.config,
                _CREATE_CAUSAL_MASK_ARG: inputs_embeds,
                "attention_mask": attention_mask,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            causal_mask = mask_function(**mask_kwargs)

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        hidden_states_pool = [] if output_hidden_states else None
        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            if output_hidden_states:
                hidden_states_pool.append(hidden_states)
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            hidden_states_pool.append(hidden_states)

        return BaseModelOutputWithPast(
            hidden_states=hidden_states_pool,
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class Mistral3Model(nn.Module):
    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}

    def __init__(self, config: Mistral3Config):
        super().__init__()
        self.language_model = MistralModel(config.text_config)
        self.config = config

    def get_input_embeddings(self):
        return self.language_model.embed_tokens

    def set_decoder(self, decoder):
        self.language_model = decoder

    def get_decoder(self):
        return self.language_model

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = 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,
        vision_feature_layer: Optional[Union[int, list[int]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        image_sizes: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Union[tuple, Mistral3ModelOutputWithPast]:
        del pixel_values, vision_feature_layer, return_dict
        output_attentions = False
        output_hidden_states = True

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You must specify exactly one of input_ids or inputs_embeds"
            )

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        outputs: BaseModelOutputWithPast = self.language_model(
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            cache_position=cache_position,
            **kwargs,
        )

        return Mistral3ModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class Mistral3ForConditionalGeneration(nn.Module, LayerwiseOffloadableModuleMixin):
    _checkpoint_conversion_mapping = {
        "^language_model.model": "model.language_model",
        "^multi_modal_projector": "model.multi_modal_projector",
        "^language_model.lm_head": "lm_head",
    }
    _tied_weights_keys = ["lm_head.weight"]
    uses_sglang_forward_context = True
    layerwise_offload_dit_group_enabled = False
    layer_names = ["model.language_model.layers"]

    def __init__(self, config: LlavaConfig):
        super().__init__()
        self.model = Mistral3Model(config.arch_config)

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_decoder(self, decoder):
        self.model.set_decoder(decoder)

    def get_decoder(self):
        return self.model.get_decoder()

    # Make modules available through conditional class for BC
    @property
    def language_model(self):
        return self.model.language_model

    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,
        output_hidden_states: Optional[bool] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        image_sizes: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Union[tuple, Mistral3CausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Example:

        """
        output_hidden_states = True

        execution_tensor = input_ids if input_ids is not None else inputs_embeds
        sdpa_context = (
            sdpa_kernel(SDPBackend.CUDNN_ATTENTION)
            if execution_tensor is not None
            and execution_tensor.device.type == "cuda"
            and current_platform.is_cuda()
            else nullcontext()
        )
        with sdpa_context:
            # FLUX.2 uses the text-only Mistral3 path but still expects the
            # same local SDPA kernel choice as the official HF implementation.
            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,
                output_hidden_states=output_hidden_states,
                return_dict=True,
                cache_position=cache_position,
                image_sizes=image_sizes,
                **kwargs,
            )

        return Mistral3CausalLMOutputWithPast(
            hidden_states=outputs.hidden_states,
        )

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # Define mapping for stacked parameters
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            name_lower = name.lower()
            if (
                "vision" in name_lower
                or "multi" in name_lower
                or "lm_head" in name_lower
            ):
                continue
            final_name = name.replace("language_model.model.", "model.language_model.")

            if final_name in params_dict:
                param = params_dict[final_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, loaded_weight)
                loaded_params.add(final_name)
            else:
                logger.warning(f"Param {name=} {final_name=} from weight is not loaded")

        return loaded_params


EntryClass = Mistral3ForConditionalGeneration