"""PyTorch Moss-VL model for SGLang - Qwen3VL Vision + Text with Cross Attention.""" from __future__ import annotations import logging from array import array from functools import partial from typing import Iterable, List, Optional, Tuple import torch import torch.nn as nn from einops import rearrange from transformers.activations import ACT2FN from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionRotaryEmbedding, ) from sglang.srt.environ import envs from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.conv import Conv3dLayer from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import ( MRotaryEmbedding, get_rope, ) from sglang.srt.layers.rotary_embedding.mrope import apply_interleaved_rope from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.managers.schedule_batch import MultimodalInputs from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) # Below this image count the per-image loop beats the vectorized path (which has a # fixed setup cost); both give the same result. _VECTORIZED_VL_POS_EMBED_MIN_IMAGES = 6 # ==================== Vision Components ==================== class MossVLVisionMLP(nn.Module): def __init__( self, in_features: int, hidden_features: int, bias: bool = True, hidden_act: str = "silu", quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.linear_fc1 = ColumnParallelLinear( in_features, hidden_features, bias=bias, quant_config=quant_config, prefix=add_prefix("linear_fc1", prefix), ) self.linear_fc2 = RowParallelLinear( hidden_features, in_features, bias=bias, quant_config=quant_config, prefix=add_prefix("linear_fc2", prefix), ) self.act = ACT2FN[hidden_act] def forward(self, x: torch.Tensor): x_fc1, _ = self.linear_fc1(x) mlp_output, _ = self.linear_fc2(self.act(x_fc1)) return mlp_output class MossVLVisionPatchEmbed(nn.Module): def __init__(self, config) -> None: super().__init__() self.patch_size = config.patch_size self.temporal_patch_size = config.temporal_patch_size self.in_channels = config.in_channels self.embed_dim = config.hidden_size kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] self.proj = Conv3dLayer( self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: target_dtype = self.proj.weight.dtype hidden_states = hidden_states.view( -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size, ) hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( -1, self.embed_dim ) return hidden_states class MossVLVisionBlock(nn.Module): def __init__( self, dim: int, num_heads: int, intermediate_dim: int, hidden_act: str = "silu", norm_layer=None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() if norm_layer is None: norm_layer = partial(nn.LayerNorm, eps=1e-6) self.norm1 = norm_layer(dim) self.norm2 = norm_layer(dim) self.attn = VisionAttention( embed_dim=dim, num_heads=num_heads, projection_size=dim, use_qkv_parallel=True, proj_bias=True, flatten_batch=True, quant_config=quant_config, prefix=add_prefix("attn", prefix), ) self.mlp = MossVLVisionMLP( dim, intermediate_dim, hidden_act=hidden_act, bias=True, quant_config=quant_config, prefix=f"{prefix}.mlp", ) def forward( self, x: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: torch.Tensor, ) -> torch.Tensor: hidden_states = self.norm1(x) hidden_states = rearrange(hidden_states, "s b ... -> b s ...") attn = self.attn( hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, ) attn = rearrange(attn, "b s ... -> s b ...") x = x + attn norm2 = self.norm2(x) mlp = self.mlp(norm2) x = x + mlp return x class MossVLVisionPatchMerger(nn.Module): """Merges spatial patches and concatenates deepstack features. Unlike Qwen3VL which uses separate merger modules per deepstack layer, Moss-VL concatenates all features and processes them through a single MLP. """ def __init__( self, config, num_deepstack_features: int = 0, norm_layer=None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() if norm_layer is None: norm_layer = partial(nn.LayerNorm, eps=1e-6) base_hidden_size = config.hidden_size * (config.spatial_merge_size**2) self.input_hidden_size = base_hidden_size * (1 + num_deepstack_features) self.hidden_size = config.hidden_size num_features = 1 + num_deepstack_features self.norms = nn.ModuleList( [norm_layer(config.hidden_size) for _ in range(num_features)] ) self.linear_fc1 = ColumnParallelLinear( self.input_hidden_size, self.input_hidden_size, bias=True, quant_config=quant_config, prefix=add_prefix("linear_fc1", prefix), ) self.act_fn = nn.GELU() self.linear_fc2 = RowParallelLinear( self.input_hidden_size, config.out_hidden_size, bias=True, quant_config=quant_config, prefix=add_prefix("linear_fc2", prefix), ) def forward( self, last_hidden_state: torch.Tensor, deepstack_features: List[torch.Tensor], ) -> torch.Tensor: all_inputs = [last_hidden_state] + deepstack_features outs = [] for i, feat in enumerate(all_inputs): outs.append(self.norms[i](feat)) x = torch.cat(outs, dim=-1) x = x.view(-1, self.input_hidden_size) x, _ = self.linear_fc1(x) x = self.act_fn(x) x, _ = self.linear_fc2(x) return x class MossVLVisionModel(nn.Module): """Moss-VL Vision Encoder (same architecture as Qwen3VL vision).""" def __init__( self, config, norm_eps: float = 1e-6, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.hidden_size = config.hidden_size self.num_heads = config.num_heads self.num_position_embeddings = config.num_position_embeddings self.patch_size = config.patch_size self.spatial_merge_size = config.spatial_merge_size self.spatial_merge_unit = self.spatial_merge_size**2 self.temporal_patch_size = config.temporal_patch_size self.deepstack_visual_indexes = config.deepstack_visual_indexes self.patch_embed = MossVLVisionPatchEmbed(config=config) self.pos_embed = nn.Embedding(self.num_position_embeddings, self.hidden_size) norm_layer = partial(nn.LayerNorm, eps=norm_eps) head_dim = self.hidden_size // self.num_heads self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) self.blocks = nn.ModuleList( [ MossVLVisionBlock( dim=self.hidden_size, num_heads=self.num_heads, intermediate_dim=config.intermediate_size, hidden_act=config.hidden_act, norm_layer=norm_layer, quant_config=quant_config, prefix=add_prefix(f"blocks.{i}", prefix), ) for i in range(config.depth) ] ) num_deepstack = len(self.deepstack_visual_indexes) self.merger = MossVLVisionPatchMerger( config=config, num_deepstack_features=num_deepstack, norm_layer=norm_layer, quant_config=quant_config, prefix=add_prefix("merger", prefix), ) @property def dtype(self) -> torch.dtype: return self.patch_embed.proj.weight.dtype @property def device(self) -> torch.device: return self.patch_embed.proj.weight.device def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: pos_ids = [] for t, h, w in grid_thw: 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).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).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 = int(grid_thw[:, 1:].max()) # transformers 5.12's rotary forward takes 1-D position_ids on the input device (grid_thw is CPU). rotary_pos_emb_full = self.rotary_pos_emb( torch.arange(max_grid_size, device=self.device) ) rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb def fast_pos_embed_interpolate(self, grid_thw: torch.Tensor) -> torch.Tensor: num_grid_per_side = int(self.num_position_embeddings**0.5) grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2] device = self.pos_embed.weight.device dtype = self.pos_embed.weight.dtype idx_parts = [[] for _ in range(4)] weight_parts = [[] for _ in range(4)] for _, h, w in zip(grid_ts, grid_hs, grid_ws): h_int, w_int = int(h.item()), int(w.item()) h_idxs = torch.linspace(0, num_grid_per_side - 1, h_int, device=device) w_idxs = torch.linspace(0, num_grid_per_side - 1, w_int, device=device) h_idxs_floor = h_idxs.int() w_idxs_floor = w_idxs.int() h_idxs_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1) w_idxs_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1) dh = h_idxs - h_idxs_floor dw = w_idxs - w_idxs_floor base_h = h_idxs_floor * num_grid_per_side base_h_ceil = h_idxs_ceil * num_grid_per_side indices = [ (base_h[None].T + w_idxs_floor[None]).flatten(), (base_h[None].T + w_idxs_ceil[None]).flatten(), (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), ] weights = [ ((1 - dh)[None].T * (1 - dw)[None]).flatten(), ((1 - dh)[None].T * dw[None]).flatten(), (dh[None].T * (1 - dw)[None]).flatten(), (dh[None].T * dw[None]).flatten(), ] for i in range(4): idx_parts[i].append(indices[i]) weight_parts[i].append(weights[i]) idx_tensor = torch.stack([torch.cat(parts) for parts in idx_parts]).to( dtype=torch.long ) weight_tensor = torch.stack([torch.cat(parts) for parts in weight_parts]).to( dtype=dtype ) pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None] patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] patch_pos_embeds = patch_pos_embeds.split( [int((h * w).item()) for h, w in zip(grid_hs, grid_ws)] ) m_size = self.spatial_merge_size patch_pos_embeds_permute = [] for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): t, h, w = int(t.item()), int(h.item()), int(w.item()) pos_embed = ( pos_embed.repeat(t, 1) .view(t, h // m_size, m_size, w // m_size, m_size, -1) .permute(0, 1, 3, 2, 4, 5) .flatten(0, 4) ) patch_pos_embeds_permute.append(pos_embed) return torch.cat(patch_pos_embeds_permute) def fast_pos_embed_interpolate_vectorized( self, grid_thw: torch.Tensor ) -> torch.Tensor: """Vectorized fast_pos_embed_interpolate (no per-image loop). Same result as the loop version; the cost no longer scales with the number of images. """ num_grid_per_side = int(self.num_position_embeddings**0.5) m = self.spatial_merge_size device = self.pos_embed.weight.device dtype = self.pos_embed.weight.dtype grid_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist() ts = [int(g[0]) for g in grid_list] hs = [int(g[1]) for g in grid_list] ws = [int(g[2]) for g in grid_list] num_images = len(grid_list) hw_list = [h * w for h, w in zip(hs, ws)] thw_list = [t * s for t, s in zip(ts, hw_list)] total_hw = sum(hw_list) total_out = sum(thw_list) def _exclusive_prefix(sizes): out, acc = [], 0 for s in sizes: out.append(acc) acc += s return torch.tensor(out, device=device, dtype=torch.long) hw_off = _exclusive_prefix(hw_list) thw_off = _exclusive_prefix(thw_list) image_arange = torch.arange(num_images, device=device) base_image_id = torch.repeat_interleave( image_arange, torch.tensor(hw_list, device=device) ) base_local = torch.arange(total_hw, device=device) - hw_off[base_image_id] w_of = torch.tensor(ws, device=device)[base_image_id] row = base_local // w_of col = base_local % w_of uniq_h, inv_h = torch.unique( torch.tensor(hs, device=device), return_inverse=True ) uniq_w, inv_w = torch.unique( torch.tensor(ws, device=device), return_inverse=True ) h_luts = [ torch.linspace(0, num_grid_per_side - 1, int(h), device=device) for h in uniq_h.tolist() ] w_luts = [ torch.linspace(0, num_grid_per_side - 1, int(w), device=device) for w in uniq_w.tolist() ] h_lut_off = _exclusive_prefix([len(x) for x in h_luts]) w_lut_off = _exclusive_prefix([len(x) for x in w_luts]) h_idxs = torch.cat(h_luts)[h_lut_off[inv_h[base_image_id]] + row] w_idxs = torch.cat(w_luts)[w_lut_off[inv_w[base_image_id]] + col] h_floor = h_idxs.int() w_floor = w_idxs.int() h_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1) w_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1) dh = h_idxs - h_floor dw = w_idxs - w_floor base_h = h_floor * num_grid_per_side base_h_ceil = h_ceil * num_grid_per_side indices = torch.stack( [ base_h + w_floor, base_h + w_ceil, base_h_ceil + w_floor, base_h_ceil + w_ceil, ], dim=0, ).to(dtype=torch.long) weights = torch.stack( [ (1 - dh) * (1 - dw), (1 - dh) * dw, dh * (1 - dw), dh * dw, ], dim=0, ).to(dtype=dtype) pe = self.pos_embed(indices) * weights[:, :, None] base_embeds = pe[0] + pe[1] + pe[2] + pe[3] # [total_hw, C] out_image_id = torch.repeat_interleave( image_arange, torch.tensor(thw_list, device=device) ) pos_in_image = torch.arange(total_out, device=device) - thw_off[out_image_id] hw_of_out = torch.tensor(hw_list, device=device)[out_image_id] frame_idx = pos_in_image // hw_of_out local_idx = pos_in_image % hw_of_out patch = base_embeds[hw_off[out_image_id] + local_idx] all_w = torch.tensor(ws, device=device)[out_image_id] rows = local_idx // all_w cols = local_idx % all_w out_within = ( frame_idx * hw_of_out + ((rows // m) * (all_w // m) + (cols // m)) * m * m + (rows % m) * m + (cols % m) ) merged = torch.empty_like(patch) merged[out_within + thw_off[out_image_id]] = patch return merged def forward( self, x: torch.Tensor, grid_thw: torch.Tensor, ) -> torch.Tensor: x = x.to(device=self.device, dtype=self.dtype) x = self.patch_embed(x) if ( envs.SGLANG_VIT_ENABLE_VECTORIZED_POS_EMBED.get() and grid_thw.shape[0] >= _VECTORIZED_VL_POS_EMBED_MIN_IMAGES ): pos_embeds = self.fast_pos_embed_interpolate_vectorized(grid_thw) else: pos_embeds = self.fast_pos_embed_interpolate(grid_thw) x = x + pos_embeds rotary_pos_emb = self.rot_pos_emb(grid_thw) seq_len, _ = x.size() rotary_pos_emb = rotary_pos_emb.to(x.device) rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) position_embeddings = (emb.cos(), emb.sin()) cu_seqlens = torch.repeat_interleave( grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] ).cumsum(dim=0) cu_seqlens = torch.cat( [ torch.zeros(1, dtype=torch.int32, device=cu_seqlens.device), cu_seqlens.to(torch.int32), ] ) x = x.unsqueeze(1) deepstack_features = [] for layer_idx, blk in enumerate(self.blocks): x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings) if layer_idx in self.deepstack_visual_indexes: deepstack_features.append(x) # Merger: concatenate last hidden state + deepstack features, then project x = self.merger(x, deepstack_features) return x def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set: stacked_params_mapping = [ ("attn.qkv.", "attn.q.", "q"), ("attn.qkv.", "attn.k.", "k"), ("attn.qkv.", "attn.v.", "v"), ] params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set = set() for name, loaded_weight in weights: 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) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params # ==================== Cross-Attention Components ==================== class MossVLTextCrossAttention(nn.Module): """Cross attention layer for Moss-VL: text queries attend to vision keys/values. Key differences from Mllama cross attention: - Uses separate q/k/v projections (q from text hidden states, k/v from vision states) - Applies RoPE to both query (text positions) and key (vision positions) - Uses QKVParallelLinear for the query projection (reusing text hidden_size) """ def __init__( self, config, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.config = config self.model_parallel_size = get_parallel().tp_size self.num_heads = config.num_attention_heads self.num_local_heads = self.num_heads // self.model_parallel_size self.num_key_value_heads = config.num_key_value_heads self.num_local_key_value_heads = ( self.num_key_value_heads // self.model_parallel_size ) self.hidden_size = config.hidden_size self.head_dim = getattr( config, "head_dim", config.hidden_size // self.num_heads ) self.layer_id = layer_id self.q_local_size = self.num_local_heads * self.head_dim self.kv_local_size = self.num_local_key_value_heads * self.head_dim self.scaling = self.head_dim**-0.5 # Query projection from text hidden states self.q_proj = ColumnParallelLinear( self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias, quant_config=quant_config, prefix=add_prefix("q_proj", prefix), ) # Key/Value projections from vision cross_attention_states self.k_proj = ColumnParallelLinear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, quant_config=quant_config, prefix=add_prefix("k_proj", prefix), ) self.v_proj = ColumnParallelLinear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, quant_config=quant_config, prefix=add_prefix("v_proj", prefix), ) self.o_proj = RowParallelLinear( self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias, input_is_parallel=True, quant_config=quant_config, prefix=add_prefix("o_proj", prefix), ) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.rope_theta = getattr(config, "rope_theta", 1000000) self.max_position_embeddings = getattr(config, "max_position_embeddings", 32768) rope_scaling = getattr(config, "rope_scaling", None) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=self.max_position_embeddings, base=self.rope_theta, rope_scaling=rope_scaling, ) self.attn = RadixAttention( self.num_local_heads, self.head_dim, self.scaling, self.num_local_key_value_heads, layer_id=layer_id, is_cross_attention=True, quant_config=quant_config, prefix=add_prefix("attn", prefix), ) def _apply_cross_attn_rotary( self, positions: torch.Tensor, states: torch.Tensor ) -> torch.Tensor: """Apply MRoPE to a single tensor (q or k) for cross-attention. Since q and k have different sequence lengths in cross-attention, we cannot use rotary_emb(positions, q, k) which assumes matching lengths. """ rotary_emb = self.rotary_emb num_tokens = positions.shape[-1] cos_sin = rotary_emb.cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) if positions.ndim == 2 and isinstance(rotary_emb, MRotaryEmbedding): if rotary_emb.mrope_section: if rotary_emb.mrope_interleaved: cos = apply_interleaved_rope(cos, rotary_emb.mrope_section) sin = apply_interleaved_rope(sin, rotary_emb.mrope_section) else: cos = torch.cat( [ m[i] for i, m in enumerate( cos.split(rotary_emb.mrope_section, dim=-1) ) ], dim=-1, ) sin = torch.cat( [ m[i] for i, m in enumerate( sin.split(rotary_emb.mrope_section, dim=-1) ) ], dim=-1, ) states_shape = states.shape states = states.view(num_tokens, -1, rotary_emb.head_size) states_rot = states[..., : rotary_emb.rotary_dim] states_pass = states[..., rotary_emb.rotary_dim :] states_rot = apply_rotary_emb(states_rot, cos, sin, rotary_emb.is_neox_style) states = torch.cat((states_rot, states_pass), dim=-1).reshape(states_shape) return states def forward( self, hidden_states: torch.Tensor, cross_attention_states: Optional[torch.Tensor], forward_batch: ForwardBatch, positions: torch.Tensor, vision_position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: # Query from text q, _ = self.q_proj(hidden_states) q = self.q_norm(q.reshape(-1, self.head_dim)).view(q.shape) if cross_attention_states is not None: # Key/Value from vision k, _ = self.k_proj(cross_attention_states) v, _ = self.v_proj(cross_attention_states) k = self.k_norm(k.reshape(-1, self.head_dim)).view(k.shape) # Apply RoPE: text positions for query, vision positions for key q = self._apply_cross_attn_rotary(positions, q) if cross_attention_states is not None and vision_position_ids is not None: k = self._apply_cross_attn_rotary(vision_position_ids, k) if cross_attention_states is None: k = None v = None output = self.attn(q, k, v, forward_batch) out, _ = self.o_proj(output) return out class MossVLCrossAttentionDecoderLayer(nn.Module): """Cross-attention transformer block with tanh-gated attention and feedforward.""" def __init__( self, config, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.layer_id = layer_id self.cross_attn = MossVLTextCrossAttention( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("cross_attn", prefix), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.cross_attn_attn_gate = nn.Parameter(torch.zeros(1)) self.mlp = MossVLTextMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) self.is_first_cross_attention_layer = ( bool(config.cross_attention_layers) and layer_id == config.cross_attention_layers[0] ) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.cross_attn_mlp_gate = nn.Parameter(torch.zeros(1)) def forward( self, hidden_states: torch.Tensor, cross_attention_states: Optional[torch.Tensor], cross_attention_mask: Optional[torch.Tensor], full_text_row_masked_out_mask: Optional[torch.Tensor], forward_batch: ForwardBatch, positions: torch.Tensor = None, vision_position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.cross_attn( hidden_states=hidden_states, cross_attention_states=cross_attention_states, forward_batch=forward_batch, positions=positions, vision_position_ids=vision_position_ids, ) hidden_states = full_text_row_masked_out_mask * hidden_states hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = full_text_row_masked_out_mask * hidden_states hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states return hidden_states class MossVLTextMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str = "silu", quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=add_prefix("gate_up_proj", prefix), ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, prefix=add_prefix("down_proj", prefix), ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. " "Only silu is supported for MossVLTextMLP." ) self.act_fn = SiluAndMul() def forward(self, x: torch.Tensor) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x # ==================== Self-Attention Decoder Layer ==================== class MossVLSelfAttention(nn.Module): """Self-attention for Moss-VL text model (same structure as Qwen3Attention).""" def __init__( self, config, layer_id: int = 0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.hidden_size = config.hidden_size self.tp_size = get_parallel().tp_size self.total_num_heads = config.num_attention_heads attn_tp_rank = get_parallel().attn_tp_rank attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size self.total_num_kv_heads = config.num_key_value_heads if self.total_num_kv_heads >= attn_tp_size: assert self.total_num_kv_heads % attn_tp_size == 0 else: assert attn_tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) self.head_dim = getattr( config, "head_dim", config.hidden_size // self.total_num_heads ) self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.rope_theta = getattr(config, "rope_theta", 1000000) self.max_position_embeddings = getattr(config, "max_position_embeddings", 32768) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.qkv_proj = QKVParallelLinear( config.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=config.attention_bias, quant_config=quant_config, tp_rank=attn_tp_rank, tp_size=attn_tp_size, prefix=add_prefix("qkv_proj", prefix), ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias, quant_config=quant_config, tp_rank=attn_tp_rank, tp_size=attn_tp_size, reduce_results=False, prefix=add_prefix("o_proj", prefix), ) rope_scaling = getattr(config, "rope_scaling", None) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=self.max_position_embeddings, base=self.rope_theta, rope_scaling=rope_scaling, ) self.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, prefix=add_prefix("attn", prefix), ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q = self.q_norm(q.reshape(-1, self.head_dim)).view(q.shape) k = self.k_norm(k.reshape(-1, self.head_dim)).view(k.shape) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v, forward_batch) output, _ = self.o_proj(attn_output) return output class MossVLSelfAttentionDecoderLayer(nn.Module): def __init__( self, config, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.layer_id = layer_id self.hidden_size = config.hidden_size self.self_attn = MossVLSelfAttention( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) self.mlp = MossVLTextMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) norm_kwargs = ( dict( weight_dtype=torch.float32, cast_x_before_out_mul=True, override_orig_dtype=torch.float32, fp32_residual=True, ) if get_server_args().rl_on_policy_target is not None else {} ) self.input_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs ) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs ) self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, is_layer_sparse=False, is_previous_layer_sparse=False, is_next_layer_sparse=False, ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: # Self Attention hidden_states, residual = self.layer_communicator.prepare_attn( hidden_states, residual, forward_batch ) if hidden_states.shape[0] != 0: hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) # MLP hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch, ) hidden_states = self.mlp(hidden_states) hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual # ==================== Text Model ==================== class MossVLTextModel(nn.Module): def __init__( self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, prefix=add_prefix("embed_tokens", prefix), ) self.cross_attention_layers = config.cross_attention_layers layers = [] for layer_id in range(config.num_hidden_layers): if layer_id in self.cross_attention_layers: layers.append( MossVLCrossAttentionDecoderLayer( config, layer_id, quant_config=quant_config, prefix=add_prefix(f"layers.{layer_id}", prefix), ) ) else: layers.append( MossVLSelfAttentionDecoderLayer( config, layer_id, quant_config=quant_config, prefix=add_prefix(f"layers.{layer_id}", prefix), ) ) self.layers = nn.ModuleList(layers) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, input_ids: torch.LongTensor, positions: torch.LongTensor, cross_attention_states: Optional[torch.Tensor], cross_attention_mask: Optional[torch.Tensor], full_text_row_masked_out_mask: Optional[torch.Tensor], forward_batch: ForwardBatch, skip_cross_attention: bool, vision_position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: hidden_states = self.embed_tokens(input_ids) residual = None for decoder_layer in self.layers: if isinstance(decoder_layer, MossVLCrossAttentionDecoderLayer): if not skip_cross_attention: # Fuse residual before cross-attention if residual is not None: hidden_states = hidden_states + residual residual = None hidden_states = decoder_layer( hidden_states=hidden_states, cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, forward_batch=forward_batch, positions=positions, vision_position_ids=vision_position_ids, ) elif isinstance(decoder_layer, MossVLSelfAttentionDecoderLayer): hidden_states, residual = decoder_layer( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, residual=residual, ) else: raise ValueError(f"Unknown decoder layer type {type(decoder_layer)}") if residual is not None: hidden_states, _ = self.norm(hidden_states, residual) else: hidden_states = self.norm(hidden_states) return hidden_states class MossVLForCausalLM(nn.Module): def __init__( self, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.vocab_size = config.vocab_size self.model = MossVLTextModel( config, quant_config, prefix=add_prefix("model", prefix) ) self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, org_num_embeddings=config.vocab_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) def forward( self, input_ids: torch.LongTensor, positions: torch.LongTensor, cross_attention_states: Optional[torch.Tensor], cross_attention_mask: Optional[torch.Tensor], full_text_row_masked_out_mask: Optional[torch.Tensor], forward_batch: ForwardBatch, skip_cross_attention: bool, vision_position_ids: Optional[torch.Tensor] = None, ) -> torch.Tensor: hidden_states = self.model( input_ids=input_ids, positions=positions, cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, forward_batch=forward_batch, skip_cross_attention=skip_cross_attention, vision_position_ids=vision_position_ids, ) return hidden_states # ==================== Main Model ==================== class MossVLForConditionalGeneration(nn.Module): def __init__(self, config, quant_config=None, prefix: str = ""): super().__init__() self.config = config self.quant_config = quant_config self.prefix = prefix vision_config = config.vision_config text_config = config.text_config self.spatial_merge_size = max( 1, int(getattr(vision_config, "spatial_merge_size", 2)) ) self.vision_seq_pad_multiple = 1 self.visual = MossVLVisionModel( vision_config, quant_config=quant_config, prefix=add_prefix("model.visual", prefix), ) self.language_model = MossVLForCausalLM( text_config, quant_config=quant_config, prefix=add_prefix("model.language_model", prefix), ) # Learnable separator token self.separator_token = nn.Parameter(torch.zeros(vision_config.out_hidden_size)) self.is_mrope_enabled = ( hasattr(text_config, "rope_scaling") and text_config.rope_scaling is not None and "mrope_section" in text_config.rope_scaling ) self.logits_processor = LogitsProcessor(text_config) def get_input_embeddings(self): return self.language_model.model.embed_tokens # ---- pad_input_ids (called at request scheduling time) ---- def _get_encoder_len(self, mm_inputs: MultimodalInputs) -> int: if not mm_inputs.mm_items: return 0 grid_thw = getattr(mm_inputs.mm_items[0], "grid_thw", None) if grid_thw is None: return 0 grid_thw = torch.as_tensor(grid_thw, dtype=torch.int64) if grid_thw.ndim == 1: grid_thw = grid_thw.unsqueeze(0) if grid_thw.numel() == 0: return 0 merge_square = self.spatial_merge_size**2 tokens_per_media = torch.prod(grid_thw, dim=1) // merge_square num_frames_per_media = grid_thw[:, 0] # Each frame contributes tokens_per_frame vision tokens + 1 separator total_len = int((tokens_per_media + num_frames_per_media).sum().item()) pad_multiple = self.vision_seq_pad_multiple if total_len % pad_multiple != 0: total_len = ((total_len + pad_multiple - 1) // pad_multiple) * pad_multiple return total_len def _build_encoder_prefix_pad_ids(self, mm_inputs: MultimodalInputs) -> array[int]: encoder_len = self._get_encoder_len(mm_inputs) if encoder_len == 0 or not mm_inputs.mm_items: return array("q") pad_value = mm_inputs.mm_items[0].pad_value return array("q", [pad_value]) * encoder_len def pad_input_ids( self, input_ids: array[int], mm_inputs: MultimodalInputs ) -> array[int]: encoder_len = self._get_encoder_len(mm_inputs) mm_inputs.num_image_tokens = encoder_len if encoder_len == 0: return input_ids return self._build_encoder_prefix_pad_ids(mm_inputs) + input_ids # ---- Collect and encode vision inputs ---- def _collect_mm_data(self, forward_batch: ForwardBatch): """Collect pixel_values, grid_thw, and vision_position_ids from uncached requests.""" if forward_batch.forward_mode.is_decode() or all(forward_batch.encoder_cached): return None, None, None pixel_values_list = [] grid_thw_list = [] vision_pos_ids_list = [] for i, mm_input in enumerate(forward_batch.mm_inputs): if forward_batch.encoder_cached[i] or mm_input is None: continue if not mm_input.mm_items: continue item = mm_input.mm_items[0] pixel_values_list.append(item.feature) grid_thw = getattr(item, "grid_thw", None) if grid_thw is not None: grid_thw_list.append(torch.as_tensor(grid_thw, dtype=torch.long)) encoder_len = forward_batch.encoder_lens_cpu[i] vp = mm_input.vision_position_ids if vp is not None: vision_pos_ids_list.append(vp[:, :encoder_len]) if not pixel_values_list: return None, None, None pixel_values = torch.cat(pixel_values_list, dim=0) grid_thw = torch.cat(grid_thw_list, dim=0) if grid_thw_list else None packed_vision_pos_ids = ( torch.cat(vision_pos_ids_list, dim=1) if vision_pos_ids_list else None ) return pixel_values, grid_thw, packed_vision_pos_ids def _get_vision_features( self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, ) -> torch.Tensor: """Run ViT encoder and insert separator tokens.""" hidden_states = self.visual(pixel_values, grid_thw=grid_thw) # hidden_states is packed: (total_vision_tokens, hidden_size) return hidden_states def _insert_separator_tokens( self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, ) -> torch.Tensor: """Insert separator token after each frame's vision tokens. Input: packed vision tokens from ViT (no separators) Output: packed vision tokens with separator tokens inserted after each frame """ merge_square = self.spatial_merge_size**2 tokens_per_media = ( grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2] ) // merge_square hidden_size = hidden_states.shape[-1] separator = self.separator_token.to(hidden_states.dtype) output_parts = [] src_offset = 0 for i in range(grid_thw.shape[0]): num_tokens = tokens_per_media[i].item() num_frames = grid_thw[i, 0].item() tokens_per_frame = num_tokens // num_frames media_hidden_states = hidden_states[ src_offset : src_offset + num_tokens ].view(num_frames, tokens_per_frame, hidden_size) separators = separator.view(1, 1, hidden_size).expand( num_frames, 1, hidden_size ) output_parts.append( torch.cat([media_hidden_states, separators], dim=1).flatten(0, 1) ) src_offset += num_tokens return torch.cat(output_parts, dim=0) # ---- prepare_forward_batch (called before attn backend init) ---- def prepare_forward_batch(self, forward_batch: ForwardBatch): """Build cross-attention custom mask before attn backend init. This hook is called by model_runner before init_forward_metadata so that the packed 1D mask is ready when FlashInfer plans cross-attention. Decode does not use a custom mask: newly generated tokens can attend to all encoder vision tokens. """ forward_batch.cross_attention_custom_mask = None if forward_batch.forward_mode.is_decode(): return if forward_batch.encoder_lens is None or forward_batch.encoder_lens.max() == 0: return custom_mask = self._build_cross_attention_custom_mask(forward_batch) if custom_mask is not None: forward_batch.cross_attention_custom_mask = custom_mask def _build_cross_attention_custom_mask( self, forward_batch: ForwardBatch ) -> Optional[torch.Tensor]: """Build packed 1D extend-stage cross-attention custom mask. The mask controls frame-level causal visibility: which vision frames each extend-stage text token can attend to during cross-attention. IMPORTANT: by the time ForwardBatch reaches the model, prepare_encoder_info_extend() has already stripped the encoder prefix from input_ids / seq_lens / extend_lens / prefix_lens. So the extend segment is purely decoder text — no encoder-prefix placeholder tokens. extend_prefix_len is the number of *cached text tokens*, and extend_seq_len is the number of *new text tokens* in this extend. Returns: 1D uint8 tensor of shape (sum_i(q_len_i * kv_len_i),) in FlashInfer packed row-major format, or None when no frame-level mask is needed. """ merge_square = self.spatial_merge_size**2 device = forward_batch.seq_lens.device mask_parts = [] need_mask = False for i in range(forward_batch.batch_size): encoder_len = forward_batch.encoder_lens_cpu[i] extend_seq_len = forward_batch.extend_seq_lens_cpu[i] extend_prefix_len = forward_batch.extend_prefix_lens_cpu[i] q_len = extend_seq_len kv_len = encoder_len if kv_len == 0 or q_len == 0: continue mm_input = forward_batch.mm_inputs[i] if forward_batch.mm_inputs else None if mm_input is None: mask_parts.append( torch.ones(q_len * kv_len, dtype=torch.uint8, device=device) ) continue visible_frame_counts = mm_input.visible_frame_counts if visible_frame_counts is None: mask_parts.append( torch.ones(q_len * kv_len, dtype=torch.uint8, device=device) ) continue item = mm_input.mm_items[0] if mm_input.mm_items else None grid_thw = getattr(item, "grid_thw", None) if item else None if grid_thw is None: mask_parts.append( torch.ones(q_len * kv_len, dtype=torch.uint8, device=device) ) continue need_mask = True grid_thw_t = torch.as_tensor(grid_thw, dtype=torch.long) if grid_thw_t.ndim == 1: grid_thw_t = grid_thw_t.unsqueeze(0) # Build frame_ranges: each frame's [start, end) in the encoder # token sequence (vision tokens + separator per frame). frame_ranges: List[Tuple[int, int]] = [] cursor = 0 for row_idx in range(grid_thw_t.shape[0]): t = grid_thw_t[row_idx, 0].item() h = grid_thw_t[row_idx, 1].item() w = grid_thw_t[row_idx, 2].item() span = (h * w) // merge_square + 1 for _ in range(t): frame_ranges.append((cursor, cursor + span)) cursor += span # The extend segment is purely text (encoder prefix already # stripped by prepare_encoder_info_extend). extend_prefix_len # is the cached-text offset into the full text sequence. text_offset = extend_prefix_len vis_counts = visible_frame_counts[text_offset : text_offset + q_len].to( device ) mask = torch.zeros(q_len, kv_len, dtype=torch.uint8, device=device) for f, (start, end) in enumerate(frame_ranges): clamped_end = min(end, kv_len) if start >= kv_len: break visible_rows = vis_counts > f if visible_rows.any(): mask[visible_rows, start:clamped_end] = 1 mask_parts.append(mask.flatten()) if not need_mask or not mask_parts: return None return torch.cat(mask_parts) # ---- full_text_row_masked_out_mask ---- def get_full_text_row_masked_out_mask( self, forward_batch: ForwardBatch ) -> torch.Tensor: """Create per-token mask that zeros cross-attn output for tokens that cannot see any vision token. HF semantics: a text token's cross-attn + cross-attn-MLP residuals are zeroed when that token has zero visible vision tokens. This is derived from the token-level cross_attention_mask, not just from whether the request has vision. For decode, HF copies the previous token's cross_attention_mask row to the new token. Since the processor's frame-level mask is prefix-causal, this reduces to copying the last prefill token's visibility. NOTE: prepare_encoder_info_extend() already strips encoder prefix tokens, so extend_seq_len / extend_prefix_len are purely text. extend_prefix_len is the cached-text offset into visible_frame_counts. """ encoder_lens_cpu = forward_batch.encoder_lens_cpu if forward_batch.forward_mode.is_decode(): device = forward_batch.encoder_lens.device full_text_row_masked_out_mask = forward_batch.encoder_lens != 0 if not forward_batch.mm_inputs: return full_text_row_masked_out_mask.reshape(-1, 1) bs = forward_batch.batch_size for i in range(bs): if not full_text_row_masked_out_mask[i]: continue mm_input = forward_batch.mm_inputs[i] visible_frame_counts = ( mm_input.visible_frame_counts if mm_input else None ) if visible_frame_counts is None: # Fall back to request-level gating only when frame-level # visibility metadata is unavailable. The request-level # encoder_lens signal already marks this row as visible. continue full_text_row_masked_out_mask[i] = visible_frame_counts[-1] > 0 else: device = forward_batch.seq_lens.device total_extend_len = int(forward_batch.extend_seq_lens.sum().item()) full_text_row_masked_out_mask = torch.zeros( total_extend_len, dtype=torch.bool, device=device ) offset = 0 for i in range(forward_batch.batch_size): encoder_len = encoder_lens_cpu[i] extend_seq_len = forward_batch.extend_seq_lens_cpu[i] extend_prefix_len = forward_batch.extend_prefix_lens_cpu[i] if extend_seq_len == 0: continue if encoder_len == 0: offset += extend_seq_len continue mm_input = ( forward_batch.mm_inputs[i] if forward_batch.mm_inputs else None ) visible_frame_counts = ( mm_input.visible_frame_counts if mm_input else None ) if visible_frame_counts is None: full_text_row_masked_out_mask[offset : offset + extend_seq_len] = ( True ) offset += extend_seq_len continue # The extend is purely text; extend_prefix_len is the # cached-text offset into the full text sequence. text_offset = extend_prefix_len vis_counts = visible_frame_counts[ text_offset : text_offset + extend_seq_len ].to(device) full_text_row_masked_out_mask[offset : offset + extend_seq_len] = ( vis_counts > 0 ) # Last prefill chunk for this request: decode will only need # visible_frame_counts[-1], so shrink the tensor to that single # element and drop the rest. .clone() detaches the view from # the original storage so the large tensor can be freed. if text_offset + extend_seq_len >= visible_frame_counts.shape[0]: mm_input.visible_frame_counts = visible_frame_counts[-1:].clone() offset += extend_seq_len return full_text_row_masked_out_mask.reshape(-1, 1) # ---- Forward ---- def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, get_embedding: bool = False, pp_proxy_tensors=None, ): if self.is_mrope_enabled: positions = forward_batch.mrope_positions # 1. Collect vision inputs for uncached requests pixel_values, grid_thw, vision_position_ids = self._collect_mm_data( forward_batch ) cross_attention_mask = None cross_attention_states = None if get_is_capture_mode(): skip_cross_attention = False else: assert len(forward_batch.encoder_lens) == len(forward_batch.seq_lens) skip_cross_attention = forward_batch.encoder_lens.max() == 0 # 2. Build full_text_row_masked_out_mask if not skip_cross_attention: full_text_row_masked_out_mask = self.get_full_text_row_masked_out_mask( forward_batch ) else: full_text_row_masked_out_mask = None # 3. Encode vision if needed if pixel_values is not None and grid_thw is not None: # Run ViT vision_hidden_states = self._get_vision_features(pixel_values, grid_thw) # Insert separator tokens after each frame. The result is already # packed (total_tokens, hidden_size) matching encoder_lens, so it # can be passed directly into the cross-attention path. cross_attention_states = self._insert_separator_tokens( vision_hidden_states, grid_thw ) # Drop heavy per-request vision tensors now that the encoder KV # has been produced and will be cached. Otherwise pixel_values and # vision_position_ids stay pinned on req.multimodal_inputs across # the entire decode phase. (visible_frame_counts is shrunk to a # single scalar element at the end of the last prefill chunk in # get_full_text_row_masked_out_mask, so decode still works.) # Note: the local `vision_position_ids` is still needed by the LM # cross-attention below, so we keep it; but we drop the per-request # copy on mm_input, which we won't read again. del pixel_values, vision_hidden_states for i, mm_input in enumerate(forward_batch.mm_inputs): if forward_batch.encoder_cached[i] or mm_input is None: continue mm_input.release_features() mm_input.vision_position_ids = None # 4. Run language model with cross attention hidden_states = self.language_model( input_ids=input_ids, positions=positions, cross_attention_states=cross_attention_states, cross_attention_mask=cross_attention_mask, full_text_row_masked_out_mask=full_text_row_masked_out_mask, forward_batch=forward_batch, skip_cross_attention=skip_cross_attention, vision_position_ids=vision_position_ids, ) return self.logits_processor( input_ids, hidden_states, self.language_model.lm_head, forward_batch, ) # ---- Weight Loading ---- def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), (".gate_up_proj", ".gate_proj", 0), (".gate_up_proj", ".up_proj", 1), ] params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue original_name = name # Map HF names to local module names. if name == "lm_head.weight": name = "language_model.lm_head.weight" elif name.startswith("model.language_model."): name = "language_model.model." + name[len("model.language_model.") :] elif name.startswith("model.visual."): name = name[len("model.") :] elif name.startswith("model.separator_token"): name = name[len("model.") :] # VisionAttention stores fused QKV weights under qkv_proj in SGLang. if "visual." in name: name = name.replace("attn.qkv.", "attn.qkv_proj.") handled = False if "visual." not in name and ".cross_attn." not in name: for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue mapped_name = name.replace(weight_name, param_name) if mapped_name.endswith(".bias") and mapped_name not in params_dict: handled = True break if mapped_name in params_dict: param = params_dict[mapped_name] param.weight_loader(param, loaded_weight, shard_id) handled = True break if handled: continue if name.endswith(".bias") and name not in params_dict: continue if name in params_dict: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) else: logger.debug(f"Skipping weight: {original_name} -> {name}") EntryClass = MossVLForConditionalGeneration