Fix: mistake sigmoid in kda (#17508)

This commit is contained in:
strgrb
2026-01-24 13:35:14 +08:00
committed by GitHub
parent 4c512a7d1d
commit 176da1bbdd
2 changed files with 125 additions and 1 deletions

View File

@@ -314,14 +314,15 @@ 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]
beta = self.b_proj(hidden_states)[0].float()
# fused_kda_gate is fused to KimiLinearAttentionBackend with decode
if not forward_batch.forward_mode.is_decode():
forget_gate = fused_kda_gate(
forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias
)
beta = beta.sigmoid()
beta = beta.unsqueeze(0)
forget_gate = forget_gate.unsqueeze(0)