[Kimi K3 Perf] Optimize `eh_proj` linear calculation, 12.9 ~ 25.2% kernel performance improvement (#53942)
Wentao Ye committed
Aug 26, 2026 at 17:08 UTC
76cfe1cd88d30d525eec8be5bff75f8b77471c88
3 files changed
+26
-2
tests/kernels/test_bf16_skinny_gemm.py
+9
-1
@@ -39,6 +39,7 @@ EXPECTED_SELECTIONS = {
39
(7168, 8448): (set(range(1, 4)), set()),
40
(8448, 7168): ({1, 2}, set()),
41
(16896, 7168): ({1, 2}, set()),
42
+ (7168, 14336): ({1, 2}, set()),
43
(20480, 7168): (set(range(1, 5)), set()),
44
(40960, 7168): (set(range(1, 5)), set()),
45
# TP16.
@@ -96,6 +97,8 @@ EXPECTED_CUTE_CONFIGS = {
97
(8448, 7168, 2): (32, 4, 4, 8),
98
(16896, 7168, 1): (224, 6, 4, 8),
99
(16896, 7168, 2): (32, 4, 4, 8),
100
+ (7168, 14336, 1): (256, 2, 1, 4),
101
+ (7168, 14336, 2): (224, 4, 2, 8),
102
(20480, 7168, 1): (224, 4, 2, 8),
103
(20480, 7168, 2): (64, 4, 2, 8),
104
(20480, 7168, 3): (64, 2, 2, 8),
@@ -218,7 +221,12 @@ def test_cute_configs_match_measured_table() -> None:
221
for n, k, num_tokens, config in configs
222
}
223
assert actual == EXPECTED_CUTE_CONFIGS
221
- assert all(config.static_k is None for *_, config in configs)
224
+ static_k_configs = {
225
+ (n, k, num_tokens, config.static_k)
226
+ for n, k, num_tokens, config in configs
227
+ if config.static_k is not None
228
+ }
229
+ assert static_k_configs == {(7168, 14336, 1, 14336)}
230
231
232
def test_residual_cute_configs_match_measured_table() -> None:
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+8
@@ -194,6 +194,14 @@ KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = {
194
),
195
name="dense_gate_up_proj",
196
),
197
+ (7168, 14336): ProjectionSpec(
198
+ 7168,
199
+ 14336,
200
+ cute_configs=(
201
+ (1, SkinnyGemmConfig(1, 256, 2, vector_width=4, static_k=14336)),
202
+ (2, _cute(2, 224, 4, 2)),
203
+ ),
204
+ ),
205
(20480, 7168): ProjectionSpec(
206
20480,
207
7168,
vllm/models/kimi_k3/nvidia/mtp.py
+9
-1
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.fused_moe import (
16
fused_moe_make_expert_params_mapping,
17
)
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
+from vllm.model_executor.layers.linear import ReplicatedLinear
20
from vllm.model_executor.layers.logits_processor import LogitsProcessor
21
from vllm.model_executor.layers.quantization import QuantizationConfig
22
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -80,7 +81,14 @@ class KimiK3MultiTokenPredictorLayer(nn.Module):
81
82
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
83
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
83
- self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
84
+ self.eh_proj = ReplicatedLinear(
85
+ config.hidden_size * 2,
86
+ config.hidden_size,
87
+ bias=False,
88
+ quant_config=None,
89
+ prefix=maybe_prefix(prefix, "eh_proj"),
90
+ return_bias=False,
91
+ )
92
93
self.shared_head = SharedHead(
94
config=config, prefix=prefix, quant_config=quant_config