[Kimi K3 Bug] Fix deepgemm support for kimi k3

[yewentao256]

Purpose

VLLM_USE_RUST_FRONTEND=0 vllm serve moonshotai/Kimi-K3 --moe-backend auto --gpu-memory-utilization 0.95 --tensor-parallel-size 8 --load-format fastsafetensors --no-enable-flashinfer-autotune --max-model-len 1048576 --kv-cache-dtype fp8 --attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"TRTLLM_RAGGED"}' --enable-prefix-caching --enable-auto-tool-choice --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 --trust_remote_code

Will raise

(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/runner/moe_runner.py", line 610, in _apply_quant_method
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     fused_out = self.routed_experts.forward_modular(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/routed_experts.py", line 1209, in forward_modular
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     return self.quant_method.apply(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/quantization/mxfp4.py", line 886, in apply
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     return self.moe_kernel.apply(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 1714, in apply
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     return self.impl.apply(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 1494, in apply
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     fused_out = self._fused_experts(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]                 ^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/modular_kernel.py", line 1342, in _fused_experts
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     self.fused_experts.apply(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py", line 628, in apply
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     deepgemm_unpermute_and_reduce(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py", line 558, in deepgemm_unpermute_and_reduce
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     return ep_gather(
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     return func(*args, **kwargs)
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]   File "/home/yewentao256/vllm-source/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py", line 428, in ep_gather
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]     assert hidden_size % BLOCK_D == 0
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP1 pid=2706766) ERROR 07-30 16:56:07 [multiproc_executor.py:1018] AssertionError

This PR fixes the issue

Test

Covered in unit test

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}