[KMCompiler] perf-opt grouped_topk for DeepSeek-v3.2

[liuxiao0909c]

PR Category

Operator

Type of Change

Performance Optimization

Description

• port kernel "grouped_topk_fused_small_expert_count_kernel" from vLLM
• replace benchmark test shapes with top-hit shapes from DeepSeek-v3.2
• add correctness test for top-hit shapes from DeepSeek-v3.2
• port torch implementation from vLLM as vllm-uninstalled fallback in correctness test
• enable tests for the metax backend
• implement grouped_topk for ASCEND backend based on common implementation, remove INT64 usage
• add torch reference implementation for benchmarking (mthreads, hygon, ascend)
• mitigate torch.topk non-stability in correctness tests

also pick following PR in FlagGems for benchmark test:

• Change test count to time (Change test count to time FlagGems#4338)
• Add cudagraph benchmark mode (Add cudagraph benchmark mode FlagGems#4321)
• [Ascend] Update benchmark do_bench to do_bench_npu on ascend ([Benchmark][Ascend] Update benchmark do_bench to do_bench_npu on ascend FlagGems#4857)
• [Ascend] Fix input param of do_bench_npu ([Ascend] Fix input param of do_bench_npu FlagGems#5451)

Performance

common config from DeepSeek-v3.2:

num_experts	n_group	topk_group	topk	scores_dtype	bias_dtype
256	8	4	8	bfloat16	float32

• backend nvidia

device: NVIDIA H800
vLLM version: 0.23.0
flagtree version: 0.6.0
test command: pytest benchmark/test_grouped_topk.py -s --mode=cudagraph
speedup vs C in vLLM (min: 1.012, max: 1.974, avg: 1.331):

num_tokens	renormalize=False, scoring_func=0	renormalize=True, scoring_func=0	renormalize=True, scoring_func=1
1	1.012	1.325	1.320
8	1.071	1.323	1.322
32	1.082	1.327	1.326
64	1.080	1.315	1.313
128	1.083	1.322	1.321
256	1.099	1.330	1.334
496	1.145	1.372	1.396
512	1.156	1.375	1.400
16384	1.890	1.974	1.926

• backend metax

device: MetaX C550
vLLM version: 0.21.0
vLLM-metax version: 0.21.0
flagtree version: 0.6.1a2+metax3.6
test command: pytest benchmark/test_grouped_topk.py -s --mode=cudagraph
speedup vs C in vLLM-metax (min: 1.432, max: 5.077, avg: 2.059):

num_tokens	renormalize=False, scoring_func=0	renormalize=True, scoring_func=0	renormalize=True, scoring_func=1
1	1.490	1.496	1.446
8	1.479	1.492	1.452
32	1.473	1.485	1.432
64	1.481	1.487	1.437
128	1.552	1.594	1.579
256	1.755	1.844	1.805
496	2.136	2.265	2.216
512	2.122	2.281	2.210
16384	4.565	4.948	5.077

• backend thead

device: PPU-ZW810E
vLLM-version: 0.19.0+cu130
flagtree version: 0.6.1+ppu3.6
test command: pytest benchmark/test_grouped_topk.py -s --mode=cudagraph
speedup vs C in vLLM (min: 1.842, max: 5.965, avg: 2.872):

num_tokens	renormalize=False, scoring_func=0	renormalize=True, scoring_func=0	renormalize=True, scoring_func=1
1	1.937	2.359	2.109
8	1.864	2.188	2.078
32	1.893	2.211	1.982
64	1.842	2.104	1.982
128	2.157	2.553	2.247
256	2.187	2.385	2.213
496	3.541	3.885	3.781
512	3.661	3.895	3.891
16384	5.387	5.965	5.266

• backend mthreads

device: MTT S5000
flagtree version: 0.6.0+mthreads.gitc64a4918
test command: pytest benchmark/test_grouped_topk.py -s --mode=cudagraph
speedup vs torch (min: 2.320, max: 6.679, avg: 4.993):

num_tokens	renormalize=False, scoring_func=0	renormalize=True, scoring_func=0	renormalize=True, scoring_func=1
1	4.382	4.875	4.434
8	4.472	4.959	4.479
32	4.567	5.099	4.644
64	4.576	5.148	4.872
128	5.067	5.797	5.005
256	5.684	5.834	6.304
496	5.896	6.450	5.982
512	6.014	6.679	6.366
16384	2.429	2.498	2.320

• backend hygon

device: bw1000
flagtree version: 0.6.1a1+hcu3.6
aiter version: 0.1.3+das.opt1.dtk2604.torch2100.2606172003.g0a655d
test command: pytest benchmark/test_grouped_topk.py -s --mode=cudagraph
speedup vs C (min: 1.142, max: 7.166, avg: 5.069):

num_tokens	renormalize=True, scoring_func=1
1	7.166
8	6.757
32	6.768
64	6.662
128	6.690
256	4.630
496	2.911
512	2.900
16384	1.142

• backend ascend

device: Ascend910_9382 (910c)
flagtree version: 0.6.0+ascend.gitca1dd135
test command: pytest benchmark/test_grouped_topk.py -s (mode=cudagraph not supported)
speedup vs torch (min: 0.506, max: 9.821, avg: 4.646):

num_tokens	renormalize=False, scoring_func=0	renormalize=True, scoring_func=0	renormalize=True, scoring_func=1
1	5.030	5.346	5.515
8	6.136	6.752	6.611
32	8.305	9.240	9.821
64	5.730	5.763	7.375
128	5.156	5.520	6.108
256	3.484	3.762	3.997
496	2.230	2.421	2.531
512	2.231	2.392	2.483
16384	0.506	0.511	0.510

[CLAassistant]

All committers have signed the CLA.