@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 85f638a2b8

skip cudagraph/DP padding in topk (#48979)

Signed-off-by: gnovack <novackgm@gmail.com>

gnovack committed Jul 21, 2026 at 15:20 UTC 85f638a2b8f9838275b911c44fb0c2abdf604476

10 files changed +302 -99

csrc/libtorch_stable/moe/moe_ops.h

+6 -3

@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,

9 torch::stable::Tensor& topk_indices,

10 torch::stable::Tensor& token_expert_indices,

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

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

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

13 + std::optional<torch::stable::Tensor> is_padding);

15 void topk_sigmoid(torch::stable::Tensor& topk_weights,

16 torch::stable::Tensor& topk_indices,

17 torch::stable::Tensor& token_expert_indices,

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

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

19 - double routed_scaling_factor);

20 + double routed_scaling_factor,

21 + std::optional<torch::stable::Tensor> is_padding);

23 void topk_softplus_sqrt(

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

@@ -25,7 +27,8 @@ void topk_softplus_sqrt(

27 double routed_scaling_factor,

28 const std::optional<torch::stable::Tensor>& correction_bias,

29 const std::optional<torch::stable::Tensor>& input_ids,

28 - const std::optional<torch::stable::Tensor>& tid2eid);

30 + const std::optional<torch::stable::Tensor>& tid2eid,

31 + const std::optional<torch::stable::Tensor>& is_padding);

33 void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,

34 std::optional<torch::stable::Tensor> topk_ids,

csrc/libtorch_stable/moe/topk_softmax_kernels.cu

+45 -24

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

174 const int end_expert,

175 const bool renormalize,

176 const float* bias,

177 - const double routed_scaling_factor)

177 + const double routed_scaling_factor,

178 + const bool* is_padding)

179 {

180

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

@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(

229 const int expert = result_kvp.key;

230 const bool node_uses_expert = expert >= start_expert && expert < end_expert;

231 const bool should_process_row = row_is_active && node_uses_expert;

232 + const bool is_pad_row = is_padding != nullptr && is_padding[block_row];

233

234 const int idx = k * block_row + k_idx;

235 // Return the unbiased scores for output weights

236 output[idx] = inputs_after_softmax[thread_read_offset + expert];

235 - indices[idx] = should_process_row ? (expert - start_expert) : num_experts;

236 - assert(indices[idx] >= 0);

237 + indices[idx] = is_pad_row ? static_cast<IndType>(-1)

238 + : (should_process_row ? (expert - start_expert) : num_experts);

239 + assert(is_pad_row || indices[idx] >= 0);

240 source_rows[idx] = k_idx * num_rows + block_row;

241 if (renormalize) {

242 selected_sum += inputs_after_softmax[thread_read_offset + expert];

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

280 __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

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

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

280 - const float* bias, const double routed_scaling_factor)

283 + const float* bias, const double routed_scaling_factor, const bool* is_padding)

284 {

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

286 std::is_same_v<InputType, __half>,

@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

548 // Add a guard to ignore experts not included by this node

549 const bool node_uses_expert = expert >= start_expert && expert < end_expert;

550 const bool should_process_row = row_is_active && node_uses_expert;

551 + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];

552

553 // The lead thread from each sub-group will write out the final results to global memory. (This will be a

554 // single) thread per row of the input/output matrices.

555 const int idx = k * thread_row + k_idx;

556 output[idx] = max_val;

553 - indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;

557 + indices[idx] = is_pad_row ? static_cast<IndType>(-1)

558 + : (should_process_row ? (expert - start_expert) : NUM_EXPERTS);

559 source_rows[idx] = k_idx * num_rows + thread_row;

560 if (renormalize) {

561 selected_sum += max_val;

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

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

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

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

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

613 + const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)

614 {

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

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

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

621

622 dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);

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

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

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

625 }

626

627 #ifndef USE_ROCM

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

632 IndType, InputType, SF>( \

633 gating_output, nullptr, topk_weights, topk_indices, \

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

630 - bias, routed_scaling_factor, stream);

635 + bias, routed_scaling_factor, stream, is_padding);

636 #else

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

638 if (WARP_SIZE == 64) { \

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

640 IndType, InputType, SF>( \

641 gating_output, nullptr, topk_weights, topk_indices, \

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

638 - bias, routed_scaling_factor, stream); \

643 + bias, routed_scaling_factor, stream, is_padding); \

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

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

