import logging import math import re from typing import Any, Dict, Iterable, List, Optional, Tuple import torch from sgl_kernel import fused_qk_norm_rope from torch import nn from transformers import PretrainedConfig from sglang_omni.models.qwen3_omni.hf_config import Qwen3OmniMoeTextConfig from sglang_omni.models.qwen3_omni.quantization import ( convert_fp8_weight_scale_inv_for_sglang, ) from sglang_omni.models.weight_loader import default_weight_loader from sglang_omni.utils import add_prefix from sglang_omni.vendor.sglang.core import ForwardBatch from sglang_omni.vendor.sglang.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang_omni.vendor.sglang.layers import ( LayerCommunicator, LayerScatterModes, MRotaryEmbedding, QKVParallelLinear, QuantizationConfig, RadixAttention, ReplicatedLinear, RMSNorm, RoutingMethodType, RowParallelLinear, TopK, VocabParallelEmbedding, get_attention_tp_rank, get_attention_tp_size, get_moe_impl_class, get_rope, should_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang_omni.vendor.sglang.models import ( apply_qk_norm, create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) from sglang_omni.vendor.sglang.server_args import get_global_server_args from sglang_omni.vendor.sglang.utils import make_layers logger = logging.getLogger(__name__) def _bind_default_weight_loaders(module: nn.Module) -> None: for param in module.parameters(): if not hasattr(param, "weight_loader"): param.weight_loader = default_weight_loader def compute_yarn_parameters( config: PretrainedConfig, ) -> tuple[float, float, float, float]: """ Refer to https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L197C1-L288C1 Computes the inverse frequencies with NTK scaling. Please refer to the [original paper](https://huggingface.co/papers/2309.00071) Args: config ([`~transformers.PretrainedConfig`]): The model configuration. Returns: factor: float, the scaling factor for the RoPE embeddings low: float, the lower bound of the dimension range high: float, the upper bound of the dimension range attention_factor: float, the post-processing scaling factor applied to the computed cos/sin """ # The config does not contain rope_scaling, which means the model is not using yarn rope_scaling = config.rope_scaling if rope_scaling is None: return 1.0, 0, 0, 1.0 base = config.rope_theta partial_rotary_factor = config.partial_rotary_factor head_dim = config.head_dim dim = int(head_dim * partial_rotary_factor) factor = rope_scaling.get("factor", 1.0) attention_factor = rope_scaling.get("attention_factor") mscale = rope_scaling.get("mscale") mscale_all_dim = rope_scaling.get("mscale_all_dim") if "original_max_position_embeddings" in rope_scaling: original_max_position_embeddings = rope_scaling[ "original_max_position_embeddings" ] factor = config.max_position_embeddings / original_max_position_embeddings else: original_max_position_embeddings = config.max_position_embeddings def get_mscale(scale, mscale=1): if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 # Sets the attention factor as suggested in the paper if attention_factor is None: if mscale and mscale_all_dim: attention_factor = float( get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim) ) else: attention_factor = get_mscale(factor) # Optional config options # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly) beta_fast = rope_scaling.get("beta_fast") or 32 beta_slow = rope_scaling.get("beta_slow") or 1 # Compute the inverse frequencies def find_correction_dim(num_rotations, dim, base, max_position_embeddings): """Inverse dimension formula to find the dimension based on the number of rotations""" return ( dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) ) / (2 * math.log(base)) def find_correction_range( low_rot, high_rot, dim, base, max_position_embeddings, truncate ): """Find dimension range bounds based on rotations""" low = find_correction_dim(low_rot, dim, base, max_position_embeddings) high = find_correction_dim(high_rot, dim, base, max_position_embeddings) if truncate: low = math.floor(low) high = math.ceil(high) return max(low, 0), min(high, dim - 1) truncate = rope_scaling.get("truncate", True) low, high = find_correction_range( beta_fast, beta_slow, dim, base, original_max_position_embeddings, truncate ) return factor, low, high, attention_factor class Qwen3OmniMoeThinkerTextAttention(nn.Module): """ Omni tailed version of qwen3_moe.py:Qwen3MoeAttention """ def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: int, layer_id: int = 0, rope_theta: float = 10000, rope_scaling: Optional[Dict[str, Any]] = None, max_position_embeddings: int = 8192, head_dim: Optional[int] = None, rms_norm_eps: float = 1e-06, attention_bias: bool = False, config: Optional[PretrainedConfig] = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", dual_chunk_attention_config: Optional[dict[str, Any]] = None, alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.hidden_size = hidden_size self.layer_id = layer_id attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() self.config = config self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= attn_tp_size: # Number of KV heads is greater than TP size, so we partition # the KV heads across multiple tensor parallel GPUs. assert self.total_num_kv_heads % attn_tp_size == 0 else: # Number of KV heads is less than TP size, so we replicate # the KV heads across multiple tensor parallel GPUs. 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 = head_dim or 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 = rope_theta self.max_position_embeddings = max_position_embeddings self.tp_rank = get_tensor_model_parallel_rank() self.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=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, hidden_size, bias=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), ) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=max_position_embeddings, base=rope_theta, rope_scaling=rope_scaling, dual_chunk_attention_config=dual_chunk_attention_config, ) self.compatible_with_fused_kv_buffer = ( False if isinstance(self.rotary_emb, MRotaryEmbedding) else True ) self.compatible_with_fused_qk_norm_rope = ( not isinstance(self.rotary_emb, MRotaryEmbedding) ) and self.head_dim in (64, 128, 256) self.use_fused_qk_norm_rope = ( get_global_server_args().enable_fused_qk_norm_rope and self.compatible_with_fused_qk_norm_rope ) self._used_fused_qk_norm_rope_last_call = False self._used_fused_set_kv_buffer_last_call = False 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), ) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.alt_stream = alt_stream def forward_prepare_native( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): qkv, _ = self.qkv_proj(hidden_states) q, k, v = self.apply_qk_norm_rope(qkv, positions, forward_batch) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state def apply_qk_norm_rope(self, qkv, positions, forward_batch): # Note:(Chenchen Hong) the talker uses a base (non-MRoPE) RotaryEmbedding # but post1 still passes MRoPE [3, seq] positions; collapse to the # temporal row so it isn't misread as 3 batches (all sections equal here). if positions.dim() == 2 and not isinstance(self.rotary_emb, MRotaryEmbedding): positions = positions[0] use_fused = self.use_fused_qk_norm_rope and qkv.dtype == torch.bfloat16 if use_fused: theta = self.config.rope_theta positions = ( positions.view(-1).to(dtype=torch.int32, device=qkv.device).contiguous() ) factor, low, high, attention_factor = compute_yarn_parameters(self.config) fused_qk_norm_rope( qkv, self.num_heads, self.num_kv_heads, self.num_kv_heads, self.head_dim, self.q_norm.variance_epsilon, self.q_norm.weight, self.k_norm.weight, theta, self.rotary_emb.is_neox_style, positions, factor, low, high, attention_factor, ) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) self._used_fused_qk_norm_rope_last_call = True self._used_fused_set_kv_buffer_last_call = False else: # Fallback to non-fused QK Norm & RoPE implementation q_linear, k_linear, v = qkv.split( [self.q_size, self.kv_size, self.kv_size], dim=-1 ) q, k = apply_qk_norm( q=q_linear, k=k_linear, q_norm=self.q_norm, k_norm=self.k_norm, head_dim=self.head_dim, alt_stream=self.alt_stream, ) use_fused_set_kv_buffer = ( enable_fused_set_kv_buffer(forward_batch) and self.compatible_with_fused_kv_buffer ) q, k = self.rotary_emb( positions, q, k, fused_set_kv_buffer_arg=( create_fused_set_kv_buffer_arg( value=v, layer=self.attn, forward_batch=forward_batch, ) if use_fused_set_kv_buffer else None ), ) self._used_fused_qk_norm_rope_last_call = False self._used_fused_set_kv_buffer_last_call = use_fused_set_kv_buffer return q, k, v def forward_prepare( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): if hidden_states.shape[0] == 0: return hidden_states, forward_batch, None return self.forward_prepare_native( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) def forward_core(self, intermediate_state): hidden_states, forward_batch, inner_state = intermediate_state if inner_state is None: return hidden_states q, k, v, fb = inner_state must_save_kv = self._used_fused_qk_norm_rope_last_call save_kv_cache = must_save_kv or not self._used_fused_set_kv_buffer_last_call attn_output = self.attn( q, k, v, fb, save_kv_cache=save_kv_cache, ) # Note:(Chenchen Hong) cast attn output to the compute dtype (v.dtype), # not o_proj.weight.dtype: for FP8 weights the latter feeds an fp8 # activation to the quantizer, which post1's sgl_kernel rejects. if attn_output.dtype != v.dtype: attn_output = attn_output.to(dtype=v.dtype) output, _ = self.o_proj(attn_output) return output def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: s = self.forward_prepare( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) return self.forward_core(s) class Qwen3OmniMoeThinkerTextSparseMoeBlock(nn.Module): """ Omni version of Qwen3MoeSparseMoeBlock """ def __init__( self, layer_id: int, config: Qwen3OmniMoeTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.layer_id = layer_id if self.tp_size > config.num_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_experts}." ) self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, use_grouped_topk=False, layer_id=layer_id, ) self.experts = get_moe_impl_class(quant_config)( num_experts=config.num_experts + get_global_server_args().ep_num_redundant_experts, top_k=config.num_experts_per_tok, layer_id=layer_id, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.Renormalize, ) self.gate = ReplicatedLinear( config.hidden_size, config.num_experts, bias=False, quant_config=None, prefix=add_prefix("gate", prefix), ) def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: return self.forward_normal( hidden_states, should_allreduce_fusion, use_reduce_scatter ) def forward_normal( self, hidden_states: torch.Tensor, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) if ( self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter and not should_use_flashinfer_cutlass_moe_fp4_allgather() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) class Qwen3OmniMoeThinkerTextDecoderLayer(nn.Module): def __init__( self, config: Qwen3OmniMoeTextConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.config = config self.hidden_size = config.hidden_size rope_theta = config.rope_theta rope_scaling = config.rope_scaling max_position_embeddings = config.max_position_embeddings head_dim = config.head_dim rms_norm_eps = config.rms_norm_eps attention_bias = config.attention_bias dual_chunk_attention_config = config.dual_chunk_attention_config self.self_attn = Qwen3OmniMoeThinkerTextAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, layer_id=layer_id, rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, head_dim=head_dim, rms_norm_eps=rms_norm_eps, attention_bias=attention_bias, config=config, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), dual_chunk_attention_config=dual_chunk_attention_config, alt_stream=alt_stream, ) self.layer_id = layer_id self.attn_tp_size = get_attention_tp_size() self.attn_tp_rank = get_attention_tp_rank() # Qwen3MoE all layers are sparse and have no nextn now self.is_layer_sparse = True is_previous_layer_sparse = True is_next_layer_sparse = True self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, is_layer_sparse=self.is_layer_sparse, is_previous_layer_sparse=is_previous_layer_sparse, is_next_layer_sparse=is_next_layer_sparse, ) if self.is_layer_sparse: self.mlp = Qwen3OmniMoeThinkerTextSparseMoeBlock( layer_id=self.layer_id, config=config, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) else: raise NotImplementedError( "Dense MLP is not implemented in Qwen3OmniMoeThinkerTextDecoderLayer yet." ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=(self.layer_id == self.config.num_hidden_layers - 1), ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = ( self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( hidden_states, residual, forward_batch, captured_last_layer_outputs=captured_last_layer_outputs, **kwargs, ) ) if hidden_states.shape[0] != 0: hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) should_allreduce_fusion = ( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( forward_batch ) ) # For DP with padding, reduce scatter can be used instead of all-reduce. use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( forward_batch ) hidden_states = self.mlp( hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter ) if should_allreduce_fusion: hidden_states._sglang_needs_allreduce_fusion = True else: hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual class Qwen3OmniMoeThinkerTextModel(nn.Module): """ Qwen3 omni text thinker only (without AuT and ViT) """ def __init__( self, config: Qwen3OmniMoeTextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.config = config self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, use_attn_tp_group=False, prefix=add_prefix(prefix, "embed_tokens"), ) alt_stream = torch.cuda.Stream() result = make_layers( config.num_hidden_layers, lambda idx, prefix: Qwen3OmniMoeThinkerTextDecoderLayer( layer_id=idx, config=config, quant_config=quant_config, prefix=prefix, alt_stream=alt_stream, ), prefix=add_prefix("layers", prefix), ) # make_layers returns (layers, start, end) with PP args, else just layers if isinstance(result, tuple): self.layers, self.start_layer, self.end_layer = result else: self.layers = result self.start_layer = 0 self.end_layer = config.num_hidden_layers self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # For EAGLE3 support self.layers_to_capture = [] _bind_default_weight_loaders(self) self._cached_params_dict = dict(self.named_parameters()) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, visual_pos_masks: Optional[torch.Tensor] = None, deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, ): if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds residual = None aux_hidden_states = [] # Capture word embeddings (before any transformer layer) if requested if "embed" in self.layers_to_capture: aux_hidden_states.append(("embed", hidden_states.clone())) for layer_idx in range(self.start_layer, self.end_layer): layer = self.layers[layer_idx] hidden_states, residual = layer( positions, hidden_states, forward_batch, residual, captured_last_layer_outputs=( aux_hidden_states if layer_idx in self.layers_to_capture else None ), ) if deepstack_visual_embeds is not None and layer_idx in range( len(deepstack_visual_embeds) ): hidden_states = self._deepstack_process( hidden_states, visual_pos_masks, deepstack_visual_embeds[layer_idx], ) if hidden_states.shape[0] != 0: if residual is None: hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states def _deepstack_process(self, hidden_states, visual_pos_masks, visual_embeds): # visual_pos_masks may be 1D boolean (SGLang path) or multi-dim (HF path) if visual_pos_masks.dim() > 1: visual_pos_masks = visual_pos_masks[..., 0] visual_pos_masks = visual_pos_masks.to(hidden_states.device) visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) local_this = hidden_states[visual_pos_masks, :] + visual_embeds hidden_states[visual_pos_masks, :] = local_this return hidden_states def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): """ " For Qwen3 omni thinker, the prefix of each part of weights are: (audio_tower): Qwen3OmniMoeAudioEncoder (visual): Qwen3OmniMoeVisionEncoder (model): Qwen3OmniMoeThinkerTextModel (lm_head): lm_head """ params_dict = self._cached_params_dict for name, loaded_weight in weights: if maybe_update_fused_qkv_proj( params_dict=params_dict, name=name, loaded_weight=loaded_weight, ): continue elif maybe_update_fused_moe_proj( params_dict=params_dict, name=name, loaded_weight=loaded_weight, config=self.config, ): continue else: if name in params_dict.keys(): param = params_dict[name] loaded_weight = convert_fp8_weight_scale_inv_for_sglang( name, loaded_weight ) param.weight_loader(param, loaded_weight) continue logger.warning(f"Parameter {name} not found in params_dict") def maybe_update_fused_qkv_proj( params_dict, name, loaded_weight, ): stacked_params_mapping = { "q_proj": ("qkv_proj", "q"), "k_proj": ("qkv_proj", "k"), "v_proj": ("qkv_proj", "v"), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } for shard_name in stacked_params_mapping: if shard_name in name: # Expert weights (mlp.experts.*) are handled by the MoE loader, # not the fused QKV/gate_up_proj path. if "mlp.experts" in name: continue fused_param_name, shard_id = stacked_params_mapping[shard_name] name = name.replace(shard_name, fused_param_name) param = params_dict[name] loaded_weight = convert_fp8_weight_scale_inv_for_sglang(name, loaded_weight) param.weight_loader(param, loaded_weight, shard_id) return True return False def maybe_update_fused_moe_proj(params_dict, name, loaded_weight, config): # replace FusedMoE.make_expert_params_mapping if res := extract_fused_experts( name=name, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=config.num_experts, ): param_name, weight_name, expert_id, shard_id = res name = name.replace(weight_name, param_name) if name in params_dict: param = params_dict[name] loaded_weight = convert_fp8_weight_scale_inv_for_sglang(name, loaded_weight) param.weight_loader( param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id, ) return True return False def extract_fused_experts( name, ckpt_gate_proj_name: str, ckpt_down_proj_name: str, ckpt_up_proj_name: str, num_experts: int, ): pattern = rf"experts\.(\d+)\.({ckpt_gate_proj_name}|{ckpt_down_proj_name}|{ckpt_up_proj_name})" match = re.search(pattern, name) if match: expert_id = int(match.group(1)) weight_type = match.group(2) if expert_id < num_experts: # Determine param_name based on weight_type param_name = ( "experts.w2_" if weight_type == ckpt_down_proj_name else "experts.w13_" ) if weight_type == ckpt_gate_proj_name: shard_id = "w1" elif weight_type == ckpt_down_proj_name: shard_id = "w2" elif weight_type == ckpt_up_proj_name: shard_id = "w3" return param_name, f"experts.{expert_id}.{weight_type}", expert_id, shard_id return None EntryClass = Qwen3OmniMoeThinkerTextModel