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

python/sglang/srt/layers/attention/linear/kda_backend.py

@@ -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)
+        q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
-        k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
+        k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
-        v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
+        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)
+        q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
-        k = rearrange(k, "n (h d) -> 1 n h d", d=layer.head_k_dim)
+        k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0)  # n (h d) -> 1 n h d
-        v = rearrange(v, "n (h d) -> 1 n h d", d=layer.head_v_dim)
+        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,

python/sglang/srt/models/kimi_linear.py

@@ -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]