From 04d918ea4f23baff064fb58225e5960b7903a7b4 Mon Sep 17 00:00:00 2001 From: apinge Date: Fri, 21 Aug 2026 03:13:58 +0000 Subject: [PATCH 1/7] support qwen3.8 on branch 0.5.15 Signed-off-by: apinge --- python/sglang/srt/layers/layernorm.py | 39 +- python/sglang/srt/models/qwen3_5.py | 334 ++++++++++++++++-- .../srt/utils/hf_transformers/common.py | 9 + 3 files changed, 337 insertions(+), 45 deletions(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 405cc19a20c4..fd3121456881 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -99,6 +99,7 @@ def layernorm( ) _has_aiter_layer_norm = False _has_vllm_rms_norm = False +_has_rocm_triton_gemma_rms_norm = False if _use_aiter: from aiter import layernorm2d_fwd as layer_norm from aiter import rmsnorm2d_fwd as rms_norm @@ -115,6 +116,19 @@ def layernorm( # Fallback: vllm not available, will use forward_native _has_vllm_rms_norm = False +if _is_hip: + try: + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm, + ) + from sglang.jit_kernel.minimax_m3.rmsnorm import ( + gemma_rmsnorm as rocm_triton_gemma_rmsnorm, + ) + + _has_rocm_triton_gemma_rms_norm = True + except ImportError: + _has_rocm_triton_gemma_rms_norm = False + if _is_cuda: # HF-semantics RMSNorm kernel (JIT-compiled). Used when `cast_x_before_out_mul=True` # (the transformers backend path) to produce outputs that are numerically identical @@ -480,14 +494,10 @@ def forward_hip( # NOTE: Remove this if aiter kernel supports discontinuous input x = x.contiguous() if residual is not None: - out = torch.empty_like(x) - residual_out = torch.empty_like(x) if post_residual_addition is not None: residual = residual + post_residual_addition - fused_add_rms_norm( - out, x, residual_out, residual, self.weight.data, self.variance_epsilon - ) - return out, residual_out + fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon) + return x, residual out = torch.empty_like(x) rms_norm(out, x, self.weight.data, self.variance_epsilon) return out @@ -788,6 +798,15 @@ def forward_hip( residual: Optional[torch.Tensor] = None, post_residual_addition: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if _has_rocm_triton_gemma_rms_norm: + if residual is not None: + if post_residual_addition is not None: + residual = residual + post_residual_addition + return rocm_triton_gemma_fused_add_rmsnorm( + x, residual, self.weight.data, self.variance_epsilon + ) + return rocm_triton_gemma_rmsnorm(x, self.weight.data, self.variance_epsilon) + if not _has_vllm_rms_norm: return self.forward_native(x, residual, post_residual_addition) @@ -811,14 +830,10 @@ def forward_hip( if not x.is_contiguous(): x = x.contiguous() if residual is not None: - out = torch.empty_like(x) - residual_out = torch.empty_like(x) if post_residual_addition is not None: residual = residual + post_residual_addition - fused_add_rms_norm( - out, x, residual_out, residual, w, self.variance_epsilon - ) - return out, residual_out + fused_add_rms_norm(x, residual, w, self.variance_epsilon) + return x, residual out = torch.empty_like(x) rms_norm(out, x, w, self.variance_epsilon) return out diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index e09080701d0c..5233e34953a0 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -30,6 +30,7 @@ from sglang.srt.configs.qwen3_5 import ( Qwen3_5Config, Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, Qwen3_5TextConfig, ) @@ -57,6 +58,7 @@ QKVParallelLinear, RowParallelLinear, ) +from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.parameter import ( BlockQuantScaleParameter, @@ -67,7 +69,10 @@ from sglang.srt.layers.radix_linear_attention import RadixLinearAttention 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 VocabParallelEmbedding +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, @@ -1103,8 +1108,17 @@ def forward( } -class Qwen3_5ForCausalLM(nn.Module): - """Qwen3.5 Model with support for dense variant.""" +def _coerce_qwen3_5_text_config( + config: Qwen3_5TextConfig, *, moe: bool = False +) -> Qwen3_5TextConfig: + cls = Qwen3_5MoeTextConfig if moe else Qwen3_5TextConfig + if isinstance(config, cls): + return config + return cls(**config.to_dict()) + + +class Qwen3_5Model(nn.Module): + """Qwen3.5 backbone (embeddings + decoder layers + norm) for the dense variant.""" packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], @@ -1183,7 +1197,7 @@ def _maybe_autodisable_shared_experts_fusion(self, config, quant_config): from sglang.srt.arg_groups.overrides import declare_load_time_override declare_load_time_override( - "Qwen3_5ForCausalLM._maybe_autodisable_shared_experts_fusion", + "Qwen3_5Model._maybe_autodisable_shared_experts_fusion", {"disable_shared_experts_fusion": True}, ) logger.info( @@ -1341,6 +1355,97 @@ def forward( return hidden_states, aux_hidden_states + +class Qwen3_5MoeModel(Qwen3_5Model): + """Qwen3.5-MoE backbone (embeddings + decoder layers + norm).""" + + def __init__( + self, + config: Qwen3_5TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config, quant_config=quant_config, prefix=prefix) + + +class Qwen3_5ForCausalLM(nn.Module): + """Qwen3.5 text-only dense causal LM.""" + + packed_modules_mapping = Qwen3_5Model.packed_modules_mapping + supported_lora_modules = Qwen3_5Model.supported_lora_modules + + def __init__( + self, + config: Qwen3_5TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + config = _coerce_qwen3_5_text_config(config, moe=False) + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.model = Qwen3_5Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_flags().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = False + + def get_hidden_dim(self, module_name: str, layer_idx: int): + return self.model.get_hidden_dim(module_name, layer_idx) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @property + def start_layer(self) -> int: + return self.model.start_layer + + @property + def end_layer(self) -> int: + 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, + ) -> Union[torch.Tensor, PPProxyTensors]: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if not self.pp_group.is_last_rank: + return hidden_states + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states + ) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1369,6 +1474,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): name = name.replace(r"model.language_model.", r"model.") if ".self_attn." in name: name = name.replace(".self_attn", "") + if ( + self.config.tie_word_embeddings + and self.pp_group.is_last_rank + and "model.embed_tokens.weight" in name + and "lm_head.weight" in params_dict + ): + lm_head_param = params_dict["lm_head.weight"] + weight_loader = getattr( + lm_head_param, "weight_loader", default_weight_loader + ) + weight_loader(lm_head_param, loaded_weight) layer_id = get_layer_id(name) if ( layer_id is not None @@ -1388,9 +1504,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue - # Skip layers on other devices. - # if is_pp_missing_parameter(name, self): - # continue if name not in params_dict: continue param = params_dict[name] @@ -1411,23 +1524,66 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): loaded_params.add(name) return loaded_params - @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, - ) - class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): + """Qwen3.5-MoE text-only causal LM.""" + def __init__( self, config: Qwen3_5TextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: - super().__init__(config=config, quant_config=quant_config, prefix=prefix) + config = _coerce_qwen3_5_text_config(config, moe=True) + nn.Module.__init__(self) + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.model = Qwen3_5MoeModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + if self.pp_group.is_last_rank: + if self.pp_group.world_size == 1 and config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_flags().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = False + + self.num_fused_shared_experts = 0 + if _use_aiter and not _disable_shared_experts_fusion(): + self.num_fused_shared_experts = self._get_num_fused_shared_experts() + self.enable_shared_expert_fusion = self.num_fused_shared_experts > 0 + + def _get_num_fused_shared_experts(self): + if not hasattr(self.model, "layers"): + return 0 + for layer_id in range(self.model.start_layer, self.model.end_layer): + mlp = getattr(self.model.layers[layer_id], "mlp", None) + if hasattr(mlp, "num_fused_shared_experts"): + return mlp.num_fused_shared_experts + return 0 + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + @classmethod + def get_model_config_for_expert_location(cls, config): + text_config = getattr(config, "text_config", config) + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.num_experts, + num_groups=None, + ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ @@ -1444,13 +1600,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("in_proj_ba.", "in_proj_a.", 1), ] + num_experts = self.config.num_experts + # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) 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, + num_experts=( + num_experts + if not self.enable_shared_expert_fusion + else num_experts + self.num_fused_shared_experts + ), ) # Skip loading extra parameters for GPTQ/modelopt models. @@ -1473,7 +1635,40 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ("experts.w2_weight", "experts.down_proj", 0, "w2"), ] - num_experts = self.config.num_experts + if self.enable_shared_expert_fusion: + """ + When shared experts are fused, we need to map the shared experts to routed experts. + + mlp.share_expert.gate_up_proj.weight --> experts.512.gate_up_proj.weight -> experts.w13_weight, expert_id = 512 + mlp.share_expert.down_proj.weight --> experts.512.down_proj.weight -> experts.w2_weight, expert_id = 512 + """ + fused_expert_params_mapping += [ + ( + "experts.w13_", + f"experts.{num_experts}.gate_up_proj.", + num_experts, + "w1", + ), + ( + "experts.w2_", + f"experts.{num_experts}.down_proj.", + num_experts, + "w2", + ), + ## shared experts may contain gate_proj and up_proj instead of gate_up_proj + ( + "experts.w13_", + f"experts.{num_experts}.gate_proj.", + num_experts, + "w1", + ), + ( + "experts.w13_", + f"experts.{num_experts}.up_proj.", + num_experts, + "w3", + ), + ] def load_fused_expert_weights( name: str, @@ -1512,6 +1707,17 @@ def load_fused_expert_weights( name = name.replace(r"model.language_model.", r"model.") if ".self_attn." in name: name = name.replace(".self_attn", "") + if ( + self.config.tie_word_embeddings + and self.pp_group.is_last_rank + and "model.embed_tokens.weight" in name + and "lm_head.weight" in params_dict + ): + lm_head_param = params_dict["lm_head.weight"] + weight_loader = getattr( + lm_head_param, "weight_loader", default_weight_loader + ) + weight_loader(lm_head_param, loaded_weight) layer_id = get_layer_id(name) if ( @@ -1521,8 +1727,18 @@ def load_fused_expert_weights( ): continue + if self.enable_shared_expert_fusion: + if "mlp.shared_expert." in name: + # Firstly map mlp.shared_expert.xx_proj to mlp.experts.512.xx_proj + name = name.replace( + "mlp.shared_expert.", + f"mlp.experts.{num_experts}.", + ) + for param_name, weight_name, shard_id in stacked_params_mapping: - if "experts.gate_up_proj" in name or "experts.down_proj" in name: + if name.endswith("experts.gate_up_proj") or name.endswith( + "experts.down_proj" + ): is_fused_expert = True expert_params_mapping = fused_expert_params_mapping @@ -1563,7 +1779,10 @@ def load_fused_expert_weights( is_expert_weight = True name_mapped = name.replace(weight_name, param_name) if is_fused_expert: + # is_fused_expert is True, the checkpoint contains gate_up_proj and down_proj for each expert if "experts.gate_up_proj" in name: + # experts.gate_up_proj contains all routed experts, excluding shared experts + # split into w1 and w3 loaded_weight = loaded_weight.chunk(2, dim=-2) load_fused_expert_weights( name_mapped, @@ -1579,7 +1798,8 @@ def load_fused_expert_weights( "w3", num_experts, ) - else: + elif "experts.down_proj" in name: + # experts.down_proj contains all routed experts, excluding shared experts load_fused_expert_weights( name_mapped, params_dict, @@ -1587,6 +1807,41 @@ def load_fused_expert_weights( shard_id, num_experts, ) + elif self.enable_shared_expert_fusion: + # shared experts should be loaded to experts.w13_weight and experts.w2_weight + param = params_dict[name_mapped] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + if f"{num_experts}.gate_up_proj" in name: + # split into w1 and w3 + loaded_weight = loaded_weight.chunk(2, dim=-2) + # load to experts.w13_weight, shard_id = w1, expert_id = num_experts + weight_loader( + param, + loaded_weight[0], + name_mapped, + "w1", + expert_id, + ) + # load to experts.w13_weight, shard_id = w3, expert_id = num_experts + weight_loader( + param, + loaded_weight[1], + name_mapped, + "w3", + expert_id, + ) + else: + # load down_proj to experts.w2_weight, shard_id = w2, expert_id = num_experts + # Or load gate_proj and up_proj to experts.w13_weight, shard_id = w1/w3, expert_id = num_experts + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id, + expert_id, + ) else: # Skip loading extra parameters for GPTQ/modelopt models. if ( @@ -1627,21 +1882,29 @@ def load_fused_expert_weights( logger.warning(f"Parameter {name} not found in params_dict") loaded_params.add(name) + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock) + } + ) + return loaded_params class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration): - packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping + packed_modules_mapping = Qwen3_5Model.packed_modules_mapping hf_to_sglang_mapper = None - supported_lora_modules = Qwen3_5ForCausalLM.supported_lora_modules + supported_lora_modules = Qwen3_5Model.supported_lora_modules def __init__( self, config: Qwen3_5Config, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", - language_model_cls=Qwen3_5ForCausalLM, + language_model_cls=Qwen3_5Model, ): super().__init__(config, quant_config, prefix, language_model_cls) @@ -1789,17 +2052,17 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): """Qwen3.5 MoE Vision-Language Model.""" - packed_modules_mapping = Qwen3_5ForCausalLM.packed_modules_mapping + packed_modules_mapping = Qwen3_5Model.packed_modules_mapping hf_to_sglang_mapper = None - supported_lora_modules = Qwen3_5ForCausalLM.supported_lora_modules + supported_lora_modules = Qwen3_5Model.supported_lora_modules def __init__( self, config: Qwen3_5MoeConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", - language_model_cls=Qwen3_5MoeForCausalLM, + language_model_cls=Qwen3_5MoeModel, ) -> None: super().__init__(config, quant_config, prefix, language_model_cls) rope_config = getattr(self.config, "rope_parameters", None) or getattr( @@ -1822,13 +2085,13 @@ def should_apply_lora(self, module_name: str) -> bool: return module_name.startswith("model.layers.") def _get_num_fused_shared_experts(self): - if not ( - hasattr(self.model, "layers") - and len(self.model.layers) > 0 - and hasattr(self.model.layers[0].mlp, "num_fused_shared_experts") - ): + if not hasattr(self.model, "layers"): return 0 - return self.model.layers[0].mlp.num_fused_shared_experts + for layer_id in range(self.model.start_layer, self.model.end_layer): + mlp = getattr(self.model.layers[layer_id], "mlp", None) + if hasattr(mlp, "num_fused_shared_experts"): + return mlp.num_fused_shared_experts + return 0 def get_embed_and_head(self): embed = self.model.embed_tokens.weight if self.pp_group.is_first_rank else None @@ -2176,4 +2439,9 @@ def get_model_config_for_expert_location(cls, config): ) -EntryClass = [Qwen3_5MoeForConditionalGeneration, Qwen3_5ForConditionalGeneration] +EntryClass = [ + Qwen3_5MoeForConditionalGeneration, + Qwen3_5ForConditionalGeneration, + Qwen3_5MoeForCausalLM, + Qwen3_5ForCausalLM, +] diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index eef80dc7a985..956f90d66594 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -51,6 +51,8 @@ Olmo3Config, Qwen3_5Config, Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, Qwen3NextConfig, Step3p5Config, Step3p7Config, @@ -104,6 +106,13 @@ DeepseekVLV2Config, Qwen3_5Config, Qwen3_5MoeConfig, + # Text-only checkpoints ship model_type "qwen3_5_text" / + # "qwen3_5_moe_text" (no multimodal wrapper). Register these so + # get_config() returns SGLang's config subclass instead of the + # transformers-native class, keeping isinstance-based dispatch + # (e.g. hybrid_gdn_config -> HybridLinearAttnBackend) working. + Qwen3_5TextConfig, + Qwen3_5MoeTextConfig, InternS2PreviewConfig, JetNemotronConfig, JetVLMConfig, From f39ec213805f8c05785adc64bbbe4b789edf7dd0 Mon Sep 17 00:00:00 2001 From: apinge Date: Fri, 21 Aug 2026 12:39:55 +0000 Subject: [PATCH 2/7] add python/sglang/srt/configs/__init__.py Signed-off-by: apinge --- python/sglang/srt/configs/__init__.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index aa05b73be721..be97ba9534ab 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -30,7 +30,12 @@ ) from sglang.srt.configs.nemotron_h import NemotronHConfig, NemotronHPuzzleConfig from sglang.srt.configs.olmo3 import Olmo3Config -from sglang.srt.configs.qwen3_5 import Qwen3_5Config, Qwen3_5MoeConfig +from sglang.srt.configs.qwen3_5 import ( + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3_5MoeTextConfig, + Qwen3_5TextConfig, +) from sglang.srt.configs.qwen3_asr import Qwen3ASRConfig from sglang.srt.configs.qwen3_next import Qwen3NextConfig from sglang.srt.configs.step3_vl import ( @@ -64,6 +69,8 @@ "Qwen3NextConfig", "Qwen3_5Config", "Qwen3_5MoeConfig", + "Qwen3_5TextConfig", + "Qwen3_5MoeTextConfig", "InternS2PreviewConfig", "DotsVLMConfig", "DotsOCRConfig", From fe8df015461e11caf43cac0e47826804a69070bb Mon Sep 17 00:00:00 2001 From: apinge Date: Fri, 21 Aug 2026 13:29:04 +0000 Subject: [PATCH 3/7] update forward_hip --- python/sglang/srt/layers/layernorm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index fd3121456881..809e856004c7 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -473,6 +473,9 @@ def forward_hip( residual: Optional[torch.Tensor] = None, post_residual_addition: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if _use_aiter: + return self.forward_aiter(x, residual, post_residual_addition) + # Fallback to native implementation if vllm is not available if not _has_vllm_rms_norm: return self.forward_native(x, residual, post_residual_addition) From d70b3259b09b1d662e5a30cafd49a23518438e39 Mon Sep 17 00:00:00 2001 From: apinge Date: Fri, 21 Aug 2026 14:04:04 +0000 Subject: [PATCH 4/7] add guard for rocm_triton_gemma_rmsnorm --- python/sglang/srt/layers/layernorm.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 809e856004c7..a2fb9ed62ad1 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -801,6 +801,16 @@ def forward_hip( residual: Optional[torch.Tensor] = None, post_residual_addition: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if x.numel() == 0: + if residual is not None: + if post_residual_addition is not None: + residual = residual + post_residual_addition + return x, residual + return x + + if is_batch_invariant_mode_enabled(): + return self.forward_native(x, residual, post_residual_addition) + if _has_rocm_triton_gemma_rms_norm: if residual is not None: if post_residual_addition is not None: From 7abdda55b323b3eb14d40ca2e9e84ccf4d0efa16 Mon Sep 17 00:00:00 2001 From: apinge Date: Sat, 22 Aug 2026 02:19:41 +0000 Subject: [PATCH 5/7] update Qwen3_5ForCausalLMMTP's member class --- python/sglang/srt/models/qwen3_5_mtp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 5cc5b6f98929..b76a5e339c97 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -33,7 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM +from sglang.srt.models.qwen3_5 import Qwen3_5Model from sglang.srt.runtime_context import get_flags, get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_npu @@ -96,7 +96,7 @@ def __init__( mtp_config = copy.deepcopy(config) mtp_config.num_hidden_layers = 1 mtp_config.full_attention_interval = 1 - self.model = Qwen3_5ForCausalLM( + self.model = Qwen3_5Model( mtp_config, quant_config, prefix=add_prefix("mtp", prefix), From a329a89c558af83d933058b3f7e8e3a81f487ae2 Mon Sep 17 00:00:00 2001 From: apinge Date: Mon, 24 Aug 2026 05:39:00 +0000 Subject: [PATCH 6/7] add workaround SGLANG_USE_ROCM_TRITON_GEMMA_RMSNORM --- python/sglang/srt/layers/layernorm.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index a2fb9ed62ad1..6d689f86a0c8 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -51,6 +51,9 @@ _is_musa = is_musa() _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_use_rocm_triton_gemma_rms_norm = ( + get_bool_env_var("SGLANG_USE_ROCM_TRITON_GEMMA_RMSNORM") and _is_hip +) _AITER_NEW_CA = get_bool_env_var("SGLANG_USE_AITER_NEW_CA", "true") _AITER_FUSED_NORM_DEFAULT = ( _use_aiter @@ -116,7 +119,7 @@ def layernorm( # Fallback: vllm not available, will use forward_native _has_vllm_rms_norm = False -if _is_hip: +if _use_rocm_triton_gemma_rms_norm: try: from sglang.jit_kernel.minimax_m3.rmsnorm import ( gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm, @@ -811,7 +814,7 @@ def forward_hip( if is_batch_invariant_mode_enabled(): return self.forward_native(x, residual, post_residual_addition) - if _has_rocm_triton_gemma_rms_norm: + if _use_rocm_triton_gemma_rms_norm and _has_rocm_triton_gemma_rms_norm: if residual is not None: if post_residual_addition is not None: residual = residual + post_residual_addition From 4bbc057eb81553194a90ca86f4fd7bb8ce3e173e Mon Sep 17 00:00:00 2001 From: apinge Date: Mon, 24 Aug 2026 13:36:38 +0000 Subject: [PATCH 7/7] remove SGLANG_USE_ROCM_TRITON_GEMMA_RMSNORM --- python/sglang/srt/layers/layernorm.py | 28 +-------------------------- 1 file changed, 1 insertion(+), 27 deletions(-) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 6d689f86a0c8..e1d6e6cf62fb 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -51,9 +51,6 @@ _is_musa = is_musa() _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -_use_rocm_triton_gemma_rms_norm = ( - get_bool_env_var("SGLANG_USE_ROCM_TRITON_GEMMA_RMSNORM") and _is_hip -) _AITER_NEW_CA = get_bool_env_var("SGLANG_USE_AITER_NEW_CA", "true") _AITER_FUSED_NORM_DEFAULT = ( _use_aiter @@ -102,7 +99,6 @@ def layernorm( ) _has_aiter_layer_norm = False _has_vllm_rms_norm = False -_has_rocm_triton_gemma_rms_norm = False if _use_aiter: from aiter import layernorm2d_fwd as layer_norm from aiter import rmsnorm2d_fwd as rms_norm @@ -119,19 +115,6 @@ def layernorm( # Fallback: vllm not available, will use forward_native _has_vllm_rms_norm = False -if _use_rocm_triton_gemma_rms_norm: - try: - from sglang.jit_kernel.minimax_m3.rmsnorm import ( - gemma_fused_add_rmsnorm as rocm_triton_gemma_fused_add_rmsnorm, - ) - from sglang.jit_kernel.minimax_m3.rmsnorm import ( - gemma_rmsnorm as rocm_triton_gemma_rmsnorm, - ) - - _has_rocm_triton_gemma_rms_norm = True - except ImportError: - _has_rocm_triton_gemma_rms_norm = False - if _is_cuda: # HF-semantics RMSNorm kernel (JIT-compiled). Used when `cast_x_before_out_mul=True` # (the transformers backend path) to produce outputs that are numerically identical @@ -814,15 +797,6 @@ def forward_hip( if is_batch_invariant_mode_enabled(): return self.forward_native(x, residual, post_residual_addition) - if _use_rocm_triton_gemma_rms_norm and _has_rocm_triton_gemma_rms_norm: - if residual is not None: - if post_residual_addition is not None: - residual = residual + post_residual_addition - return rocm_triton_gemma_fused_add_rmsnorm( - x, residual, self.weight.data, self.variance_epsilon - ) - return rocm_triton_gemma_rmsnorm(x, self.weight.data, self.variance_epsilon) - if not _has_vllm_rms_norm: return self.forward_native(x, residual, post_residual_addition) @@ -842,7 +816,7 @@ def forward_hip( return rms_norm(x, w, self.variance_epsilon) else: # vllm API: rms_norm(out, input, weight, eps) -> None (in-place) - # fused_add_rms_norm(out, input, residual_out, residual, weight, eps) + # fused_add_rms_norm(input, residual, weight, eps) -> None if not x.is_contiguous(): x = x.contiguous() if residual is not None: