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

perf(kimi_linear): replace einops rearrange with native torch ops in Kimi-Linear KDA path (#20396)

Browse Source
This commit is contained in:
Vedant V Jhaveri
2026-03-19 19:38:12 -07:00
committed by GitHub
parent fa0d8f6629
commit db995fba47
2 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/layers/attention/linear/kda_backend.py
View File

@@ -1,7 +1,6 @@
from typing import Tuple, Union
import torch
from einops import rearrange
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
@@ -151,9 +150,9 @@ class KDAAttnBackend(MambaAttnBackendBase):
conv_state_indices=cache_indices,
)
q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) # n (h d) -> 1 n h d
k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) # n (h d) -> 1 n h d
v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) # n (h d) -> 1 n h d
return self.kernel_dispatcher.decode(
q=q,
@@ -232,9 +231,9 @@ class KDAAttnBackend(MambaAttnBackendBase):
seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
).transpose(0, 1)
q = rearrange(q, "n (h d) -> 1 n h d", d=layer.head_q_dim)
k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) # n (h d) -> 1 n h d
k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) # n (h d) -> 1 n h d
v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) # n (h d) -> 1 n h d
core_attn_out = self.kernel_dispatcher.extend(
q=q,

7
python/sglang/srt/models/kimi_linear.py
View File

@@ -4,7 +4,6 @@ from collections.abc import Iterable
from typing import Optional
import torch
from einops import rearrange
from torch import nn
from sglang.srt.configs.kimi_linear import KimiLinearConfig
@@ -399,9 +398,11 @@ class KimiDeltaAttention(nn.Module):
b=beta,
)
norm_gate = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
norm_gate = g_proj_states.unflatten(
-1, (-1, self.head_dim)
) # ... (h d) -> ... h d
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)")
core_attn_out = core_attn_out.squeeze(0).flatten(-2) # 1 n h d -> n (h d)
return self.o_proj(core_attn_out)[0]