@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5d23ca47ab

[Kernel] Applies routed_scaling_factor internally (#47408)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

Jee Jee Li committed Jul 7, 2026 at 17:00 UTC 5d23ca47ab597d9a725df1dbf14c9e7877072407

8 files changed +56 -44

csrc/libtorch_stable/moe/moe_ops.h

+2 -1

@@ -15,7 +15,8 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights,

15 torch::stable::Tensor& topk_indices,

16 torch::stable::Tensor& token_expert_indices,

17 torch::stable::Tensor& gating_output, bool renormalize,

18 - std::optional<torch::stable::Tensor> bias);

18 + std::optional<torch::stable::Tensor> bias,

19 + double routed_scaling_factor);

21 void topk_softplus_sqrt(

22 torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

+42 -36

@@ -173,7 +173,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(

173 const int start_expert,

174 const int end_expert,

175 const bool renormalize,

176 - const float* bias)

176 + const float* bias,

177 + const double routed_scaling_factor)

178 {

179

180 using cub_kvp = cub::KeyValuePair<int, float>;

@@ -241,14 +242,16 @@ __launch_bounds__(TPB) __global__ void moeTopK(

242 __syncthreads();

243 }

244

244 - // Renormalize the k weights for this row to sum to 1, if requested.

245 - if (renormalize) {

246 - if (threadIdx.x == 0) {

245 + // Apply renormalization and routed scaling factor to final weights.

246 + if (threadIdx.x == 0) {

247 + float scale = static_cast<float>(routed_scaling_factor);

248 + if (renormalize) {

249 const float denom = selected_sum > 0.f ? selected_sum : 1.f;

248 - for (int k_idx = 0; k_idx < k; ++k_idx) {

249 - const int idx = k * block_row + k_idx;

250 - output[idx] = output[idx] / denom;

251 - }

250 + scale /= denom;

251 + }

252 + for (int k_idx = 0; k_idx < k; ++k_idx) {

253 + const int idx = k * block_row + k_idx;

254 + output[idx] = output[idx] * scale;

255 }

256 }

257 }

@@ -274,7 +277,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA

277 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

278 void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,

279 int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,

277 - const float* bias)

280 + const float* bias, const double routed_scaling_factor)

281 {

282 static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||

283 std::is_same_v<InputType, __half>,

@@ -570,17 +573,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

573 }

574 }

575

573 - // Renormalize the k weights for this row to sum to 1, if requested.

574 - if (renormalize) {

575 - if (thread_group_idx == 0)

576 - {

577 - const float denom = selected_sum > 0.f ? selected_sum : 1.f;

578 - for (int k_idx = 0; k_idx < k; ++k_idx)

579 - {

580 - const int idx = k * thread_row + k_idx;

581 - output[idx] = output[idx] / denom;

582 - }

583 - }

576 + // Apply renormalization and routed scaling factor to final weights.

577 + if (thread_group_idx == 0) {

578 + float scale = static_cast<float>(routed_scaling_factor);

579 + if (renormalize) {

580 + const float denom = selected_sum > 0.f ? selected_sum : 1.f;

581 + scale /= denom;

582 + }

583 + for (int k_idx = 0; k_idx < k; ++k_idx) {

584 + const int idx = k * thread_row + k_idx;

585 + output[idx] = output[idx] * scale;

586 + }

587 }

588 }

589

@@ -602,7 +605,7 @@ struct TopkConstants

605 template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>

606 void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,

607 int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,

605 - const float* bias, cudaStream_t stream)

608 + const float* bias, const double routed_scaling_factor, cudaStream_t stream)

609 {

610 static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);

611 using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;

@@ -613,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa

616

617 dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);

618 topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(

616 - input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);

619 + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);

620 }

621

622 #ifndef USE_ROCM

@@ -624,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa

627 IndType, InputType, SF>( \

628 gating_output, nullptr, topk_weights, topk_indices, \

629 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \

627 - bias, stream);

630 + bias, routed_scaling_factor, stream);

631 #else

632 #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \

633 if (WARP_SIZE == 64) { \

@@ -632,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa

635 IndType, InputType, SF>( \

636 gating_output, nullptr, topk_weights, topk_indices, \

637 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \

635 - bias, stream); \

638 + bias, routed_scaling_factor, stream); \

