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

[Kimi-Linear] Refactor kimi-linear gate calculation to avoid duplicated code (#17160)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-01-20 14:29:24 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent a3addd6203
commit e6b7c04947
2 changed files with 18 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/kimi_linear.py
+13 -9
View File
@@ -15,7 +15,7 @@ from sglang.srt.distributed import (
tensor_model_parallel_all_reduce,
)
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.attention.fla.kda import FusedRMSNormGated
from sglang.srt.layers.attention.fla.kda import FusedRMSNormGated, fused_kda_gate
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
ColumnParallelLinear,
@@ -314,6 +314,14 @@ class KimiDeltaAttention(nn.Module):
self.v_conv1d.weight.size(0), self.v_conv1d.weight.size(2)
)
beta = self.b_proj(hidden_states)[0].float().sigmoid()
forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
forget_gate = fused_kda_gate(
forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias
)
beta = beta.unsqueeze(0)
forget_gate = forget_gate.unsqueeze(0)
kwargs = {
"q_proj_states": q_proj_states,
"k_proj_states": k_proj_states,
@@ -324,14 +332,10 @@ class KimiDeltaAttention(nn.Module):
"q_conv_bias": self.q_conv1d.bias,
"k_conv_bias": self.k_conv1d.bias,
"v_conv_bias": self.v_conv1d.bias,
"dt_bias": self.dt_bias,
"b_proj": self.b_proj,
"f_a_proj": self.f_a_proj,
"f_b_proj": self.f_b_proj,
"A_log": self.A_log,
"head_dim": self.head_dim,
"hidden_states": hidden_states,
"layer_id": self.layer_idx,
"beta": beta,
"gate": forget_gate,
}
core_attn_out = forward_batch.attn_backend.forward(
@@ -344,8 +348,8 @@ class KimiDeltaAttention(nn.Module):
)
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
g = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
core_attn_out = self.o_norm(core_attn_out, g)
norm_gate = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
core_attn_out = self.o_norm(core_attn_out, norm_gate)
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
return self.o_proj(core_attn_out)[0]