646 IndType, InputType, SF>( \

647 gating_output, nullptr, topk_weights, topk_indices, \

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

644 - bias, routed_scaling_factor, stream); \

649 + bias, routed_scaling_factor, stream, is_padding); \

650 } else { \

651 assert(false && \

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

@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(

666 const bool renormalize,

667 const float* bias,

668 const double routed_scaling_factor,

664 - cudaStream_t stream) {

669 + cudaStream_t stream,

670 + const bool* is_padding) {

671 static constexpr int WARPS_PER_TB = 4;

672 static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;

673 #ifndef USE_ROCM

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

742 }

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

744 workspace, nullptr, topk_weights, topk_indices, token_expert_indices,

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

745 + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);

746 }

747 }

748 }

@@ -755,7 +761,8 @@ void dispatch_topk_launch(

761 int num_tokens, int num_experts, int topk, bool renormalize,

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

763 double routed_scaling_factor,

758 - cudaStream_t stream)

764 + cudaStream_t stream,

765 + std::optional<torch::stable::Tensor> is_padding)

766 {

767 const float* bias_ptr = nullptr;

768 if (bias.has_value()) {

@@ -769,6 +776,18 @@ void dispatch_topk_launch(

776 bias_ptr = bias_tensor.const_data_ptr<float>();

777 }

778

779 + const bool* is_padding_ptr = nullptr;

780 + if (is_padding.has_value()) {

781 + const torch::stable::Tensor& is_padding_tensor = is_padding.value();

782 + STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,

783 + "is_padding tensor must be bool");

784 + STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");

785 + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,

786 + "is_padding size mismatch, expected: ", num_tokens);

787 + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");

788 + is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();

789 + }

790 +

791 if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {

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

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

@@ -777,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,

780 - bias_ptr, routed_scaling_factor, stream);

799 + bias_ptr, routed_scaling_factor, stream, is_padding_ptr);

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

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

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

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

805 token_expert_indices.mutable_data_ptr<int>(),

806 softmax_workspace.mutable_data_ptr<float>(),

807 num_tokens, num_experts, topk, renormalize,

789 - bias_ptr, routed_scaling_factor, stream);

808 + bias_ptr, routed_scaling_factor, stream, is_padding_ptr);

809 } else {

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

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

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

815 token_expert_indices.mutable_data_ptr<int>(),

816 softmax_workspace.mutable_data_ptr<float>(),

817 num_tokens, num_experts, topk, renormalize,

799 - bias_ptr, routed_scaling_factor, stream);

818 + bias_ptr, routed_scaling_factor, stream, is_padding_ptr);

819 }

820 }

821

@@ -806,7 +825,8 @@ void topk_softmax(

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

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

827 bool renormalize,

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

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

829 + std::optional<torch::stable::Tensor> is_padding)

830 {

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

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

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

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

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

847 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

828 - bias, 1.0, stream);

848 + bias, 1.0, stream, is_padding);

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

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

851 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

832 - bias, 1.0, stream);

852 + bias, 1.0, stream, is_padding);

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

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

855 token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,

836 - bias, 1.0, stream);

856 + bias, 1.0, stream, is_padding);

857 } else {

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

859 }

@@ -846,7 +866,8 @@ void topk_sigmoid(

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

867 bool renormalize,

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

849 - double routed_scaling_factor)

869 + double routed_scaling_factor,

870 + std::optional<torch::stable::Tensor> is_padding)

871 {

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

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

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

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

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

888 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

868 - bias, routed_scaling_factor, stream);

889 + bias, routed_scaling_factor, stream, is_padding);

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

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

892 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

872 - bias, routed_scaling_factor, stream);

893 + bias, routed_scaling_factor, stream, is_padding);

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

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

896 token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,

876 - bias, routed_scaling_factor, stream);

897 + bias, routed_scaling_factor, stream, is_padding);

898 } else {

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

900 }

csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu

+100 -49

@@ -80,22 +80,27 @@ __launch_bounds__(128) __global__

80 OutIndType* indices, int num_rows,

81 int num_experts, float routed_scaling_factor,

82 const HashIndType* input_ids,

83 - const HashIndType* tid2eid) {

83 + const HashIndType* tid2eid,

84 + const bool* is_padding) {

85 const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;

86 const int lane = threadIdx.x % 32;

87 if (warp >= num_rows) return;

88 const int64_t token_id = load_index_as_int64(input_ids, warp);

89 + const bool is_pad_row = is_padding != nullptr && is_padding[warp];

91 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)