639 } else if (WARP_SIZE == 32) { \

640 topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \

641 IndType, InputType, SF>( \

642 gating_output, nullptr, topk_weights, topk_indices, \

643 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \

641 - bias, stream); \

644 + bias, routed_scaling_factor, stream); \

645 } else { \

646 assert(false && \

647 "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \

@@ -657,6 +660,7 @@ void topkGatingKernelLauncher(

660 const int topk,

661 const bool renormalize,

662 const float* bias,

663 + const double routed_scaling_factor,

664 cudaStream_t stream) {

665 static constexpr int WARPS_PER_TB = 4;

666 static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;

@@ -732,7 +736,7 @@ void topkGatingKernelLauncher(

736 }

737 moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(

738 workspace, nullptr, topk_weights, topk_indices, token_expert_indices,

735 - num_experts, topk, 0, num_experts, renormalize, bias);

739 + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);

740 }

741 }

742 }

@@ -750,6 +754,7 @@ void dispatch_topk_launch(

754 torch::stable::Tensor& softmax_workspace,

755 int num_tokens, int num_experts, int topk, bool renormalize,

756 std::optional<torch::stable::Tensor> bias,

757 + double routed_scaling_factor,

758 cudaStream_t stream)

759 {

760 const float* bias_ptr = nullptr;

@@ -772,7 +777,7 @@ void dispatch_topk_launch(

777 token_expert_indices.mutable_data_ptr<int>(),

778 softmax_workspace.mutable_data_ptr<float>(),

779 num_tokens, num_experts, topk, renormalize,

775 - bias_ptr, stream);

780 + bias_ptr, routed_scaling_factor, stream);

781 } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {

782 vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(

783 reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),

@@ -781,7 +786,7 @@ void dispatch_topk_launch(

786 token_expert_indices.mutable_data_ptr<int>(),

787 softmax_workspace.mutable_data_ptr<float>(),

788 num_tokens, num_experts, topk, renormalize,

784 - bias_ptr, stream);

789 + bias_ptr, routed_scaling_factor, stream);

790 } else {

791 STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);

792 vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(

@@ -791,7 +796,7 @@ void dispatch_topk_launch(

796 token_expert_indices.mutable_data_ptr<int>(),

797 softmax_workspace.mutable_data_ptr<float>(),

798 num_tokens, num_experts, topk, renormalize,

794 - bias_ptr, stream);

799 + bias_ptr, routed_scaling_factor, stream);

800 }

801 }

802

@@ -820,15 +825,15 @@ void topk_softmax(

825 if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {

826 dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,

827 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

823 - bias, stream);

828 + bias, 1.0, stream);

829 } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {

830 dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,

831 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

827 - bias, stream);

832 + bias, 1.0, stream);

833 } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {

834 dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,

835 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

831 - bias, stream);

836 + bias, 1.0, stream);

837 } else {

838 STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());

839 }

@@ -840,7 +845,8 @@ void topk_sigmoid(

845 torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]

846 torch::stable::Tensor& gating_output, // [num_tokens, num_experts]

847 bool renormalize,

843 - std::optional<torch::stable::Tensor> bias)

848 + std::optional<torch::stable::Tensor> bias,

849 + double routed_scaling_factor)

850 {

851 const int num_experts = gating_output.size(-1);

852 const auto num_tokens = gating_output.numel() / num_experts;

@@ -859,15 +865,15 @@ void topk_sigmoid(

865 if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {

866 dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,

867 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

862 - bias, stream);

868 + bias, routed_scaling_factor, stream);

869 } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {

870 dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,

871 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

866 - bias, stream);

872 + bias, routed_scaling_factor, stream);

873 } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {

874 dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,

875 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

870 - bias, stream);

876 + bias, routed_scaling_factor, stream);

877 } else {

878 STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());

879 }

csrc/libtorch_stable/moe/torch_bindings.cpp

+2 -2

@@ -13,8 +13,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {

13 // Apply topk sigmoid to the gating outputs.

14 m.def(

15 "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "

16 - "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "

17 - "bias) -> ()");

16 + "token_expert_indices, Tensor gating_output, bool renormalize, "

17 + "Tensor? bias, float routed_scaling_factor) -> ()");

