[NPU] support kimi k3 on A5 and improve performance by zhaozx-cn · Pull Request #39589 · sgl-project/sglang

zhaozx-cn · 2026-09-15T09:48:40Z

Co-Authored-By: hanwlax
Co-Authored-By: Hexq0210
Co-Authored-By: McZyWu
Co-Authored-By: qybnb
Co-Authored-By: sherdavincl9

Motivation

support kimi k3 on A5 and improve performance.

Modifications

1.add compressed w4a8 mxfp4 moe.
2.shared expert: add fine-grained dual steam and support shared expert specified tp size.
3.add fused qkvg proj.
4.add situ mx quant kernel.
5.add kv nz for k3 mla.
6.add fia v2 for mtp branch.
7.add chunk kda kernel for prefill branch.
8.fix kimi k3 dspark acc pd disaggregation.

Accuracy Tests

#!/usr/bin/env bash
set -eo pipefail

NODE_RANK="$1"
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

export SGLANG_NPU_USE_FIAS_V2_BSND=True
export SGLANG_NPU_FINE_GRAINED_MOE_DUAL_STREAM=True

export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1
export SGLANG_ENABLE_SPEC_V2=1
export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_DSPARK_FOLDED_PROPOSAL=0
export SGLANG_DSPARK_FOLDED_SAMPLING=0
export SGLANG_DSPARK_STACKED_CTX_KV=0
export SGLANG_DSPARK_EMBED_IN_GRAPH=0
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export HCCL_SOCKET_IFNAME=enp194s0f0
export GLOO_SOCKET_IFNAME=enp194s0f0
export STREAMS_PER_DEVICE=32
export DEEP_NORMAL_MODE_USE_INT8_QUANT=1
export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128
export HCCL_BUFFSIZE=2000
export DEEPEP_NORMAL_LONG_SEQ_ROUND=64
export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=512
export HCCL_OP_EXPANSION_MODE=AIV

unset ASCEND_CUSTOM_OPP_PATH
unset SGLANG_NPU_FUSED_MOE_MODE
unset ENABLE_PROFILING
unset SGLANG_K3_TRACE_STATE_FILE
unset SGLANG_K3_TRACE_HIDDEN_FILE

sglang serve \
  --model-loader-extra-config '{"enable_multithread_load": true}' \
  --dist-init-addr <<first node>>:30143 \
  --nnodes 4 \
  --node-rank "$NODE_RANK" \
  --model-path /home/weights/Kimi-K3 \
  --tokenizer-path /home/weights/Kimi-K3 \
  --trust-remote-code \
  --attention-backend ascend \
  --device npu \
  --quantization modelslim \
  --dtype bfloat16 \
  --tp-size 32 \
  --enable-dp-attention \
  --enable-dp-lm-head \
  --mem-fraction-static 0.9 \
  --chunked-prefill-size 8192 \
  --cuda-graph-bs 32 \
  --max-running-requests 32 \
  --host 0.0.0.0 \
  --enable-shared-experts-attn-tp \
  --enable-dense-mlp-attn-tp \
  --shared-experts-tp-size 4 \
  --port 30543 \
  --reasoning-parser kimi_k3 \
  --moe-a2a-backend deepep \
  --deepep-mode auto \
  --speculative-algorithm DSPARK \
  --speculative-draft-model-path /home/weights/Kimi-K3-DSpark \
  --speculative-draft-model-quantization unquant \
  --speculative-dspark-block-size 7 \
  --speculative-draft-attention-backend ascend \
  --linear-attn-verify-backend triton \
  --speculative-eagle-topk 1 \
  --disable-radix-cache \
  --disable-custom-all-reduce \
  --watchdog-timeout 9000

DeprecationWarning: sglang.test.few_shot_gsm8k is deprecated. Use sglang.test.run_eval with eval_name='gsm8k' instead.
  metrics = run_eval(args)
100%|████████████████████████████████████████████████████████| 200/200 [01:54<00:00,  1.74it/s]
Accuracy: 0.985
Invalid: 0.005
Latency: 116.842 s
Output throughput: 199.406 token/s
metrics={'accuracy': np.float64(0.98), 'invalid': np.float64(0.005), 'latency': 116.84225159999914, 'output_throughput': 199.405606113809}
metrics['accuracy']=np.float64(0.98)

Speed Tests and Profiling

128k in 1k out dspark 7 tp32 ep 32

Metric	Value
Mean TFTT	5604.54 ms
Mean TPOT	23.02 ms
Throughput	47 tok/s

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.