From 295bc175760a1d690ca3f8e7c963dd5221df3815 Mon Sep 17 00:00:00 2001 From: chengshuang18 <60605542+chengshuang18@users.noreply.github.com> Date: Fri, 20 Feb 2026 13:58:15 +0800 Subject: [PATCH] Feature/sdar support (#19044) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: root Co-authored-by: chengshuang Co-authored-by: 赵晨阳 --- .../diffusion_language_models.md | 2 + python/sglang/srt/dllm/config.py | 17 +- python/sglang/srt/models/sdar.py | 589 ++++++++++++++ python/sglang/srt/models/sdar_moe.py | 746 ++++++++++++++++++ 4 files changed, 1348 insertions(+), 6 deletions(-) create mode 100644 python/sglang/srt/models/sdar.py create mode 100644 python/sglang/srt/models/sdar_moe.py diff --git a/docs/supported_models/text_generation/diffusion_language_models.md b/docs/supported_models/text_generation/diffusion_language_models.md index 2faa0206e..7dbb4828b 100644 --- a/docs/supported_models/text_generation/diffusion_language_models.md +++ b/docs/supported_models/text_generation/diffusion_language_models.md @@ -107,3 +107,5 @@ Below the supported models are summarized in a table. | Model Family | Example Model | Description | | -------------------------- | ---------------------------- | ---------------------------------------------------------------------------------------------------- | | **LLaDA2.0 (mini, flash)** | `inclusionAI/LLaDA2.0-flash` | LLaDA2.0-flash is a diffusion language model featuring a 100B Mixture-of-Experts (MoE) architecture. | +| **SDAR (JetLM)** | `JetLM/SDAR-8B-Chat` | SDAR series diffusion language model (Chat), dense architecture. | +| **SDAR (JetLM)** | `JetLM/SDAR-30B-A3B-Chat` | SDAR series diffusion language model (Chat), MoE architecture. | diff --git a/python/sglang/srt/dllm/config.py b/python/sglang/srt/dllm/config.py index ce601cbdc..edd204926 100644 --- a/python/sglang/srt/dllm/config.py +++ b/python/sglang/srt/dllm/config.py @@ -31,14 +31,19 @@ class DllmConfig: model_path=server_args.model_path, model_revision=server_args.revision, ) + DLLM_PARAMS = { + "LLaDA2MoeModelLM": {"block_size": 32, "mask_id": 156895}, + "SDARForCausalLM": {"block_size": 4, "mask_id": 151669}, + "SDARMoeForCausalLM": {"block_size": 4, "mask_id": 151669}, + } - if model_config.hf_config.architectures[0] == "LLaDA2MoeModelLM": - block_size = 32 - mask_id = 156895 + arch = model_config.hf_config.architectures[0] + if arch in DLLM_PARAMS: + params = DLLM_PARAMS[arch] + block_size = params["block_size"] + mask_id = params["mask_id"] else: - raise RuntimeError( - f"Unknown diffusion LLM: {model_config.hf_config.architectures[0]}" - ) + raise RuntimeError(f"Unknown diffusion LLM: {arch}") max_running_requests = ( 1 diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py new file mode 100644 index 000000000..70ab59a48 --- /dev/null +++ b/python/sglang/srt/models/sdar.py @@ -0,0 +1,589 @@ +# coding=utf-8 +""" +SGLang SDARModelLM (block diffusion / dLLM-style forward). +""" + +import logging +from typing import Iterable, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +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 ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import AttentionType, RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +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, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.utils import ( + apply_qk_norm, + create_fused_set_kv_buffer_arg, + enable_fused_set_kv_buffer, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix, is_cuda, make_layers + +logger = logging.getLogger(__name__) +_is_cuda = is_cuda() + + +class SDARMLP(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config=None, + reduce_results: bool = True, + prefix: str = "", + ): + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + config.hidden_size, + [config.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=False, + reduce_results=reduce_results, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + self.act_fn = SiluAndMul() + + def forward(self, hidden_states: torch.Tensor, use_reduce_scatter: bool = False): + gate_up, _ = self.gate_up_proj(hidden_states) + hidden_states = self.act_fn(gate_up) + hidden_states, _ = self.down_proj( + hidden_states, skip_all_reduce=use_reduce_scatter + ) + return hidden_states + + +class SDARAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config=None, + reduce_results: bool = True, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.tp_size = get_tensor_model_parallel_world_size() + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + + self.head_dim = getattr( + config, "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.scale = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=getattr(config, "attention_bias", False), + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=getattr(config, "attention_bias", False), + quant_config=quant_config, + reduce_results=reduce_results, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("o_proj", prefix), + ) + + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + rope_theta = getattr(config, "rope_theta", 10000.0) + rope_scaling = getattr(config, "rope_scaling", None) + max_pos = getattr(config, "max_position_embeddings", 32768) + self.rotary_dim = self.head_dim + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_pos, + base=rope_theta, + rope_scaling=rope_scaling, + ) + + # RadixAttention: ENCODER_ONLY lets ForwardBatch provide non-causal / block masks (dLLM) + # NOTE: this is the key change vs AR Llama-style DECODER self-attn. + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scale, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + attn_type=AttentionType.ENCODER_ONLY, + prefix=add_prefix("attn", prefix), + ) + self.alt_stream = alt_stream + + def forward_prepare_native(self, positions, hidden_states, forward_batch): + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = apply_qk_norm( + q=q, + k=k, + q_norm=self.q_norm, + k_norm=self.k_norm, + head_dim=self.head_dim, + alt_stream=self.alt_stream, + ) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=v, + layer=self.attn, + forward_batch=forward_batch, + ) + if enable_fused_set_kv_buffer(forward_batch) + else None + ), + ) + return q, k, v + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ): + if get_global_server_args().rl_on_policy_target is not None: + hidden_states = hidden_states.bfloat16() + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = apply_qk_norm( + q=q, + k=k, + q_norm=self.q_norm, + k_norm=self.k_norm, + head_dim=self.head_dim, + alt_stream=self.alt_stream, + ) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=v, + layer=self.attn, + forward_batch=forward_batch, + ) + if enable_fused_set_kv_buffer(forward_batch) + else None + ), + ) + + if get_global_server_args().rl_on_policy_target is not None: + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + + context_layer = self.attn( + q, + k, + v, + forward_batch, + save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), + ) + out, _ = self.o_proj(context_layer) + return out + + +class SDARBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_id = layer_id + + norm_kwargs = ( + dict( + weight_dtype=torch.float32, + cast_x_before_out_mul=True, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + if get_global_server_args().rl_on_policy_target is not None + else {} + ) + self.input_layernorm = RMSNorm( + self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) + self.post_attention_layernorm = RMSNorm( + self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) + + self.self_attn = SDARAttention( + layer_id=layer_id, + config=config, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + ) + + self.mlp = SDARMLP( + config=config, + quant_config=quant_config, + reduce_results=True, + prefix=add_prefix("mlp", prefix), + ) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=False, + is_previous_layer_sparse=False, + is_next_layer_sparse=False, + ) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, + residual, + forward_batch, + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, + residual, + forward_batch, + ) + + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + hidden_states = self.mlp(hidden_states, use_reduce_scatter=use_reduce_scatter) + + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + +class SDARModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + quant_config=quant_config, + use_attn_tp_group=is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: SDARBlock( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + alt_stream=alt_stream, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + norm_kwargs = ( + dict( + weight_dtype=torch.float32, + cast_x_before_out_mul=True, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + if get_global_server_args().rl_on_policy_target is not None + else {} + ) + self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + hidden_states = ( + self.embed_tokens(input_ids) if input_embeds is None else input_embeds + ) + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors.get("residual", None) + + for i in range(self.start_layer, self.end_layer): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + + 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(): + hidden_states, residual = self.norm(hidden_states, residual) + return hidden_states + + +class SDARForCausalLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.pp_group = get_pp_group() + assert self.pp_group.world_size == 1, ( + f"SDARMoeForCausalLM does not support pipeline parallel (pp_size={self.pp_group.world_size}). " + "Please set pp_size=1." + ) + + self.config = config + self.quant_config = quant_config + alt_stream = torch.cuda.Stream() if _is_cuda else None + + self.model = SDARModel( + config, + quant_config=quant_config, + prefix=add_prefix("model", ""), + alt_stream=alt_stream, + ) + + if self.pp_group.is_last_rank: + tp_size = get_tensor_model_parallel_world_size() + if ( + self.pp_group.world_size == 1 + and config.tie_word_embeddings + and tp_size == 1 + ): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + prefix=add_prefix("lm_head", prefix), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config, return_full_logits=True) + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + input_embeds=input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + else: + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if not name.startswith("model.") and ( + name.startswith("layers.") + or name.startswith("embed_tokens.") + or name.startswith("norm.") + ): + name = add_prefix(name, "model") + + if name == "model.embed_tokens.weight": + if self.pp_group.is_last_rank and self.config.tie_word_embeddings: + param = params_dict["lm_head.weight"] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + if "scale" in name: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + +EntryClass = SDARForCausalLM diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py new file mode 100644 index 000000000..a1b20f67c --- /dev/null +++ b/python/sglang/srt/models/sdar_moe.py @@ -0,0 +1,746 @@ +# coding=utf-8 +""" +SGLang SDARMoeModelLM (block diffusion / dLLM-style forward) with MoE MLP. +""" + +import logging +from typing import Iterable, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_moe_expert_parallel_world_size, + get_pp_group, + 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.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +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 ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_use_flashinfer_cutlass_moe_fp4_allgather, +) +from sglang.srt.layers.moe.ep_moe.layer import 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.moe.utils import ( + RoutingMethodType, + filter_moe_weight_param_global_expert, +) +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import AttentionType, RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +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, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.utils import ( + apply_qk_norm, + create_fused_set_kv_buffer_arg, + enable_fused_set_kv_buffer, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers + +logger = logging.getLogger(__name__) +_is_cuda = is_cuda() + + +class SDARMoeSparseMoeBlock(nn.Module): + """ + Qwen3MoE-style sparse MoE block: + - gate: ReplicatedLinear(hidden, num_experts) + - topk routing: TopK + - experts: get_moe_impl_class(quant_config)(...) + """ + + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.layer_id = layer_id + self.tp_size = get_tensor_model_parallel_world_size() + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} > num_experts {config.num_experts}." + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + use_grouped_topk=False, + layer_id=layer_id, + ) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=config.num_experts + + get_global_server_args().ep_num_redundant_experts, + top_k=config.num_experts_per_tok, + layer_id=layer_id, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + prefix=add_prefix("experts", prefix), + routing_method_type=RoutingMethodType.Renormalize, + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + # Deepep / FuseEP support + if get_moe_a2a_backend().is_deepep(): + self.ep_size = get_moe_expert_parallel_world_size() + self.num_experts = ( + config.num_experts + get_global_server_args().ep_num_redundant_experts + ) + self.top_k = config.num_experts_per_tok + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + if ( + not get_moe_a2a_backend().is_deepep() + and not get_moe_a2a_backend().is_ascend_fuseep() + ): + return self.forward_normal( + hidden_states, + should_allreduce_fusion=should_allreduce_fusion, + use_reduce_scatter=use_reduce_scatter, + ) + else: + assert forward_batch is not None, "deepep/fuseep MoE needs forward_batch" + return self.forward_deepep(hidden_states, forward_batch) + + def forward_normal( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + router_logits, _ = self.gate(hidden_states) # (T, E) + topk_output = self.topk(hidden_states, router_logits) + out = self.experts(hidden_states, topk_output) # (T, H) + + # TP all-reduce (unless fused / reduce_scatter / fp4 allgather path) + if ( + self.tp_size > 1 + and not should_allreduce_fusion + and not use_reduce_scatter + and not should_use_flashinfer_cutlass_moe_fp4_allgather() + ): + out = tensor_model_parallel_all_reduce(out) + + return out.view(num_tokens, hidden_dim) + + def forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + if hidden_states.shape[0] > 0: + router_logits, _ = self.gate(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id + ), + ) + else: + topk_output = self.topk.empty_topk_output(hidden_states.device) + + out = self.experts(hidden_states=hidden_states, topk_output=topk_output) + return out + + def get_moe_weights(self): + return [ + p.data + for name, p in self.experts.named_parameters() + if name not in ["correction_bias"] + and filter_moe_weight_param_global_expert( + name, p, self.experts.num_local_experts + ) + ] + + +class SDARMoeAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= attn_tp_size: + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + + self.head_dim = getattr( + config, "head_dim", 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.scale = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=getattr(config, "attention_bias", False), + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=getattr(config, "attention_bias", False), + quant_config=quant_config, + reduce_results=reduce_results, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("o_proj", prefix), + ) + + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + rope_theta = getattr(config, "rope_theta", 10000.0) + rope_scaling = getattr(config, "rope_scaling", None) + max_pos = getattr(config, "max_position_embeddings", 32768) + self.rotary_dim = self.head_dim + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_pos, + base=rope_theta, + rope_scaling=rope_scaling, + ) + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scale, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + attn_type=AttentionType.ENCODER_ONLY, + prefix=add_prefix("attn", prefix), + ) + self.alt_stream = alt_stream + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if get_global_server_args().rl_on_policy_target is not None: + hidden_states = hidden_states.bfloat16() + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = apply_qk_norm( + q=q, + k=k, + q_norm=self.q_norm, + k_norm=self.k_norm, + head_dim=self.head_dim, + alt_stream=self.alt_stream, + ) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=v, + layer=self.attn, + forward_batch=forward_batch, + ) + if enable_fused_set_kv_buffer(forward_batch) + else None + ), + ) + + if get_global_server_args().rl_on_policy_target is not None: + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + + context = self.attn( + q, + k, + v, + forward_batch, + save_kv_cache=not enable_fused_set_kv_buffer(forward_batch), + ) + out, _ = self.o_proj(context) + return out + + +class SDARMoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.layer_id = layer_id + + norm_kwargs = ( + dict( + weight_dtype=torch.float32, + cast_x_before_out_mul=True, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + if get_global_server_args().rl_on_policy_target is not None + else {} + ) + self.input_layernorm = RMSNorm( + self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) + self.post_attention_layernorm = RMSNorm( + self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs + ) + + self.self_attn = SDARMoeAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, + ) + + self.mlp = SDARMoeSparseMoeBlock( + layer_id=layer_id, + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=True, + is_previous_layer_sparse=True, + is_next_layer_sparse=True, + ) + + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=(layer_id == config.num_hidden_layers - 1), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + hidden_states = self.mlp( + hidden_states, + forward_batch=forward_batch, + should_allreduce_fusion=should_allreduce_fusion, + use_reduce_scatter=use_reduce_scatter, + ) + + if should_allreduce_fusion: + hidden_states._sglang_needs_allreduce_fusion = True + else: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + +class SDARMoeModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + quant_config=quant_config, + use_attn_tp_group=is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: SDARMoeBlock( + config=config, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=alt_stream, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + + if self.pp_group.is_last_rank: + norm_kwargs = ( + dict( + weight_dtype=torch.float32, + cast_x_before_out_mul=True, + override_orig_dtype=torch.float32, + fp32_residual=True, + ) + if get_global_server_args().rl_on_policy_target is not None + else {} + ) + self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + hidden_states = ( + self.embed_tokens(input_ids) if input_embeds is None else input_embeds + ) + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors.get("residual", None) + + for i in range(self.start_layer, self.end_layer): + layer = self.layers[i] + with get_global_expert_distribution_recorder().with_current_layer(i): + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + if not forward_batch.forward_mode.is_idle(): + hidden_states, residual = self.norm(hidden_states, residual) + return hidden_states + + +class SDARMoeForCausalLM(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.pp_group = get_pp_group() + assert self.pp_group.world_size == 1, ( + f"SDARMoeForCausalLM does not support pipeline parallel (pp_size={self.pp_group.world_size}). " + "Please set pp_size=1." + ) + + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + alt_stream = torch.cuda.Stream() if _is_cuda else None + + self.model = SDARMoeModel( + config, + quant_config=quant_config, + prefix=add_prefix("model", ""), + alt_stream=alt_stream, + ) + + if self.pp_group.is_last_rank: + tp_size = get_tensor_model_parallel_world_size() + if ( + self.pp_group.world_size == 1 + and getattr(config, "tie_word_embeddings", False) + and tp_size == 1 + ): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + prefix=add_prefix("lm_head", prefix), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config, return_full_logits=True) + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + input_embeds=input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("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, + ) + + if not hasattr(self, "_cached_params_dict"): + self._cached_params_dict = dict(self.named_parameters()) + params_dict = self._cached_params_dict + + for name, loaded_weight in weights: + if not name.startswith("model.") and ( + name.startswith("layers.") + or name.startswith("embed_tokens.") + or name.startswith("norm.") + ): + name = add_prefix(name, "model") + + if name == "model.embed_tokens.weight": + if self.pp_group.is_last_rank and getattr( + self.config, "tie_word_embeddings", False + ): + if "lm_head.weight" in params_dict: + param = params_dict["lm_head.weight"] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + + if "scale" in name: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + + name2 = name.replace(weight_name, param_name) + if name2.endswith(".bias") and name2 not in params_dict: + continue + if name2 not in params_dict: + continue + + param = params_dict[name2] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + break + else: + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + is_expert_weight = True + + name2 = name.replace(weight_name, param_name) + if name2 not in params_dict: + continue + + param = params_dict[name2] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader( + param, + loaded_weight, + name2, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if is_expert_weight: + continue + + # 3) regular params + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + if not hasattr(self, "routed_experts_weights_of_layer"): + self.routed_experts_weights_of_layer = LazyValue( + lambda: { + lid: self.model.layers[lid].mlp.get_moe_weights() + for lid in range(self.start_layer, self.end_layer) + if isinstance(self.model.layers[lid].mlp, SDARMoeSparseMoeBlock) + } + ) + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None, + ) + + +EntryClass = SDARMoeForCausalLM