Fix warp illegal instruction in kimi k2 thinking PCG

[BBuf]

Motivation

The kimi_k2_moe_fused_gate kernels fail to initialize indices_ptr when no valid expert is found (when expert_id is out of bounds or max_expert == -1), leaving uninitialized garbage values that propagate to moe_align_block_size_kernel and cause out-of-bounds memory access in atomicAdd(&shared_counts[expert_id], 1), triggering a "Warp Illegal Instruction" error.

Kimi K2 Thinking piecewise cuda graph

python3 -m sglang.launch_server --model-path moonshotai/Kimi-K2-Thinking --tp 8 --trust-remote-code  --tool-call-parser kimi_k2 --reasoning-parser kimi_k2 --model-loader-extra-config='{"enable_multithread_load": "true","num_threads": 64}' --enable-piecewise-cuda-graph

➜  sglang git:(pcg-marlin-moe) ✗ python3 benchmark/gsm8k/bench_sglang.py --num-questions 1400 --parallel 1400
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [00:24<00:00, 53.88it/s]
Accuracy: 0.938
Invalid: 0.001
Latency: 24.583 s
Output throughput: 5294.458 token/s

Kimi K2 Thinking without piecewise cuda graph

python3 -m sglang.launch_server --model-path moonshotai/Kimi-K2-Thinking --tp 8 --trust-remote-code  --tool-call-parser kimi_k2 --reasoning-parser kimi_k2 --model-loader-extra-config='{"enable_multithread_load": "true","num_threads": 64}' 

➜  sglang git:(pcg-marlin-moe) ✗ python3 benchmark/gsm8k/bench_sglang.py --num-questions 1400 --parallel 1400
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [00:24<00:00, 54.09it/s]
Accuracy: 0.941
Invalid: 0.001
Latency: 24.558 s
Output throughput: 5289.763 token/s

Kimi K2 Thinking EP

python3 -m sglang.launch_server --model-path moonshotai/Kimi-K2-Thinking --tp 8 --ep 4 --trust-remote-code  --tool-call-parser kimi_k2 --reasoning-parser kimi_k2 --model-loader-extra-config='{"enable_multithread_load": "true","num_threads": 64}'

➜  sglang git:(pcg-marlin-moe) ✗ python3 benchmark/gsm8k/bench_sglang.py --num-questions 1400 --parallel 1400
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [00:23<00:00, 55.14it/s]
Accuracy: 0.939
Invalid: 0.000
Latency: 24.119 s
Output throughput: 5519.985 token/s

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ispobock]

/tag-and-rerun-ci