Repository navigation
[Perf] Optimize grouped topk kernel, 1.2%~2% E2E Throughput improvement - #32058
Conversation
Signed-off-by: yewentao256 <zhyanwentao@126.com>
There was a problem hiding this comment.
Code Review
This pull request introduces a significant performance optimization by fusing two kernels in the grouped top-k operation into a single grouped_topk_fused_kernel. This change effectively reduces global memory traffic by using shared memory to pass intermediate results (group_scores) between stages. The overall approach is solid and should lead to performance gains as intended. My review includes one high-severity comment regarding shared memory allocation, which could be tightened to potentially further improve performance by increasing kernel occupancy.
|
nice job wentao |
|
I ran The inference starts without any issue, however when doing the benchmark: the server crashed with this error: I tried tweaking the vllm args such as removing |
@raihan0824 What's the full error log? Seems like an OOM issue I met before, H100 is usually not enough to hold the GLM4.7 FP8 |
yeah its OOM issue, the error vanished after i put anw, in my opinion, the server shouldn't be even started if it's gonna be crashed when handling heavy task |
I think the kv cache / activations count for memory as well, not just for the model weights. There might be some cases that you can setup model but OOM in requests. Feel free to submit an issue and discuss there, seems out of this PR's scope |
mgoin
left a comment
There was a problem hiding this comment.
Took a relatively careful look through and seems solid, I guess there is space to improve over TRTLLM after all :)
I think the two kernel design caused SM underutilization for small numbers of tokens. So while phase 2 only uses a single warp per block, it is only dealing with a single token. This has me a bit concerned about regression for prefill/large numbers of tokens, so I would like to see a prefill-heavy benchmark to confirm.
| # Kimi uses `num_expert_group` instead of `n_group`. | ||
| if n_group is None: | ||
| n_group = getattr(text_config, "num_expert_group", None) |
There was a problem hiding this comment.
| switch (n_group) { | ||
| case 4: { | ||
| launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 4>); | ||
| break; | ||
| } | ||
| case 8: { | ||
| launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 8>); | ||
| break; | ||
| } | ||
| case 16: { | ||
| launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 16>); | ||
| break; | ||
| } | ||
| case 32: { | ||
| launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 32>); | ||
| break; | ||
| } | ||
| default: { | ||
| launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF>); | ||
| break; | ||
| } | ||
| } |
There was a problem hiding this comment.
nit: it could still be valuable to template on num groups for loop unrolling
There was a problem hiding this comment.
Originally we have for (g = 0; g < n_group; ++g), which could get perf from unrolling n_group.
Now for (g = 0; g < topk_group; ++g) maybe I can explore unrolling for topk_group in following up PRs.
Great suggestion, thanks!
There was a problem hiding this comment.
Give it a try and seems no improvement, topk_group usually small I suppose
Added the benchmark of prefill in the PR description, also faster for prefill |
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
…nt (vllm-project#32058) Signed-off-by: yewentao256 <zhyanwentao@126.com>
Purpose
Part of the #31755
Here we combine the two kernels together, reducing global write and read
Test
Test
export MODEL="zai-org/GLM-4.7-FP8"vllm serve $MODEL -tp 8 --port 9256 --enable-expert-parallel --max_num_seqs 128Acc
lm_eval --model local-completions --model_args "base_url=http://127.0.0.1:9256/v1/completions,model=$MODEL,num_concurrent=1024" --tasks gsm8k --limit 100Perf
Decode
vllm bench serve --model $MODEL --dataset-name random --host 127.0.0.1 --port 9256 --random-input-len 2 --random-output-len 128 --request-rate inf --num-prompts 128Prefill
vllm bench serve --model $MODEL --dataset-name random --host 127.0.0.1 --port 9256 --random-input-len 1024 --random-output-len 0 --request-rate inf --num-prompts 128Note
Cursor Bugbot is generating a summary for commit 3a5eda1. Configure here.
Note
Fuses grouped top‑k into a single CUDA kernel and removes the intermediate buffer and second pass.
grouped_topk_fused_kernelthat computes per‑group scores and final experttopkin one pass; eliminates previoustopk_with_k2_kernel+group_idx_and_topk_idx_kernelpipelineinvokeNoAuxTcto launch the fused kernel (one block per token, one warp per group) with dynamic shared memory; dropsgroup_scoresparameter and associated allocations/callsWarpSortand minor sync cleanup; adjusts shared‑memory staging for WarpSelectgrouped_topk(positivity checks,topk_group <= n_group, andtopk <= topk_group * (num_experts / n_group)), and removes the temporarygroup_scorestensorfloat32values,int32indices) and existing scoring options; fills uniform defaults when selection is not possibleWritten by Cursor Bugbot for commit 3a5eda1. This will update automatically on new commits. Configure here.
Note
Fuses grouped top‑k into a single CUDA kernel and simplifies API/launch.
grouped_topk_fused_kernelto compute per‑group scores and final experttopkin one pass; removestopk_with_k2_kernelandgroup_idx_and_topk_idx_kerneland their launch plumbinginvokeNoAuxTc(one block per token, one warp per group) and drops thegroup_scoresparameter; sets dynamic shared‑memory layout for WarpSelect staging + group scoresWarpSort, cleans up a sync, and adjusts shared‑memory staginggrouped_topkvalidation (positivity, bounds) and removes temporarygroup_scorestensor; outputs remainfloat32values andint32indicesWritten by Cursor Bugbot for commit c6a40fb. This will update automatically on new commits. Configure here.
Note
Fuses grouped top‑k into a single CUDA pass
grouped_topk_fused_kernelthat computes per‑group scores and final experttopkin one block (one warp per group); removestopk_with_k2_kernelandgroup_idx_and_topk_idx_kernelinvokeNoAuxTc(dropsgroup_scores, sets dynamic smem layout for WarpSelect staging + group scores) and updates instantiationsWarpSort, removes an unnecessary sync, and adjusts shared‑memory staginggrouped_topkinput checks (positivity/bounds) and removes temporarygroup_scorestensor; outputs remainfloat32values andint32indicesnum_expert_groupandnum_experts_per_tokenin dummy config to ensure MoE setupWritten by Cursor Bugbot for commit 9eba6cc. This will update automatically on new commits. Configure here.
Note
Combines grouped top‑k into a single pass to reduce memory traffic and simplify launches.
grouped_topk_fused_kernel(one block per token, one warp per group); removes prior two‑kernel pipeline andgroup_scoresinvokeNoAuxTcsignature and launch (dynamic smem layout for WarpSelect + per‑group scores); updates explicit instantiationsWarpSortaccessors, adjusts shared‑memory staging, and removes an unnecessary syncgrouped_topkinput checks (positivity/bounds) and eliminates temporary tensor; outputs remainfloat32values andint32indicesdummy_hf_overridesnow recognizesnum_expert_groupandnum_experts_per_tokenfor MoE setupWritten by Cursor Bugbot for commit af26c8b. This will update automatically on new commits. Configure here.