19 m.def(

20 "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "

vllm/_custom_ops.py

@@ -2244,6 +2244,7 @@ def topk_sigmoid(

2244 gating_output: torch.Tensor,

2245 renormalize: bool = False,

2246 e_score_correction_bias: torch.Tensor | None = None,

2247 + routed_scaling_factor: float = 1.0,

2248 ) -> None:

2249 torch.ops._moe_C.topk_sigmoid(

2250 topk_weights,

@@ -2252,6 +2253,7 @@ def topk_sigmoid(

2253 gating_output,

2254 renormalize,

2255 e_score_correction_bias,

2256 + routed_scaling_factor,

2257 )

2258

2259

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

+3 -2

@@ -43,6 +43,7 @@ def vllm_topk_sigmoid(

43 gating_output: torch.Tensor,

44 renormalize: bool = False,

45 e_score_correction_bias: torch.Tensor | None = None,

46 + routed_scaling_factor: float = 1.0,

47 ) -> tuple[torch.Tensor, ...]:

48 ops.topk_sigmoid(

49 topk_weights,

@@ -51,6 +52,7 @@ def vllm_topk_sigmoid(

52 gating_output,

53 renormalize,

54 e_score_correction_bias,

55 + routed_scaling_factor,

56 )

58 return topk_weights, topk_indices

@@ -216,9 +218,8 @@ def fused_topk_bias(

218 gating_output,

219 renormalize,

220 e_score_correction_bias,

221 + routed_scaling_factor,

222 )

220 - if routed_scaling_factor != 1.0:

221 - topk_weights *= routed_scaling_factor

223 return topk_weights, topk_ids

224 elif scoring_func == "sqrtsoftplus":

225 return vllm_topk_softplus_sqrt(

vllm/model_executor/models/deepseek_v2.py

+4 -2

@@ -280,6 +280,7 @@ class DeepseekV2MoE(nn.Module):

280 parallel_config: ParallelConfig,

281 quant_config: QuantizationConfig | None = None,

282 prefix: str = "",

283 + apply_routed_scale_to_output: bool = False,

284 ):

285 super().__init__()

286 self.tp_size = get_tensor_model_parallel_world_size()

@@ -372,9 +373,8 @@ class DeepseekV2MoE(nn.Module):

373 topk_group=getattr(config, "topk_group", 1),

374 prefix=f"{prefix}.experts",

375 scoring_func=getattr(config, "scoring_func", "softmax"),

375 - # aiter applies routed_scaling_factor internally

376 routed_scaling_factor=self.routed_scaling_factor,

377 - apply_routed_scale_to_output=not self.is_rocm_aiter_moe_enabled,

377 + apply_routed_scale_to_output=apply_routed_scale_to_output,

378 e_score_correction_bias=self.gate.e_score_correction_bias,

379 enable_eplb=self.enable_eplb,

380 num_redundant_experts=self.n_redundant_experts,

@@ -1246,6 +1246,8 @@ class DeepseekV2DecoderLayer(nn.Module):

1246 parallel_config=parallel_config,

1247 quant_config=quant_config,

1248 prefix=f"{prefix}.mlp",

1249 + # aiter applies routed_scaling_factor internally

1250 + apply_routed_scale_to_output=not rocm_aiter_ops.is_fused_moe_enabled(),

1251 )

1252 else:

1253 self.mlp = DeepseekV2MLP(

vllm/models/deepseek_v32/nvidia/model.py

@@ -77,6 +77,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):

77 parallel_config=parallel_config,

78 quant_config=quant_config,

79 prefix=f"{prefix}.mlp",

80 + apply_routed_scale_to_output=False,

81 )

82 # Defer the MoE cross-rank all-reduce; it is fused into the next

83 # layer's input_layernorm (or the final norm) via

vllm/models/minimax_m3/nvidia/model.py

-1

@@ -258,7 +258,6 @@ class MiniMaxM3MoE(nn.Module):

258 swiglu_alpha=config.swiglu_alpha,

259 swiglu_beta=config.swiglu_beta,

260 routed_scaling_factor=self.routed_scaling_factor,

261 - apply_routed_scale_to_output=True,

261 router_logits_dtype=self.gate.out_dtype,

262 shared_experts=self.shared_experts,

263 quant_config=quant_config,