92 cudaGridDependencySynchronize();

93 #endif

94 int expert = 0;

95 float weight = 0.f;

94 - if (lane < 6) {

96 + if (lane < 6 && !is_pad_row) {

97 // only load and calculate for 6 experts

98 expert = static_cast<int>(tid2eid[token_id * 6 + lane]);

99 const float x = input[warp * num_experts + expert];

100 weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));

101 + if (isnan(weight)) {

102 + weight = 0.f;

103 + }

104 }

105 float weight_sum = weight;

106 #pragma unroll

@@ -111,7 +116,8 @@ __launch_bounds__(128) __global__

116 const int offset = warp * 6 + lane;

117 output[offset] =

118 weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);

114 - indices[offset] = static_cast<OutIndType>(expert);

119 + indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)

120 + : static_cast<OutIndType>(-1);

121 }

122 }

123

@@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,

126 int num_rows, int num_experts,

127 double routed_scaling_factor,

128 const HashIndType* input_ids,

123 - const HashIndType* tid2eid, cudaStream_t stream) {

129 + const HashIndType* tid2eid, cudaStream_t stream,

130 + const bool* is_padding) {

131 if (num_rows == 0) return;

132 auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;

133 cudaLaunchConfig_t config = {};

@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,

141 config.numAttrs = 1;

142 const float scale = static_cast<float>(routed_scaling_factor);

143 cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,

137 - num_experts, scale, input_ids, tid2eid);

144 + num_experts, scale, input_ids, tid2eid, is_padding);

145 }

146 #endif

147

@@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

173 const int num_rows, IndType* indices, int* source_rows, const int k,

174 const int start_expert, const int end_expert, const bool renormalize,

175 double routed_scaling_factor, const float* correction_bias,

169 - const HashIndType* input_ids, const HashIndType* tid2eid) {

176 + const HashIndType* input_ids, const HashIndType* tid2eid,

177 + const bool* is_padding) {

178 static_assert(std::is_same_v<InputType, float> ||

179 std::is_same_v<InputType, __nv_bfloat16> ||

180 std::is_same_v<InputType, __half>,

@@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

239 return;

240 }

241 const bool row_is_active = finished ? !finished[thread_row] : true;

242 + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];

243

244 // We finally start setting up the read pointers for each thread. First, each

245 // thread jumps to the start of the row it will read.

@@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

258 cudaGridDependencySynchronize();

259 #endif

260

252 - // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert

253 - // to float

