[Kimi K3 Bug] Fix deepgemm support for kimi k3 (#50458)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 30, 2026 at 20:10 UTC
0bff0ce5d39249d5119b8d688071671736f3f01f
2 files changed
+4
-1
tests/kernels/moe/test_deepgemm.py
+1
@@ -381,6 +381,7 @@ def run_single_fp4_case(m, n, k, topk, num_experts):
381
FP4_MNKs = [
382
(128, 4096, 4096), # DeepSeek V4 shape
383
(256, 2048, 2048), # Half-size variant
384
+ (128, 384, 3584), # Kimi-K3 TP8 latent-MoE shape
385
]
386
387
FP4_TOPKS = [2]
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+3
-1
@@ -5,6 +5,8 @@ Taken from https://github.com/ModelTC/LightLLM/blob/8ed97c74c18f11505b048b1ba00b
5
and updated to fit vllm needs and terminology.
6
"""
7
8
+import math
9
+
10
import torch
11
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -424,7 +426,7 @@ def ep_gather(
426
num_warps = 2
427
num_tokens = output_tensor.shape[0]
428
hidden_size = input_tensor.shape[1]
427
- BLOCK_D = min(hidden_size, 1024)
429
+ BLOCK_D = math.gcd(hidden_size, 1024)
430
assert hidden_size % BLOCK_D == 0
431
grid = (triton.cdiv(hidden_size, BLOCK_D), min(num_tokens, 1024))
432