From 2c94fd2f2294d08c91f5f677d45ec4dcf7e8671f Mon Sep 17 00:00:00 2001 From: Yash Akhauri Date: Thu, 13 Aug 2026 00:59:37 +0000 Subject: [PATCH 1/5] Add exact MoVA serving to xLLM Serve the 36B MoVA checkpoint with packed Q/K/gate projections, output-sharded routed values, native router and norm semantics, ordinary RadixAttention caching, and live-update mappings. Reuse the fused-MoE first GEMM behind a compile-safe custom-op boundary and preserve the legacy K2 path. --- python/sglang/srt/layers/mova.py | 303 ++++++++++++ python/sglang/srt/models/xllm.py | 578 +++++++++++++++++++++-- test/registered/unit/layers/test_mova.py | 404 ++++++++++++++++ 3 files changed, 1237 insertions(+), 48 deletions(-) create mode 100644 python/sglang/srt/layers/mova.py create mode 100644 test/registered/unit/layers/test_mova.py diff --git a/python/sglang/srt/layers/mova.py b/python/sglang/srt/layers/mova.py new file mode 100644 index 000000000000..1b4edf6ea355 --- /dev/null +++ b/python/sglang/srt/layers/mova.py @@ -0,0 +1,303 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference primitives for mixture-of-value attention (MoVA).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.srt.utils import set_weight_attrs +from sglang.srt.utils.custom_op import register_custom_op + +_ROUTED_LINEAR_CHUNK_SIZE = 64 * 1024 + + +def mova_router_topk( + router_logits: torch.Tensor, + router_bias: Optional[torch.Tensor], + *, + score_func: str, + top_k: int, + scaling_factor: float, + renormalize: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply xLLM's selection-only router-bias semantics. + + Scores are computed in fp32. ``router_bias`` changes which value experts + are selected, but the mixture coefficients are gathered from the unbiased + scores. Scaling happens after optional top-k renormalization. + """ + + if top_k <= 0 or top_k > router_logits.shape[-1]: + raise ValueError( + f"top_k must be in [1, {router_logits.shape[-1]}], got {top_k}" + ) + if router_logits.is_cuda and (score_func == "sigmoid" or router_bias is None): + # Reuse SGLang's fused sigmoid/softmax top-k kernels. They implement + # the same selection-only correction-bias contract and return fp32 + # mixture weights plus int32 expert ids. + from sglang.srt.layers.moe.topk import fused_topk + + weights, selected = fused_topk( + hidden_states=router_logits, + gating_output=router_logits, + topk=top_k, + # Native xLLM leaves a top-1 route at its raw probability. + renormalize=renormalize and top_k > 1, + correction_bias=router_bias, + scoring_func=score_func, + ) + return (weights * scaling_factor).to(router_logits.dtype), selected + if score_func == "sigmoid": + scores = torch.sigmoid(router_logits.float()) + elif score_func == "softmax": + scores = F.softmax(router_logits, dim=-1, dtype=torch.float32) + else: + raise ValueError(f"Unsupported MoVA router score function: {score_func}") + + selection_scores = scores + if router_bias is not None: + selection_scores = selection_scores + router_bias.to(selection_scores) + + selected = torch.topk(selection_scores, top_k, dim=-1).indices + weights = torch.gather(scores, dim=-1, index=selected) + if renormalize and top_k > 1: + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * scaling_factor + return weights.to(router_logits.dtype), selected.to(torch.int32) + + +def routed_linear_reference( + hidden_states: torch.Tensor, + expert_weights: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, +) -> torch.Tensor: + """Straightforward MoVA value projection used as the correctness oracle.""" + + if hidden_states.ndim != 2: + raise ValueError("MoVA routed linear expects [tokens, hidden] inputs") + if expert_weights.ndim != 3: + raise ValueError("MoVA expert weights must be [experts, output, hidden]") + if routing_weights.shape != selected_experts.shape: + raise ValueError("MoVA routing weights and expert ids must have equal shape") + if hidden_states.shape[0] != selected_experts.shape[0]: + raise ValueError("MoVA route count must match the token count") + if hidden_states.shape[1] != expert_weights.shape[2]: + raise ValueError("MoVA input and expert hidden dimensions differ") + if hidden_states.shape[0] == 0: + return hidden_states.new_empty((0, expert_weights.shape[1])) + + # This deliberately favors clarity over memory use. Production CUDA paths + # use ``routed_linear`` below and never materialize selected expert weights. + selected_weights = expert_weights[selected_experts.to(torch.long)] + projected = torch.einsum("mknh,mh->mkn", selected_weights, hidden_states) + projected = F.silu(projected) + return (projected * routing_weights.to(projected).unsqueeze(-1)).sum(dim=1) + + +def _fake_mova_routed_linear_cuda( + hidden_states: torch.Tensor, + expert_weights: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, +) -> torch.Tensor: + return hidden_states.new_empty((hidden_states.shape[0], expert_weights.shape[1])) + + +@register_custom_op( + op_name="mova_routed_linear_cuda", + fake_impl=_fake_mova_routed_linear_cuda, +) +def _routed_linear_cuda_chunk( + hidden_states: torch.Tensor, + expert_weights: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, +) -> torch.Tensor: + # Keep these imports local: CPU model inspection and mapping tests should + # not initialize Triton or require the CUDA extension. + import triton.language as tl + + from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + invoke_fused_moe_kernel, + ) + from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_config import ( + get_config_dtype_str, + try_get_optimal_moe_config, + ) + from sglang.srt.layers.moe.fused_moe_triton.moe_align_block_size import ( + moe_align_block_size, + ) + + num_tokens = hidden_states.shape[0] + top_k = selected_experts.shape[1] + num_experts, output_size, input_size = expert_weights.shape + dtype_name = get_config_dtype_str(dtype=hidden_states.dtype) + # ``try_get_optimal_moe_config`` uses the last dimension of its synthetic + # second-GEMM shape as N. MoVA has no second GEMM, so describe the desired + # routed projection output explicitly. + config = try_get_optimal_moe_config( + expert_weights.shape, + (num_experts, input_size, output_size), + top_k, + dtype_name, + num_tokens, + ) + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + selected_experts, config["BLOCK_SIZE_M"], num_experts + ) + projected = torch.empty( + (num_tokens * top_k, output_size), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16 + invoke_fused_moe_kernel( + hidden_states, + expert_weights, + None, + projected, + None, + None, + None, + routing_weights, + selected_experts, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + False, # Routing weights are applied after SiLU. + top_k, + config, + compute_type=compute_type, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + per_channel_quant=False, + filter_expert=False, + ) + projected = F.silu(projected.view(num_tokens, top_k, output_size)) + return (projected * routing_weights.to(projected).unsqueeze(-1)).sum(dim=1) + + +def routed_linear( + hidden_states: torch.Tensor, + expert_weights: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, +) -> torch.Tensor: + """Run routed value projections using SGLang's fused-MoE first GEMM.""" + + if not hidden_states.is_cuda: + return routed_linear_reference( + hidden_states, expert_weights, routing_weights, selected_experts + ) + if hidden_states.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("Fused MoVA routed linear supports fp16 and bf16 only") + if hidden_states.shape[0] == 0: + return hidden_states.new_empty((0, expert_weights.shape[1])) + if not hidden_states.is_contiguous() or not expert_weights.is_contiguous(): + raise ValueError("Fused MoVA inputs and expert weights must be contiguous") + + outputs = [] + for begin in range(0, hidden_states.shape[0], _ROUTED_LINEAR_CHUNK_SIZE): + end = min(begin + _ROUTED_LINEAR_CHUNK_SIZE, hidden_states.shape[0]) + outputs.append( + _routed_linear_cuda_chunk( + hidden_states[begin:end], + expert_weights, + routing_weights[begin:end], + selected_experts[begin:end], + ) + ) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + + +class RoutedValueExperts(nn.Module): + """Persistent output-sharded MoVA value-expert weights.""" + + def __init__( + self, + num_experts: int, + input_size: int, + output_size: int, + *, + tp_rank: int, + tp_size: int, + ) -> None: + super().__init__() + if output_size % tp_size: + raise ValueError( + f"MoVA value width {output_size} is not divisible by TP={tp_size}" + ) + self.num_experts = num_experts + self.input_size = input_size + self.output_size = output_size + self.output_size_per_partition = output_size // tp_size + self.tp_rank = tp_rank + self.tp_size = tp_size + self.weight = nn.Parameter( + torch.empty(num_experts, self.output_size_per_partition, input_size), + requires_grad=False, + ) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: Optional[int] = None, + ) -> None: + output_begin = self.tp_rank * self.output_size_per_partition + if loaded_shard_id is None: + expected = (self.num_experts, self.output_size, self.input_size) + if tuple(loaded_weight.shape) != expected: + raise ValueError( + f"Packed MoVA value weight must be {expected}, got " + f"{tuple(loaded_weight.shape)}" + ) + local_weight = loaded_weight.narrow( + 1, output_begin, self.output_size_per_partition + ) + param.data.copy_(local_weight) + return + + if not 0 <= loaded_shard_id < self.num_experts: + raise ValueError(f"Invalid MoVA value expert id: {loaded_shard_id}") + expected = (self.output_size, self.input_size) + if tuple(loaded_weight.shape) != expected: + raise ValueError( + f"MoVA value expert must be {expected}, got {tuple(loaded_weight.shape)}" + ) + local_weight = loaded_weight.narrow( + 0, output_begin, self.output_size_per_partition + ) + param.data[loaded_shard_id].copy_(local_weight) + + def forward( + self, + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + ) -> torch.Tensor: + return routed_linear( + hidden_states, + self.weight, + routing_weights, + selected_experts, + ) diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index f0ac97e61df1..1d91e887923c 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -20,9 +20,10 @@ # - No shared_expert_gate (shared expert output added directly) # - Dense layers specified via mlp_only_layers config # - Partial RoPE (rope_head_dim < head_dim) -"""Inference-only Xllm K2MoE model compatible with HuggingFace weights.""" +"""Inference-only xLLM K2MoE and MoVA models compatible with HF weights.""" import logging +import math from contextlib import nullcontext from typing import Any, Dict, Iterable, List, Optional, Tuple, Union @@ -51,9 +52,11 @@ try: from sglang.srt.layers.communicator import enable_moe_dense_fully_dp except ImportError: + def enable_moe_dense_fully_dp(): return getattr(get_global_server_args(), "moe_dense_tp_size", -1) == 1 + from sglang.srt.layers.dp_attention import ( get_attention_tp_rank, get_attention_tp_size, @@ -69,13 +72,22 @@ class XllmGroupRMSNorm(nn.Module): Matches the HF XllmRMSNorm implementation used by the xllm model family. """ - def __init__(self, hidden_size: int, n_groups: int = 1, eps: float = 1e-6): + def __init__( + self, + hidden_size: int, + n_groups: int = 1, + eps: float = 1e-6, + zero_centered: bool = False, + ): super().__init__() self.n_groups = n_groups self.hidden_size = hidden_size assert hidden_size % n_groups == 0 self.variance_epsilon = eps - self.weight = nn.Parameter(torch.ones(hidden_size)) + self.zero_centered = zero_centered + self.weight = nn.Parameter( + torch.zeros(hidden_size) if zero_centered else torch.ones(hidden_size) + ) def forward(self, hidden_states, residual=None, post_residual_addition=None): if residual is not None: @@ -93,7 +105,8 @@ def forward(self, hidden_states, residual=None, post_residual_addition=None): hidden_states.pow(2).mean(-1, keepdim=True) + self.variance_epsilon ) hidden_states = hidden_states.reshape(*hidden_states.shape[:-2], -1) - hidden_states = (self.weight * hidden_states).to(orig_dtype) + weight = self.weight + 1.0 if self.zero_centered else self.weight + hidden_states = (weight * hidden_states).to(orig_dtype) if residual is not None: return hidden_states, residual return hidden_states @@ -102,14 +115,19 @@ def forward(self, hidden_states, residual=None, post_residual_addition=None): def _make_norm(config): """Create the appropriate RMSNorm for this config.""" n_groups = getattr(config, "layernorm_num_groups", 1) - if n_groups is None or n_groups <= 1: + is_mova = getattr(config, "num_values", 0) > 0 + if (n_groups is None or n_groups <= 1) and not is_mova: return RMSNorm(config.hidden_size, eps=config.rms_norm_eps) return XllmGroupRMSNorm( - config.hidden_size, n_groups=n_groups, eps=config.rms_norm_eps + config.hidden_size, + n_groups=n_groups or 1, + eps=config.rms_norm_eps, + zero_centered=is_mova, ) from sglang.srt.layers.linear import ( + ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, ReplicatedLinear, @@ -120,6 +138,8 @@ def _make_norm(config): from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.mova import RoutedValueExperts, mova_router_topk + try: from sglang.srt.layers.moe.utils import RoutingMethodType except ImportError: @@ -134,6 +154,7 @@ class RoutingMethodType(IntEnum): TopK = 5 Unspecified = 6 + try: from sglang.srt.layers.moe.utils import filter_moe_weight_param_global_expert except ImportError: @@ -144,6 +165,8 @@ def filter_moe_weight_param_global_expert(name, x, num_local_experts): and x.data.ndim > 0 and x.data.shape[0] == num_local_experts ) + + from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope @@ -162,7 +185,7 @@ def filter_moe_weight_param_global_expert(name, x, num_local_experts): is_cpu, is_cuda, make_layers, - use_intel_amx_backend, + set_weight_attrs, ) logger = logging.getLogger(__name__) @@ -172,6 +195,115 @@ def filter_moe_weight_param_global_expert(name, x, num_local_experts): _is_cpu_amx_available = cpu_has_amx_support() +def _validate_mova_config( + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], +) -> None: + """Fail early for phase-one combinations that cannot be served exactly.""" + + if getattr(config, "num_values", 0) <= 0: + return + if torch.get_default_dtype() != torch.bfloat16: + raise ValueError( + "MoVA phase 1 requires --dtype bfloat16. The converted 36B HF " + "artifact reports float32, so SGLang dtype=auto would select fp16." + ) + if quant_config is not None: + raise ValueError("MoVA phase 1 does not support quantized model weights") + if getattr(config, "attention_bias", False): + raise ValueError("MoVA phase 1 requires bias-free Q/K/V/O projections") + if getattr(config, "query_key_norm", False): + raise ValueError("MoVA phase 1 does not support query/key normalization") + if not getattr(config, "apply_attn_gate", False): + raise ValueError("MoVA phase 1 requires the xLLM attention gate") + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + if head_dim % 2: + raise ValueError(f"MoVA requires an even RoPE head dimension, got {head_dim}") + if config.num_attention_heads % config.num_key_value_heads: + raise ValueError("MoVA requires query heads to be divisible by KV heads") + if getattr(config, "rope_head_dim", head_dim) != head_dim: + raise ValueError("MoVA phase 1 requires full-head interleaved RoPE") + if getattr(config, "rope_scaling", None) is not None: + raise ValueError("MoVA phase 1 does not support non-default RoPE scaling") + if getattr(config, "sliding_window", None) is not None or getattr( + config, "use_sliding_window", False + ): + raise ValueError("MoVA phase 1 uses full causal RadixAttention only") + if getattr(config, "attn_gate_func", "silu") not in ("silu", "softplus"): + raise ValueError("MoVA supports only silu and softplus attention gates") + if getattr(config, "router_score_func", "sigmoid") not in ("sigmoid", "softmax"): + raise ValueError("MoVA supports only sigmoid and softmax value routing") + router_scale = getattr(config, "router_scaling_factor", 1.0) + if router_scale is None or not math.isfinite(router_scale) or router_scale <= 0: + raise ValueError( + f"MoVA requires a positive finite router scaling factor, got {router_scale}" + ) + num_dense_layers = getattr(config, "num_dense_layers", None) + if ( + num_dense_layers is None + or not 0 <= num_dense_layers <= config.num_hidden_layers + ): + raise ValueError( + "MoVA requires num_dense_layers in [0, num_hidden_layers], got " + f"{num_dense_layers}" + ) + expected_dense_layers = list(range(num_dense_layers)) + if list(getattr(config, "mlp_only_layers", [])) != expected_dense_layers: + raise ValueError( + "MoVA phase 1 requires dense attention and dense FFN prefix layers to " + f"match exactly; expected mlp_only_layers={expected_dense_layers}" + ) + if getattr(config, "decoder_sparse_step", 1) != 1: + raise ValueError("MoVA phase 1 requires decoder_sparse_step=1") + num_values = config.num_values + top_k = getattr(config, "num_values_per_tok", 0) + if not 0 < top_k <= num_values: + raise ValueError( + f"num_values_per_tok must be in [1, {num_values}], got {top_k}" + ) + if getattr(config, "num_experts", 0) <= 0: + raise ValueError("MoVA requires sparse MoE feed-forward layers") + n_groups = getattr(config, "layernorm_num_groups", 1) or 1 + if config.hidden_size % n_groups: + raise ValueError( + f"hidden size {config.hidden_size} is not divisible by {n_groups} norm groups" + ) + attn_tp_size = get_attention_tp_size() + if config.num_attention_heads % attn_tp_size: + raise ValueError(f"MoVA query heads must be divisible by TP={attn_tp_size}") + if config.num_key_value_heads % attn_tp_size: + raise ValueError( + "MoVA phase 1 requires TP <= KV heads and KV heads divisible by TP; " + f"got TP={attn_tp_size}, KV heads={config.num_key_value_heads}" + ) + + +def _xllm_stacked_params_mapping(config: PretrainedConfig): + if getattr(config, "num_values", 0) <= 0: + return [ + (".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), + ] + + mapping = [ + (".qkg_proj", ".q_proj", "q"), + (".qkg_proj", ".attn_gate_proj", "gate"), + (".qkg_proj", ".k_proj", "k"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + mapping.extend( + (".v_experts.weight", f".v_experts.{expert_id}.weight", expert_id) + for expert_id in range(config.num_values) + ) + return mapping + + def permute_to_xllm(x): """Interleave first half and second half: [0,1,...,63,64,...,127] -> [0,64,1,65,...,63,127]""" return x.reshape(*x.shape[:-1], 2, -1).transpose(-1, -2).reshape(*x.shape[:-1], -1) @@ -182,6 +314,149 @@ def permute_to_hf(x): return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1) +def _interleave_rope_weight(weight: torch.Tensor, num_heads: int) -> torch.Tensor: + """Convert HF/NeoX Q or K rows to native interleaved RoPE order.""" + + if weight.shape[0] % num_heads: + raise ValueError( + f"Projection width {weight.shape[0]} is not divisible by {num_heads} heads" + ) + head_dim = weight.shape[0] // num_heads + if head_dim % 2: + raise ValueError(f"RoPE head dimension must be even, got {head_dim}") + rows = weight.reshape(num_heads, head_dim, weight.shape[1]) + rows = rows.reshape(num_heads, 2, head_dim // 2, weight.shape[1]) + return rows.transpose(1, 2).reshape_as(weight).contiguous() + + +class XllmQKGParallelLinear(nn.Module): + """One TP-sharded Q/K/gate projection packed by local GQA group. + + The persistent layout matches Megatron's ``linear_qkg`` contract: + ``[Q heads in group, gate heads in group, K head]`` for every KV group. + Canonical HF Q/K tensors are converted to interleaved RoPE row order while + loading, removing a permutation from every forward pass. + """ + + _SHARD_IDS = {"q", "gate", "k"} + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + *, + tp_rank: int, + tp_size: int, + ) -> None: + super().__init__() + if num_heads % num_kv_heads: + raise ValueError( + f"Query heads {num_heads} must be divisible by KV heads {num_kv_heads}" + ) + if num_kv_heads % tp_size: + raise ValueError( + "MoVA requires attention TP to divide the number of KV heads; " + f"got TP={tp_size}, KV heads={num_kv_heads}" + ) + self.hidden_size = hidden_size + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.tp_rank = tp_rank + self.tp_size = tp_size + self.num_kv_heads = num_kv_heads // tp_size + self.queries_per_kv = num_heads // num_kv_heads + self.group_width = 2 * self.queries_per_kv + 1 + self.weight = nn.Parameter( + torch.empty( + self.num_kv_heads * self.group_width * head_dim, + hidden_size, + ), + requires_grad=False, + ) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + @property + def q_size(self) -> int: + return self.num_kv_heads * self.queries_per_kv * self.head_dim + + @property + def kv_size(self) -> int: + return self.num_kv_heads * self.head_dim + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str, + ) -> None: + if loaded_shard_id not in self._SHARD_IDS: + raise ValueError(f"Invalid QKG shard id: {loaded_shard_id}") + param_data = param.data.view( + self.num_kv_heads, + self.group_width, + self.head_dim, + self.hidden_size, + ) + kv_begin = self.tp_rank * self.num_kv_heads + + if loaded_shard_id in ("q", "gate"): + expected_heads = self.total_num_heads + expected = (expected_heads * self.head_dim, self.hidden_size) + if tuple(loaded_weight.shape) != expected: + raise ValueError( + f"{loaded_shard_id} projection must be {expected}, got " + f"{tuple(loaded_weight.shape)}" + ) + if loaded_shard_id == "q": + loaded_weight = _interleave_rope_weight( + loaded_weight, self.total_num_heads + ) + local = loaded_weight.view( + self.total_num_kv_heads, + self.queries_per_kv, + self.head_dim, + self.hidden_size, + ).narrow(0, kv_begin, self.num_kv_heads) + offset = 0 if loaded_shard_id == "q" else self.queries_per_kv + param_data[:, offset : offset + self.queries_per_kv].copy_(local) + return + + expected = (self.total_num_kv_heads * self.head_dim, self.hidden_size) + if tuple(loaded_weight.shape) != expected: + raise ValueError( + f"k projection must be {expected}, got {tuple(loaded_weight.shape)}" + ) + loaded_weight = _interleave_rope_weight(loaded_weight, self.total_num_kv_heads) + local = loaded_weight.view( + self.total_num_kv_heads, 1, self.head_dim, self.hidden_size + ).narrow(0, kv_begin, self.num_kv_heads) + param_data[:, -1:].copy_(local) + + def forward( + self, hidden_states: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qkg = F.linear(hidden_states, self.weight) + qkg = qkg.view( + hidden_states.shape[0], + self.num_kv_heads, + self.group_width, + self.head_dim, + ) + q = qkg[:, :, : self.queries_per_kv].reshape( + hidden_states.shape[0], self.q_size + ) + gate = qkg[ + :, + :, + self.queries_per_kv : 2 * self.queries_per_kv, + ].reshape(hidden_states.shape[0], self.q_size) + k = qkg[:, :, -1].reshape(hidden_states.shape[0], self.kv_size) + return q, k, gate + + class XllmMLP(nn.Module): def __init__( self, @@ -549,9 +824,7 @@ def _apply_partial_rope( k_rope_flat = permute_to_hf(k_rope).reshape( -1, self.num_kv_heads * self.rope_head_dim ) - q_rope_flat, k_rope_flat = self.rotary_emb( - positions, q_rope_flat, k_rope_flat - ) + q_rope_flat, k_rope_flat = self.rotary_emb(positions, q_rope_flat, k_rope_flat) q_rope = permute_to_xllm( q_rope_flat.reshape(-1, self.num_heads, self.rope_head_dim) @@ -587,6 +860,193 @@ def forward( return output +class _XllmMoVAAttentionBase(nn.Module): + """Shared gated-GQA path for dense and routed-value MoVA layers.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig], + prefix: str, + ) -> None: + super().__init__() + if quant_config is not None: + raise ValueError("MoVA phase 1 supports unquantized bf16/fp16 weights only") + + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + self.rope_head_dim = getattr(config, "rope_head_dim", self.head_dim) + self.apply_attn_gate = getattr(config, "apply_attn_gate", False) + self.attn_gate_func = getattr(config, "attn_gate_func", "silu") + self.scaling = self.head_dim**-0.5 + + self.tp_rank = get_attention_tp_rank() + self.tp_size = get_attention_tp_size() + if self.total_num_heads % self.tp_size: + raise ValueError( + f"Attention heads {self.total_num_heads} are not divisible by TP={self.tp_size}" + ) + if self.total_num_kv_heads % self.tp_size: + raise ValueError( + "MoVA phase 1 requires TP <= KV heads and KV heads divisible by TP; " + f"got TP={self.tp_size}, KV heads={self.total_num_kv_heads}" + ) + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = self.total_num_kv_heads // self.tp_size + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.qkg_proj = XllmQKGParallelLinear( + config.hidden_size, + self.total_num_heads, + self.total_num_kv_heads, + self.head_dim, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + reduce_results=False, + prefix=add_prefix("o_proj", prefix), + ) + # HF Q/K rows are converted to this interleaved convention by the QKG + # loader, so no per-token permutation is needed here. + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rope_head_dim, + max_position=getattr(config, "max_position_embeddings", 8192), + base=getattr(config, "rope_theta", 10000), + rope_scaling=getattr(config, "rope_scaling", None), + is_neox_style=False, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=None, + prefix=add_prefix("attn", prefix), + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def _activate_gate(self, gate: torch.Tensor) -> torch.Tensor: + if self.attn_gate_func == "silu": + return F.silu(gate) + if self.attn_gate_func == "softplus": + return F.softplus(gate, beta=math.log(2)) + raise ValueError( + f"Unsupported xLLM attention gate function: {self.attn_gate_func}" + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + q, k, gate = self.qkg_proj(hidden_states) + value = self._project_value(hidden_states) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, value, forward_batch) + if self.apply_attn_gate: + attn_output = attn_output * self._activate_gate(gate) + output, _ = self.o_proj(attn_output) + return output + + +class XllmGatedAttention(_XllmMoVAAttentionBase): + """Dense GQA used by the prefix layers of a MoVA checkpoint.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + self.v_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=False, + quant_config=None, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + prefix=add_prefix("v_proj", prefix), + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + value, _ = self.v_proj(hidden_states) + return value + + +class XllmMoVAAttention(_XllmMoVAAttentionBase): + """Sparse MoVA attention with output-sharded routed value experts.""" + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + self.num_values = config.num_values + self.num_values_per_tok = config.num_values_per_tok + self.router_score_func = getattr(config, "router_score_func", "sigmoid") + self.router_scaling_factor = getattr(config, "router_scaling_factor", 1.0) + self.renormalize = getattr(config, "norm_topk_prob", True) + self.v_router = ReplicatedLinear( + config.hidden_size, + self.num_values, + bias=False, + quant_config=None, + prefix=add_prefix("v_router", prefix), + ) + if getattr(config, "moe_gate_bias", False): + # SGLang's fused sigmoid top-k requires correction bias in fp32. + # It remains a loadable parameter for Miles weight updates, but is + # never included in the router logits matmul. + self.v_router.bias = nn.Parameter( + torch.empty(self.num_values, dtype=torch.float32), + requires_grad=False, + ) + self.v_experts = RoutedValueExperts( + self.num_values, + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) + + def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Router bias is deliberately omitted from the logits matmul. It only + # changes route selection inside ``mova_router_topk``. + router_logits = F.linear(hidden_states, self.v_router.weight) + routing_weights, selected_values = mova_router_topk( + router_logits, + self.v_router.bias, + score_func=self.router_score_func, + top_k=self.num_values_per_tok, + scaling_factor=self.router_scaling_factor, + renormalize=self.renormalize, + ) + return self.v_experts(hidden_states, routing_weights, selected_values) + + class XllmDecoderLayer(nn.Module): def __init__( self, @@ -603,23 +1063,10 @@ def __init__( rope_scaling = getattr(config, "rope_scaling", None) max_position_embeddings = getattr(config, "max_position_embeddings", 8192) qkv_bias = getattr(config, "attention_bias", False) - head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - rope_head_dim = getattr(config, "rope_head_dim", head_dim) - - self.self_attn = XllmAttention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - head_dim=head_dim, - rope_head_dim=rope_head_dim, - layer_id=layer_id, - rope_theta=rope_theta, - rope_scaling=rope_scaling, - max_position_embeddings=max_position_embeddings, - qkv_bias=qkv_bias, - quant_config=quant_config, - prefix=add_prefix("self_attn", prefix), + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads ) + rope_head_dim = getattr(config, "rope_head_dim", head_dim) self.layer_id = layer_id @@ -630,20 +1077,50 @@ def __init__( mlp_only_layers = getattr(config, "mlp_only_layers", []) decoder_sparse_step = getattr(config, "decoder_sparse_step", 1) if (layer_id not in mlp_only_layers) and ( - config.num_experts > 0 - and (layer_id + 1) % decoder_sparse_step == 0 + config.num_experts > 0 and (layer_id + 1) % decoder_sparse_step == 0 ): self.is_layer_sparse = True else: self.is_layer_sparse = False + is_mova_config = getattr(config, "num_values", 0) > 0 + is_mova_attention = is_mova_config and layer_id >= config.num_dense_layers + if is_mova_attention: + self.self_attn = XllmMoVAAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + elif is_mova_config: + self.self_attn = XllmGatedAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + else: + self.self_attn = XllmAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + qkv_bias=qkv_bias, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + # Check neighbors for scatter modes def _is_sparse(lid): if lid < 0 or lid >= config.num_hidden_layers: return False return (lid not in mlp_only_layers) and ( - config.num_experts > 0 - and (lid + 1) % decoder_sparse_step == 0 + config.num_experts > 0 and (lid + 1) % decoder_sparse_step == 0 ) is_previous_layer_sparse = _is_sparse(layer_id - 1) @@ -726,7 +1203,9 @@ def forward( ) if isinstance(self.mlp, XllmMLP): - hidden_states = self.mlp(hidden_states, use_reduce_scatter=use_reduce_scatter) + hidden_states = self.mlp( + hidden_states, use_reduce_scatter=use_reduce_scatter + ) else: hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) @@ -821,7 +1300,8 @@ def forward( for i in range(self.start_layer, self.end_layer): _server_args = get_global_server_args() _disable_pcg = getattr( - _server_args, "disable_piecewise_cuda_graph", + _server_args, + "disable_piecewise_cuda_graph", not getattr(_server_args, "enable_piecewise_cuda_graph", True), ) ctx = ( @@ -875,6 +1355,7 @@ def __init__( self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config + _validate_mova_config(config, quant_config) # Keep the 375B MoE path single-stream. CUDA graph capture is still # allowed, but the capture-mode dual-stream path adds async # shared-expert/router risk without being required for correctness. @@ -895,6 +1376,17 @@ def __init__( self.logits_processor = LogitsProcessor(config) self.capture_aux_hidden_states = False + # Value experts are shards of one attention-TP parameter, not FFN/EP + # experts. ParameterMapper therefore stages all 64 canonical HF shards + # before writing the persistent packed tensor during live updates. + self.stacked_params_mapping = _xllm_stacked_params_mapping(config) + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + @torch.no_grad() def forward( self, @@ -920,7 +1412,10 @@ def forward( ) if logits_output.next_token_logits is not None: logits_output.next_token_logits = torch.nan_to_num( - logits_output.next_token_logits, nan=0.0, posinf=65504.0, neginf=-65504.0 + logits_output.next_token_logits, + nan=0.0, + posinf=65504.0, + neginf=-65504.0, ) return logits_output else: @@ -935,21 +1430,8 @@ def end_layer(self): return self.model.end_layer 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), - ] - - expert_params_mapping = FusedMoE.make_expert_params_mapping( - ckpt_gate_proj_name="gate_proj", - ckpt_down_proj_name="down_proj", - ckpt_up_proj_name="up_proj", - num_experts=self.config.num_experts, - ) + stacked_params_mapping = self.stacked_params_mapping + expert_params_mapping = self.expert_params_mapping params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: diff --git a/test/registered/unit/layers/test_mova.py b/test/registered/unit/layers/test_mova.py new file mode 100644 index 000000000000..3fb08fbed4fa --- /dev/null +++ b/test/registered/unit/layers/test_mova.py @@ -0,0 +1,404 @@ +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from sglang.srt.layers.mova import ( + RoutedValueExperts, + mova_router_topk, + routed_linear, + routed_linear_reference, +) +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.model_loader.parameter_mapper import ParameterMapper +from sglang.srt.models.xllm import ( + XllmGroupRMSNorm, + XllmQKGParallelLinear, + _interleave_rope_weight, + _validate_mova_config, + _xllm_stacked_params_mapping, + _XllmMoVAAttentionBase, +) + + +@pytest.fixture(autouse=True) +def _server_args_for_kernel_helpers(monkeypatch): + """ModelRunner normally installs these globals before model creation.""" + + monkeypatch.setattr( + "sglang.srt.server_args._global_server_args", + SimpleNamespace( + enable_deterministic_inference=False, + rl_on_policy_target=None, + ), + ) + + +def _valid_mova_config(**overrides): + values = dict( + num_values=64, + num_values_per_tok=4, + num_hidden_layers=48, + num_dense_layers=3, + mlp_only_layers=[0, 1, 2], + decoder_sparse_step=1, + num_experts=100, + hidden_size=2560, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + rope_head_dim=128, + attention_bias=False, + query_key_norm=False, + apply_attn_gate=True, + attn_gate_func="softplus", + rope_scaling=None, + sliding_window=None, + use_sliding_window=False, + router_score_func="sigmoid", + router_scaling_factor=2.5, + layernorm_num_groups=2, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def test_mova_router_bias_changes_selection_not_weight(): + logits = torch.tensor([[2.0, 1.0, -1.0]], dtype=torch.float32) + bias = torch.tensor([0.0, 0.0, 10.0], dtype=torch.float32) + + weights, selected = mova_router_topk( + logits, + bias, + score_func="sigmoid", + top_k=1, + scaling_factor=2.5, + ) + + assert selected.tolist() == [[2]] + expected = torch.sigmoid(logits)[0, 2] * 2.5 + torch.testing.assert_close(weights[0, 0], expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("top_k", [1, 4]) +def test_fused_mova_router_matches_selection_only_reference(top_k): + torch.manual_seed(5) + logits = torch.randn(19, 64, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(64, device="cuda", dtype=torch.float32) + + weights, selected = mova_router_topk( + logits, + bias, + score_func="sigmoid", + top_k=top_k, + scaling_factor=2.5, + ) + scores = torch.sigmoid(logits.float()) + expected_selected = torch.topk(scores + bias, top_k, dim=-1).indices + expected_weights = torch.gather(scores, 1, expected_selected) + if top_k > 1: + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + expected_weights = (expected_weights * 2.5).to(logits.dtype) + + torch.testing.assert_close(selected.long(), expected_selected) + torch.testing.assert_close(weights, expected_weights, rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("top_k", [1, 2, 4]) +def test_routed_linear_reference_matches_explicit_mixture(top_k): + torch.manual_seed(7) + hidden = torch.randn(5, 6) + experts = torch.randn(4, 3, 6) + selected = torch.stack([torch.randperm(4)[:top_k] for _ in range(hidden.shape[0])]) + routing = torch.rand(hidden.shape[0], top_k) + + actual = routed_linear_reference(hidden, experts, routing, selected) + expected = torch.zeros_like(actual) + for token in range(hidden.shape[0]): + for slot in range(top_k): + expert = selected[token, slot] + expected[token] += routing[token, slot] * F.silu( + F.linear(hidden[token], experts[expert]) + ) + + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("num_tokens", [1, 17, 257]) +def test_fused_routed_linear_matches_reference(num_tokens): + torch.manual_seed(11) + device = torch.device("cuda") + hidden = torch.randn(num_tokens, 64, device=device, dtype=torch.bfloat16) + experts = torch.randn(8, 32, 64, device=device, dtype=torch.bfloat16) + # Include a hot expert while leaving at least one expert empty. + selected = torch.randint(0, 7, (num_tokens, 4), device=device, dtype=torch.int32) + selected[:, 0] = 0 + routing = torch.rand(num_tokens, 4, device=device, dtype=torch.bfloat16) + + expected = routed_linear_reference(hidden, experts, routing, selected) + actual = routed_linear(hidden, experts, routing, selected) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_fused_routed_linear_matches_36b_tp8_shape(): + torch.manual_seed(17) + hidden = torch.randn(9, 2560, device="cuda", dtype=torch.bfloat16) + experts = torch.randn(64, 128, 2560, device="cuda", dtype=torch.bfloat16) + selected = torch.randint(0, 64, (9, 4), device="cuda", dtype=torch.int32) + routing = torch.rand(9, 4, device="cuda", dtype=torch.bfloat16) + + expected = routed_linear_reference(hidden, experts, routing, selected) + actual = routed_linear(hidden, experts, routing, selected) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("num_tokens", [1, 128]) +def test_fused_routed_linear_cuda_graph_capture_and_replay(num_tokens): + torch.manual_seed(23) + hidden = torch.randn(num_tokens, 64, device="cuda", dtype=torch.bfloat16) + experts = torch.randn(8, 32, 64, device="cuda", dtype=torch.bfloat16) + selected = torch.randint(0, 8, (num_tokens, 4), device="cuda", dtype=torch.int32) + routing = torch.rand(num_tokens, 4, device="cuda", dtype=torch.bfloat16) + + # Warm up Triton compilation and allocator state outside capture. + routed_linear(hidden, experts, routing, selected) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = routed_linear(hidden, experts, routing, selected) + + hidden.copy_(torch.randn_like(hidden)) + selected.copy_(torch.randint_like(selected, 0, 8)) + routing.copy_(torch.rand_like(routing)) + expected = routed_linear_reference(hidden, experts, routing, selected) + graph.replay() + torch.testing.assert_close(captured, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("num_tokens", [1, 128]) +def test_fused_routed_linear_torch_compile_fullgraph(num_tokens): + torch.manual_seed(29) + hidden = torch.randn(num_tokens, 64, device="cuda", dtype=torch.bfloat16) + experts = torch.randn(8, 32, 64, device="cuda", dtype=torch.bfloat16) + selected = torch.randint(0, 8, (num_tokens, 4), device="cuda", dtype=torch.int32) + routing = torch.rand(num_tokens, 4, device="cuda", dtype=torch.bfloat16) + compiled = torch.compile(routed_linear, backend="eager", fullgraph=True) + + expected = routed_linear_reference(hidden, experts, routing, selected) + actual = compiled(hidden, experts, routing, selected) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +def test_value_experts_load_output_shards_and_forward(): + layer = RoutedValueExperts( + num_experts=3, + input_size=4, + output_size=6, + tp_rank=1, + tp_size=2, + ) + full = torch.arange(3 * 6 * 4, dtype=torch.float32).view(3, 6, 4) + layer.weight_loader(layer.weight, full) + torch.testing.assert_close(layer.weight, full[:, 3:]) + + replacement = torch.full((6, 4), -2.0) + layer.weight_loader(layer.weight, replacement, 1) + torch.testing.assert_close(layer.weight[1], replacement[3:]) + + hidden = torch.randn(2, 4) + selected = torch.tensor([[0, 1], [2, 1]]) + routing = torch.tensor([[0.25, 0.75], [0.6, 0.4]]) + torch.testing.assert_close( + layer(hidden, routing, selected), + routed_linear_reference(hidden, layer.weight, routing, selected), + ) + + +def test_qkg_loader_packs_one_local_gqa_group_and_interleaves_qk(): + hidden_size = 3 + num_heads = 4 + num_kv_heads = 2 + head_dim = 4 + qkg = XllmQKGParallelLinear( + hidden_size, + num_heads, + num_kv_heads, + head_dim, + tp_rank=1, + tp_size=2, + ) + q = torch.arange(num_heads * head_dim * hidden_size, dtype=torch.float32).view( + num_heads * head_dim, hidden_size + ) + gate = q + 1000 + k = ( + torch.arange(num_kv_heads * head_dim * hidden_size, dtype=torch.float32).view( + num_kv_heads * head_dim, hidden_size + ) + + 2000 + ) + + qkg.weight_loader(qkg.weight, q, "q") + qkg.weight_loader(qkg.weight, gate, "gate") + qkg.weight_loader(qkg.weight, k, "k") + + packed = qkg.weight.view(1, 5, head_dim, hidden_size) + expected_q = _interleave_rope_weight(q, num_heads).view( + num_kv_heads, 2, head_dim, hidden_size + )[1] + expected_gate = gate.view(num_kv_heads, 2, head_dim, hidden_size)[1] + expected_k = _interleave_rope_weight(k, num_kv_heads).view( + num_kv_heads, 1, head_dim, hidden_size + )[1] + torch.testing.assert_close(packed[0, :2], expected_q) + torch.testing.assert_close(packed[0, 2:4], expected_gate) + torch.testing.assert_close(packed[0, 4:], expected_k) + + x = torch.randn(2, hidden_size) + q_out, k_out, gate_out = qkg(x) + torch.testing.assert_close(q_out, F.linear(x, expected_q.reshape(-1, hidden_size))) + torch.testing.assert_close(k_out, F.linear(x, expected_k.reshape(-1, hidden_size))) + torch.testing.assert_close( + gate_out, F.linear(x, expected_gate.reshape(-1, hidden_size)) + ) + + +def test_zero_centered_group_norm_uses_weight_plus_one(): + norm = XllmGroupRMSNorm(4, n_groups=2, eps=0.0, zero_centered=True) + x = torch.tensor([[3.0, 4.0, 5.0, 12.0]]) + expected = torch.tensor([[3.0, 4.0, 5.0, 12.0]]) / torch.tensor( + [[12.5**0.5, 12.5**0.5, 84.5**0.5, 84.5**0.5]] + ) + torch.testing.assert_close(norm(x), expected) + + norm.weight.data.fill_(1.0) + torch.testing.assert_close(norm(x), 2 * expected) + + +def test_group_norm_preserves_legacy_direct_weight_semantics(): + norm = XllmGroupRMSNorm(4, n_groups=2, eps=0.0, zero_centered=False) + norm.weight.data.fill_(2.0) + x = torch.tensor([[3.0, 4.0, 5.0, 12.0]]) + expected = ( + 2 + * torch.tensor([[3.0, 4.0, 5.0, 12.0]]) + / torch.tensor([[12.5**0.5, 12.5**0.5, 84.5**0.5, 84.5**0.5]]) + ) + torch.testing.assert_close(norm(x), expected) + + +def test_mova_config_rejects_misaligned_attention_and_ffn_layout(monkeypatch): + config = _valid_mova_config(mlp_only_layers=[0, 1]) + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.bfloat16 + ) + monkeypatch.setattr("sglang.srt.models.xllm.get_attention_tp_size", lambda: 1) + with pytest.raises(ValueError, match="mlp_only_layers"): + _validate_mova_config(config, quant_config=None) + + +def test_mova_config_accepts_36b_contract(monkeypatch): + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.bfloat16 + ) + monkeypatch.setattr("sglang.srt.models.xllm.get_attention_tp_size", lambda: 8) + _validate_mova_config(_valid_mova_config(), quant_config=None) + + +def test_mova_config_rejects_accidental_auto_fp16(monkeypatch): + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.float16 + ) + with pytest.raises(ValueError, match="--dtype bfloat16"): + _validate_mova_config(_valid_mova_config(), quant_config=None) + + +def test_softplus_attention_gate_uses_ln2_beta(): + gate = torch.tensor([-2.0, 0.0, 2.0]) + module = SimpleNamespace(attn_gate_func="softplus") + actual = _XllmMoVAAttentionBase._activate_gate(module, gate) + expected = F.softplus(gate, beta=math.log(2)) + torch.testing.assert_close(actual, expected) + + +def test_interleaved_rope_matches_hf_neox_after_weight_permutation(): + torch.manual_seed(13) + head_dim = 8 + hidden_size = 5 + positions = torch.tensor([1, 7, 31], dtype=torch.long) + hidden = torch.randn(3, hidden_size) + hf_weight = torch.randn(head_dim, hidden_size) + hf_projection = F.linear(hidden, hf_weight) + native_projection = F.linear( + hidden, _interleave_rope_weight(hf_weight, num_heads=1) + ) + + hf_rope = get_rope( + head_dim, + rotary_dim=head_dim, + max_position=64, + base=10000, + is_neox_style=True, + ) + native_rope = get_rope( + head_dim, + rotary_dim=head_dim, + max_position=64, + base=10000, + is_neox_style=False, + ) + hf_rotated, _ = hf_rope.forward_native(positions, hf_projection, hf_projection) + native_rotated, _ = native_rope.forward_native( + positions, native_projection, native_projection + ) + torch.testing.assert_close( + native_rotated, + _interleave_rope_weight(hf_rotated.transpose(0, 1), num_heads=1).transpose( + 0, 1 + ), + ) + + +def test_mova_parameter_mapper_stages_qkg_and_all_value_experts(): + config = SimpleNamespace(num_values=64) + model = SimpleNamespace( + stacked_params_mapping=_xllm_stacked_params_mapping(config), + expert_params_mapping=[], + ) + mapper = ParameterMapper.from_model(model) + + q = mapper.map("model.layers.3.self_attn.q_proj.weight") + assert q.sglang_name == "model.layers.3.self_attn.qkg_proj.weight" + assert q.shard_id == "q" + assert q.num_shards == 3 + + gate = mapper.map("model.layers.3.self_attn.attn_gate_proj.weight") + assert gate.sglang_name == "model.layers.3.self_attn.qkg_proj.weight" + assert gate.shard_id == "gate" + assert gate.num_shards == 3 + + value = mapper.map("model.layers.3.self_attn.v_experts.63.weight") + assert value.sglang_name == "model.layers.3.self_attn.v_experts.weight" + assert value.shard_id == 63 + assert value.num_shards == 64 + + router = mapper.map("model.layers.3.self_attn.v_router.bias") + assert router.sglang_name == "model.layers.3.self_attn.v_router.bias" + assert router.num_shards == 1 + + +def test_legacy_xllm_mapping_is_unchanged(): + assert _xllm_stacked_params_mapping(SimpleNamespace(num_values=0)) == [ + (".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), + ] From d39803581ed22cb8fea08e5544c5ba7ef23145cc Mon Sep 17 00:00:00 2001 From: Yash Akhauri Date: Thu, 13 Aug 2026 01:09:51 +0000 Subject: [PATCH 2/5] Accept normalized default RoPE for MoVA Transformers 5 materializes a null rope_scaling field as an explicit default-RoPE dictionary. Treat that normalized representation as unscaled while retaining the fail-fast guard for actual scaling modes. --- python/sglang/srt/models/xllm.py | 9 ++++++- test/registered/unit/layers/test_mova.py | 33 ++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index 1d91e887923c..3c228e85677f 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -225,7 +225,14 @@ def _validate_mova_config( raise ValueError("MoVA requires query heads to be divisible by KV heads") if getattr(config, "rope_head_dim", head_dim) != head_dim: raise ValueError("MoVA phase 1 requires full-head interleaved RoPE") - if getattr(config, "rope_scaling", None) is not None: + rope_scaling = getattr(config, "rope_scaling", None) + # Transformers 5 normalizes a JSON ``rope_scaling: null`` into an + # explicit default-RoPE dictionary. That representation does not change + # the rotary math and must not be confused with linear/dynamic scaling. + if rope_scaling is not None and not ( + isinstance(rope_scaling, dict) + and rope_scaling.get("rope_type", rope_scaling.get("type")) == "default" + ): raise ValueError("MoVA phase 1 does not support non-default RoPE scaling") if getattr(config, "sliding_window", None) is not None or getattr( config, "use_sliding_window", False diff --git a/test/registered/unit/layers/test_mova.py b/test/registered/unit/layers/test_mova.py index 3fb08fbed4fa..39c93a3ecd9c 100644 --- a/test/registered/unit/layers/test_mova.py +++ b/test/registered/unit/layers/test_mova.py @@ -312,6 +312,39 @@ def test_mova_config_accepts_36b_contract(monkeypatch): _validate_mova_config(_valid_mova_config(), quant_config=None) +@pytest.mark.parametrize( + "rope_scaling", + [ + {"rope_type": "default", "rope_theta": 10_000_000.0}, + {"type": "default"}, + ], +) +def test_mova_config_accepts_transformers_normalized_default_rope( + monkeypatch, rope_scaling +): + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.bfloat16 + ) + monkeypatch.setattr("sglang.srt.models.xllm.get_attention_tp_size", lambda: 8) + _validate_mova_config( + _valid_mova_config(rope_scaling=rope_scaling), quant_config=None + ) + + +def test_mova_config_rejects_actual_rope_scaling(monkeypatch): + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.bfloat16 + ) + monkeypatch.setattr("sglang.srt.models.xllm.get_attention_tp_size", lambda: 8) + with pytest.raises(ValueError, match="non-default RoPE scaling"): + _validate_mova_config( + _valid_mova_config( + rope_scaling={"rope_type": "linear", "factor": 2.0} + ), + quant_config=None, + ) + + def test_mova_config_rejects_accidental_auto_fp16(monkeypatch): monkeypatch.setattr( "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.float16 From 687ee776387a962da8b133764d4186c3c0ff449d Mon Sep 17 00:00:00 2001 From: Yash Akhauri Date: Thu, 13 Aug 2026 02:09:59 +0000 Subject: [PATCH 3/5] Preserve FP32 xLLM router semantics xLLM applies sigmoid, correction bias, expert selection, and top-k normalization in FP32 after the BF16 router projection. Match that contract in the native fallback so near-boundary routes do not change during serving. --- python/sglang/srt/layers/moe/topk.py | 6 +++++- test/srt/cpu/test_topk.py | 27 +++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 892dcebaea81..a4dd21f1cd91 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -422,7 +422,11 @@ def scoring_func_impl(gating_output: torch.Tensor) -> torch.Tensor: if correction_bias is not None: n_routed_experts = gating_output.shape[-1] - scores = scoring_func_impl(gating_output) + # Keep routing scores in FP32. In particular, xLLM computes its BF16 + # router GEMM first, then applies sigmoid, correction bias, top-k, and + # renormalization in FP32. Applying sigmoid in BF16 can change both + # mixture weights and expert selection near a top-k boundary. + scores = scoring_func_impl(gating_output.float()) scores_for_choice = scores.view( -1, n_routed_experts ) + correction_bias.unsqueeze(0) diff --git a/test/srt/cpu/test_topk.py b/test/srt/cpu/test_topk.py index 9f3dfc1b4163..7b4975e573cf 100644 --- a/test/srt/cpu/test_topk.py +++ b/test/srt/cpu/test_topk.py @@ -112,6 +112,33 @@ def test_biased_grouped_topk(self): class TestTopK(CustomTestCase): + def test_native_biased_sigmoid_routes_in_fp32(self): + hidden_states = torch.zeros(1, 4, dtype=torch.bfloat16) + gating_output = torch.tensor( + [[0.0, 0.00390625]], dtype=torch.bfloat16 + ) + correction_bias = torch.tensor([0.0005, 0.0], dtype=torch.float32) + + topk_weights, topk_ids = native_fused_topk( + hidden_states, + gating_output, + topk=1, + renormalize=True, + correction_bias=correction_bias, + scoring_func="sigmoid", + ) + + scores = torch.sigmoid(gating_output.float()) + expected_ids = torch.topk( + scores + correction_bias.unsqueeze(0), k=1, dim=-1, sorted=False + )[1] + expected_weights = scores.gather(1, expected_ids) + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + + self.assertEqual(topk_weights.dtype, torch.float32) + torch.testing.assert_close(topk_ids, expected_ids) + torch.testing.assert_close(topk_weights, expected_weights) + def _run_single_test(self, M, E, topk, renormalize, dtype): torch.manual_seed(1998) From de220dcc66ba28218427b2a8e50f291c8555dc90 Mon Sep 17 00:00:00 2001 From: Yash Akhauri Date: Thu, 13 Aug 2026 17:52:35 +0000 Subject: [PATCH 4/5] Fix remote seed transfer engine IP lookup The remote-instance seed path re-imported get_local_ip_auto from the utils package root after that API moved to utils.network. This caused every TP scheduler to fail before Mooncake initialization. Use the existing canonical module-level import and cover the transfer-engine and bootstrap contracts. --- .../sglang/srt/model_executor/model_runner.py | 2 - .../test_remote_instance_transfer_engine.py | 85 +++++++++++++++++++ 2 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/model_executor/test_remote_instance_transfer_engine.py diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e9b7604aded5..5792e3991857 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -720,8 +720,6 @@ def remote_instance_init_transfer_engine(self): ) return - from sglang.srt.utils import get_local_ip_auto - self.remote_instance_transfer_engine = TransferEngine() local_ip = get_local_ip_auto() self.remote_instance_transfer_engine.initialize( diff --git a/test/registered/unit/model_executor/test_remote_instance_transfer_engine.py b/test/registered/unit/model_executor/test_remote_instance_transfer_engine.py new file mode 100644 index 000000000000..0aa9271f72b7 --- /dev/null +++ b/test/registered/unit/model_executor/test_remote_instance_transfer_engine.py @@ -0,0 +1,85 @@ +import sys +import types +import unittest +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.model_executor.model_runner import ModelRunner + +register_cpu_ci(est_time=1, suite="stage-a-test-cpu") + + +class TestRemoteInstanceTransferEngine(CustomTestCase): + def test_initializes_with_detected_local_ip(self): + class FakeTransferEngine: + def __init__(self): + self.initialize_args = None + + def initialize(self, *args): + self.initialize_args = args + + def get_rpc_port(self): + return 12345 + + mooncake = types.ModuleType("mooncake") + mooncake.__path__ = [] + mooncake_engine = types.ModuleType("mooncake.engine") + mooncake_engine.TransferEngine = FakeTransferEngine + + runner = ModelRunner.__new__(ModelRunner) + with patch.dict( + sys.modules, + {"mooncake": mooncake, "mooncake.engine": mooncake_engine}, + ), patch( + "sglang.srt.model_executor.model_runner.get_local_ip_auto", + return_value="10.20.30.40", + ), patch( + "sglang.srt.model_executor.model_runner.envs.MOONCAKE_DEVICE.get", + return_value="mlx5_0", + ): + runner.remote_instance_init_transfer_engine() + + self.assertEqual( + runner.remote_instance_transfer_engine.initialize_args, + ("10.20.30.40", "P2PHANDSHAKE", "rdma", "mlx5_0"), + ) + self.assertEqual( + runner.remote_instance_transfer_engine_session_id, + "10.20.30.40:12345", + ) + + def test_registers_initialized_engine_info_with_bootstrap(self): + runner = ModelRunner.__new__(ModelRunner) + runner.server_args = types.SimpleNamespace( + dist_init_addr="10.20.30.1:17503", + engine_info_bootstrap_port=17502, + ) + runner.tp_rank = 3 + runner.remote_instance_transfer_engine_session_id = "10.20.30.40:12345" + runner.remote_instance_transfer_engine_weight_info = { + "model.weight": (4096, 128, 2) + } + + with patch("requests.put") as put: + put.return_value.status_code = 200 + runner._register_to_engine_info_bootstrap() + + put.assert_called_once_with( + "http://10.20.30.1:17502/register_transfer_engine_info", + json={ + "tp_rank": 3, + "transfer_engine_info": { + "session_id": "10.20.30.40:12345", + "weights_info_dict": {"model.weight": (4096, 128, 2)}, + }, + }, + timeout=5, + ) + + +if __name__ == "__main__": + unittest.main() From 198d611c005b789879ac73732e67d0d55d6362fd Mon Sep 17 00:00:00 2001 From: Yash Akhauri Date: Fri, 14 Aug 2026 02:13:51 +0000 Subject: [PATCH 5/5] Match native xLLM router GEMM provenance Preserve legacy router behavior when source topology is absent. For explicit MP1 or MP2 provenance, reproduce native BF16 local GEMM rounding and FP32 logits at both FFN and value routers. --- python/sglang/srt/models/xllm.py | 110 +++++- test/registered/unit/layers/test_mova.py | 413 +++++++++++++++++++++++ 2 files changed, 521 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index 3c228e85677f..bcf5106633ae 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -194,6 +194,99 @@ def filter_moe_weight_param_global_expert(name, x, num_local_experts): _is_cpu = is_cpu() _is_cpu_amx_available = cpu_has_amx_support() +_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY = "xllm_source_router_gemm_partitions" +_XLLM_SOURCE_ROUTER_PARTITIONS_MISSING = object() + + +def _get_xllm_source_router_gemm_partitions( + config: PretrainedConfig, +) -> Optional[int]: + """Read optional source-router provenance without inferring it from TP.""" + + partitions = getattr( + config, + _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY, + _XLLM_SOURCE_ROUTER_PARTITIONS_MISSING, + ) + if partitions is _XLLM_SOURCE_ROUTER_PARTITIONS_MISSING: + return None + if ( + isinstance(partitions, bool) + or not isinstance(partitions, int) + or partitions not in (1, 2) + ): + raise ValueError( + f"{_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY}={partitions!r} is " + f"invalid (type={type(partitions).__name__}); when present it must " + "be the integer 1 or 2. Omit the key to preserve legacy router " + "GEMM behavior." + ) + if config.hidden_size % partitions: + raise ValueError( + f"explicit {_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY}={partitions} " + f"requires hidden_size divisible by {partitions}; got " + f"hidden_size={config.hidden_size}" + ) + return partitions + + +def _xllm_router_gemm( + hidden_states: torch.Tensor, + weight: torch.Tensor, + source_partitions: Optional[int], +) -> torch.Tensor: + """Reproduce the source xLLM router GEMM's partition rounding contract.""" + + # Old xLLM artifacts have no source-topology provenance. Preserve their + # exact pre-contract behavior instead of guessing how the router was run. + if source_partitions is None: + return F.linear(hidden_states, weight) + if isinstance(source_partitions, bool) or not isinstance(source_partitions, int): + raise ValueError( + "explicit xLLM router source partitions must be the integer 1 or " + f"2; got {source_partitions!r} " + f"(type={type(source_partitions).__name__})" + ) + if source_partitions not in (1, 2): + raise ValueError( + "explicit xLLM router source partitions must be 1 or 2, got " + f"{source_partitions}" + ) + if hidden_states.ndim < 1 or weight.ndim != 2: + raise ValueError( + "xLLM router GEMM expects input [..., hidden] and weight " + f"[routes, hidden]; got input={tuple(hidden_states.shape)}, " + f"weight={tuple(weight.shape)}" + ) + if hidden_states.shape[-1] != weight.shape[-1]: + raise ValueError( + "xLLM router input and weight hidden dimensions differ; got " + f"input={tuple(hidden_states.shape)}, weight={tuple(weight.shape)}" + ) + if hidden_states.shape[-1] % source_partitions: + raise ValueError( + f"explicit xLLM router partitions={source_partitions} requires " + f"hidden size divisible by {source_partitions}; got hidden_size=" + f"{hidden_states.shape[-1]}" + ) + if hidden_states.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + raise ValueError( + "Explicit xLLM source router GEMM provenance requires BF16 input " + f"and weight; got input={hidden_states.dtype}, weight={weight.dtype}" + ) + + if source_partitions == 1: + return F.linear(hidden_states, weight).float() + + # Native xLLM row-shards each router across MP2. Each rank performs a BF16 + # partial GEMM, rounds that result to BF16, casts it to FP32, and then the + # FP32 all-reduce adds the two partials. Emulate that ordering locally. + input_parts = hidden_states.chunk(source_partitions, dim=-1) + weight_parts = weight.chunk(source_partitions, dim=-1) + first = F.linear(input_parts[0].contiguous(), weight_parts[0].contiguous()) + second = F.linear(input_parts[1].contiguous(), weight_parts[1].contiguous()) + return first.float() + second.float() + def _validate_mova_config( config: PretrainedConfig, @@ -208,6 +301,7 @@ def _validate_mova_config( "MoVA phase 1 requires --dtype bfloat16. The converted 36B HF " "artifact reports float32, so SGLang dtype=auto would select fp16." ) + _get_xllm_source_router_gemm_partitions(config) if quant_config is not None: raise ValueError("MoVA phase 1 does not support quantized model weights") if getattr(config, "attention_bias", False): @@ -526,6 +620,9 @@ class XllmMoEGate(nn.Module): def __init__(self, config: PretrainedConfig): super().__init__() + self.source_router_gemm_partitions = ( + _get_xllm_source_router_gemm_partitions(config) + ) self.weight = nn.Parameter( torch.empty((config.num_experts, config.hidden_size)) ) @@ -538,7 +635,9 @@ def __init__(self, config: PretrainedConfig): self.bias = None def forward(self, hidden_states: torch.Tensor): - return F.linear(hidden_states, self.weight) + return _xllm_router_gemm( + hidden_states, self.weight, self.source_router_gemm_partitions + ) class XllmSparseMoeBlock(nn.Module): @@ -1016,6 +1115,9 @@ def __init__( self.router_score_func = getattr(config, "router_score_func", "sigmoid") self.router_scaling_factor = getattr(config, "router_scaling_factor", 1.0) self.renormalize = getattr(config, "norm_topk_prob", True) + self.source_router_gemm_partitions = ( + _get_xllm_source_router_gemm_partitions(config) + ) self.v_router = ReplicatedLinear( config.hidden_size, self.num_values, @@ -1042,7 +1144,11 @@ def __init__( def _project_value(self, hidden_states: torch.Tensor) -> torch.Tensor: # Router bias is deliberately omitted from the logits matmul. It only # changes route selection inside ``mova_router_topk``. - router_logits = F.linear(hidden_states, self.v_router.weight) + router_logits = _xllm_router_gemm( + hidden_states, + self.v_router.weight, + self.source_router_gemm_partitions, + ) routing_weights, selected_values = mova_router_topk( router_logits, self.v_router.bias, diff --git a/test/registered/unit/layers/test_mova.py b/test/registered/unit/layers/test_mova.py index 39c93a3ecd9c..713fbe991604 100644 --- a/test/registered/unit/layers/test_mova.py +++ b/test/registered/unit/layers/test_mova.py @@ -4,7 +4,9 @@ import pytest import torch import torch.nn.functional as F +from transformers import PretrainedConfig +from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.mova import ( RoutedValueExperts, mova_router_topk, @@ -15,9 +17,14 @@ from sglang.srt.model_loader.parameter_mapper import ParameterMapper from sglang.srt.models.xllm import ( XllmGroupRMSNorm, + XllmMoEGate, + XllmMoVAAttention, XllmQKGParallelLinear, + _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY, + _get_xllm_source_router_gemm_partitions, _interleave_rope_weight, _validate_mova_config, + _xllm_router_gemm, _xllm_stacked_params_mapping, _XllmMoVAAttentionBase, ) @@ -60,11 +67,395 @@ def _valid_mova_config(**overrides): router_score_func="sigmoid", router_scaling_factor=2.5, layernorm_num_groups=2, + xllm_source_router_gemm_partitions=2, ) values.update(overrides) return SimpleNamespace(**values) +def _native_mp2_router_gemm_reference(hidden, weight): + hidden_parts = hidden.chunk(2, dim=-1) + weight_parts = weight.chunk(2, dim=-1) + return sum( + F.linear(hidden_part.contiguous(), weight_part.contiguous()).float() + for hidden_part, weight_part in zip(hidden_parts, weight_parts) + ) + + +def _native_router_topk_reference( + router_logits, + router_bias, + *, + top_k, + scaling_factor, + output_dtype=torch.float32, +): + scores = torch.sigmoid(router_logits.float()) + selected = torch.topk(scores + router_bias.float(), top_k, dim=-1).indices + weights = torch.gather(scores, dim=-1, index=selected) + if top_k > 1: + weights = weights / weights.sum(dim=-1, keepdim=True) + return ( + (weights * scaling_factor).to(output_dtype), + selected.to(torch.int32), + ) + + +def _canonical_routes(weights, selected): + order = torch.argsort(selected.long(), dim=-1) + return ( + torch.gather(weights, dim=-1, index=order), + torch.gather(selected.long(), dim=-1, index=order), + ) + + +def _real_shape_boundary_router_case(*, num_routes, top_k, device): + """Build a literal BF16 rounding boundary where MP2 flips the last route.""" + + hidden_size = 2560 + split = hidden_size // 2 + native_candidate = top_k - 1 + full_gemm_candidate = top_k + + hidden = torch.zeros(1, hidden_size, device=device, dtype=torch.bfloat16) + weight = torch.zeros( + num_routes, hidden_size, device=device, dtype=torch.bfloat16 + ) + hidden[0, 0] = 1.0 + hidden[0, split] = 1.0 + weight[native_candidate, 0] = 1.0 + weight[native_candidate, split] = 2**-8 + weight[full_gemm_candidate, 0] = 1.0 + + native_logits = _native_mp2_router_gemm_reference(hidden, weight) + full_gemm_logits = F.linear(hidden, weight).float() + # Analytic literal golden for native xLLM 5494c84 MP2 ordering: each BF16 + # partial is rounded before its FP32 cast and the FP32 all-reduce sum. + expected_native_pair = torch.tensor( + [[1.0 + 2**-8, 1.0]], device=device, dtype=torch.float32 + ) + expected_full_pair = torch.ones(1, 2, device=device, dtype=torch.float32) + torch.testing.assert_close( + native_logits[:, native_candidate : full_gemm_candidate + 1], + expected_native_pair, + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + full_gemm_logits[:, native_candidate : full_gemm_candidate + 1], + expected_full_pair, + rtol=0.0, + atol=0.0, + ) + + native_scores = torch.sigmoid(native_logits) + full_gemm_scores = torch.sigmoid(full_gemm_logits) + native_margin = ( + native_scores[0, native_candidate] + - native_scores[0, full_gemm_candidate] + ) + full_gemm_margin = ( + full_gemm_scores[0, native_candidate] + - full_gemm_scores[0, full_gemm_candidate] + ) + assert native_margin > full_gemm_margin + + bias = torch.full((num_routes,), -10.0, device=device, dtype=torch.float32) + bias[: top_k - 1] = 10.0 + bias[native_candidate] = 0.0 + bias[full_gemm_candidate] = (native_margin + full_gemm_margin) / 2 + return hidden, weight, bias, native_logits, full_gemm_logits + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" + ), + ), + ], +) +def test_xllm_router_gemm_matches_native_mp2_rounding_reference(device): + hidden = ( + torch.sin(torch.arange(24, device=device, dtype=torch.float32) * 0.173) + .view(3, 8) + .to(torch.bfloat16) + ) + weight = ( + torch.cos(torch.arange(80, device=device, dtype=torch.float32) * 0.097) + .view(10, 8) + .to(torch.bfloat16) + ) + + actual = _xllm_router_gemm(hidden, weight, source_partitions=2) + expected = _native_mp2_router_gemm_reference(hidden, weight) + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" + ), + ), + ], +) +def test_xllm_router_gemm_explicit_mp1_returns_fp32_native_logits(device): + hidden = torch.randn(3, 8, device=device, dtype=torch.bfloat16) + weight = torch.randn(10, 8, device=device, dtype=torch.bfloat16) + + actual = _xllm_router_gemm(hidden, weight, source_partitions=1) + expected = F.linear(hidden, weight).float() + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_xllm_router_gemm_missing_provenance_preserves_legacy_behavior(): + hidden = torch.randn(3, 8, dtype=torch.bfloat16) + weight = torch.randn(10, 8, dtype=torch.bfloat16) + + actual = _xllm_router_gemm(hidden, weight, source_partitions=None) + expected = F.linear(hidden, weight) + + assert actual.dtype == torch.bfloat16 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +_MISSING_ROUTER_PROVENANCE = object() + + +@pytest.mark.parametrize( + "partitions,expected", + [ + pytest.param(_MISSING_ROUTER_PROVENANCE, None, id="missing-legacy"), + pytest.param(1, 1, id="explicit-mp1"), + pytest.param(2, 2, id="explicit-mp2"), + ], +) +def test_xllm_source_router_partition_config_roundtrip(partitions, expected): + kwargs = {"hidden_size": 8} + if partitions is not _MISSING_ROUTER_PROVENANCE: + kwargs[_XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY] = partitions + config = PretrainedConfig(**kwargs) + restored = PretrainedConfig.from_dict(config.to_dict()) + + assert _get_xllm_source_router_gemm_partitions(restored) == expected + assert ( + _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY in restored.to_dict() + ) == (partitions is not _MISSING_ROUTER_PROVENANCE) + + +@pytest.mark.parametrize("partitions", [None, True, 0, -1, 3, 1.5, "2"]) +def test_xllm_source_router_partition_config_rejects_bad_values(partitions): + config = SimpleNamespace( + hidden_size=8, + xllm_source_router_gemm_partitions=partitions, + ) + + with pytest.raises( + ValueError, + match="source_router_gemm_partitions.*Omit the key", + ): + _get_xllm_source_router_gemm_partitions(config) + + +def test_xllm_source_router_partition_config_rejects_odd_hidden_size(): + config = SimpleNamespace( + hidden_size=7, + xllm_source_router_gemm_partitions=2, + ) + + with pytest.raises(ValueError, match="requires hidden_size divisible"): + _get_xllm_source_router_gemm_partitions(config) + + +@pytest.mark.parametrize( + "hidden,weight,partitions,error", + [ + ( + torch.randn(2, 7, dtype=torch.bfloat16), + torch.randn(4, 7, dtype=torch.bfloat16), + 2, + "divisible", + ), + (torch.randn(2, 8), torch.randn(4, 8), 1, "BF16"), + (torch.randn(2, 8), torch.randn(4, 8), 2, "BF16"), + ( + torch.randn(2, 8, dtype=torch.bfloat16), + torch.randn(4, 6, dtype=torch.bfloat16), + 2, + "differ", + ), + ], +) +def test_xllm_router_gemm_rejects_invalid_runtime_contract( + hidden, weight, partitions, error +): + with pytest.raises(ValueError, match=error): + _xllm_router_gemm(hidden, weight, source_partitions=partitions) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" + ), + ), + ], +) +def test_ffn_top8_routes_match_native_mp2_at_real_shape_boundary(device): + top_k = 8 + scaling_factor = 2.5 + hidden, weight, bias, native_logits, full_gemm_logits = ( + _real_shape_boundary_router_case( + num_routes=100, + top_k=top_k, + device=device, + ) + ) + config = SimpleNamespace( + hidden_size=2560, + num_experts=100, + num_experts_per_tok=top_k, + moe_gate_bias=True, + xllm_source_router_gemm_partitions=2, + ) + gate = XllmMoEGate(config).to(device=device) + gate.weight.data = weight + with torch.no_grad(): + gate.bias.copy_(bias) + topk = TopK( + top_k=top_k, + renormalize=True, + scoring_func="sigmoid", + correction_bias=gate.bias, + ) + + actual_logits = gate(hidden) + actual_topk = topk.forward_native(hidden, actual_logits) + actual_weights = actual_topk.topk_weights * scaling_factor + expected_weights, expected_ids = _native_router_topk_reference( + native_logits, + bias, + top_k=top_k, + scaling_factor=scaling_factor, + ) + _, full_gemm_ids = _native_router_topk_reference( + full_gemm_logits, + bias, + top_k=top_k, + scaling_factor=scaling_factor, + ) + + torch.testing.assert_close(actual_logits, native_logits, rtol=0.0, atol=0.0) + actual_weights, actual_ids = _canonical_routes( + actual_weights, actual_topk.topk_ids + ) + expected_weights, expected_ids = _canonical_routes( + expected_weights, expected_ids + ) + _, full_gemm_ids = _canonical_routes(expected_weights, full_gemm_ids) + assert actual_ids.tolist() == [list(range(8))] + assert full_gemm_ids.tolist() == [[0, 1, 2, 3, 4, 5, 6, 8]] + torch.testing.assert_close(actual_ids, expected_ids, rtol=0.0, atol=0.0) + torch.testing.assert_close(actual_weights, expected_weights, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" + ), + ), + ], +) +def test_value_top4_routes_match_native_mp2_at_real_shape_boundary(device): + top_k = 4 + scaling_factor = 2.5 + hidden, weight, bias, native_logits, full_gemm_logits = ( + _real_shape_boundary_router_case( + num_routes=64, + top_k=top_k, + device=device, + ) + ) + attention = object.__new__(XllmMoVAAttention) + torch.nn.Module.__init__(attention) + attention.source_router_gemm_partitions = 2 + attention.router_score_func = "sigmoid" + attention.router_scaling_factor = scaling_factor + attention.renormalize = True + attention.num_values_per_tok = top_k + attention.v_router = torch.nn.Linear( + 2560, 64, bias=False, device=device, dtype=torch.bfloat16 + ) + attention.v_router.bias = torch.nn.Parameter( + bias.clone(), requires_grad=False + ) + with torch.no_grad(): + attention.v_router.weight.copy_(weight) + + seen = {} + + class RecordingValueExperts(torch.nn.Module): + def forward(self, hidden_states, routing_weights, selected_values): + seen["weights"] = routing_weights + seen["ids"] = selected_values + # The real CPU and CUDA value-expert paths cast these coefficients + # to the BF16 projected activation immediately before multiplying. + seen["native_weights"] = routing_weights.to(hidden_states.dtype) + return hidden_states + + attention.v_experts = RecordingValueExperts() + output = attention._project_value(hidden) + expected_weights, expected_ids = _native_router_topk_reference( + native_logits, + bias, + top_k=top_k, + scaling_factor=scaling_factor, + output_dtype=torch.bfloat16, + ) + _, full_gemm_ids = _native_router_topk_reference( + full_gemm_logits, + bias, + top_k=top_k, + scaling_factor=scaling_factor, + ) + + assert seen["weights"].dtype == torch.float32 + actual_weights, actual_ids = _canonical_routes( + seen["native_weights"], seen["ids"] + ) + expected_weights, expected_ids = _canonical_routes( + expected_weights, expected_ids + ) + _, full_gemm_ids = _canonical_routes(expected_weights, full_gemm_ids) + assert actual_ids.tolist() == [[0, 1, 2, 3]] + assert full_gemm_ids.tolist() == [[0, 1, 2, 4]] + torch.testing.assert_close(actual_ids, expected_ids, rtol=0.0, atol=0.0) + torch.testing.assert_close(actual_weights, expected_weights, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(output, hidden) + + def test_mova_router_bias_changes_selection_not_weight(): logits = torch.tensor([[2.0, 1.0, -1.0]], dtype=torch.float32) bias = torch.tensor([0.0, 0.0, 10.0], dtype=torch.float32) @@ -312,6 +703,28 @@ def test_mova_config_accepts_36b_contract(monkeypatch): _validate_mova_config(_valid_mova_config(), quant_config=None) +@pytest.mark.parametrize( + "partitions", + [ + pytest.param(_MISSING_ROUTER_PROVENANCE, id="missing-legacy"), + pytest.param(1, id="explicit-mp1"), + pytest.param(2, id="explicit-mp2"), + ], +) +def test_mova_config_accepts_supported_source_router_modes(monkeypatch, partitions): + config = _valid_mova_config() + if partitions is _MISSING_ROUTER_PROVENANCE: + delattr(config, _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY) + else: + setattr(config, _XLLM_SOURCE_ROUTER_PARTITIONS_CONFIG_KEY, partitions) + monkeypatch.setattr( + "sglang.srt.models.xllm.torch.get_default_dtype", lambda: torch.bfloat16 + ) + monkeypatch.setattr("sglang.srt.models.xllm.get_attention_tp_size", lambda: 8) + + _validate_mova_config(config, quant_config=None) + + @pytest.mark.parametrize( "rope_scaling", [