chore: import upstream snapshot with attribution

python/sglang/srt/models/whisper.py

@@ -0,0 +1,493 @@
+from __future__ import annotations
+
+from array import array
+from typing import Any, Iterable, Optional, Tuple
+
+import torch
+from transformers import WhisperConfig
+
+from sglang.srt.layers.activation import get_act_fn
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput
+from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
+from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
+from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
+
+
+class WhisperAttention(torch.nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        bias: bool = True,
+        layer_id: Optional[int] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        is_cross_attention: bool = False,
+        is_encoder=False,
+    ):
+        super().__init__()
+        self.total_num_heads = num_heads
+        head_dim = embed_dim // num_heads
+        self.is_cross_attention = is_cross_attention
+        self.is_encoder = is_encoder
+
+        tp_size = get_parallel().tp_size
+        assert (
+            num_heads % tp_size == 0
+        ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})"
+        self.num_heads = num_heads // tp_size
+
+        if (head_dim * num_heads) != embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {embed_dim}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.scaling = head_dim**-0.5
+        self.head_dim = head_dim
+        self.kv_size = self.num_heads * head_dim
+
+        if is_cross_attention:
+            self.q_proj = ColumnParallelLinear(
+                embed_dim, embed_dim, quant_config=quant_config
+            )
+            self.kv_proj = QKVParallelLinear(
+                hidden_size=embed_dim,
+                head_size=head_dim,
+                total_num_heads=0,
+                total_num_kv_heads=num_heads,
+                bias=bias,
+                quant_config=quant_config,
+            )
+        else:
+            self.qkv_proj = QKVParallelLinear(
+                embed_dim, head_dim, num_heads, quant_config=quant_config
+            )
+        self.out_proj = RowParallelLinear(
+            embed_dim, embed_dim, bias=bias, quant_config=quant_config
+        )
+        self.attn = RadixAttention(
+            self.num_heads,
+            head_dim,
+            scaling=1.0,
+            num_kv_heads=self.num_heads,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            is_cross_attention=is_cross_attention,
+            attn_type=(
+                AttentionType.ENCODER_ONLY if is_encoder else AttentionType.DECODER
+            ),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+        cross_hidden_states: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Input shape: Batch x Time x Channel"""
+
+        if self.is_cross_attention:
+            # Cross-attention: KV cached during prefill, read from pool during decode.
+            q, _ = self.q_proj(hidden_states)
+            q = q * self.scaling
+            if cross_hidden_states is not None:
+                kv, _ = self.kv_proj(cross_hidden_states)
+                k, v = kv.split([self.kv_size, self.kv_size], dim=-1)
+            else:
+                k = None
+                v = None
+            attn_output = self.attn(q, k, v, forward_batch)
+        else:
+            qkv, _ = self.qkv_proj(hidden_states)
+            q, k, v = qkv.chunk(chunks=3, dim=-1)
+            q = q * self.scaling
+
+            if self.is_encoder:
+                num_heads = self.attn.tp_q_head_num
+                head_dim = self.attn.head_dim
+                batch_size, seq_len, _ = hidden_states.shape
+
+                q = q.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
+                k = k.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
+                v = v.view(batch_size, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
+
+                attn_output = torch.nn.functional.scaled_dot_product_attention(
+                    q, k, v, scale=1.0
+                )
+                attn_output = attn_output.permute(0, 2, 1, 3).reshape(
+                    batch_size, seq_len, num_heads * head_dim
+                )
+            else:
+                attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=True)
+
+        attn_output, _ = self.out_proj(attn_output)
+
+        return attn_output
+
+
+class WhisperEncoderLayer(torch.nn.Module):
+    def __init__(
+        self,
+        config: WhisperConfig,
+        layer_id: Optional[int] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__()
+        self.embed_dim = config.d_model
+
+        self.self_attn = WhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.encoder_attention_heads,
+            layer_id=layer_id,
+            quant_config=quant_config,
+            is_encoder=True,
+        )
+        self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+
+        self.activation_fn = get_act_fn(
+            config.activation_function, quant_config=quant_config
+        )
+
+        self.fc1 = ColumnParallelLinear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = RowParallelLinear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, forward_batch)
+
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states, _ = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+
+        hidden_states, _ = self.fc2(hidden_states)
+
+        hidden_states = residual + hidden_states
+
+        if hidden_states.dtype == torch.float16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(
+                hidden_states, min=-clamp_value, max=clamp_value
+            )
+        return hidden_states
+
+
+class WhisperDecoderLayer(torch.nn.Module):
+    def __init__(
+        self,
+        config: WhisperConfig,
+        layer_id: Optional[int] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+    ):
+        super().__init__()
+        self.embed_dim = config.d_model
+
+        # Offset decoder layer IDs to avoid overlap with encoder layers
+        decoder_self_attn_layer_id = config.encoder_layers + layer_id
+        decoder_cross_attn_layer_id = (
+            config.encoder_layers + config.decoder_layers + layer_id
+        )
+
+        self.self_attn = WhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.decoder_attention_heads,
+            layer_id=decoder_self_attn_layer_id,
+            quant_config=quant_config,
+        )
+
+        self.activation_fn = get_act_fn(
+            config.activation_function, quant_config=quant_config
+        )
+
+        self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+        self.encoder_attn = WhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.decoder_attention_heads,
+            layer_id=decoder_cross_attn_layer_id,
+            quant_config=quant_config,
+            is_cross_attention=True,
+        )
+        self.encoder_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+        self.fc1 = ColumnParallelLinear(self.embed_dim, config.decoder_ffn_dim)
+        self.fc2 = RowParallelLinear(config.decoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        decoder_hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+    ) -> torch.Tensor:
+
+        residual = decoder_hidden_states
+        decoder_hidden_states = self.self_attn_layer_norm(decoder_hidden_states)
+        decoder_hidden_states = self.self_attn(decoder_hidden_states, forward_batch)
+        decoder_hidden_states = residual + decoder_hidden_states
+
+        residual = decoder_hidden_states
+        decoder_hidden_states = self.encoder_attn_layer_norm(decoder_hidden_states)
+        decoder_hidden_states = self.encoder_attn(
+            decoder_hidden_states, forward_batch, encoder_hidden_states
+        )
+        decoder_hidden_states = residual + decoder_hidden_states
+
+        residual = decoder_hidden_states
+        decoder_hidden_states = self.final_layer_norm(decoder_hidden_states)
+        decoder_hidden_states, _ = self.fc1(decoder_hidden_states)
+        decoder_hidden_states = self.activation_fn(decoder_hidden_states)
+        decoder_hidden_states, _ = self.fc2(decoder_hidden_states)
+
+        decoder_hidden_states = residual + decoder_hidden_states
+
+        return decoder_hidden_states
+
+
+class WhisperEncoder(torch.nn.Module):
+
+    def __init__(
+        self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
+    ):
+        super().__init__()
+
+        embed_dim = config.d_model
+        self.embed_scale = embed_dim**-0.5 if config.scale_embedding else 1.0
+
+        self.conv1 = torch.nn.Conv1d(
+            config.num_mel_bins, embed_dim, kernel_size=3, padding=1
+        )
+        self.conv2 = torch.nn.Conv1d(
+            embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
+        )
+        self.embed_positions = torch.nn.Embedding(
+            config.max_source_positions, embed_dim
+        )
+
+        self.layers = torch.nn.ModuleList(
+            [
+                WhisperEncoderLayer(config, id, quant_config)
+                for id in range(config.encoder_layers)
+            ]
+        )
+        self.layer_norm = torch.nn.LayerNorm(config.d_model)
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        position_ids: torch.Tensor,
+        forward_batch: ForwardBatch,
+    ):
+        device = self.conv1.weight.device
+        input_features = input_features.to(device=device)
+        position_ids = position_ids.to(device=device)
+
+        inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
+        inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))
+
+        inputs_embeds = inputs_embeds.mT
+
+        hidden_states = inputs_embeds + self.embed_positions(position_ids)
+
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(hidden_states, forward_batch)
+
+        hidden_states = self.layer_norm(hidden_states)
+        return hidden_states
+
+
+class WhisperDecoder(torch.nn.Module):
+
+    def __init__(
+        self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
+    ):
+        super().__init__()
+        self.max_target_positions = config.max_target_positions
+        self.max_source_positions = config.max_source_positions
+        self.embed_scale = config.d_model**-0.5 if config.scale_embedding else 1.0
+
+        self.embed_tokens = torch.nn.Embedding(
+            config.vocab_size, config.d_model, padding_idx=config.pad_token_id
+        )
+        self.embed_positions = torch.nn.Embedding(
+            self.max_target_positions, config.d_model
+        )
+
+        self.layers = torch.nn.ModuleList(
+            [
+                WhisperDecoderLayer(config, layer_idx, quant_config)
+                for layer_idx in range(config.decoder_layers)
+            ]
+        )
+
+        self.layer_norm = torch.nn.LayerNorm(config.d_model)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+        position_ids=None,
+    ):
+        inputs_embeds = self.embed_tokens(input_ids)
+        position_ids = position_ids.clamp(max=self.max_target_positions - 1)
+        positions = self.embed_positions(position_ids)
+        hidden_states = inputs_embeds + positions.to(inputs_embeds.device)
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states, encoder_hidden_states, forward_batch
+            )
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        return hidden_states
+
+
+class WhisperForConditionalGeneration(torch.nn.Module):
+
+    def __init__(
+        self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
+    ):
+        super().__init__()
+        self.encoder = WhisperEncoder(config, quant_config)
+        self.decoder = WhisperDecoder(config, quant_config)
+        self.proj_out = ParallelLMHead(
+            config.vocab_size, config.d_model, quant_config=quant_config
+        )
+        self.logits_processor = LogitsProcessor(config)
+        self.config = config
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
+        stacked_params_mapping = [
+            (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
+            (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
+            (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
+            (".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
+            (".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        weights_dict = dict(weights)
+
+        # Whisper has no k_proj bias, create zeros
+        for layer_idx in range(self.config.decoder_layers):
+            layer_prefix = f"model.decoder.layers.{layer_idx}.encoder_attn."
+            k_proj_key = layer_prefix + "k_proj.weight"
+            if k_proj_key in weights_dict:
+                k_proj_weight = weights_dict[k_proj_key]
+                bias_key = layer_prefix + "k_proj.bias"
+                if bias_key not in weights_dict:
+                    weights_dict[bias_key] = torch.zeros(k_proj_weight.size(0))
+
+        weights_dict["proj_out.weight"] = weights_dict[
+            "model.decoder.embed_tokens.weight"
+        ]
+
+        for name, loaded_weight in weights_dict.items():
+            name = name.replace("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:
+                    break
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                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)
+
+    def pad_input_ids(
+        self, input_ids: array[int], mm_inputs: MultimodalInputs
+    ) -> array[int]:
+        # Prepend dummy encoder tokens so that prepare_encoder_info_extend
+        # correctly allocates encoder KV cache locations in the KV pool.
+        # These dummy tokens are stripped before the model forward receives input_ids.
+        encoder_len = self.config.max_source_positions
+        mm_inputs.num_image_tokens = encoder_len
+        return array("q", [0]) * encoder_len + input_ids
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        **kwargs: Any,
+    ) -> LogitsProcessorOutput:
+        dtype = self.encoder.conv1.weight.dtype
+
+        # Run encoder for requests that haven't cached encoder output yet.
+        # During decode or when encoder is already cached, encoder_hidden_states
+        # is None and cross-attention reads KV from the pool via RadixAttention.
+        encoder_hidden_states = None
+        if not forward_batch.forward_mode.is_decode():
+            mm_inputs_list = forward_batch.mm_inputs if forward_batch.mm_inputs else []
+            encoder_cached_list = (
+                forward_batch.encoder_cached if forward_batch.encoder_cached else []
+            )
+
+            # Collect features from all uncached requests for batched encoding
+            features_to_encode = []
+            for mm_input, cached in zip(mm_inputs_list, encoder_cached_list):
+                if cached or mm_input is None or not mm_input.mm_items:
+                    continue
+                features = mm_input.mm_items[0].feature
+                if features.ndim == 2:
+                    features = features.unsqueeze(0)
+                features_to_encode.append(features.to(dtype))
+
+            if features_to_encode:
+                # Batch all features and run encoder once instead of sequentially
+                features_batch = torch.cat(features_to_encode, dim=0)
+                encoder_len = features_batch.shape[-1] // 2
+                encoder_position_ids = torch.arange(
+                    encoder_len, device=features_batch.device
+                )
+
+                batched_output = self.encoder(
+                    features_batch, encoder_position_ids, forward_batch
+                )
+                # Flatten [N, seq_len, dim] → [N*seq_len, dim] for cross-attention
+                encoder_hidden_states = batched_output.reshape(
+                    -1, batched_output.shape[-1]
+                )
+
+        decoder_outputs = self.decoder(
+            input_ids, encoder_hidden_states, forward_batch, positions
+        )
+
+        logits = self.logits_processor(
+            input_ids=input_ids,
+            lm_head=self.proj_out,
+            hidden_states=decoder_outputs,
+            logits_metadata=forward_batch,
+        )
+
+        return logits
+
+
+EntryClass = [WhisperForConditionalGeneration]