254 - if constexpr (std::is_same_v<InputType, float>) {

261 + if (is_pad_row) {

262 +#pragma unroll

263 + for (int ii = 0; ii < VPT; ++ii) {

264 + row_chunk[ii] = 0.f;

265 + }

266 + } else if constexpr (std::is_same_v<InputType, float>) {

267 using VecType = AlignedArray<float, ELTS_PER_LDG>;

268 VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);

269 const VecType* vec_thread_read_ptr =

@@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

327 if constexpr (USE_HASH) {

328 const int64_t token_id = load_index_as_int64(input_ids, thread_row);

329 const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);

330 + if (!is_pad_row) {

331 #pragma unroll

319 - for (int ii = 0; ii < VPT; ++ii) {

320 - float val = row_chunk[ii];

321 - float val_b = val * beta;

322 - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;

323 - row_chunk[ii] = sqrtf(val);

332 + for (int ii = 0; ii < VPT; ++ii) {

333 + float val = row_chunk[ii];

334 + float val_b = val * beta;

335 + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;

336 + val = sqrtf(val);

337 +

338 + // Dummy/padding tokens can result in NaN values, so

339 + // clamp them to 0.0. Note: this clamp could likely be removed if

340 + // 'is_padding' is made mandatory

341 + if (isnan(val)) {

342 + val = 0.f;

343 + }

344 + row_chunk[ii] = val;

345 + }

346 }

347 float selected_sum = 0.f;

348 #pragma unroll

@@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

357 group_id * THREADS_PER_ROW * ELTS_PER_LDG +

358 local_id;

359 if (expert == expert_idx) {

338 - indices[idx] = static_cast<IndType>(expert);

360 + indices[idx] = !is_pad_row ? static_cast<IndType>(expert)

361 + : static_cast<IndType>(-1);

362 selected_sum += row_chunk[ii];

363 break;

364 }

@@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

402 #endif

403 return;

404 } else {

405 + if (!is_pad_row) {

406 #pragma unroll

383 - for (int ii = 0; ii < VPT; ++ii) {

384 - float val = row_chunk[ii];

385 - float val_b = val * beta;

386 - // Compute softplus: log(1 + exp(val)) with numerical stability

387 - // When val > threshold, softplus(x) ≈ x to avoid exp overflow

388 - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;

389 - val = sqrtf(val);

390 - if (correction_bias) {

391 - const int group_id = ii / ELTS_PER_LDG;

392 - const int local_id = ii % ELTS_PER_LDG;

393 - const int expert_idx = first_elt_read_by_thread +

394 - group_id * THREADS_PER_ROW * ELTS_PER_LDG +

395 - local_id;

396 - val = val + correction_bias[expert_idx];

407 + for (int ii = 0; ii < VPT; ++ii) {

408 + float val = row_chunk[ii];

409 + float val_b = val * beta;

410 + // Compute softplus: log(1 + exp(val)) with numerical stability

411 + // When val > threshold, softplus(x) ≈ x to avoid exp overflow

412 + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;

413 + val = sqrtf(val);

414 + // Dummy/padding tokens can result in NaN values, so

415 + // clamp them to 0.0. Note: this clamp could likely be removed if

416 + // 'is_padding' is made mandatory

417 + if (isnan(val)) {

418 + val = 0.f;

419 + }

420 + if (correction_bias) {

421 + const int group_id = ii / ELTS_PER_LDG;

422 + const int local_id = ii % ELTS_PER_LDG;

423 + const int expert_idx = first_elt_read_by_thread +

424 + group_id * THREADS_PER_ROW * ELTS_PER_LDG +

425 + local_id;

426 + val = val + correction_bias[expert_idx];

427 + }

428 + row_chunk[ii] = val;

429 }

398 - row_chunk[ii] = val;

430 }

431

432 // Original TopK path: find top-k experts by score

@@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__

481 // Add a guard to ignore experts not included by this node

482 const bool node_uses_expert =

483 expert >= start_expert && expert < end_expert;

453 - const bool should_process_row = row_is_active && node_uses_expert;

484 + const bool should_process_row =

485 + row_is_active && node_uses_expert && !is_pad_row;

486

487 // The lead thread from each sub-group will write out the final results

488 // to global memory. (This will be a single) thread per row of the

489 // input/output matrices.

490 const int idx = k * thread_row + k_idx;

459 - if (correction_bias != nullptr) {

491 + if (correction_bias != nullptr && should_process_row) {

492 max_val -= correction_bias[expert];

493 }

494 output[idx] = max_val;

495 indices[idx] =

464 - should_process_row ? (expert - start_expert) : NUM_EXPERTS;

496 + !is_pad_row ? expert - start_expert : static_cast<IndType>(-1);

497 source_rows[idx] = k_idx * num_rows + thread_row;

498 if (renormalize) {

499 selected_sum += max_val;

@@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(

576 const int start_expert, const int end_expert, const bool renormalize,

577 double routed_scaling_factor, const float* correction_bias,

578 const bool use_hash, const HashIndType* input_ids,

547 - const HashIndType* tid2eid, cudaStream_t stream) {

579 + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {

580 static constexpr int BYTES_PER_LDG =

581 MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);

582 using Constants =

@@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(

605 cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,

606 indices, source_row, k, start_expert, end_expert,

607 renormalize, routed_scaling_factor, correction_bias,

576 - input_ids, tid2eid);

608 + input_ids, tid2eid, is_padding);

609 #else

610 kernel<<<num_blocks, block_dim, 0, stream>>>(

611 input, finished, output, num_rows, indices, source_row, k, start_expert,

612 end_expert, renormalize, routed_scaling_factor, correction_bias,

581 - input_ids, tid2eid);

613 + input_ids, tid2eid, is_padding);

614 #endif

615 })

616 }

@@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(

624 gating_output, nullptr, topk_weights, topk_indices, \

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

626 routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \

595 - stream);

627 + stream, is_padding);

628 #else

629 #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \

630 if (WARP_SIZE == 64) { \

@@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(

633 gating_output, nullptr, topk_weights, topk_indices, \

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

635 routed_scaling_factor, correction_bias, use_hash, input_ids, \

604 - tid2eid, stream); \

636 + tid2eid, stream, is_padding); \

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

638 topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \

639 MAX_BYTES>( \

640 gating_output, nullptr, topk_weights, topk_indices, \

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

642 routed_scaling_factor, correction_bias, use_hash, input_ids, \

611 - tid2eid, stream); \

643 + tid2eid, stream, is_padding); \

