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

Clean up GDN Init (#14855)

Browse Source
This commit is contained in:
Stefan He
2025-12-13 02:56:54 -08:00
committed by GitHub
parent d36299ad77
commit f1bbd26ff7
1 changed files with 5 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/models/qwen3_next.py
View File

@@ -7,7 +7,7 @@ from torch import nn
from sglang.srt.compilation.piecewise_context_manager import get_forward_context
from sglang.srt.configs.qwen3_next import Qwen3NextConfig
from sglang.srt.distributed import divide, get_pp_group
from sglang.srt.distributed import get_pp_group
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
@@ -221,7 +221,6 @@ class Qwen3GatedDeltaNet(nn.Module):
self.activation = config.hidden_act
self.layer_norm_epsilon = config.rms_norm_eps
# QKV
self.conv_dim = self.key_dim * 2 + self.value_dim
self.conv1d = ColumnParallelLinear(
input_size=self.conv_kernel_size,
@@ -232,7 +231,6 @@ class Qwen3GatedDeltaNet(nn.Module):
tp_size=self.attn_tp_size,
)
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
# projection of the input hidden states
projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
projection_size_ba = self.num_v_heads * 2
@@ -272,17 +270,11 @@ class Qwen3GatedDeltaNet(nn.Module):
},
)
# selective projection used to make dt, B and C input dependent
self.dt_bias = nn.Parameter(torch.zeros(self.num_v_heads // self.attn_tp_size))
# time step projection (discretization)
# instantiate once and copy inv_dt in init_weights of PretrainedModel
self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads // self.attn_tp_size))
A = torch.empty(
divide(self.num_v_heads, self.attn_tp_size), dtype=torch.float32
).uniform_(0, 16)
self.A_log = nn.Parameter(torch.log(A))
self.A_log._no_weight_decay = True
self.A_log = nn.Parameter(
torch.zeros(self.num_v_heads // self.attn_tp_size, dtype=torch.float32)
)
set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})