minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

GLM-4.7-Flash: standalone MLA impl and MLA NextN/MTP (#26088)

Browse Source
This commit is contained in:
Yuxuan Zhang
2026-05-26 13:17:39 +08:00
committed by GitHub
parent 59cad671e2
commit 7ef06bfc06
4 changed files with 798 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+7 -4
View File
@@ -426,11 +426,13 @@ class ModelConfig:
self.hf_config.architectures[0] = "DeepseekV4ForCausalLMNextN"
self.hf_config.num_nextn_predict_layers = 1
if is_draft_model and self.hf_config.architectures[0] in [
"Glm4MoeForCausalLM",
"Glm4MoeLiteForCausalLM",
]:
if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM":
self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN"
if (
is_draft_model
and self.hf_config.architectures[0] == "Glm4MoeLiteForCausalLM"
):
self.hf_config.architectures[0] = "Glm4MoeLiteForCausalLMNextN"
if is_draft_model and self.hf_config.architectures[0] in [
"GlmOcrForConditionalGeneration",
@@ -601,6 +603,7 @@ class ModelConfig:
or "DeepseekV3ForCausalLM" in self.hf_config.architectures
or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures
or "Glm4MoeLiteForCausalLM" in self.hf_config.architectures
or "Glm4MoeLiteForCausalLMNextN" in self.hf_config.architectures
or "GlmMoeDsaForCausalLM" in self.hf_config.architectures
or "LongcatFlashForCausalLM" in self.hf_config.architectures
or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
python/sglang/srt/model_loader/weight_utils.py
+7 -1
View File
@@ -685,7 +685,13 @@ def maybe_add_mtp_safetensors(
getattr(hf_config, "num_nextn_predict_layers", 0),
)
if not (
arch in ["Glm4MoeForCausalLM", "Glm4MoeForCausalLMNextN"]
arch
in [
"Glm4MoeForCausalLM",
"Glm4MoeForCausalLMNextN",
"Glm4MoeLiteForCausalLM",
"Glm4MoeLiteForCausalLMNextN",
]
and num_nextn_layers > 0
):
return hf_weights_files
python/sglang/srt/models/glm4_moe_lite.py
+602 -80
View File
@@ -1,4 +1,4 @@
# Copyright 2025-2026 SGLang Team
# Copyright 2026-2027 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
@@ -12,11 +12,11 @@
# limitations under the License.
# ==============================================================================
"""Inference-only GLM-4.7-Flash model compatible with HuggingFace weights"""
"""Inference-only GLM-4.7-Flash model compatible with HuggingFace weights."""
import logging
import re
from typing import Iterable, Optional, Tuple
from typing import Iterable, List, Optional, Tuple, Union
import torch
import torch.nn.functional as F
@@ -24,21 +24,29 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import (
get_moe_expert_parallel_world_size,
get_pp_group,
get_tensor_model_parallel_world_size,
parallel_state,
tensor_model_parallel_all_reduce,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
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.activation import SiluAndMul
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.communicator import (
LayerCommunicator,
LayerScatterModes,
enable_moe_dense_fully_dp,
get_attn_tp_context,
)
from sglang.srt.layers.dp_attention import (
get_attention_tp_rank,
get_attention_tp_size,
is_allocation_symmetric,
is_dp_attention_enabled,
)
from sglang.srt.layers.layernorm import RMSNorm
@@ -46,43 +54,39 @@ from sglang.srt.layers.linear import MergedColumnParallelLinear, RowParallelLine
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import (
get_moe_a2a_backend,
should_skip_post_experts_all_reduce,
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.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
from sglang.srt.layers.moe.utils import filter_moe_weight_param_global_expert
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
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,
DeepseekV2DecoderLayer,
DeepseekV2ForCausalLM,
DeepseekV2Model,
DeepseekV2MoE,
from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
DeepseekV2WeightLoaderMixin,
)
from sglang.srt.models.deepseek_common.utils import _is_cuda, _use_aiter
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
BumpAllocator,
LazyValue,
add_prefix,
get_device_sm,
is_cuda,
is_non_idle_and_non_empty,
log_info_on_rank0,
make_layers,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config
_is_cuda = is_cuda()
_device_sm = get_device_sm()
if _is_cuda:
from sgl_kernel import dsv3_router_gemm
logger = logging.getLogger(__name__)
@@ -132,34 +136,14 @@ class Glm4MoeLiteMLP(nn.Module):
forward_batch=None,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
gemm_output_zero_allocator: BumpAllocator = None,
):
# Keep parity with DeepseekV2MLP.forward signature since DeepseekV2DecoderLayer
# invokes MLP modules with these extra arguments.
if (self.tp_size == 1) and x.shape[0] == 0:
return x
# Some quantization wrappers store the underlying parameter as `weight_packed`.
if not hasattr(self.gate_up_proj, "weight"):
self.gate_up_proj.weight = getattr(self.gate_up_proj, "weight_packed")
if not hasattr(self.down_proj, "weight"):
self.down_proj.weight = getattr(self.down_proj, "weight_packed")
if (
gemm_output_zero_allocator is not None
and x.shape[0] <= 256
and self.gate_up_proj.weight.dtype == torch.uint8
):
y = gemm_output_zero_allocator.allocate(
x.shape[0] * self.gate_up_proj.output_size_per_partition
).view(x.shape[0], self.gate_up_proj.output_size_per_partition)
x = (x, None, y)
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(
x,
skip_all_reduce=should_allreduce_fusion or use_reduce_scatter,
x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
)
return x
@@ -179,28 +163,18 @@ class Glm4MoeLiteGate(nn.Module):
self.e_score_correction_bias = nn.Parameter(
torch.empty((config.n_routed_experts), dtype=torch.float32)
)
# GLM requires FP32 gate projection; cache to avoid per-forward cast.
# FIXME: if gate weight is updated at runtime (e.g. expert rebalancing), _weight_fp32 must be invalidated.
self.register_buffer("_weight_fp32", None, persistent=False)
def forward(self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None):
# NOTE: For some unknown reason, router_gemm seems degrade accept length.
if (
_is_cuda
and not self.is_nextn
and hidden_states.shape[0] < 4
and hidden_states.shape[1] == 7168
and self.weight.shape[0] == 256
and _device_sm >= 90
):
logits = dsv3_router_gemm(hidden_states, self.weight).to(
hidden_states.dtype
)
else:
logits = F.linear(hidden_states, self.weight, None)
def forward(self, hidden_states):
if self._weight_fp32 is None:
self._weight_fp32 = self.weight.data.to(torch.float32)
logits = F.linear(hidden_states.to(torch.float32), self._weight_fp32, None)
return logits
class Glm4MoeLiteSparseMoeBlock(DeepseekV2MoE):
class Glm4MoeLiteSparseMoeBlock(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -210,7 +184,7 @@ class Glm4MoeLiteSparseMoeBlock(DeepseekV2MoE):
alt_stream: Optional[torch.cuda.Stream] = None,
is_nextn: bool = False,
):
nn.Module.__init__(self)
super().__init__()
self.tp_size = get_tensor_model_parallel_world_size()
self.routed_scaling_factor = config.routed_scaling_factor
self.n_shared_experts = config.n_shared_experts
@@ -273,7 +247,8 @@ class Glm4MoeLiteSparseMoeBlock(DeepseekV2MoE):
self.shared_experts_is_int8 = False
self.shared_experts_is_fp8 = False
# self.shared_experts_weight_block_size = None
self.shared_experts_weight_block_size = None
self._shared_expert_tp1 = False
if config.n_shared_experts is not None and self.num_fused_shared_experts == 0:
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
# disable tp for shared experts when enable deepep moe, or with fp4 allgather
@@ -327,8 +302,241 @@ class Glm4MoeLiteSparseMoeBlock(DeepseekV2MoE):
)
self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
def get_moe_weights(self):
return [
x.data
for name, x in self.experts.named_parameters()
if name not in ["correction_bias"]
and filter_moe_weight_param_global_expert(
name, x, self.experts.num_local_experts
)
]
class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
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 self._enable_a2a_moe:
if (
self.alt_stream is not None
and self.num_fused_shared_experts == 0
and hidden_states.shape[0] > 0
and get_is_capture_mode()
):
return self.forward_normal_dual_stream(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else:
return self.forward_normal(
hidden_states, should_allreduce_fusion, use_reduce_scatter
)
else:
return self.forward_deepep(hidden_states, forward_batch)
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
shared_output = self._forward_shared_experts(hidden_states)
with torch.cuda.stream(self.alt_stream):
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output)
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
final_hidden_states *= self.routed_scaling_factor
current_stream.wait_stream(self.alt_stream)
final_hidden_states += shared_output
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
def forward_normal(
self,
hidden_states: torch.Tensor,
should_allreduce_fusion: bool = False,
use_reduce_scatter: bool = False,
) -> torch.Tensor:
if hidden_states.shape[0] > 0:
shared_output = self._forward_shared_experts(hidden_states)
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
else:
shared_output = None
topk_output = self.topk.empty_topk_output(hidden_states.device)
final_hidden_states = self.experts(hidden_states, topk_output)
if not _is_cuda and not _use_aiter:
final_hidden_states *= self.routed_scaling_factor
if shared_output is not None:
with use_symmetric_memory(
parallel_state.get_tp_group(), disabled=not is_allocation_symmetric()
):
final_hidden_states_out = torch.empty_like(final_hidden_states)
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
final_hidden_states = final_hidden_states_out
if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
is_tp_path=True,
use_reduce_scatter=use_reduce_scatter,
should_allreduce_fusion=should_allreduce_fusion,
):
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
return final_hidden_states
def forward_deepep(
self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
) -> torch.Tensor:
shared_output = None
if hidden_states.shape[0] > 0:
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states)
shared_output = self._forward_shared_experts(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)
final_hidden_states = self.experts(
hidden_states=hidden_states,
topk_output=topk_output,
)
if shared_output is not None:
x = shared_output
if self.experts.should_fuse_routed_scaling_factor_in_topk:
x.add_(final_hidden_states)
else:
x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
final_hidden_states = x
else:
if not self.experts.should_fuse_routed_scaling_factor_in_topk:
final_hidden_states *= self.routed_scaling_factor
return final_hidden_states
def _forward_shared_experts(self, hidden_states: torch.Tensor):
if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0):
return self.shared_experts(hidden_states)
else:
return None
def op_gate(self, state):
if is_non_idle_and_non_empty(
state.forward_batch.forward_mode, state.hidden_states_mlp_input
):
# router_logits: (num_tokens, n_experts)
state.router_logits = self.gate(state.hidden_states_mlp_input)
else:
state.router_logits = None
def op_shared_experts(self, state):
hidden_states_mlp_input = state.pop("hidden_states_mlp_input")
if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty(
state.forward_batch.forward_mode, hidden_states_mlp_input
):
state.shared_output = self.shared_experts(hidden_states_mlp_input)
else:
state.shared_output = None
def op_select_experts(self, state):
router_logits = state.pop("router_logits")
hidden_states = state.hidden_states_mlp_input
if router_logits is not None:
with get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
):
state.topk_output = self.topk(
hidden_states=hidden_states,
router_logits=router_logits,
num_token_non_padded=state.forward_batch.num_token_non_padded,
expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
layer_id=self.layer_id,
),
)
else:
state.topk_output = self.topk.empty_topk_output(hidden_states.device)
def op_dispatch_a(self, state):
if self.ep_size > 1:
self.experts.dispatcher.dispatch_a(
hidden_states=state.hidden_states_mlp_input,
topk_output=state.pop("topk_output"),
tbo_subbatch_index=state.get("tbo_subbatch_index"),
)
def op_dispatch_b(self, state):
if self.ep_size > 1:
with get_global_expert_distribution_recorder().with_current_layer(
self.layer_id
):
state.dispatch_output = self.experts.dispatcher.dispatch_b(
tbo_subbatch_index=state.get("tbo_subbatch_index"),
)
def op_experts(self, state):
state.combine_input = self.experts.run_moe_core(
dispatch_output=state.dispatch_output,
)
def op_combine_a(self, state):
if self.ep_size > 1:
self.experts.dispatcher.combine_a(
combine_input=state.pop("combine_input"),
tbo_subbatch_index=state.get("tbo_subbatch_index"),
)
state.pop("dispatch_output")
def op_combine_b(self, state):
if self.ep_size > 1:
state.hidden_states_after_combine = self.experts.dispatcher.combine_b(
tbo_subbatch_index=state.get("tbo_subbatch_index"),
)
def op_output(self, state):
final_hidden_states = state.pop("hidden_states_after_combine")
if get_moe_a2a_backend().is_mori():
num_tokens = state.pop("num_tokens")
final_hidden_states = final_hidden_states[:num_tokens]
if (shared_output := state.pop("shared_output")) is not None:
x = shared_output
if _use_aiter:
x.add_(final_hidden_states)
else:
x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
final_hidden_states = x
elif _use_aiter:
# fused in aiter_biased_grouped_topk so we can skip here
pass
else:
final_hidden_states *= self.routed_scaling_factor
state.hidden_states_mlp_output = final_hidden_states
class Glm4MoeLiteDecoderLayer(nn.Module):
def __init__(
self,
config: PretrainedConfig,
@@ -338,13 +546,14 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
prefix: str = "",
alt_stream: Optional[torch.cuda.Stream] = None,
) -> None:
nn.Module.__init__(self)
super().__init__()
self.hidden_size = config.hidden_size
self.config = config
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
rope_theta, rope_scaling = get_rope_config(config)
max_position_embeddings = getattr(config, "max_position_embeddings", 202752)
self.layer_id = layer_id
self.is_nextn = is_nextn
self.self_attn = DeepseekV2AttentionMLA(
config=config,
@@ -418,26 +627,171 @@ class Glm4MoeLiteDecoderLayer(DeepseekV2DecoderLayer):
qkv_latent_func=self.self_attn.prepare_qkv_latent,
)
def _detect_gfx95_quant_format(self) -> str:
from sglang.srt.models.deepseek_common.utils import _is_gfx95_supported
if not _is_gfx95_supported:
return ""
weight = getattr(
getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None), "weight", None
)
if weight is None:
return ""
if weight.dtype == torch.uint8:
return "mxfp4"
if weight.dtype == getattr(torch, "float8_e4m3fn", None):
return "fp8"
return ""
def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
return is_nextn or (
self.config.n_routed_experts is not None
and layer_id >= self.config.first_k_dense_replace
and layer_id % self.config.moe_layer_freq == 0
)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor],
zero_allocator: BumpAllocator,
) -> torch.Tensor:
hidden_states, residual = self.layer_communicator.prepare_attn(
hidden_states,
residual,
forward_batch,
getattr(self, "_gfx95_quant_format", ""),
)
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
zero_allocator=zero_allocator,
layer_scatter_modes=self.layer_scatter_modes,
)
if isinstance(hidden_states, tuple):
hidden_states = hidden_states[0]
get_attn_tp_context().clear_attn_inputs()
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
)
)
# For DP with padding, reduce scatter can be used instead of all-reduce.
use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
forward_batch
)
hidden_states = self.mlp(
hidden_states, forward_batch, should_allreduce_fusion, 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
def op_comm_prepare_attn(
self,
state,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
residual: Optional[torch.Tensor],
zero_allocator: BumpAllocator,
tbo_subbatch_index: Optional[int] = None,
):
state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = (
self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch)
)
if get_moe_a2a_backend().is_mori():
state.num_tokens = hidden_states.shape[0]
state.update(
dict(
forward_batch=forward_batch,
positions=positions,
zero_allocator=zero_allocator,
tbo_subbatch_index=tbo_subbatch_index,
)
)
def op_comm_prepare_mlp(self, state):
state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = (
self.layer_communicator.prepare_mlp(
state.pop("hidden_states_after_attn"),
state.pop("residual_after_input_ln"),
state.forward_batch,
)
)
def op_mlp(self, state):
hidden_states = state.pop("hidden_states_mlp_input")
if not (
enable_moe_dense_fully_dp()
and (not self.is_layer_sparse)
and hidden_states.shape[0] == 0
):
state.hidden_states_mlp_output = self.mlp(
hidden_states, state.forward_batch
)
else:
state.hidden_states_mlp_output = hidden_states
def op_comm_postprocess_layer(self, state):
hidden_states, residual = self.layer_communicator.postprocess_layer(
state.pop("hidden_states_mlp_output"),
state.pop("residual_after_comm_pre_mlp"),
state.forward_batch,
)
output = dict(
positions=state.positions,
hidden_states=hidden_states,
residual=residual,
forward_batch=state.forward_batch,
zero_allocator=state.zero_allocator,
tbo_subbatch_index=state.tbo_subbatch_index,
)
state.clear(
expect_keys={
"positions",
"forward_batch",
"zero_allocator",
"tbo_subbatch_index",
}
)
return output
class Glm4MoeLiteModel(nn.Module):
fall_back_to_pt_during_load = False
class Glm4MoeLiteModel(DeepseekV2Model):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
):
nn.Module.__init__(self)
super().__init__()
self.padding_id = config.pad_token_id
self.vocab_size = config.vocab_size
self.first_k_dense_replace = config.first_k_dense_replace
self.pp_group = get_pp_group()
# DeepseekV2Model.forward expects these attributes to exist.
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.cp_size = get_attention_tp_size() if self.dsa_enable_prefill_cp else None
self.gemm_output_zero_allocator_size = 0
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
@@ -467,15 +821,103 @@ class Glm4MoeLiteModel(DeepseekV2Model):
self.norm = PPMissingLayer(return_tuple=True)
self.layers_to_capture = []
def get_input_embeddings(self) -> torch.Tensor:
return self.embed_tokens
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]:
total_num_layers = self.end_layer - self.start_layer
if self.pp_group.is_first_rank:
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
residual = None
else:
assert pp_proxy_tensors is not None
hidden_states = pp_proxy_tensors["hidden_states"]
residual = pp_proxy_tensors["residual"]
device = hidden_states.device
zero_allocator = BumpAllocator(
buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1),
dtype=torch.float32,
device=device,
)
normal_start_layer = self.start_layer
normal_end_layer = self.end_layer
if forward_batch.can_run_tbo:
if (
self.first_k_dense_replace > normal_start_layer
and self.first_k_dense_replace < normal_end_layer
):
normal_end_layer = self.first_k_dense_replace
elif self.first_k_dense_replace < normal_start_layer:
normal_end_layer = normal_start_layer = 0
aux_hidden_states = []
for i in range(normal_start_layer, normal_end_layer):
with get_global_expert_distribution_recorder().with_current_layer(i):
if i in self.layers_to_capture:
aux_hidden_states.append(hidden_states + residual)
layer = self.layers[i]
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
zero_allocator,
)
if normal_end_layer != self.end_layer:
hidden_states, residual = model_forward_maybe_tbo(
layers=self.layers[normal_end_layer : self.end_layer],
enable_tbo=True,
positions=positions,
forward_batch=forward_batch,
hidden_states=hidden_states,
residual=residual,
input_data_scatter_mode=self.layers[
normal_end_layer - 1
].layer_scatter_modes.layer_output_mode,
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)
if len(aux_hidden_states) == 0:
return hidden_states
return hidden_states, aux_hidden_states
class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
# for quark model load
packed_modules_mapping = {}
class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
nn.Module.__init__(self)
super().__init__()
config.moe_layer_freq = 1
self.config = config
self.tp_size = get_tensor_model_parallel_world_size()
@@ -503,12 +945,9 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
)
self.capture_aux_hidden_states = False
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_attention_tp_rank()
self.cp_size = get_attention_tp_size()
else:
self.cp_rank = self.cp_size = None
@property
def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value
def determine_num_fused_shared_experts(
self, architecture: str = "Glm4MoeLiteForCausalLM"
@@ -539,6 +978,89 @@ class Glm4MoeLiteForCausalLM(DeepseekV2ForCausalLM):
self.num_fused_shared_experts = self.config.n_shared_experts
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
with get_attn_tp_context().maybe_input_scattered(forward_batch):
hidden_states = self.model(
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
)
aux_hidden_states = None
if self.capture_aux_hidden_states:
hidden_states, aux_hidden_states = hidden_states
if self.pp_group.is_last_rank:
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
)
else:
return hidden_states
@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):
return self.model.embed_tokens.weight, self.lm_head.weight
def set_embed_and_head(self, embed, head):
del self.model.embed_tokens.weight
del self.lm_head.weight
self.model.embed_tokens.weight = embed
self.lm_head.weight = head
torch.cuda.empty_cache()
torch.cuda.synchronize()
@classmethod
def get_model_config_for_expert_location(cls, config):
return ModelConfigForExpertLocation(
num_layers=config.num_hidden_layers,
num_logical_experts=config.n_routed_experts,
num_groups=config.n_group,
)
def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
if not self.pp_group.is_last_rank:
return
if layer_ids is None:
self.capture_aux_hidden_states = True
num_layers = self.config.num_hidden_layers
self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
else:
self.capture_aux_hidden_states = True
# TODO (Qiaolin-Yu): check if other draft models need similar layer id
# adjustment
if layer_ids and layer_ids[0] == 1:
self.model.layers_to_capture = [val + 1 for val in layer_ids]
else:
self.model.layers_to_capture = list(layer_ids)
def set_dflash_layers_to_capture(self, layer_ids: List[int]):
if not self.pp_group.is_last_rank:
return
if layer_ids is None:
raise ValueError(
"DFLASH requires explicit layer_ids for aux hidden capture."
)
self.capture_aux_hidden_states = True
self.model.layers_to_capture = [val + 1 for val in layer_ids]
def load_weights(
self,
weights: Iterable[Tuple[str, torch.Tensor]],
python/sglang/srt/models/glm4_moe_lite_nextn.py
+182
View File
@@ -0,0 +1,182 @@
# Copyright 2026-2027 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Inference-only GLM-4.7-Flash Speculative Decoding (NextN) compatible with HuggingFace weights."""
import logging
from typing import Iterable, Optional, Tuple
import torch
from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_tensor_model_parallel_world_size
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.glm4_moe_lite import (
Glm4MoeLiteDecoderLayer,
Glm4MoeLiteForCausalLM,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import BumpAllocator, add_prefix, is_npu
logger = logging.getLogger(__name__)
class Glm4MoeLiteModelNextN(nn.Module):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
logger.warning(
"Overriding Glm4MoeLiteForCausalLMNextN quant config for modelopt_fp4 "
"GLM-4.7-Flash model."
)
quant_config = None
self.vocab_size = config.vocab_size
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
use_attn_tp_group=is_dp_attention_enabled(),
prefix=add_prefix("embed_tokens", prefix),
)
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
self.decoder = Glm4MoeLiteDecoderLayer(
config,
0,
quant_config=quant_config,
is_nextn=True,
prefix=add_prefix("decoder", prefix),
)
self.shared_head = nn.Module()
self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
# Glm4MoeLiteDecoderLayer uses DeepseekV2AttentionMLA, which requires a
# zero_allocator (the GQA glm4_moe_nextn path does not pass one).
zero_allocator = BumpAllocator(
buffer_size=2,
dtype=torch.float32,
device=(
input_embeds.device if input_embeds is not None else input_ids.device
),
)
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
if hidden_states.shape[0] > 0:
hidden_states = self.eh_proj(
torch.cat(
(
self.enorm(hidden_states),
self.hnorm(forward_batch.spec_info.hidden_states),
),
dim=-1,
)
)
residual = None
with get_global_expert_distribution_recorder().disable_this_region():
hidden_states, residual = self.decoder(
positions, hidden_states, forward_batch, residual, zero_allocator
)
if not forward_batch.forward_mode.is_idle():
if residual is not None:
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
else:
hidden_states = self.shared_head.norm(hidden_states)
return hidden_states
class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
nn.Module.__init__(self)
self.config = config
self.tp_size = get_tensor_model_parallel_world_size()
if (
is_npu()
and get_global_server_args().speculative_draft_model_quantization is None
):
quant_config = None
self.quant_config = quant_config
self.model = Glm4MoeLiteModelNextN(
config, quant_config, prefix=add_prefix("model", prefix)
)
self.lm_head = ParallelLMHead(
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=add_prefix("model.shared_head.head", prefix),
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
)
self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_global_server_args().disable_shared_experts_fusion else 1
)
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
hidden_states = self.model(input_ids, positions, forward_batch)
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
super().load_weights(weights, is_nextn=True)
EntryClass = [Glm4MoeLiteForCausalLMNextN]