644 } else { \

645 assert(false && \

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

@@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher(

654 const int topk, const bool renormalize, double routed_scaling_factor,

655 const float* correction_bias, const bool use_hash,

656 const HashIndType* input_ids, const HashIndType* tid2eid,

625 - cudaStream_t stream) {

657 + cudaStream_t stream, const bool* is_padding) {

658 #ifndef USE_ROCM

659 if constexpr (std::is_same_v<InputType, float>) {

660 if (use_hash && topk == 6 && renormalize &&

661 (num_experts == 256 || num_experts == 384)) {

662 launchDsv4HashTopk<IndType, HashIndType>(

663 gating_output, topk_weights, topk_indices, num_tokens, num_experts,

632 - routed_scaling_factor, input_ids, tid2eid, stream);

664 + routed_scaling_factor, input_ids, tid2eid, stream, is_padding);

665 return;

666 }

667 }

@@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(

760 int num_experts, int topk, bool renormalize, double routed_scaling_factor,

761 const std::optional<torch::stable::Tensor>& correction_bias,

762 const std::optional<torch::stable::Tensor>& input_ids,

731 - const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {

763 + const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,

764 + const std::optional<torch::stable::Tensor>& is_padding) {

765 const float* bias_ptr = nullptr;

766 if (correction_bias.has_value()) {

767 bias_ptr = correction_bias.value().const_data_ptr<float>();

@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(

770 auto launch = [&](auto* topk_indices_ptr) {

771 using OutIndType =

772 typename std::remove_pointer<decltype(topk_indices_ptr)>::type;

773 +

774 + const bool* is_padding_ptr = nullptr;

775 + if (is_padding.has_value()) {

776 + const torch::stable::Tensor& is_padding_tensor = is_padding.value();

777 + STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==

778 + torch::headeronly::ScalarType::Bool,

779 + "is_padding tensor must be bool");

780 + STD_TORCH_CHECK(is_padding_tensor.dim() == 1,

781 + "is_padding tensor must be 1D");

782 + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,

783 + "is_padding size mismatch, expected: ", num_tokens);

784 + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),

785 + "is_padding tensor must be contiguous");

786 + is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();

787 + }

788 +

789 if (tid2eid.has_value()) {

790 STD_TORCH_CHECK(input_ids.has_value(),

791 "input_ids is required for hash MoE");

@@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(

800 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),

801 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,

802 bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),

754 - tid2eid.value().const_data_ptr<int64_t>(), stream);

803 + tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);

804 } else {

805 STD_TORCH_CHECK(tid2eid.value().scalar_type() ==

806 torch::headeronly::ScalarType::Int);

@@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(

810 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),

811 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,

812 bias_ptr, true, input_ids.value().const_data_ptr<int>(),

764 - tid2eid.value().const_data_ptr<int>(), stream);

813 + tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);

814 }

815 } else {

816 vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(

@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(

818 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),

819 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,

820 bias_ptr, false, static_cast<const OutIndType*>(nullptr),

772 - static_cast<const OutIndType*>(nullptr), stream);

821 + static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);

822 }

823 };

824

@@ -793,7 +842,8 @@ void topk_softplus_sqrt(

842 bool renormalize, double routed_scaling_factor,

843 const std::optional<torch::stable::Tensor>& correction_bias,

844 const std::optional<torch::stable::Tensor>& input_ids,

796 - const std::optional<torch::stable::Tensor>& tid2eid) {

845 + const std::optional<torch::stable::Tensor>& tid2eid,

846 + const std::optional<torch::stable::Tensor>& is_padding) {

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

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

849 const int topk = topk_weights.size(-1);

@@ -806,21 +856,22 @@ void topk_softplus_sqrt(

856 dispatch_topk_softplus_sqrt_launch<float>(

857 gating_output.const_data_ptr<float>(), topk_weights, topk_indices,

858 token_expert_indices, num_tokens, num_experts, topk, renormalize,

809 - routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);

859 + routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,

860 + is_padding);

861 } else if (gating_output.scalar_type() ==

862 torch::headeronly::ScalarType::Half) {

863 dispatch_topk_softplus_sqrt_launch<__half>(

864 reinterpret_cast<const __half*>(gating_output.const_data_ptr()),

865 topk_weights, topk_indices, token_expert_indices, num_tokens,

866 num_experts, topk, renormalize, routed_scaling_factor, correction_bias,

816 - input_ids, tid2eid, stream);

867 + input_ids, tid2eid, stream, is_padding);

