A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 08:45 UTC
[Refactor] Small refactor for group topk (#30562)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Wentao Ye committed
Dec 16, 2025 at 14:50 UTC
f21f5ea38c6fa0e824bc00d5762d17e049199cd3
2 files changed
+10
-4
csrc/moe/grouped_topk_kernels.cu
+10
-3
@@ -446,9 +446,13 @@ __device__ inline T apply_sigmoid(T val) {
446
447
template <ScoringFunc SF, typename T>
448
__device__ inline T apply_scoring(T val) {
449
- if constexpr (SF == SCORING_SIGMOID) {
449
+ if constexpr (SF == SCORING_NONE) {
450
+ return val;
451
+ } else if constexpr (SF == SCORING_SIGMOID) {
452
return apply_sigmoid(val);
453
} else {
454
+ static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID,
455
+ "Unsupported ScoringFunc in apply_scoring");
456
return val;
457
}
458
}
@@ -670,10 +674,13 @@ __global__ void group_idx_and_topk_idx_kernel(
674
675
if (case_id < num_tokens) {
676
if (if_proceed_next_topk) {
677
+ float scale = routed_scaling_factor;
678
+ if (renormalize) {
679
+ scale /= topk_sum;
680
+ }
681
for (int i = lane_id; i < topk; i += WARP_SIZE) {
682
float base = cuda_cast<float, T>(s_topk_value[i]);
675
- float value = renormalize ? (base / topk_sum * routed_scaling_factor)
676
- : (base * routed_scaling_factor);
683
+ float value = base * scale;
684
topk_indices[i] = s_topk_idx[i];
685
topk_values[i] = value;
686
}
tests/v1/determinism/test_batch_invariance.py
-1
@@ -188,7 +188,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
188
llm = LLM(
189
model=model_name,
190
tensor_parallel_size=tp_size,
191
- # enable_prefix_caching=False,
191
max_num_seqs=32,
192
max_model_len=8192,
193
dtype="bfloat16", # not everything is supported