minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] Xiaomi MiMo-V2.5 day0 support (#23811)

Browse Source 
Co-authored-by: 张袁 <zhangyuan36@xiaomi.com>
Co-authored-by: 刘安岐 <liuanqi6@xiaomi.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
Zhonghua Deng
2026-05-01 00:02:26 +08:00
committed by GitHub
co-authored by 张袁 刘安岐 Xinyuan Tong Shangming Cai Xinyuan Tong
parent cf4f462094
commit 651af06a0b
16 changed files with 4369 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+36 -23
View File
@@ -39,6 +39,36 @@ from sglang.utils import is_in_ci
logger = logging.getLogger(__name__)
MIMO_V2_MODEL_ARCHS = (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
)
MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",)
def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
layout = getattr(hf_config, "attention_projection_layout", None)
if layout is None:
return None
if layout != "fused_qkv":
raise ValueError(
"MiMoV2 hf_config has unsupported "
f"attention_projection_layout={layout!r}; expected 'fused_qkv' "
"or unset."
)
num_key_value_heads = getattr(hf_config, "num_key_value_heads", None)
text_config = getattr(hf_config, "text_config", None)
if num_key_value_heads is None and text_config is not None:
num_key_value_heads = getattr(text_config, "num_key_value_heads", None)
if num_key_value_heads is None:
raise ValueError(
"MiMoV2 hf_config has attention_projection_layout='fused_qkv' "
"but num_key_value_heads is missing; this value is required to "
"derive the fused qkv_proj TP size."
)
return num_key_value_heads
class AttentionArch(IntEnum):
MLA = auto()
@@ -337,10 +367,7 @@ class ModelConfig:
if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM":
self.hf_config.architectures[0] = "MiMoMTP"
if is_draft_model and self.hf_config.architectures[0] in (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
):
if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS:
self.hf_config.architectures[0] = "MiMoV2MTP"
if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM":
self.hf_config.architectures[0] = "Step3p5MTP"
@@ -397,8 +424,7 @@ class ModelConfig:
self.has_attention_sinks = self._detect_attention_sinks()
self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
*MIMO_V2_MODEL_ARCHS,
"MiMoV2MTP",
"Gemma4ForCausalLM",
"Gemma4ForConditionalGeneration",
@@ -417,14 +443,7 @@ class ModelConfig:
return True
# MiMoV2 creates sinks only when the config flags are set.
if any(
a in archs
for a in (
"MiMoV2FlashForCausalLM",
"MiMoV2ForCausalLM",
"MiMoV2MTP",
)
):
if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "MiMoV2MTP")):
return getattr(
self.hf_text_config, "add_swa_attention_sink_bias", False
) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
@@ -1383,6 +1402,7 @@ multimodal_model_archs = [
"LlavaVidForCausalLM",
"Lfm2VlForConditionalGeneration",
"LightOnOCRForConditionalGeneration",
*MIMO_V2_MULTIMODAL_ARCHS,
"MiniCPMO",
"MiniCPMV",
"Mistral3ForConditionalGeneration",
@@ -1528,8 +1548,7 @@ def is_hybrid_swa_model(model_architectures: List[str]):
hybrid_swa_archs = {
"Llama4ForConditionalGeneration",
"GptOssForCausalLM",
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
*MIMO_V2_MODEL_ARCHS,
"MiMoV2MTP",
"Step3p5ForCausalLM",
"Step3p5MTP",
@@ -1559,13 +1578,7 @@ def get_hybrid_layer_ids(
full_attention_layer_ids = [
i for i, x in enumerate(layer_types) if x == "full_attention"
]
elif any(
x in model_architectures
for x in (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
)
):
elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures):
hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None)
swa_attention_layer_ids = [
i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1
python/sglang/srt/layers/attention/vision.py
+61 -17
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import dataclasses
import functools
import math
import warnings
from functools import lru_cache, partial
from typing import Any, Callable, Optional, Tuple
@@ -406,34 +407,39 @@ class VisionFlash3Attention(nn.Module):
Returns:
[b * s, h, head_size]
"""
window_size = kwargs.get("window_size", (-1, -1))
s_aux = kwargs.get("s_aux", None)
if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
max_seqlen = cu_seqlens[1]
output = flash_attn_varlen_func(
q,
k,
v,
fa_kwargs = dict(
cu_seqlens_q=cu_seqlens[0],
cu_seqlens_k=cu_seqlens[0],
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
softmax_scale=softmax_scale,
window_size=window_size,
)
if s_aux is not None:
fa_kwargs["sinks"] = s_aux
output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
else:
cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
output = flash_attn_varlen_func(
q,
k,
v,
fa_kwargs = dict(
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
softmax_scale=softmax_scale,
window_size=window_size,
)
if s_aux is not None:
fa_kwargs["sinks"] = s_aux
output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
return output
@@ -749,7 +755,8 @@ class VisionAttention(nn.Module):
num_heads: int,
projection_size: int,
use_qkv_parallel: bool,
head_size: Optional[int] = None,
num_kv_heads: Optional[int] = None,
head_dim: Optional[int] = None,
qkv_backend: Optional[str] = None,
quant_config: Optional[QuantizationConfig] = None,
dropout: float = 0.0,
@@ -770,13 +777,23 @@ class VisionAttention(nn.Module):
use_dp_attention_reduce: bool = False,
aux_stream: Optional[torch.cuda.Stream] = None,
workspace_buffer: Optional[torch.Tensor] = None,
use_sink: bool = False,
window_size: Tuple[int, int] = (-1, -1),
**kwargs,
):
super().__init__()
if head_dim is None and "head_size" in kwargs:
head_dim = kwargs.pop("head_size")
warnings.warn(
"VisionAttention(head_size=...) is deprecated; use head_dim=...",
DeprecationWarning,
stacklevel=2,
)
self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
self.dropout = dropout
self.head_size = head_size if head_size is not None else embed_dim // num_heads
num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
self.head_size = head_dim if head_dim is not None else embed_dim // num_heads
self.hidden_size_per_attention_head = dist_utils.divide(
projection_size, num_heads
)
@@ -784,7 +801,7 @@ class VisionAttention(nn.Module):
num_dummy_heads + num_heads, self.tp_size
)
self.num_attention_kv_heads_per_partition = dist_utils.divide(
num_dummy_heads + num_heads, self.tp_size
num_dummy_heads + num_kv_heads, self.tp_size
)
self.q_size = self.num_attention_heads_per_partition * self.head_size
@@ -838,7 +855,7 @@ class VisionAttention(nn.Module):
hidden_size=embed_dim,
head_size=self.head_size,
total_num_heads=num_dummy_heads + num_heads,
total_num_kv_heads=num_dummy_heads + num_heads,
total_num_kv_heads=num_dummy_heads + num_kv_heads,
bias=qkv_bias,
quant_config=quant_config,
tp_rank=self.tp_rank,
@@ -870,6 +887,20 @@ class VisionAttention(nn.Module):
self.aux_stream = aux_stream
self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else []
self.window_size = window_size
if use_sink:
# Allocate the full (unsharded) sink tensor for weight loading;
# only the local TP slice is used in forward.
self.sinks = nn.Parameter(
torch.empty(
self.num_attention_heads_per_partition * self.tp_size,
dtype=torch.bfloat16,
),
requires_grad=False,
)
else:
self.sinks = None
def _init_qk_norm(
self, norm_dim: int, eps: float, var_hidden_size: Optional[int] = None
):
@@ -989,6 +1020,7 @@ class VisionAttention(nn.Module):
rotary_pos_emb_cos: Optional[torch.Tensor] = None,
rotary_pos_emb_sin: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
full_attn: bool = True,
**kwargs,
) -> torch.Tensor:
r"""
@@ -1070,19 +1102,20 @@ class VisionAttention(nn.Module):
sin = rotary_pos_emb_sin
if cos is not None and sin is not None:
original_shape = q.shape
original_q_shape = q.shape
original_k_shape = k.shape
# [total_tokens, head, head_size]
# [total_tokens, head, head_size] for q / [total_tokens, kv_head, head_size] for k
q = q.view(-1, head, self.head_size)
k = k.view(-1, head, self.head_size)
k = k.view(-1, kv_head, self.head_size)
if cos.size(-1) * 2 == self.head_size:
cos = torch.cat([cos, cos], dim=-1)
sin = torch.cat([sin, sin], dim=-1)
q, k = apply_rotary_pos_emb(q, k, cos, sin)
q = q.view(original_shape)
k = k.view(original_shape)
q = q.view(original_q_shape)
k = k.view(original_k_shape)
if q.dim() == 4:
# [b, s, head, head_size] --> [b * s, head, head_size]
@@ -1118,6 +1151,15 @@ class VisionAttention(nn.Module):
else:
q, k = self._apply_qk_norm(q, k)
if full_attn or self.sinks is None:
effective_window_size = (-1, -1)
s_aux = None
else:
effective_window_size = self.window_size
q_head_start = self.tp_rank * self.num_attention_heads_per_partition
q_head_end = (self.tp_rank + 1) * self.num_attention_heads_per_partition
s_aux = self.sinks[q_head_start:q_head_end]
output = self.qkv_backend.forward(
q=q,
k=k,
@@ -1130,6 +1172,8 @@ class VisionAttention(nn.Module):
max_seqlen=max_seqlen,
output_ws=attn_output_ws,
softmax_scale=self.softmax_scale,
window_size=effective_window_size,
s_aux=s_aux,
)
assert output.dim() == 3, output.shape
python/sglang/srt/layers/linear.py
+1 -1
View File
@@ -1047,7 +1047,7 @@ class QKVParallelLinear(ColumnParallelLinear):
block_n, _ = self.quant_method.quant_config.weight_block_size
q_size = self.total_num_heads * self.head_size // block_n
k_size = self.total_num_kv_heads * self.head_size // block_n
v_size = self.total_num_kv_heads * self.head_size // block_n
v_size = self.total_num_kv_heads * self.v_head_size // block_n
shard_offsets = [
# (shard_id, shard_offset, shard_size)
("q", 0, q_size),
python/sglang/srt/managers/io_struct.py
+2 -2
View File
@@ -38,7 +38,7 @@ from sglang.srt.observability.req_time_stats import (
SchedulerReqTimeStats,
)
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.utils import ImageData
from sglang.srt.utils import ImageData, VideoData
# Handle serialization of Image for pydantic
if TYPE_CHECKING:
@@ -118,7 +118,7 @@ class SessionParams:
# Individual data item types for each modality
ImageDataInputItem = Union[Image, str, ImageData, Dict]
AudioDataInputItem = Union[str, Dict]
VideoDataInputItem = Union[str, Dict]
VideoDataInputItem = Union[str, VideoData, Dict]
# Union type for any multimodal data item
MultimodalDataInputItem = Union[
ImageDataInputItem, VideoDataInputItem, AudioDataInputItem
python/sglang/srt/managers/mm_utils.py
+38 -9
View File
@@ -1608,15 +1608,35 @@ def wrap_shm_features(obj):
if hasattr(obj, "mm_inputs") and obj.mm_inputs:
for item in obj.mm_inputs.mm_items:
if (
hasattr(item, "feature")
and isinstance(item.feature, torch.Tensor)
and item.feature.is_cpu
):
item.feature = ShmPointerMMData(item.feature)
if not hasattr(item, "feature"):
continue
feat = item.feature
if isinstance(feat, torch.Tensor) and feat.is_cpu:
item.feature = ShmPointerMMData(feat)
elif isinstance(feat, (list, tuple)):
wrapped = [
(
ShmPointerMMData(t)
if isinstance(t, torch.Tensor) and t.is_cpu
else t
)
for t in feat
]
item.feature = (
type(feat)(wrapped) if isinstance(feat, tuple) else wrapped
)
return obj
def _feature_has_shm(feat) -> bool:
"""Check whether a single feature (tensor, ShmPointer, or list) contains ShmPointerMMData."""
if isinstance(feat, ShmPointerMMData):
return True
if isinstance(feat, (list, tuple)):
return any(isinstance(t, ShmPointerMMData) for t in feat)
return False
def has_shm_features(recv_reqs):
"""Return True if any request in the list contains ShmPointerMMData."""
for req in recv_reqs:
@@ -1625,7 +1645,7 @@ def has_shm_features(recv_reqs):
return True
elif hasattr(req, "mm_inputs") and req.mm_inputs:
for item in req.mm_inputs.mm_items:
if isinstance(item.feature, ShmPointerMMData):
if _feature_has_shm(item.feature):
return True
return False
@@ -1646,6 +1666,15 @@ def unwrap_shm_features(obj):
if hasattr(obj, "mm_inputs") and obj.mm_inputs:
mm_items = obj.mm_inputs.mm_items
for item in mm_items:
if isinstance(item.feature, ShmPointerMMData):
item.feature = item.feature.materialize()
feat = item.feature
if isinstance(feat, ShmPointerMMData):
item.feature = feat.materialize()
elif isinstance(feat, (list, tuple)):
unwrapped = [
t.materialize() if isinstance(t, ShmPointerMMData) else t
for t in feat
]
item.feature = (
type(feat)(unwrapped) if isinstance(feat, tuple) else unwrapped
)
return obj
python/sglang/srt/models/mimo_audio.py
+1350
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/models/mimo_v2.py
+222 -13
View File
@@ -13,13 +13,14 @@
# ==============================================================================
import logging
from typing import Any, Dict, Iterable, Optional, Tuple, Union
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
from torch import nn
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size
from sglang.srt.distributed import (
get_moe_expert_parallel_world_size,
get_pp_group,
@@ -63,11 +64,18 @@ from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.managers.mm_utils import (
MultiModalityDataPaddingPatternMultimodalTokens,
general_mm_embed_routine,
)
from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
kv_cache_scales_loader,
)
from sglang.srt.models.mimo_audio import MiMoAudioEncoder, MiMoAudioEncoderConfig
from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
LazyValue,
@@ -81,6 +89,38 @@ MiMoV2Config = None
logger = logging.getLogger(__name__)
def load_mimo_v2_qkv_proj_weight(
name, param, loaded_weight, expected_fused_tp_size: Optional[int] = None
):
if loaded_weight.shape == param.shape:
# The checkpoint already stores this rank's qkv_proj shard.
default_weight_loader(param, loaded_weight)
return
if loaded_weight.ndim != param.ndim or loaded_weight.shape[1:] != param.shape[1:]:
raise ValueError(
f"qkv_proj weight {name}: unexpected shape {tuple(loaded_weight.shape)}; "
f"expected sharded {tuple(param.shape)}"
)
tp_size = get_attention_tp_size()
tp_rank = get_attention_tp_rank()
if expected_fused_tp_size is not None and tp_size != expected_fused_tp_size:
raise ValueError(
f"MiMoV2 fused qkv_proj checkpoint is TP={expected_fused_tp_size}-"
f"interleaved; got attention tp_size={tp_size} while loading {name}."
)
fused_shape = (param.shape[0] * tp_size, *param.shape[1:])
if tuple(loaded_weight.shape) != fused_shape:
raise ValueError(
f"qkv_proj weight {name}: unexpected shape {tuple(loaded_weight.shape)}; "
f"expected fused {fused_shape} or sharded {tuple(param.shape)}"
)
default_weight_loader(param, loaded_weight.chunk(tp_size, dim=0)[tp_rank])
class MiMoV2MLP(nn.Module):
def __init__(
self,
@@ -995,6 +1035,24 @@ class MiMoV2ForCausalLM(nn.Module):
self.logits_processor = LogitsProcessor(config)
vision_config = getattr(config, "vision_config", None)
audio_config = getattr(config, "audio_config", None)
self._is_multimodal = vision_config is not None and audio_config is not None
if self._is_multimodal:
if hasattr(vision_config, "to_dict"):
vision_config = vision_config.to_dict()
if hasattr(audio_config, "to_dict"):
audio_config = audio_config.to_dict()
self.visual = MiMoVisionTransformer(
MiMoVLVisionConfig.from_dict(vision_config),
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
quant_config=None,
prefix=add_prefix("visual", prefix),
)
self.audio_config = MiMoAudioEncoderConfig(**audio_config)
self.audio_encoder = MiMoAudioEncoder(self.audio_config)
self._routed_experts_weights_of_layer = LazyValue(
lambda: {
layer_id: layer.mlp.get_moe_weights()
@@ -1010,6 +1068,31 @@ class MiMoV2ForCausalLM(nn.Module):
def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embedding(input_ids)
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
pattern = MultiModalityDataPaddingPatternMultimodalTokens()
return pattern.pad_input_tokens(input_ids, mm_inputs)
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
pixel_values = torch.cat([item.feature for item in items], dim=0).type(
self.visual.dtype
)
image_grid_thw = torch.cat([item.image_grid_thw for item in items], dim=0)
assert pixel_values.dim() == 2, pixel_values.dim()
assert image_grid_thw.dim() == 2, image_grid_thw.dim()
return self.visual(pixel_values, grid_thw=image_grid_thw)
def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
pixel_values = torch.cat([item.feature for item in items], dim=0).type(
self.visual.dtype
)
video_grid_thw = torch.cat([item.video_grid_thw for item in items], dim=0)
assert pixel_values.dim() == 2, pixel_values.dim()
assert video_grid_thw.dim() == 2, video_grid_thw.dim()
return self.visual(pixel_values, grid_thw=video_grid_thw)
def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
return self.audio_encoder.get_audio_feature(items)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
@@ -1022,13 +1105,23 @@ class MiMoV2ForCausalLM(nn.Module):
input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
hidden_states, hidden_states_before_norm = self.model(
input_ids,
positions,
forward_batch,
input_embeds,
pp_proxy_tensors=pp_proxy_tensors,
)
if self._is_multimodal:
hidden_states, hidden_states_before_norm = general_mm_embed_routine(
input_ids=input_ids,
forward_batch=forward_batch,
language_model=self.model,
multimodal_model=self,
positions=positions,
pp_proxy_tensors=pp_proxy_tensors,
)
else:
hidden_states, hidden_states_before_norm = self.model(
input_ids,
positions,
forward_batch,
input_embeds,
pp_proxy_tensors=pp_proxy_tensors,
)
if self.pp_group.is_last_rank:
return self.logits_processor(
@@ -1058,6 +1151,11 @@ class MiMoV2ForCausalLM(nn.Module):
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
stacked_params_mapping_vit = [
# (param_name, shard_name, shard_id)
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
# (param_name, weight_name, expert_id, shard_id)
expert_params_mapping = DeepEPMoE.make_expert_params_mapping(
@@ -1068,8 +1166,105 @@ class MiMoV2ForCausalLM(nn.Module):
)
params_dict = dict(self.named_parameters())
skipped_mtp_weights = False
for name, loaded_weight in weights:
if not self._is_multimodal and (
name.startswith(("visual.", "vision_model.", "audio_encoder."))
or name.startswith("audio_")
or "speech_embeddings" in name
):
continue
if self._is_multimodal and "audio" in name:
if "projection" in name:
if (
"audio_encoder.audio_projection" in name
and "audio_encoder.projection" not in name
):
name = name.replace(
"audio_encoder.audio_projection", "audio_encoder.projection"
)
elif (
"audio_projection" in name
and "audio_encoder.projection" not in name
):
name = name.replace(
"audio_projection", "audio_encoder.projection"
)
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
continue
if "input_local_transformer" in name:
if (
"audio_input_local_transformer" in name
and "audio_encoder.input_local_transformer" not in name
):
name = name.replace(
"audio_input_local_transformer",
"audio_encoder.input_local_transformer",
)
if name not in params_dict:
logger.warning(
f"Parameter {name} not found in params_dict, skipping"
)
continue
param = params_dict[name]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
continue
if self._is_multimodal and "speech_embeddings" in name:
if (
"speech_embeddings" in name
and "audio_encoder.speech_embeddings" not in name
):
name = name.replace(
"speech_embeddings", "audio_encoder.speech_embeddings"
)
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight[: param.shape[0], :])
continue
if self._is_multimodal and "visual" in name:
name = name.replace("vision_model.", "")
name = name.replace(r"attn.qkv.", r"attn.qkv_proj.")
match_stacked_vit = False
for param_name, weight_name, shard_id in stacked_params_mapping_vit:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
match_stacked_vit = True
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
match_stacked_vit = True
break
if match_stacked_vit:
continue
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
if name.endswith("patch_embed.proj.weight"):
patch_embed = self.get_submodule(name.rsplit(".", 2)[0])
if hasattr(patch_embed, "sync_proj_weight_linear_format"):
patch_embed.sync_proj_weight_linear_format()
continue
layer_id = get_layer_id(name)
if (
layer_id is not None
@@ -1099,21 +1294,35 @@ class MiMoV2ForCausalLM(nn.Module):
else:
continue
# TODO: skip mtp weights for now, need to implement mtp
if "mtp" in name:
if not skipped_mtp_weights:
logger.info(
"Skipping draft-only MiMo-V2 MTP weights while loading the "
"target model; MiMoV2MTP loads these weights in the draft "
"model runner."
)
skipped_mtp_weights = True
continue
# Support fused qkv_proj checkpoint (Pro format)
if "qkv_proj" in name:
if name in params_dict:
tp_size = get_attention_tp_size()
tp_rank = get_attention_tp_rank()
param = params_dict[name]
loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
default_weight_loader(param, loaded_weight)
expected_fused_tp_size = get_mimo_v2_fused_qkv_expected_tp_size(
self.config
)
load_mimo_v2_qkv_proj_weight(
name, param, loaded_weight, expected_fused_tp_size
)
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if (
"compression_attention" in name
or "hybrid_softmax_attention" in name
or "compressed_softmax_attn" in name
):
continue
if weight_name not in name:
continue
if ("mlp.experts." in name) and name not in params_dict:
python/sglang/srt/models/mimo_v2_nextn.py
+12 -7
View File
@@ -19,6 +19,7 @@ import torch
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.communicator import (
@@ -28,7 +29,6 @@ from sglang.srt.layers.communicator import (
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
is_dp_attention_enabled,
)
from sglang.srt.layers.layernorm import RMSNorm
@@ -44,6 +44,7 @@ from sglang.srt.models.mimo_v2 import (
MiMoV2Attention,
MiMoV2ForCausalLM,
MiMoV2MLP,
load_mimo_v2_qkv_proj_weight,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
@@ -304,20 +305,24 @@ class MiMoV2MTP(MiMoV2ForCausalLM):
# Support fused qkv_proj checkpoint (Pro format)
if "qkv_proj" in name:
if name in params_dict:
tp_size = get_attention_tp_size()
tp_rank = get_attention_tp_rank()
param = params_dict[name]
loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
default_weight_loader(param, loaded_weight)
load_mimo_v2_qkv_proj_weight(
name,
param,
loaded_weight,
expected_fused_tp_size=get_mimo_v2_fused_qkv_expected_tp_size(
self.config
),
)
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
if f".{weight_name}." not in name:
continue
if "mtp_block" not in name:
break
name = name.replace(weight_name, param_name)
name = name.replace(f".{weight_name}.", f".{param_name}.")
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
python/sglang/srt/models/mimo_vl.py
+507
View File
@@ -0,0 +1,507 @@
"""Inference-only MiMo vision model: attention + ViT."""
from __future__ import annotations
from functools import partial
from typing import Optional, Tuple, Type
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from transformers.configuration_utils import PretrainedConfig
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VisionRotaryEmbedding,
)
from sglang.srt.layers.attention.vision import VisionAttention
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.models.qwen2_5_vl import Qwen2_5_VisionPatchMerger, Qwen2_5_VLMLP
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix
class MiMoVLVisionConfig(PretrainedConfig):
model_type = "mimovl"
base_config_key = "vision_config"
def __init__(
self,
depth=28,
hidden_size=1280,
hidden_act="silu",
intermediate_size=4608,
num_heads=32,
in_channels=3,
patch_size=16,
spatial_merge_size=2,
temporal_patch_size=2,
tokens_per_second=2,
window_size=128,
out_hidden_size=2048,
fullatt_block_indexes=[7, 15, 23, 31],
initializer_range=0.02,
kv_channels=64,
qk_channels=64,
num_query_groups=4,
num_key_value_heads=8,
vit_window_attn_types=None,
visual_token_window_size=64,
**kwargs,
):
super().__init__(**kwargs)
self.depth = depth
self.hidden_size = hidden_size
self.hidden_act = hidden_act
self.intermediate_size = intermediate_size
self.num_heads = num_heads
if num_key_value_heads is None:
num_key_value_heads = num_heads
self.num_key_value_heads = num_key_value_heads
self.in_channels = in_channels
self.patch_size = patch_size
self.spatial_merge_size = spatial_merge_size
self.temporal_patch_size = temporal_patch_size
self.tokens_per_second = tokens_per_second
self.window_size = window_size
self.fullatt_block_indexes = fullatt_block_indexes
self.out_hidden_size = out_hidden_size
self.initializer_range = initializer_range
self.kv_channels = kv_channels
self.qk_channels = qk_channels
self.num_query_groups = num_query_groups
self.vit_window_attn_types = vit_window_attn_types or [-1] * depth
self.visual_token_window_size = visual_token_window_size
class MiMoVisionPatchEmbed(nn.Module):
def __init__(
self,
patch_size: int = 16,
temporal_patch_size: int = 2,
in_channels: int = 3,
embed_dim: int = 1536,
) -> None:
super().__init__()
self.patch_size = patch_size
self.temporal_patch_size = temporal_patch_size
self.in_channels = in_channels
self.embed_dim = embed_dim
kernel_size = [temporal_patch_size, patch_size, patch_size]
self.proj = nn.Conv3d(
in_channels,
embed_dim,
kernel_size=kernel_size,
stride=kernel_size,
bias=False,
)
self.proj_weight_linear_format = None
@torch.no_grad()
def sync_proj_weight_linear_format(self):
self.proj_weight_linear_format = self.proj.weight.view(self.embed_dim, -1)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
target_dtype = self.proj.weight.dtype
hidden_states = F.linear(
hidden_states.to(dtype=target_dtype), self.proj_weight_linear_format
)
return hidden_states
class MiMoVisionBlock(nn.Module):
def __init__(
self,
dim: int,
intermediate_dim: int,
num_heads: int,
hidden_act="silu",
norm_layer: Type[nn.Module] = None,
attn_implementation: Optional[str] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
num_dummy_heads: int = 0,
rms_norm_eps: float = 1e-6,
use_sink: bool = False,
window_size: Tuple[int, int] = (-1, -1),
num_kv_heads: Optional[int] = None,
head_dim: Optional[int] = None,
use_data_parallel: bool = False,
) -> None:
super().__init__()
if norm_layer is None:
norm_layer = partial(nn.LayerNorm, eps=1e-6)
self.norm1 = RMSNorm(dim, eps=rms_norm_eps)
self.norm2 = RMSNorm(dim, eps=rms_norm_eps)
self.use_data_parallel = use_data_parallel
if attn_implementation is None:
softmax_in_single_precision = False
qkv_backend = None
flatten_batch = True
elif attn_implementation == "sdpa":
softmax_in_single_precision = False
qkv_backend = "sdpa"
flatten_batch = True
elif attn_implementation == "flash_attention_2":
softmax_in_single_precision = False
qkv_backend = "triton_attn"
flatten_batch = True
elif attn_implementation == "eager":
softmax_in_single_precision = True
qkv_backend = "sdpa"
flatten_batch = True
elif attn_implementation == "flash_attention_3":
softmax_in_single_precision = False
qkv_backend = "fa3"
flatten_batch = True
self.attn = VisionAttention(
embed_dim=dim,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
head_dim=head_dim,
projection_size=dim,
use_qkv_parallel=True,
proj_bias=True,
qkv_bias=True,
qkv_backend=qkv_backend,
softmax_in_single_precision=softmax_in_single_precision,
flatten_batch=flatten_batch,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
num_dummy_heads=num_dummy_heads,
use_sink=use_sink,
window_size=window_size,
use_data_parallel=use_data_parallel,
)
self.mlp = Qwen2_5_VLMLP(
dim,
intermediate_dim,
hidden_act=hidden_act,
quant_config=quant_config,
prefix=add_prefix("mlp", prefix),
use_data_parallel=use_data_parallel,
)
def forward(
self,
x: torch.Tensor,
cu_seqlens: torch.Tensor,
max_seqlen: int,
position_embeddings: torch.Tensor,
full_attn: bool = True,
) -> torch.Tensor:
S, B, H = x.shape
# norm1: flatten to 2D -> [S*B, H], then reshape back
x2d = x.reshape(-1, H)
hidden_states = self.norm1(x2d).reshape(S, B, H)
# Attention expects [B, S, H]
hidden_states = rearrange(hidden_states, "s b h -> b s h")
attn = self.attn(
hidden_states,
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
position_embeddings=position_embeddings,
full_attn=full_attn,
)
attn = rearrange(attn, "b s h -> s b h")
# norm2 with fused residual-add: also 2D
attn2d = attn.reshape(-1, H)
x_norm_2d, x_after_add_2d = self.norm2(x2d, residual=attn2d)
x_norm = x_norm_2d.reshape(S, B, H)
x_after_add = x_after_add_2d.reshape(S, B, H)
# MLP and final residual
mlp_out = self.mlp(x_norm)
x = x_after_add + mlp_out
return x
class MiMoVisionTransformer(nn.Module):
def __init__(
self,
vision_config: MiMoVLVisionConfig,
norm_eps: float = 1e-6,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.server_args = get_global_server_args()
self.vit_window_attn_types = vision_config.vit_window_attn_types
patch_size: int = vision_config.patch_size
temporal_patch_size: int = vision_config.temporal_patch_size
spatial_merge_size: int = vision_config.spatial_merge_size
self.spatial_merge_size = spatial_merge_size
self.spatial_merge_unit: int = spatial_merge_size * spatial_merge_size
in_channels: int = vision_config.in_channels
hidden_size: int = vision_config.hidden_size
depth: int = vision_config.depth
num_heads: int = vision_config.num_heads
num_kv_heads = getattr(vision_config, "num_key_value_heads", None)
if num_kv_heads is None:
num_kv_heads = num_heads
self.num_kv_heads = num_kv_heads
self.qk_channels = getattr(vision_config, "qk_channels", None)
self.kv_channels = getattr(vision_config, "kv_channels", None)
self.fullatt_block_indexes = vision_config.fullatt_block_indexes
self.window_size = vision_config.window_size
self.patch_size = vision_config.patch_size
self.use_data_parallel = self.server_args.mm_enable_dp_encoder
mlp_hidden_size: int = vision_config.intermediate_size
self.patch_embed = MiMoVisionPatchEmbed(
patch_size=patch_size,
temporal_patch_size=temporal_patch_size,
in_channels=in_channels,
embed_dim=hidden_size,
)
self.use_sink = getattr(vision_config, "use_sink", False)
norm_layer = partial(nn.LayerNorm, eps=norm_eps)
head_dim = (
self.qk_channels
if self.qk_channels is not None
else hidden_size // num_heads
)
self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
self.visual_token_window_size = getattr(
vision_config, "visual_token_window_size", -1
)
self.blocks = nn.ModuleList(
[
MiMoVisionBlock(
dim=hidden_size,
intermediate_dim=mlp_hidden_size,
num_heads=num_heads,
hidden_act=vision_config.hidden_act,
norm_layer=norm_layer,
attn_implementation="flash_attention_3",
quant_config=quant_config,
prefix=add_prefix(f"blocks.{i}", prefix),
use_sink=(
self.use_sink if i not in self.fullatt_block_indexes else False
),
window_size=(
self.visual_token_window_size,
self.visual_token_window_size,
),
num_kv_heads=num_kv_heads,
head_dim=self.qk_channels,
use_data_parallel=self.use_data_parallel,
)
for i in range(depth)
]
)
self.vision_config = vision_config
self.merger = Qwen2_5_VisionPatchMerger(
dim=vision_config.out_hidden_size,
context_dim=hidden_size,
spatial_merge_size=spatial_merge_size,
quant_config=quant_config,
prefix=add_prefix("merger", prefix),
use_data_parallel=self.use_data_parallel,
)
self._post_init()
def apply_index(self, tensor: torch.Tensor, index: torch.Tensor):
tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))
tensor = tensor[index]
tensor = tensor.flatten(0, 1)
return tensor
def _post_init(self):
for name, param in self.named_parameters():
if "bias" in name:
param.data.zero_()
def get_window_index_1d(self, grid_thw, col=True):
window_index: list = []
window_index_id = 0
for grid_t, grid_h, grid_w in grid_thw:
llm_grid_h, llm_grid_w = (
grid_h // self.spatial_merge_size,
grid_w // self.spatial_merge_size,
)
index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(
grid_t, llm_grid_h, llm_grid_w
)
if col:
index_new = index.transpose(1, 2).reshape(-1)
else:
index_new = index.reshape(-1)
window_index.append(index_new + window_index_id)
window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
window_index = torch.cat(
window_index,
dim=0,
)
return window_index
@property
def dtype(self) -> torch.dtype:
return self.patch_embed.proj.weight.dtype
@property
def device(self) -> torch.device:
return self.blocks[0].mlp.gate_up_proj.weight.device
def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
pos_ids = []
for i in range(grid_thw.size(0)):
t, h, w = grid_thw[i].tolist()
hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
hpos_ids = hpos_ids.reshape(
h // self.spatial_merge_size,
self.spatial_merge_size,
w // self.spatial_merge_size,
self.spatial_merge_size,
)
hpos_ids = hpos_ids.permute(0, 2, 1, 3)
hpos_ids = hpos_ids.flatten()
wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
wpos_ids = wpos_ids.reshape(
h // self.spatial_merge_size,
self.spatial_merge_size,
w // self.spatial_merge_size,
self.spatial_merge_size,
)
wpos_ids = wpos_ids.permute(0, 2, 1, 3)
wpos_ids = wpos_ids.flatten()
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
pos_ids = torch.cat(pos_ids, dim=0)
max_grid_size = grid_thw[:, 1:].max()
rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
return rotary_pos_emb
def _prepare_forward(
self,
x: torch.Tensor,
grid_thw: torch.Tensor,
):
# patchify
x = x.to(device=self.device, dtype=self.dtype)
x = self.patch_embed(x)
# compute position embedding
rotary_pos_emb = self.rot_pos_emb(grid_thw)
window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(
device=x.device
)
reverse_window_index_1d_col = torch.argsort(window_index_1d_col).to(
device=x.device
)
rotary_pos_emb = rotary_pos_emb.to(device=x.device)
emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
def get_position_embeddings(emb, x):
position_embeddings = (emb.cos(), emb.sin())
position_embeddings = (
position_embeddings[0].to(x.device),
position_embeddings[1].to(x.device),
)
return position_embeddings
seqlens = torch.repeat_interleave(
grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
)
cu_seqlens = torch.cat(
[
torch.tensor([0], device=x.device, dtype=torch.int32),
seqlens.cumsum(dim=0).to(device=x.device, dtype=torch.int32),
]
)
max_seqlen = seqlens.max().item()
row_based_embeddings = get_position_embeddings(emb, x)
col_based_embeddings = get_position_embeddings(
self.apply_index(emb, window_index_1d_col), x
)
# transformers
x = x.unsqueeze(1) # [S, 1, H]
return (
x,
row_based_embeddings,
col_based_embeddings,
window_index_1d_col,
reverse_window_index_1d_col,
cu_seqlens,
max_seqlen,
)
def run_blocks(
self,
x: torch.Tensor,
row_based_embeddings: Tuple[torch.Tensor, torch.Tensor],
col_based_embeddings: Tuple[torch.Tensor, torch.Tensor],
window_index_1d_col: torch.Tensor,
reverse_window_index_1d_col: torch.Tensor,
cu_seqlens: torch.Tensor,
max_seqlen: int,
) -> torch.Tensor:
for layer_num, blk in enumerate(self.blocks):
window_attn_type = self.vit_window_attn_types[layer_num]
# window_attn_type = 1: col-based SWA
if window_attn_type == 1 and (
layer_num == 0 or self.vit_window_attn_types[layer_num - 1] != 1
):
x = self.apply_index(x, window_index_1d_col)
if (
layer_num > 0
and window_attn_type != 1
and self.vit_window_attn_types[layer_num - 1] == 1
):
x = self.apply_index(x, reverse_window_index_1d_col)
position_embeddings = (
col_based_embeddings if window_attn_type == 1 else row_based_embeddings
)
full_attn = layer_num in self.fullatt_block_indexes
x = blk(
x,
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
position_embeddings=position_embeddings,
full_attn=full_attn,
)
x = self.merger(x)
return x
def forward(
self,
x: torch.Tensor,
grid_thw: torch.Tensor,
) -> torch.Tensor:
(
x,
row_based_embeddings,
col_based_embeddings,
window_index_1d_col,
reverse_window_index_1d_col,
cu_seqlens,
max_seqlen,
) = self._prepare_forward(x, grid_thw)
return self.run_blocks(
x,
row_based_embeddings,
col_based_embeddings,
window_index_1d_col,
reverse_window_index_1d_col,
cu_seqlens,
max_seqlen,
)
python/sglang/srt/multimodal/processors/mimo_v2.py
+2039
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/parser/conversation.py
+9 -4
View File
@@ -35,7 +35,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
from typing_extensions import Literal
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
from sglang.srt.utils import ImageData, read_system_prompt_from_file
from sglang.srt.utils import ImageData, VideoData, read_system_prompt_from_file
class SeparatorStyle(IntEnum):
@@ -97,7 +97,7 @@ class Conversation:
audio_token: str = "<audio>"
image_data: Optional[List[ImageData]] = None
video_data: Optional[List[str]] = None
video_data: Optional[List[Union[str, VideoData]]] = None
modalities: Optional[List[str]] = None
stop_token_ids: Optional[int] = None
@@ -413,9 +413,14 @@ class Conversation:
"""Append a new image."""
self.image_data.append(ImageData(url=image, detail=detail))
def append_video(self, video: str):
def append_video(self, video: str, preprocess_kwargs: Optional[Dict] = None):
"""Append a new video."""
self.video_data.append(video)
if preprocess_kwargs:
self.video_data.append(
VideoData(video, preprocess_kwargs=preprocess_kwargs)
)
else:
self.video_data.append(video)
def append_audio(self, audio: str):
"""Append a new audio."""
python/sglang/srt/server_args.py
+35 -8
View File
@@ -77,6 +77,10 @@ logger = logging.getLogger(__name__)
# Define constants
DEFAULT_UVICORN_ACCESS_LOG_EXCLUDE_PREFIXES = ()
MIMO_V2_MODEL_ARCHS = (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
)
SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"}
@@ -1631,7 +1635,10 @@ class ServerArgs:
)
def _handle_model_specific_adjustments(self):
from sglang.srt.configs.model_config import is_deepseek_nsa
from sglang.srt.configs.model_config import (
get_mimo_v2_fused_qkv_expected_tp_size,
is_deepseek_nsa,
)
if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
return
@@ -1967,13 +1974,33 @@ class ServerArgs:
self.ep_size == 1
), "Triton kernel MoE is only supported when ep_size == 1"
elif any(
x in model_arch
for x in (
"MiMoV2ForCausalLM",
"MiMoV2FlashForCausalLM",
)
):
elif model_arch in MIMO_V2_MODEL_ARCHS:
if model_arch == "MiMoV2ForCausalLM":
expected_attn_tp_size = get_mimo_v2_fused_qkv_expected_tp_size(
hf_config
)
attn_dp_size = self.dp_size if self.enable_dp_attention else 1
effective_attn_tp_size = (
self.tp_size // attn_dp_size // self.attn_cp_size
)
if (
expected_attn_tp_size is not None
and effective_attn_tp_size != expected_attn_tp_size
):
raise ValueError(
"MiMoV2ForCausalLM requires effective attention TP "
f"size {expected_attn_tp_size} because its fused "
"qkv_proj weights are "
f"TP={expected_attn_tp_size}-interleaved; got "
f"{effective_attn_tp_size} "
f"(tp_size={self.tp_size}, dp_size={self.dp_size}, "
f"enable_dp_attention={self.enable_dp_attention}, "
f"attn_cp_size={self.attn_cp_size}). "
"Set --tp, --dp, --enable-dp-attention, and "
"--attention-context-parallel-size so the effective "
f"attention TP size is {expected_attn_tp_size}."
)
if self.speculative_algorithm == "EAGLE":
self.enable_multi_layer_eagle = True
logger.info(
python/sglang/srt/speculative/eagle_info.py
+3 -1
View File
@@ -74,10 +74,12 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
grammar: BaseGrammarObject = None
# Shape info for padding
num_tokens_per_req: int = -1
num_tokens_per_req: int = -1 # -1 auto-fills from draft_token_num.
def __post_init__(self):
super().__init__(SpecInputType.EAGLE_VERIFY)
if self.num_tokens_per_req < 0:
self.num_tokens_per_req = self.draft_token_num
def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
return self.draft_token_num, self.draft_token_num
python/sglang/srt/utils/common.py
+12 -1
View File
@@ -787,6 +787,13 @@ class ImageData:
url: str
detail: Optional[Literal["auto", "low", "high"]] = "auto"
max_dynamic_patch: Optional[int] = None
preprocess_kwargs: Optional[Dict] = None
@dataclass
class VideoData:
url: str
preprocess_kwargs: Optional[Dict] = None
image_extension_names = (".png", ".jpg", ".jpeg", ".webp", ".gif")
@@ -924,7 +931,11 @@ def _normalize_video_input(
return None
def load_video(video_file: Union[str, bytes], use_gpu: bool = True):
def load_video(video_file: Union[str, bytes, VideoData], use_gpu: bool = True):
if isinstance(video_file, VideoData):
# preprocess_kwargs is consumed by the multimodal processor, not here.
video_file = video_file.url
if isinstance(video_file, (list, tuple, torch.Tensor, np.ndarray)):
return video_file
python/sglang/test/server_fixtures/mmmu_fixture.py
+4 -1
View File
@@ -22,6 +22,8 @@ class MMMUServerBase(CustomTestCase):
This fixture handles server lifecycle for single-model MMMU tests.
For multi-model tests that need to start/stop servers within test methods,
use MMMUMultiModelTestBase instead.
Set server_api_key = None to launch without auth when sharing the server
with mixins whose clients do not send API keys.
"""
model = None
@@ -29,6 +31,7 @@ class MMMUServerBase(CustomTestCase):
timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
other_args: list[str] = []
mem_fraction_static: float = DEFAULT_MEM_FRACTION_STATIC
server_api_key = "sk-123456"
@classmethod
def setUpClass(cls):
@@ -53,7 +56,7 @@ class MMMUServerBase(CustomTestCase):
cls.model,
cls.base_url,
timeout=cls.timeout,
api_key=cls.api_key,
api_key=cls.server_api_key,
other_args=server_args,
env=process_env,
)