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/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/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/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index f0ac97e61df1..bcf5106633ae 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__) @@ -171,6 +194,216 @@ 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, + 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." + ) + _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): + 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") + 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 + ): + 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]""" @@ -182,6 +415,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, @@ -244,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)) ) @@ -256,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): @@ -549,9 +930,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 +966,200 @@ 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.source_router_gemm_partitions = ( + _get_xllm_source_router_gemm_partitions(config) + ) + 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 = _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, + 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 +1176,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 +1190,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 +1316,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 +1413,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 +1468,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 +1489,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 +1525,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 +1543,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..713fbe991604 --- /dev/null +++ b/test/registered/unit/layers/test_mova.py @@ -0,0 +1,850 @@ +import math +from types import SimpleNamespace + +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, + 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, + 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, +) + + +@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, + 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) + + 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) + + +@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", + [ + {"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 + ) + 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), + ] 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() 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)