868 } else if (gating_output.scalar_type() ==

869 torch::headeronly::ScalarType::BFloat16) {

870 dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(

871 reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),

872 topk_weights, topk_indices, token_expert_indices, num_tokens,

873 num_experts, topk, renormalize, routed_scaling_factor, correction_bias,

823 - input_ids, tid2eid, stream);

874 + input_ids, tid2eid, stream, is_padding);

875 } else {

876 STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",

877 gating_output.scalar_type());

csrc/libtorch_stable/moe/torch_bindings.cpp

+3 -3

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

8 m.def(

9 "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "

10 "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "

11 - "bias) -> ()");

11 + "bias, Tensor? is_padding) -> ()");

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, "

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

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

19 m.def(

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

21 "token_expert_indices, Tensor gating_output, bool renormalize, float "

22 "routed_scaling_factor, Tensor? "

23 - "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");

23 + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");

25 // Calculate the result of moe by summing up the partial results

26 // from all selected experts. topk_ids/expert_map are optional and, when

tests/kernels/moe/test_topk_softplus_sqrt.py

+117

@@ -6,6 +6,7 @@ import pytest

6 import torch

7 import torch.nn.functional as F

9 +import vllm._custom_ops as ops

10 from vllm.model_executor.layers.fused_moe.config import (

11 RoutingMethodType,

12 get_routing_method_type,

@@ -231,3 +232,119 @@ def test_dsv4_fast_topk(

232 atol=2e-5,

233 rtol=2e-5,

234 )

235 +

236 +

237 +@pytest.mark.skipif(

238 + not current_platform.is_cuda_alike(),

239 + reason="This test is skipped on non-CUDA platform.",

240 +)

241 +@pytest.mark.parametrize("use_hash", [False, True])

242 +@pytest.mark.parametrize("use_bias", [False, True])

243 +@pytest.mark.parametrize("use_padding_mask", [False, True])

244 +@pytest.mark.parametrize("pad_with_nan", [False, True])

245 +@pytest.mark.parametrize("num_experts", [128, 256, 384])

246 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])

247 +def test_fused_topk_softplus_sqrt_padding(

248 + use_hash: bool,

249 + use_bias: bool,

250 + use_padding_mask: bool,

251 + pad_with_nan: bool,

252 + num_experts: int,

253 + dtype: torch.dtype,

254 +):

255 + """Verify explicit padding and NaN-padded rows do not affect real rows."""

256 + torch.manual_seed(0)

257 + num_tokens = 8

258 + topk = 6

259 + indices_dtype = torch.int32

260 +

261 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")

262 +

263 + padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")

264 + padding_rows[1::2] = True

265 + if pad_with_nan:

266 + gating_output[padding_rows] = float("nan")

267 + is_padding = padding_rows if use_padding_mask else None

268 +

269 + # A negative correction bias makes explicit pad rows look selectable unless

270 + # the kernel uses the is_padding guard.

271 + e_score_correction_bias = None

272 + if use_bias:

273 + e_score_correction_bias = (

274 + -torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0

275 + )

276 +

277 + input_ids = None

278 + hash_indices_table = None

279 + if use_hash:

280 + vocab_size = 64

281 + hash_indices_table = torch.stack(

282 + [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]

283 + ).to(device="cuda", dtype=indices_dtype)

284 + input_ids = torch.randint(

285 + 0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda"

286 + )

287 +

288 + topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda")

289 + topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda")

290 + token_expert_indices = torch.empty(

291 + num_tokens, topk, dtype=torch.int32, device="cuda"

292 + )

293 +

294 + ops.topk_hash_softplus_sqrt(

295 + topk_weights,

296 + topk_ids,

297 + token_expert_indices,

298 + gating_output,

299 + renormalize=True,

300 + routed_scaling_factor=1.0,

301 + e_score_correction_bias=e_score_correction_bias,

302 + input_tokens=input_ids,

303 + hash_indices_table=hash_indices_table,

304 + is_padding=is_padding,

305 + )

