# coding=utf-8 # Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved. import copy import logging from typing import Callable, Iterable, Optional, Set, Tuple, Union import torch import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import ( get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.dp_attention import ( get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.layers.quantization.fp8_utils import ( block_quant_dequant, block_quant_to_tensor_quant, channel_quant_to_tensor_quant, normalize_e4m3fn_to_e4m3fnuz, requant_weight_ue8m0_inplace, ) from sglang.srt.layers.quantization.int8_utils import ( block_dequant as int8_block_dequant, ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip from sglang.srt.models.utils import WeightsMapper from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, add_prefix, bind_or_assign, cpu_has_amx_support, get_bool_env_var, get_device_sm, is_cpu, is_cuda, is_flashinfer_available, is_gfx95_supported, is_hip, is_npu, is_sm100_supported, make_layers, ) from sglang.srt.utils.common import rank0_log _is_hip = is_hip() _is_cuda = is_cuda() _is_npu = is_npu() _is_fp8_fnuz = is_fp8_fnuz() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() _device_sm = get_device_sm() _is_gfx95_supported = is_gfx95_supported() _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported if _use_aiter_gfx95: pass if _is_cuda: from sgl_kernel import awq_dequantize elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: from sglang.srt.layers.quantization.awq_triton import ( awq_dequantize_triton as awq_dequantize, ) else: from vllm._custom_ops import awq_dequantize if _is_hip: pass _is_flashinfer_available = is_flashinfer_available() _is_sm100_supported = is_cuda() and is_sm100_supported() class DsV3MLA(DeepseekV2AttentionMLA): def __init__(self, **kwargs): super().__init__(**kwargs) if kwargs["rope_scaling"]: self.rotary_emb.forward = self.rotary_emb.forward_cuda LoraConfig = None logger = logging.getLogger(__name__) _is_cpu = is_cpu() def is_linear_layer(layer_idx, layer_group_size): if layer_idx is None: return False if layer_group_size > 0: return (layer_idx + 1) % layer_group_size != 0 else: return False def is_pp_missing_parameter( name: str, model: torch.nn.Module, ) -> bool: if isinstance(model, PPMissingLayer): return True return False def weight_loader_with_alias(alias: str): def wrapper(func: Callable): def inner_func( param: torch.Tensor, loaded_weight: torch.Tensor, *args, prefix: str = None, **kwargs, ): # pf = "[vLLM][load]" + " " if prefix is None else f"[{prefix}] " value = func(param, loaded_weight, *args, **kwargs) return value return inner_func return wrapper class BailingMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, reduce_results=True, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() def forward( self, x, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ): x, _ = self.gate_up_proj(x) x = self.act_fn(x) x, _ = self.down_proj( x, skip_all_reduce=use_reduce_scatter or should_allreduce_fusion, ) return x class BailingMoEGate(nn.Module): def __init__( self, config, params_dtype: Optional[torch.dtype] = None, prefix: str = "", ): super().__init__() if params_dtype is None: params_dtype = torch.get_default_dtype() self.params_dtype = params_dtype self.weight = nn.Parameter( torch.empty( (config.num_experts, config.hidden_size), dtype=self.params_dtype, ), ) if getattr(config, "moe_router_enable_expert_bias", False): self.expert_bias = nn.Parameter( torch.empty((config.num_experts,), dtype=torch.float32), ) else: self.expert_bias = None def forward(self, hidden_states): logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None).to( hidden_states.dtype ) return logits class BailingMoE(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, layer_id: int = 0, prefix: str = "moe", ): super().__init__() self.layer_id = layer_id self.tp_size = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() self.top_k = config.num_experts_per_tok self.norm_expert_prob = getattr(config, "norm_topk_prob", False) self.hidden_size = config.hidden_size self.intermediate_size = config.moe_intermediate_size self.num_shared_experts = getattr(config, "num_shared_experts", 0) self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.score_function = getattr(config, "score_function", None) # Gate always runs at half / full precision for now. router_dtype = getattr(config, "router_dtype", None) if router_dtype is None: self.router_dtype = torch.float32 elif router_dtype == "fp32": self.router_dtype = torch.float32 else: self.router_dtype = torch.bfloat16 # check group topk self.num_expert_group = getattr(config, "n_group", 0) self.topk_group = getattr(config, "topk_group", 0) if self.num_expert_group > 0 or self.topk_group > 0: assert ( self.num_expert_group > 0 and 0 < self.topk_group <= self.num_expert_group ) self.use_grouped_topk = True else: self.num_expert_group = self.topk_group = None self.use_grouped_topk = False self.num_experts = config.num_experts self.gate = BailingMoEGate( config=config, params_dtype=self.router_dtype, prefix=add_prefix("gate", prefix), ) self.correction_bias = ( self.gate.expert_bias.data if self.gate.expert_bias is not None else None ) if self.score_function is not None: assert ( self.score_function == "softmax" and self.correction_bias is None ) or ( self.score_function == "sigmoid" and self.correction_bias is not None ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" self.topk = TopK( top_k=self.top_k, use_grouped_topk=self.use_grouped_topk, renormalize=self.norm_expert_prob, num_expert_group=self.num_expert_group, topk_group=self.topk_group, correction_bias=self.correction_bias, routed_scaling_factor=self.routed_scaling_factor, ) moe_cls = get_moe_impl_class(quant_config) self.experts = moe_cls( num_experts=self.num_experts, top_k=self.top_k, layer_id=self.layer_id, hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, prefix=f"{prefix}.experts", ) if self.num_shared_experts > 0: intermediate_size = self.intermediate_size * self.num_shared_experts self.shared_experts = BailingMLP( hidden_size=self.hidden_size, intermediate_size=intermediate_size, reduce_results=False, prefix=f"{prefix}.shared_experts", quant_config=quant_config, ) def forward( self, hidden_states: torch.Tensor, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ) -> torch.Tensor: num_tokens, hidden_size = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_size) if self.num_shared_experts > 0: shared_output = self.shared_experts(hidden_states) router_logits = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) if self.num_shared_experts > 0: final_hidden_states = final_hidden_states + shared_output if self.tp_size > 1 and not use_reduce_scatter and not should_allreduce_fusion: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states class BailingGroupRMSNormGate(RMSNormGated): def __init__( self, hidden_size, eps=1e-5, group_size=None, norm_before_gate=True, device=None, dtype=None, ): super().__init__( hidden_size, eps=eps, group_size=group_size, norm_before_gate=norm_before_gate, device=device, dtype=dtype, activation="sigmoid", ) self.weight.weight_loader = self.weight_loader @staticmethod def weight_loader( param: torch.nn.Parameter, loaded_weight: torch.Tensor, ) -> None: tp_size = get_attention_tp_size() tp_rank = get_attention_tp_rank() shard_size = loaded_weight.shape[0] // tp_size shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) param.data.copy_(loaded_weight[shard].contiguous()) return class BailingMoELinearAttention(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, layer_id: int = 0, prefix: str = "linear_attn", ): super().__init__() self.layer_id = layer_id self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads self.total_kv_heads = config.num_attention_heads # MHA self.head_dim = getattr(config, "head_dim", None) if self.head_dim is None: self.head_dim = config.hidden_size // self.total_num_heads self.hidden_inner_size = self.head_dim * self.total_num_heads self.scaling = self.head_dim**-0.5 self.tp_size = get_attention_tp_size() self.tp_rank = get_attention_tp_rank() assert self.total_num_heads % self.tp_size == 0 self.tp_heads = self.total_num_heads // self.tp_size self.max_position_embeddings = config.max_position_embeddings self.rope_theta = getattr(config, "rope_theta", 600000) self.tp_kv_heads = self.total_kv_heads // self.tp_size self.q_size_per_rank = self.head_dim * self.tp_heads self.kv_size_per_rank = self.head_dim * self.tp_kv_heads self.use_qk_norm = getattr(config, "use_qk_norm", False) # minimax / seg_la / fla # TODO support fla self.linear_backend = getattr(config, "linear_backend", "seg_la") logger.debug(f"linear_backend in bailing_moe_linear: {self.linear_backend}") self.linear_scale = True if self.linear_backend == "minimax" else False self.linear_rope = getattr(config, "linear_rope", True) if hasattr(config, "use_linear_silu"): self.linear_silu = config.use_linear_silu elif hasattr(config, "linear_silu"): self.linear_silu = config.linear_silu else: self.linear_silu = False self.query_key_value = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_kv_heads, bias=(config.use_bias or config.use_qkv_bias), quant_config=quant_config, prefix=f"{prefix}.qkv_proj", tp_rank=self.tp_rank, tp_size=self.tp_size, ) if self.use_qk_norm: self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.g_proj = ColumnParallelLinear( self.hidden_size, self.hidden_inner_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.output_gate", tp_rank=self.tp_rank, tp_size=self.tp_size, ) self.dense = RowParallelLinear( self.hidden_inner_size, self.hidden_size, bias=config.use_bias, quant_config=quant_config, prefix=f"{prefix}.out_proj", tp_rank=self.tp_rank, tp_size=self.tp_size, reduce_results=False, ) self.attn = RadixAttention( self.tp_heads, self.head_dim, self.scaling, num_kv_heads=self.tp_kv_heads, layer_id=layer_id, quant_config=quant_config, prefix=f"{prefix}.attn", ) self.group_norm_size = getattr(config, "group_norm_size", 1) self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) assert ( self.tp_size <= self.group_norm_size ), "tp_size must be less than or equal to group_norm_size that can use local rms norm" assert ( self.group_norm_size % self.tp_size == 0 ), "group_norm_size must be divisible by tp_size" self.g_norm = BailingGroupRMSNormGate( hidden_size=self.hidden_inner_size // self.tp_size, eps=self.rms_norm_eps, group_size=self.hidden_inner_size // self.group_norm_size, ) # use fp32 rotary embedding if hasattr(config, "rotary_dim"): rotary_dim = config.rotary_dim elif hasattr(config, "partial_rotary_factor"): rotary_dim = int(self.head_dim * config.partial_rotary_factor) else: rotary_dim = self.head_dim self.rotary_emb = get_rope_wrapper( self.head_dim, rotary_dim=rotary_dim, max_position=self.max_position_embeddings, base=self.rope_theta, rope_scaling=config.rope_scaling, is_neox_style=True, device=get_global_server_args().device, dtype=torch.float32, ) @staticmethod def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight) return def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, **kwargs, ) -> torch.Tensor: qkv, _ = self.query_key_value(hidden_states) # logger.warning(f"===={self.layer_id=}, 1-1 {qkv.shape=}") # use rotary_emb support fp32 qkv = qkv.to(torch.float32) if self.linear_silu: qkv = F.silu(qkv) q, k, v = torch.split( qkv, [self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank], dim=-1, ) if self.use_qk_norm: q = q.reshape(-1, self.tp_heads, self.head_dim) k = k.reshape(-1, self.tp_kv_heads, self.head_dim) q = layernorm_fn( q, self.query_layernorm.weight.data, bias=None, eps=self.rms_norm_eps, is_rms_norm=True, ) k = layernorm_fn( k, self.key_layernorm.weight.data, bias=None, eps=self.rms_norm_eps, is_rms_norm=True, ) q = q.reshape(-1, self.q_size_per_rank) k = k.reshape(-1, self.kv_size_per_rank) if self.linear_rope: q, k = self.rotary_emb(positions, q, k) q = q.view((qkv.shape[0], self.tp_heads, self.head_dim)) k = k.view((qkv.shape[0], self.tp_kv_heads, self.head_dim)) v = v.view((qkv.shape[0], self.tp_kv_heads, self.head_dim)) # logger.warning(f"===={self.layer_id=}, 1-2 {q.shape=}, {k.shape=}, {v.shape=}") if self.linear_scale: q = q * self.scaling # q = q.to(torch.float32) # k = k.to(torch.float32) # v = v.to(torch.float32) hidden = self.attn(q, k, v, forward_batch).to(hidden_states.dtype) gate, _ = self.g_proj(hidden_states) # logger.warning( # f"===={self.layer_id=}, 1-3 {gate.shape=}, {hidden.shape=}, {gate.dtype=}, {hidden_states.dtype=}, {hidden.dtype=}" # ) if self.group_norm_size > 1: hidden = self.g_norm(hidden, gate) else: hidden = self.g_norm(hidden) hidden = F.sigmoid(gate) * hidden # logger.warning(f"===={self.layer_id=}, 1-4 {hidden.shape=}") hidden = hidden.data.to(hidden_states.dtype) hidden, _ = self.dense(hidden) # logger.warning(f"===={self.layer_id=}, 1-5 {hidden.shape=}") return hidden class BailingMoEAttention(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, layer_id: int = None, prefix: str = "mha", ) -> None: super().__init__() self.layer_id = layer_id self.hidden_size = config.hidden_size tp_size = get_attention_tp_size() self.total_num_heads = config.num_attention_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = config.num_key_value_heads if self.total_num_kv_heads >= tp_size: assert self.total_num_kv_heads % tp_size == 0 else: assert tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) self.head_dim = getattr(config, "head_dim", None) if self.head_dim is None: self.head_dim = self.hidden_size // self.total_num_heads self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.split_qkv = getattr(config, "using_split_qkv_in_self_attention", False) assert not self.split_qkv, "split_qkv is not supported for now" self.use_qk_norm = getattr(config, "use_qk_norm", False) self.query_key_value = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=(config.use_bias or config.use_qkv_bias), quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) if self.use_qk_norm: self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.dense = RowParallelLinear( self.total_num_heads * self.head_dim, self.hidden_size, bias=config.use_bias, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) if hasattr(config, "rotary_dim"): self.rotary_dim = config.rotary_dim elif hasattr(config, "partial_rotary_factor"): self.rotary_dim = int(self.head_dim * config.partial_rotary_factor) else: self.rotary_dim = self.head_dim self.max_position_embeddings = config.max_position_embeddings self.rope_theta = getattr(config, "rope_theta", 600000) self.rotary_emb = get_rope_wrapper( self.head_dim, rotary_dim=self.rotary_dim, max_position=self.max_position_embeddings, base=self.rope_theta, rope_scaling=config.rope_scaling, device=get_global_server_args().device, ) self.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, quant_config=quant_config, prefix=f"{prefix}.attn", ) def _apply_qk_norm( self, q: torch.Tensor, k: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: q_by_head = q.reshape(-1, self.head_dim) q_by_head = self.query_layernorm(q_by_head) q = q_by_head.view(q.shape) k_by_head = k.reshape(-1, self.head_dim) k_by_head = self.key_layernorm(k_by_head) k = k_by_head.view(k.shape) return q, k def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, **kwargs, ) -> torch.Tensor: qkv, _ = self.query_key_value(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: q, k = self._apply_qk_norm(q, k) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v, forward_batch) output, _ = self.dense(attn_output) return output class BailingMoELinearDecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, layer_id: int = 0, prefix: str = "layer", is_nextn: bool = False, ) -> None: super().__init__() self.layer_id = layer_id self.use_mla = getattr(config, "full_attention_type", "mla") == "mla" alt_stream = None # tptest # todo nextn if config.attention_type == 0: # Linear layer self.attention = BailingMoELinearAttention( config, quant_config=quant_config, layer_id=self.layer_id, prefix=prefix + ".attention", ) elif config.attention_type == 1: # softmax layer if self.use_mla: self.attention = DsV3MLA( config=config, hidden_size=config.hidden_size, num_heads=config.num_attention_heads, qk_nope_head_dim=config.qk_nope_head_dim, qk_rope_head_dim=config.qk_rope_head_dim, v_head_dim=config.v_head_dim, q_lora_rank=( config.q_lora_rank if hasattr(config, "q_lora_rank") else None ), kv_lora_rank=config.kv_lora_rank, rope_theta=getattr(config, "rope_theta", 600000), rope_scaling=config.rope_scaling, max_position_embeddings=262144, quant_config=quant_config, layer_id=layer_id, reduce_results=False, prefix=add_prefix("attention", prefix), alt_stream=alt_stream, ) else: logger.debug(f"layer {layer_id} use gqa") self.attention = BailingMoEAttention( config, quant_config=quant_config, layer_id=self.layer_id, prefix=prefix + ".attention", ) else: raise ValueError(f"Unsupported attention type: {config.attention_type}") self.expert_num = config.num_experts self.hidden_size = config.hidden_size is_moe_layer = self._is_layer_sparse(config, self.layer_id) is_previous_moe_layer = self._is_layer_sparse(config, self.layer_id - 1) is_next_layer_moe_layer = self._is_layer_sparse(config, self.layer_id + 1) if self.expert_num == 1: self.mlp = BailingMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) else: if is_nextn or self.layer_id >= config.first_k_dense_replace: # MoE layer self.mlp = BailingMoE( config, quant_config=quant_config, layer_id=self.layer_id, prefix=add_prefix("mlp", prefix), ) else: # dense layer self.mlp = BailingMLP( hidden_size=self.hidden_size, intermediate_size=config.intermediate_size, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) self.input_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, is_layer_sparse=is_moe_layer, is_previous_layer_sparse=is_previous_moe_layer, is_next_layer_sparse=is_next_layer_moe_layer, ) qkv_latent_func = ( self.attention.prepare_qkv_latent if config.attention_type == 1 and self.use_mla else None ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=False, qkv_latent_func=qkv_latent_func, ) def _is_layer_sparse( self, config: PretrainedConfig, layer_id: int, is_nextn: bool = False ) -> bool: return is_nextn or ( config.num_experts is not None and layer_id >= config.first_k_dense_replace ) def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], zero_allocator: BumpAllocator, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = self.layer_communicator.prepare_attn( hidden_states, residual, forward_batch ) # logger.warning( # f"===={self.layer_id=}, 1 shape= {hidden_states.shape}, {residual.shape}" # ) if not forward_batch.forward_mode.is_idle(): if self.use_mla: hidden_states = self.attention( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, ) else: hidden_states = self.attention( hidden_states=hidden_states, positions=positions, forward_batch=forward_batch, ) # logger.warning( # f"===={self.layer_id=}, 2 shape= {hidden_states.shape}, {residual.shape}" # ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) # logger.warning( # f"===={self.layer_id=}, 3 shape= {hidden_states.shape}, {residual.shape}" # ) should_allreduce_fusion = ( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( forward_batch ) ) use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( forward_batch ) hidden_states = self.mlp( hidden_states, should_allreduce_fusion, use_reduce_scatter ) hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual @staticmethod def shared_moe_coefficient_loader( param: torch.Tensor, loaded_weight: torch.Tensor ) -> None: assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight.to(torch.float32)) return class BailingMoELinearModel(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.pp_group = get_pp_group() self.config = config self.vocab_size = config.vocab_size self.embed_dim = config.hidden_size self.num_layers = config.num_hidden_layers self.layer_group_size = getattr(config, "layer_group_size", 1) self.decoder_attention_types = [ 0 if is_linear_layer(i, self.layer_group_size) else 1 for i in range(self.num_layers) ] num_linear = sum(1 for t in self.decoder_attention_types if t == 0) num_full = sum(1 for t in self.decoder_attention_types if t == 1) rank0_log( f"Layer config: {num_linear} linear attention layers, {num_full} full attention layers" ) assert ( self.num_layers % self.layer_group_size == 0 ), f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" if self.pp_group.is_first_rank: self.word_embeddings = VocabParallelEmbedding( self.vocab_size, self.embed_dim, enable_tp=not is_dp_attention_enabled(), org_num_embeddings=self.vocab_size, ) else: self.word_embeddings = PPMissingLayer() def layer_fn(idx, prefix): layer_idx = idx layer_config = copy.deepcopy(config) layer_config.attention_type = self.decoder_attention_types[layer_idx] decoder_kwargs = {"quant_config": quant_config, "layer_id": layer_idx} return BailingMoELinearDecoderLayer( layer_config, **decoder_kwargs, prefix=prefix ) self.layers, self.start_layer, self.end_layer = make_layers( self.num_layers, layer_fn, pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, prefix=f"{prefix}.layers", ) norm_kwargs = {} if hasattr(config, "rms_norm_eps"): norm_kwargs["eps"] = config.rms_norm_eps if self.pp_group.is_last_rank: self.norm = RMSNorm(config.hidden_size, **norm_kwargs) else: self.norm = PPMissingLayer() self.embed_scale = 1.0 return def forward( self, input_ids: Optional[torch.Tensor], positions: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, inputs_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: if self.pp_group.is_first_rank: if inputs_embeds is None: hidden_states = self.word_embeddings(input_ids) else: hidden_states = inputs_embeds residual = None else: assert pp_proxy_tensors is not None hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] total_num_layers = self.end_layer - self.start_layer device = inputs_embeds.device if inputs_embeds is not None else input_ids.device zero_allocator = BumpAllocator( buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), dtype=torch.float32, device=device, ) for i in range(self.start_layer, self.end_layer): with get_global_expert_distribution_recorder().with_current_layer(i): layer = self.layers[i] hidden_states, residual = layer( hidden_states=hidden_states, positions=positions, forward_batch=forward_batch, residual=residual, zero_allocator=zero_allocator, ) if not self.pp_group.is_last_rank: return PPProxyTensors( {"hidden_states": hidden_states, "residual": residual} ) else: if not forward_batch.forward_mode.is_idle(): if residual is None: hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) return hidden_states class BailingMoELinearForCausalLM(nn.Module): packed_modules_mapping = { "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"], "gate_up_proj": ["gate_proj", "up_proj"], } # To ensure correct weight loading and mapping. hf_to_sglang_mapper = WeightsMapper( orig_to_new_substr={ "attention.dense": "attention.out_proj", "layers.7.attention.out_proj": "layers.7.attention.o_proj", "layers.15.attention.out_proj": "layers.15.attention.o_proj", "layers.23.attention.out_proj": "layers.23.attention.o_proj", "layers.31.attention.out_proj": "layers.31.attention.o_proj", "layers.39.attention.out_proj": "layers.39.attention.o_proj", "layers.47.attention.out_proj": "layers.47.attention.o_proj", "layers.55.attention.out_proj": "layers.55.attention.o_proj", "layers.63.attention.out_proj": "layers.63.attention.o_proj", "layers.71.attention.out_proj": "layers.71.attention.o_proj", "layers.79.attention.out_proj": "layers.79.attention.o_proj", "attention.query_key_value": "attention.qkv_proj", "attention.g_proj": "attention.output_gate", }, ) def __init__( self, *, config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config self.model = BailingMoELinearModel( self.config, quant_config, prefix=add_prefix("model", prefix) ) if self.pp_group.is_last_rank: self.lm_head = ( self.word_embeddings if config.tie_word_embeddings else ParallelLMHead( config.vocab_size, config.hidden_size, params_dtype=torch.float32, quant_config=quant_config, use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) ) self.logits_processor = LogitsProcessor(config) else: self.lm_head = PPMissingLayer() @property def start_layer(self): return self.model.start_layer @property def end_layer(self): return self.model.end_layer def get_embed_and_head(self): """Used by the eagle_worker.""" return self.model.word_embeddings.weight, self.lm_head.weight def post_load_weights(self, is_nextn=False, weight_names=None): # Perform post-processing after loading weights if is_nextn: layer_ids = [self.config.num_hidden_layers] else: if weight_names is None: layer_ids = range(self.model.start_layer, self.model.end_layer) else: layer_ids = set() for name in weight_names: if "kv_b_proj" in name: layer_id = int(name.split(".")[2]) if ( layer_id < self.model.end_layer and layer_id >= self.model.start_layer ): layer_ids.add(layer_id) logger.debug(f"weight loading layer_ids: {layer_ids}") for layer_id in layer_ids: self_attn = ( self.model.layers[layer_id].attention if not is_nextn else self.model.decoder.attention ) if not hasattr(self_attn, "kv_b_proj"): continue if hasattr(self_attn.kv_b_proj, "qweight"): # AWQ compatible if _is_cuda or _is_hip: w = awq_dequantize( self_attn.kv_b_proj.qweight, self_attn.kv_b_proj.scales, self_attn.kv_b_proj.qzeros, ).T else: w = awq_dequantize( self_attn.kv_b_proj.qweight, self_attn.kv_b_proj.scales, self_attn.kv_b_proj.qzeros, 0, 0, 0, ).T else: w = self_attn.kv_b_proj.weight # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`. # This may affect the accuracy of fp8 model. # Fix deepseek v3 blockwise bmm by using deep_gemm use_deep_gemm_bmm = False if w.dtype in ( torch.float8_e4m3fn, torch.float8_e4m3fnuz, ): if ( hasattr(self.quant_config, "weight_block_size") and self.quant_config.weight_block_size is not None ): weight_block_size = self.quant_config.weight_block_size assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") if _is_fp8_fnuz: weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( weight=w, weight_scale=self_attn.kv_b_proj.weight_scale_inv, input_scale=None, ) else: weight = w weight_scale = self_attn.kv_b_proj.weight_scale_inv if ( _is_cuda and weight_block_size[0] == 128 and weight_block_size[1] == 128 ): if ( deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not deep_gemm_wrapper.DEEPGEMM_BLACKWELL and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false") ): block_scale = weight_scale use_deep_gemm_bmm = True else: w = block_quant_dequant( weight, weight_scale, weight_block_size, torch.bfloat16, ) else: w, scale = block_quant_to_tensor_quant( weight, weight_scale, weight_block_size ) self_attn.w_scale = scale else: if _is_fp8_fnuz: weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( weight=w, weight_scale=self_attn.kv_b_proj.weight_scale, input_scale=None, ) else: weight = w weight_scale = self_attn.kv_b_proj.weight_scale w, scale = channel_quant_to_tensor_quant(weight, weight_scale) self_attn.w_scale = scale if w.dtype == torch.int8: if hasattr(self.quant_config, "weight_block_size"): # block-wise int8 need it weight_block_size = self.quant_config.weight_block_size if weight_block_size is not None: assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") weight = w weight_scale = self_attn.kv_b_proj.weight_scale_inv w = int8_block_dequant( weight, weight_scale, weight_block_size ).to(torch.bfloat16) else: # channel-wise int8 need it w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( torch.bfloat16 ) w_kc, w_vc = w.unflatten( 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) if not use_deep_gemm_bmm: self_attn.w_kc = bind_or_assign( self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2) ) self_attn.w_vc = bind_or_assign( self_attn.w_vc, w_vc.contiguous().transpose(1, 2) ) if ( hasattr(self_attn.kv_b_proj, "weight_scale") and self_attn.w_scale is None ): self_attn.w_scale = bind_or_assign( self_attn.w_scale, self_attn.kv_b_proj.weight_scale ) if _is_hip: self_attn.w_scale *= 2.0 # TODO: remove this after adding FP8 support in bmm cpu kernel if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: self_attn.w_kc = ( self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale ) self_attn.w_vc = ( self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale ) else: num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] num_tiles_n = self_attn.v_head_dim // weight_block_size[0] ws_kc, ws_vc = block_scale.unflatten( 0, (-1, (num_tiles_k + num_tiles_n)) ).split([num_tiles_k, num_tiles_n], dim=1) self_attn.w_scale_k = bind_or_assign( self_attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() ) self_attn.w_scale_v = bind_or_assign( self_attn.w_scale_v, ws_vc.contiguous() ) self_attn.w_kc = bind_or_assign( self_attn.w_kc, w_kc.transpose(1, 2).contiguous() ) self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) self_attn.use_deep_gemm_bmm = True if ( deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 and hasattr(self.quant_config, "weight_block_size") and self.quant_config.weight_block_size is not None ): self._weight_requant_ue8m0(is_nextn) @classmethod def get_model_config_for_expert_location(cls, config): num_groups = getattr(config, "n_group", 0) from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation return ModelConfigForExpertLocation( num_layers=config.num_hidden_layers, num_logical_experts=config.num_experts, num_groups=None if num_groups == 0 else num_groups, ) def _weight_requant_ue8m0(self, is_nextn=False): weight_block_size = self.quant_config.weight_block_size moe_layers = list( range( self.config.first_k_dense_replace, self.config.num_hidden_layers, self.config.moe_layer_freq, ) ) num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers for layer_id in range(num_hidden_layers): if is_nextn: layer = self.model.decoder else: layer = self.model.layers[layer_id] module_list = [ layer.self_attn.kv_b_proj, layer.self_attn.o_proj, ] if self.config.q_lora_rank is not None: module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa) module_list.append(layer.self_attn.q_b_proj) else: module_list.append(layer.self_attn.kv_a_proj_with_mqa) module_list.append(layer.self_attn.q_proj) for module in module_list: requant_weight_ue8m0_inplace( module.weight, module.weight_scale_inv, weight_block_size ) if layer_id in moe_layers or is_nextn: shared_experts = getattr(layer.mlp, "shared_experts", None) if shared_experts is not None: for module in [ shared_experts.gate_up_proj, shared_experts.down_proj, ]: requant_weight_ue8m0_inplace( module.weight, module.weight_scale_inv, weight_block_size ) experts = layer.mlp.experts if isinstance(experts, DeepEPMoE): for w in [ experts.w13_weight_fp8, experts.w2_weight_fp8, ]: requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size) else: mlp = layer.mlp assert isinstance(mlp, DeepseekV2MLP) for module in [ mlp.gate_up_proj, mlp.down_proj, ]: requant_weight_ue8m0_inplace( module.weight, module.weight_scale_inv, weight_block_size ) def get_decoder_attention_types(self): return self.model.decoder_attention_types def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, inputs_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: hidden_states = self.model( input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, forward_batch=forward_batch, pp_proxy_tensors=pp_proxy_tensors, ) if self.pp_group.is_last_rank: return self.logits_processor( input_ids, hidden_states.float(), self.lm_head, forward_batch ) else: return hidden_states def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False ) -> Set[str]: def load_linear_attn_weight( name: str, loaded_weight: torch.Tensor, self ) -> None: if is_pp_missing_parameter(name, self): return param = params_dict[name] weight_loader = getattr( param, "weight_loader", BailingMoELinearAttention.weight_direct_load ) weight_loader = weight_loader_with_alias(name)(weight_loader) weight_loader(param, loaded_weight) return if is_nextn: if hasattr(self.config, "num_nextn_predict_layers"): num_nextn_layers = self.config.num_nextn_predict_layers assert num_nextn_layers == 1, "Only 1 nextn layer is supported" # compatible with old design nextn_layer_id = ( 0 if self.config.num_hidden_layers == 1 else self.config.num_hidden_layers ) else: raise ValueError("num nextn_predict_layers is not in the config") stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("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, ) if is_nextn: nextn_layer_prefix = f"model.layers.{nextn_layer_id}" nextn_spec_weight_names = [ "final_layernorm", "eh_proj", "enorm", "hnorm", ] params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() weight_names = [] fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( self.config.q_lora_rank is not None ) cached_a_proj = {} if fuse_qkv_a_proj else None for name, loaded_weight in weights: if name.startswith("model.mtp"): continue layer_idx = None if "model.layers." in name: layer_idx = int(name.split(".")[2]) if ( ("v_head" in name) or ("inv_freq" in name) or (self.config.tie_word_embeddings and "lm_head" in name) ): continue weight_names.append(name) if is_nextn: if not name.startswith(nextn_layer_prefix): continue # Use shared head and embed weights from target model if "shared_head.head" in name or "embed_tokens" in name: continue is_decoder = True # For nextn specific weights for weight_name in nextn_spec_weight_names: if weight_name in name: name = name.replace(nextn_layer_prefix, "model") is_decoder = False break # For decoder layer weights if is_decoder: name = name.replace(nextn_layer_prefix, "model.decoder") for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue if "mlp.experts" in name: continue name = name.replace(weight_name, param_name) if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue if is_pp_missing_parameter(name, self): continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) if name not in params_dict: continue if is_pp_missing_parameter(name, self): continue param = params_dict[name] weight_loader = param.weight_loader weight_loader( param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id, ) break else: if name.endswith(".bias") and name not in params_dict: continue if "slope" in name: continue if fuse_qkv_a_proj and ( "q_a_proj" in name or "kv_a_proj_with_mqa" in name ): cached_a_proj[name] = loaded_weight q_a_proj_name = ( name if "q_a_proj" in name else name.replace("kv_a_proj_with_mqa", "q_a_proj") ) kv_a_proj_name = ( name if "kv_a_proj_with_mqa" in name else name.replace("q_a_proj", "kv_a_proj_with_mqa") ) # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter if ( q_a_proj_name in cached_a_proj and kv_a_proj_name in cached_a_proj ): q_a_proj_weight = cached_a_proj[q_a_proj_name] kv_a_proj_weight = cached_a_proj[kv_a_proj_name] cat_dim = 0 if self.quant_config is not None and ( self.quant_config.get_name() == "awq" or self.quant_config.get_name() == "awq_marlin" or self.quant_config.get_name() == "moe_wna16" ): cat_dim = 1 fused_weight = torch.cat( [q_a_proj_weight, kv_a_proj_weight], dim=cat_dim ) param_name = ( name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") if "q_a_proj" in name else name.replace( "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa", ) ) if param_name not in params_dict: continue param = params_dict[param_name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, fused_weight) cached_a_proj.pop(q_a_proj_name) cached_a_proj.pop(kv_a_proj_name) else: if name not in params_dict: name = name.replace(".dense.", ".o_proj.") if name not in params_dict: continue if is_pp_missing_parameter(name, self): continue if ( "attention" in name and "slope" not in name and is_linear_layer(layer_idx, self.model.layer_group_size) ): load_linear_attn_weight(name, loaded_weight, self) loaded_params.add(name) continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) return loaded_params class BailingMoeV2_5ForCausalLM(BailingMoELinearForCausalLM): pass EntryClass = [ BailingMoeV2_5ForCausalLM, ]