Kimi K2 Fused MoE kernels Optimization configs

[samanamp]

Purpose

Optimize Kimi K2 Fused MoE kernels Optimization configs

Execution improvement: We see 17.8% in Output token throughput (tok/s). The bulk of improvement is in Inter-token Latency.

Test Plan

1. Single benchmark

python benchmarks/kernels/benchmark_moe_old.py --model $MODEL --dtype "fp8_w8a8" --tp 4

2. Model E2E
   Server side:

MODEL=moonshotai/Kimi-K2-Instruct
python -m vllm.entrypoints.openai.api_server --model $MODEL --enable-expert-parallel --max-model-len=65000 -tp 4 --port 8101 --gpu_memory_utilization=0.95

Bench:

python benchmarks/benchmark_serving.py --backend vllm --model $MODEL --dataset-name random --random-input-len 7500 --random-output-len 7500 --max-concurrency 128 --num-prompts 512 --port 8101`

Test Result

1. Single benchmark

2. Model E2E

BEFORE:

============ Serving Benchmark Result ============
Successful requests:                     200       
Maximum request concurrency:             64        
Benchmark duration (s):                  499.10    
Total input tokens:                      1499747   
Total generated tokens:                  276363    
Request throughput (req/s):              0.40      
Output token throughput (tok/s):         553.73    
Total Token throughput (tok/s):          3558.65   
---------------Time to First Token----------------
Mean TTFT (ms):                          12740.87  
Median TTFT (ms):                        13355.67  
P99 TTFT (ms):                           21087.93  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          272.12    
Median TPOT (ms):                        360.46    
P99 TPOT (ms):                           387.18    
---------------Inter-token Latency----------------
Mean ITL (ms):                           62.89     
Median ITL (ms):                         57.45     
P99 ITL (ms):                            363.68    
==================================================

AFTER:

============ Serving Benchmark Result ============
Successful requests:                     200       
Maximum request concurrency:             64        
Benchmark duration (s):                  424.97    
Total input tokens:                      1499747   
Total generated tokens:                  277351    
Request throughput (req/s):              0.47      
Output token throughput (tok/s):         652.64    
Total Token throughput (tok/s):          4181.70   
---------------Time to First Token----------------
Mean TTFT (ms):                          13136.72  
Median TTFT (ms):                        13603.26  
P99 TTFT (ms):                           22379.70  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          269.90    
Median TPOT (ms):                        359.65    
P99 TPOT (ms):                           386.72    
---------------Inter-token Latency----------------
Mean ITL (ms):                           52.96     
Median ITL (ms):                         47.41     
P99 ITL (ms):                            363.21    
==================================================

We see 17.8% in Output token throughput (tok/s). The bulk of improvement is in Inter-token Latency.

---

• [x ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• [x ] The test plan, such as providing test command.
• [x ] The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces optimized configurations for Fused MoE kernels for the Kimi K2 model on various NVIDIA GPUs (B200, GB200, H200). The changes consist of adding new JSON configuration files with tuned parameters for different batch sizes, which results in a significant performance improvement of 17.8% in output token throughput, primarily by reducing inter-token latency, as demonstrated by the provided benchmark results. The changes are well-documented and appear to be correct. I have reviewed the new configuration files and found no issues.