306 +

307 + if use_padding_mask:

308 + pad_ids = topk_ids[padding_rows]

309 + pad_weights = topk_weights[padding_rows]

310 + assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), (

311 + f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}"

312 + )

313 + assert (pad_weights == 0).all(), (

314 + "Explicit pad rows should have all-zero weights, "

315 + f"got {pad_weights.tolist()}"

316 + )

317 +

318 + if pad_with_nan:

319 + nan_pad_weights = topk_weights[padding_rows]

320 + assert torch.isfinite(nan_pad_weights).all(), (

321 + f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}"

322 + )

323 + assert (nan_pad_weights == 0).all(), (

324 + "NaN-padded rows should have all-zero weights, "

325 + f"got {nan_pad_weights.tolist()}"

326 + )

327 +

328 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(

329 + gating_output=gating_output,

330 + topk=topk,

331 + renormalize=True,

332 + routed_scaling_factor=1.0,

333 + e_score_correction_bias=e_score_correction_bias,

334 + input_ids=input_ids,

335 + hash_indices_table=hash_indices_table,

336 + )

337 +

338 + rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda")

339 + if use_padding_mask or pad_with_nan:

340 + rows_to_compare = ~padding_rows

341 +

342 + sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1)

343 + sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1)

344 + torch.testing.assert_close(

345 + sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0

346 + )

347 +

348 + sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)

349 + sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)

350 + torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)

vllm/_custom_ops.py

@@ -2389,6 +2389,7 @@ def topk_softmax(

2389 gating_output: torch.Tensor,

2390 renormalize: bool = False,

2391 e_score_correction_bias: torch.Tensor | None = None,

2392 + is_padding: torch.Tensor | None = None,

2393 ) -> None:

2394 torch.ops._moe_C.topk_softmax(

2395 topk_weights,

@@ -2397,6 +2398,7 @@ def topk_softmax(

2398 gating_output,

2399 renormalize,

2400 e_score_correction_bias,

2401 + is_padding,

2402 )

2403

2404

@@ -2408,6 +2410,7 @@ def topk_sigmoid(

2410 renormalize: bool = False,

2411 e_score_correction_bias: torch.Tensor | None = None,

2412 routed_scaling_factor: float = 1.0,

2413 + is_padding: torch.Tensor | None = None,

2414 ) -> None:

2415 torch.ops._moe_C.topk_sigmoid(

2416 topk_weights,

@@ -2417,6 +2420,7 @@ def topk_sigmoid(

2420 renormalize,

2421 e_score_correction_bias,

2422 routed_scaling_factor,

2423 + is_padding,

2424 )

2425

2426

@@ -2430,6 +2434,7 @@ def topk_hash_softplus_sqrt(

2434 e_score_correction_bias: torch.Tensor | None = None,

2435 input_tokens: torch.Tensor | None = None,

2436 hash_indices_table: torch.Tensor | None = None,

2437 + is_padding: torch.Tensor | None = None,

2438 ) -> None:

2439 torch.ops._moe_C.topk_softplus_sqrt(

2440 topk_weights,

@@ -2441,6 +2446,7 @@ def topk_hash_softplus_sqrt(

2446 e_score_correction_bias,

2447 input_tokens,

2448 hash_indices_table,

2449 + is_padding,

2450 )

2451

2452

vllm/envs.py

+3 -4

@@ -193,7 +193,7 @@ if TYPE_CHECKING:

193 "relax",

194 ] = "relax"

195 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True

196 - VLLM_MOE_SKIP_PADDING: bool = False

196 + VLLM_MOE_SKIP_PADDING: bool = True

197 VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True

198 VLLM_USE_FLASHINFER_MOE_INT4: bool = False

199 VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None

@@ -1519,9 +1519,8 @@ environment_variables: dict[str, Callable[[], Any]] = {

1519 ),

1520 # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert

1521 # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that

1522 - # treats topk_id == -1 as a skip sentinel; off by default because not all

1523 - # kernels support it yet.

1524 - "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))),

1522 + # treats topk_id == -1 as a skip sentinel

1523 + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))),

1524 # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.

1525 # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).

1526 "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(

vllm/model_executor/layers/fused_moe/modular_kernel.py

-16

@@ -10,7 +10,6 @@ from typing import final

10 import torch

12 import vllm.envs as envs

13 -from vllm.forward_context import get_forward_context, is_forward_context_available

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.fused_moe.activation import (

15 MoEActivation,

@@ -1194,21 +1193,6 @@ class FusedMoEKernelModularImpl:

1193 The _prepare method is a wrapper around self.prepare_finalize.prepare

1194 that handles DBO and async.

1195 """

1197 - # Skip cudagraph/DP padding tokens uniformly across all a2a backends:

1198 - # forcing padded rows' expert ids to -1 makes every prepare_finalize drop

1199 - # them (not dispatched / not computed by the experts). The V2 model runner

1200 - # marks them in forward_context.is_padding; it is None for runners that do

1201 - # not populate it, leaving topk_ids unchanged.

1202 - # Gated by VLLM_MOE_SKIP_PADDING (off by default) because this requires the

1203 - # experts kernel to treat topk_id == -1 as a skip sentinel, which not all

1204 - # MoE backends support yet.

1205 - is_padding = None

1206 - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

1207 - is_padding = get_forward_context().is_padding

1208 - if is_padding is not None:

1209 - n = topk_ids.shape[0]

1210 - # TODO: Properly support DBO (padding lives at the batch tail).

1211 - topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids)

1196

1197 if not self.prepare_finalize.supports_async():

1198 # We shouldn't be running an a2a kernel that doesn't

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

+11

@@ -9,6 +9,7 @@ import vllm._custom_ops as ops

9 import vllm.envs as envs

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.distributed.eplb.eplb_state import EplbLayerState

12 +from vllm.forward_context import get_forward_context, is_forward_context_available

13 from vllm.model_executor.layers.fused_moe.config import (

14 RoutingMethodType,

15 get_routing_method_type,

@@ -20,6 +21,13 @@ from vllm.model_executor.layers.fused_moe.router.dsv4_topk import (

21 )

24 +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None:

25 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

26 + is_padding = get_forward_context().is_padding

27 + return is_padding[:num_tokens] if is_padding is not None else None

28 + return None

29 +

30 +

31 def vllm_topk_softmax(

32 topk_weights: torch.Tensor,

33 topk_indices: torch.Tensor,

@@ -35,6 +43,7 @@ def vllm_topk_softmax(

43 gating_output,

44 renormalize,

45 e_score_correction_bias,

46 + is_padding=_get_padding_mask(topk_indices.shape[0]),

47 )

49 return topk_weights, topk_indices

@@ -57,6 +66,7 @@ def vllm_topk_sigmoid(

66 renormalize,

67 e_score_correction_bias,

68 routed_scaling_factor,

69 + is_padding=_get_padding_mask(topk_indices.shape[0]),

70 )

72 return topk_weights, topk_indices

@@ -144,6 +154,7 @@ def vllm_topk_softplus_sqrt(

154 e_score_correction_bias,

155 input_tokens,

156 hash_indices_table,

157 + is_padding=_get_padding_mask(topk_indices.shape[0]),

158 )

159

160 return topk_weights, topk_indices

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

+11

@@ -5,8 +5,10 @@ from collections.abc import Callable

5 import torch

7 import vllm._custom_ops as ops

8 +import vllm.envs as envs

9 from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.distributed.eplb.eplb_state import EplbLayerState

11 +from vllm.forward_context import get_forward_context, is_forward_context_available

12 from vllm.model_executor.layers.fused_moe.config import (

13 RoutingMethodType,

14 get_routing_method_type,

@@ -14,6 +16,13 @@ from vllm.model_executor.layers.fused_moe.config import (

16 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter

19 +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None:

20 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

21 + is_padding = get_forward_context().is_padding

22 + return is_padding[:num_tokens] if is_padding is not None else None

23 + return None

24 +

25 +

26 def vllm_topk_softmax(

27 topk_weights: torch.Tensor,

28 topk_indices: torch.Tensor,

@@ -27,6 +36,7 @@ def vllm_topk_softmax(

36 token_expert_indices,

37 gating_output,

38 renormalize,

39 + is_padding=_get_padding_mask(topk_indices.shape[0]),

40 )

42 return topk_weights, topk_indices

@@ -45,6 +55,7 @@ def vllm_topk_sigmoid(

55 token_expert_indices,

56 gating_output,

57 renormalize,

58 + is_padding=_get_padding_mask(topk_indices.shape[0]),

59 )

61 return topk_weights, topk_indices