@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0852527647

Fork 0

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 13, 14:30 UTC

[Perf][DeepSeek] Add sigmoid+bias fusion to fused_grouped_topk from TRTLLM (#28124)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Michael Goin committed Nov 8, 2025 at 10:20 UTC 085252764710fdb42ac180983ef4c37da0ed72d3

5 files changed +149 -75

csrc/moe/grouped_topk_kernels.cu

+99 -54

@@ -427,11 +427,29 @@ __device__ inline bool is_finite(const T val) {

427 #endif

428 }

429

430 +// Scoring function enums

431 +enum ScoringFunc {

432 + SCORING_NONE = 0, // no activation function

433 + SCORING_SIGMOID = 1 // apply sigmoid

434 +};

435 +

436 +// Efficient sigmoid approximation from TensorRT-LLM

437 +__device__ inline float sigmoid_accurate(float x) {

438 + return 0.5f * tanhf(0.5f * x) + 0.5f;

439 +}

440 +

441 template <typename T>

431 -__device__ void topk_with_k2(T* output, T const* input,

442 +__device__ inline T apply_sigmoid(T val) {

443 + float f = cuda_cast<float, T>(val);

444 + return cuda_cast<T, float>(sigmoid_accurate(f));

445 +}

446 +

447 +template <typename T>

448 +__device__ void topk_with_k2(T* output, T const* input, T const* bias,

449 cg::thread_block_tile<32> const& tile,

450 int32_t const lane_id,

434 - int const num_experts_per_group) {

451 + int const num_experts_per_group,

452 + int const scoring_func) {

453 // Get the top2 per thread

454 T largest = neg_inf<T>();

455 T second_largest = neg_inf<T>();

@@ -439,6 +457,12 @@ __device__ void topk_with_k2(T* output, T const* input,

457 if (num_experts_per_group > WARP_SIZE) {

458 for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {

459 T value = input[i];

460 + // Apply scoring function if needed

461 + if (scoring_func == SCORING_SIGMOID) {

462 + value = apply_sigmoid(value);

463 + }

464 + value = value + bias[i];

465 +

466 if (value > largest) {

467 second_largest = largest;

468 largest = value;

@@ -448,7 +472,13 @@ __device__ void topk_with_k2(T* output, T const* input,

472 }

473 } else {

474 for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {

451 - largest = input[i];

475 + T value = input[i];

476 + // Apply scoring function if needed

477 + if (scoring_func == SCORING_SIGMOID) {

478 + value = apply_sigmoid(value);

479 + }

480 + value = value + bias[i];

481 + largest = value;

482 }

483 }

484

@@ -472,17 +502,21 @@ __device__ void topk_with_k2(T* output, T const* input,

502 }

503

504 template <typename T>

475 -__global__ void topk_with_k2_kernel(T* output, T* input,

505 +__global__ void topk_with_k2_kernel(T* output, T* input, T const* bias,

506 int64_t const num_tokens,

507 int64_t const num_cases,

508 int64_t const n_group,

479 - int64_t const num_experts_per_group) {

509 + int64_t const num_experts_per_group,

510 + int const scoring_func) {

511 int32_t warp_id = threadIdx.x / WARP_SIZE;

512 int32_t lane_id = threadIdx.x % WARP_SIZE;

513

514 int32_t case_id = blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;

515 if (case_id < num_cases) {

516 input += case_id * num_experts_per_group;

517 + // bias is per expert group, offset to current group

518 + int32_t group_id = case_id % n_group;

519 + T const* group_bias = bias + group_id * num_experts_per_group;

520 output += case_id;

521

522 cg::thread_block block = cg::this_thread_block();

@@ -491,7 +525,8 @@ __global__ void topk_with_k2_kernel(T* output, T* input,

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

526 asm volatile("griddepcontrol.wait;");

527 #endif

494 - topk_with_k2(output, input, tile, lane_id, num_experts_per_group);

528 + topk_with_k2(output, input, group_bias, tile, lane_id,

529 + num_experts_per_group, scoring_func);

530 }

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

532 asm volatile("griddepcontrol.launch_dependents;");

@@ -500,16 +535,15 @@ __global__ void topk_with_k2_kernel(T* output, T* input,

535

536 template <typename T, typename IdxT>

537 __global__ void group_idx_and_topk_idx_kernel(

503 - T* scores, T const* group_scores, T* topk_values, IdxT* topk_indices,

504 - T* scores_with_bias, int64_t const num_tokens, int64_t const n_group,

538 + T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,

539 + T const* bias, int64_t const num_tokens, int64_t const n_group,

540 int64_t const topk_group, int64_t const topk, int64_t const num_experts,

541 int64_t const num_experts_per_group, bool renormalize,

507 - double routed_scaling_factor) {

542 + double routed_scaling_factor, int scoring_func) {

543 int32_t warp_id = threadIdx.x / WARP_SIZE;

544 int32_t lane_id = threadIdx.x % WARP_SIZE;

545 int32_t case_id =

546 blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id; // one per token

512 - scores_with_bias += case_id * num_experts;

547 scores += case_id * num_experts;

548 group_scores += case_id * n_group;

549 topk_values += case_id * topk;

@@ -577,10 +611,16 @@ __global__ void group_idx_and_topk_idx_kernel(

611 int32_t offset = i_group * num_experts_per_group;

612 for (int32_t i = lane_id; i < align_num_experts_per_group;

613 i += WARP_SIZE) {

580 - T candidates = (i < num_experts_per_group) &&

581 - is_finite(scores_with_bias[offset + i])

582 - ? scores_with_bias[offset + i]

583 - : neg_inf<T>();

614 + T candidates = neg_inf<T>();

615 + if (i < num_experts_per_group) {

616 + // Apply scoring function (if any) and add bias

617 + T input = scores[offset + i];

618 + if (is_finite(input)) {

619 + T score = (scoring_func == SCORING_SIGMOID) ? apply_sigmoid(input)

620 + : input;

621 + candidates = score + bias[offset + i];

622 + }

623 + }

624 queue.add(candidates, offset + i);

625 }

626 if (group_scores[i_group] == topk_group_value) {

@@ -602,11 +642,12 @@ __global__ void group_idx_and_topk_idx_kernel(

642 for (int i = lane_id;

643 i < warp_topk::round_up_to_multiple_of<WARP_SIZE>(topk);

644 i += WARP_SIZE) {

605 - T value =

606 - i < topk

607 - ? scores[s_topk_idx[i]]

608 - : cuda_cast<T, float>(0.0f); // Load the valid value of expert

645 + T value = cuda_cast<T, float>(0.0f);

646 if (i < topk) {

647 + // Load the score value (without bias) for normalization

648 + T input = scores[s_topk_idx[i]];

649 + value =

650 + (scoring_func == SCORING_SIGMOID) ? apply_sigmoid(input) : input;

651 s_topk_value[i] = value;

652 }

653 topk_sum +=

@@ -627,12 +668,12 @@ __global__ void group_idx_and_topk_idx_kernel(

668 value = cuda_cast<float, T>(s_topk_value[i]) * routed_scaling_factor;

669 }

670 topk_indices[i] = s_topk_idx[i];

630 - topk_values[i] = cuda_cast<T, float>(value);

671 + topk_values[i] = value;

672 }

673 } else {

674 for (int i = lane_id; i < topk; i += WARP_SIZE) {

675 topk_indices[i] = i;

635 - topk_values[i] = cuda_cast<T, float>(1.0f / topk);

676 + topk_values[i] = 1.0f / topk;

677 }

678 }

679 // Note: when if_proceed_next_topk==false, choose the first 8 experts as the

@@ -644,12 +685,12 @@ __global__ void group_idx_and_topk_idx_kernel(

685 }

686

687 template <typename T, typename IdxT>

647 -void invokeNoAuxTc(T* scores, T* group_scores, T* topk_values,

648 - IdxT* topk_indices, T* scores_with_bias,

649 - int64_t const num_tokens, int64_t const num_experts,

650 - int64_t const n_group, int64_t const topk_group,

651 - int64_t const topk, bool const renormalize,

652 - double const routed_scaling_factor, bool enable_pdl = false,

688 +void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,

689 + IdxT* topk_indices, T const* bias, int64_t const num_tokens,

690 + int64_t const num_experts, int64_t const n_group,

691 + int64_t const topk_group, int64_t const topk,

692 + bool const renormalize, double const routed_scaling_factor,

693 + int const scoring_func, bool enable_pdl = false,

694 cudaStream_t const stream = 0) {

695 int64_t num_cases = num_tokens * n_group;

696 int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;

@@ -664,8 +705,9 @@ void invokeNoAuxTc(T* scores, T* group_scores, T* topk_values,

705 attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;

706 config.numAttrs = 1;

707 config.attrs = attrs;

667 - cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores_with_bias,

668 - num_tokens, num_cases, n_group, num_experts / n_group);

708 + cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores, bias,

709 + num_tokens, num_cases, n_group, num_experts / n_group,

710 + scoring_func);

711

712 int64_t topk_with_k_group_num_blocks =

713 (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;

@@ -682,19 +724,18 @@ void invokeNoAuxTc(T* scores, T* group_scores, T* topk_values,

724 config.numAttrs = 1;

725 config.attrs = attrs;

726 cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,

685 - topk_values, topk_indices, scores_with_bias, num_tokens,

686 - n_group, topk_group, topk, num_experts,

687 - num_experts / n_group, renormalize, routed_scaling_factor);

727 + topk_values, topk_indices, bias, num_tokens, n_group,

728 + topk_group, topk, num_experts, num_experts / n_group,

729 + renormalize, routed_scaling_factor, scoring_func);

730 }

731

732 #define INSTANTIATE_NOAUX_TC(T, IdxT) \

733 template void invokeNoAuxTc<T, IdxT>( \

692 - T * scores, T * group_scores, T * topk_values, IdxT * topk_indices, \

693 - T * scores_with_bias, int64_t const num_tokens, \

694 - int64_t const num_experts, int64_t const n_group, \

695 - int64_t const topk_group, int64_t const topk, bool const renormalize, \

696 - double const routed_scaling_factor, bool enable_pdl, \

697 - cudaStream_t const stream);

734 + T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \

735 + T const* bias, int64_t const num_tokens, int64_t const num_experts, \

736 + int64_t const n_group, int64_t const topk_group, int64_t const topk, \

737 + bool const renormalize, double const routed_scaling_factor, \

738 + int const scoring_func, bool enable_pdl, cudaStream_t const stream);

739

740 INSTANTIATE_NOAUX_TC(float, int32_t);

741 INSTANTIATE_NOAUX_TC(half, int32_t);

@@ -703,28 +744,32 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, int32_t);

744 } // namespace vllm

745

746 std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

706 - torch::Tensor const& scores, torch::Tensor const& scores_with_bias,

707 - int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize,

708 - double routed_scaling_factor) {

709 - auto data_type = scores_with_bias.scalar_type();

710 - auto input_size = scores_with_bias.sizes();

747 + torch::Tensor const& scores, int64_t n_group, int64_t topk_group,

748 + int64_t topk, bool renormalize, double routed_scaling_factor,

749 + torch::Tensor const& bias, int64_t scoring_func = 0) {

750 + auto data_type = scores.scalar_type();

751 + auto input_size = scores.sizes();

752 int64_t num_tokens = input_size[0];

753 int64_t num_experts = input_size[1];

713 - TORCH_CHECK(input_size.size() == 2, "scores_with_bias must be a 2D Tensor");

754 + TORCH_CHECK(input_size.size() == 2, "scores must be a 2D Tensor");

755 TORCH_CHECK(num_experts % n_group == 0,

756 "num_experts should be divisible by n_group");

757 TORCH_CHECK(n_group <= 32,

758 "n_group should be smaller than or equal to 32 for now");

759 TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now");

760 + TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE ||

761 + scoring_func == vllm::moe::SCORING_SIGMOID,

762 + "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)");

763

764 torch::Tensor group_scores = torch::empty(

765 {num_tokens, n_group}, torch::dtype(data_type).device(torch::kCUDA));

766 + // Always output float32 for topk_values (eliminates Python-side conversion)

767 torch::Tensor topk_values = torch::empty(

723 - {num_tokens, topk}, torch::dtype(data_type).device(torch::kCUDA));

768 + {num_tokens, topk}, torch::dtype(torch::kFloat32).device(torch::kCUDA));

769 torch::Tensor topk_indices = torch::empty(

770 {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));

771

727 - auto stream = c10::cuda::getCurrentCUDAStream(scores_with_bias.get_device());

772 + auto stream = c10::cuda::getCurrentCUDAStream(scores.get_device());

773

774 switch (data_type) {

775 case torch::kFloat16:

@@ -732,11 +777,11 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

777 vllm::moe::invokeNoAuxTc<half, int32_t>(

778 reinterpret_cast<half*>(scores.mutable_data_ptr()),

779 reinterpret_cast<half*>(group_scores.mutable_data_ptr()),

735 - reinterpret_cast<half*>(topk_values.mutable_data_ptr()),

780 + reinterpret_cast<float*>(topk_values.mutable_data_ptr()),

781 reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),

737 - reinterpret_cast<half*>(scores_with_bias.data_ptr()), num_tokens,

782 + reinterpret_cast<half const*>(bias.data_ptr()), num_tokens,

783 num_experts, n_group, topk_group, topk, renormalize,

739 - routed_scaling_factor, false, stream);

784 + routed_scaling_factor, static_cast<int>(scoring_func), false, stream);

785 break;

786 case torch::kFloat32:

787 // Handle Float32

@@ -745,20 +790,20 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

790 reinterpret_cast<float*>(group_scores.mutable_data_ptr()),

791 reinterpret_cast<float*>(topk_values.mutable_data_ptr()),

792 reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),

748 - reinterpret_cast<float*>(scores_with_bias.data_ptr()), num_tokens,

793 + reinterpret_cast<float const*>(bias.data_ptr()), num_tokens,

794 num_experts, n_group, topk_group, topk, renormalize,

750 - routed_scaling_factor, false, stream);

795 + routed_scaling_factor, static_cast<int>(scoring_func), false, stream);

796 break;

797 case torch::kBFloat16:

798 // Handle BFloat16

799 vllm::moe::invokeNoAuxTc<__nv_bfloat16, int32_t>(

800 reinterpret_cast<__nv_bfloat16*>(scores.mutable_data_ptr()),

801 reinterpret_cast<__nv_bfloat16*>(group_scores.mutable_data_ptr()),

757 - reinterpret_cast<__nv_bfloat16*>(topk_values.mutable_data_ptr()),

802 + reinterpret_cast<float*>(topk_values.mutable_data_ptr()),

803 reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),

759 - reinterpret_cast<__nv_bfloat16*>(scores_with_bias.data_ptr()),

760 - num_tokens, num_experts, n_group, topk_group, topk, renormalize,

761 - routed_scaling_factor, false, stream);

804 + reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), num_tokens,

805 + num_experts, n_group, topk_group, topk, renormalize,

806 + routed_scaling_factor, static_cast<int>(scoring_func), false, stream);

807 break;

808 default:

809 // Handle other data types

csrc/moe/moe_ops.h

+3 -3

@@ -39,9 +39,9 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,

39 int64_t BLOCK_SIZE_K, int64_t bit);

41 std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

42 - torch::Tensor const& scores, torch::Tensor const& scores_with_bias,

43 - int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize,

44 - double routed_scaling_factor);

42 + torch::Tensor const& scores, int64_t n_group, int64_t topk_group,

43 + int64_t topk, bool renormalize, double routed_scaling_factor,

44 + torch::Tensor const& bias, int64_t scoring_func);

45 #endif

47 bool moe_permute_unpermute_supported();

csrc/moe/torch_bindings.cpp

+3 -2

@@ -107,9 +107,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {

107

108 // Apply grouped topk routing to select experts.

109 m.def(

110 - "grouped_topk(Tensor scores, Tensor scores_with_bias, int n_group, int "

110 + "grouped_topk(Tensor scores, int n_group, int "

111 "topk_group, int topk, bool renormalize, float "

112 - "routed_scaling_factor) -> (Tensor, Tensor)");

112 + "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, "

113 + "Tensor)");

114 m.impl("grouped_topk", torch::kCUDA, &grouped_topk);

115 #endif

116 }

vllm/_custom_ops.py

+17 -2

@@ -1898,25 +1898,40 @@ def topk_softmax(

1898

1899 def grouped_topk(

1900 scores: torch.Tensor,

1901 - scores_with_bias: torch.Tensor,

1901 num_expert_group: int,

1902 topk_group: int,

1903 topk: int,

1904 renormalize: bool,

1905 routed_scaling_factor: float,

1906 + bias: torch.Tensor,

1907 + scoring_func: int = 0,

1908 ):

1909 + """

1910 + Perform grouped top-k routing for mixture of experts.

1911 +

1912 + Args:

1913 + scores: Raw inputs (logits if scoring_func=1, scores if scoring_func=0)

1914 + num_expert_group: Number of expert groups

1915 + topk_group: Number of groups to select

1916 + topk: Number of experts to select per token

1917 + renormalize: Whether to renormalize the output weights

1918 + routed_scaling_factor: Scaling factor for routing weights

1919 + bias: Bias tensor (e_score_correction_bias). Always fused in kernel.

1920 + scoring_func: 0=none (no activation), 1=sigmoid

1921 + """

1922 if not current_platform.is_cuda():

1923 raise NotImplementedError(

1924 "The fused grouped_topk kernel is only available on CUDA platforms"

1925 )

1926 return torch.ops._moe_C.grouped_topk(

1927 scores,

1914 - scores_with_bias,

1928 num_expert_group,

1929 topk_group,

1930 topk,

1931 renormalize,

1932 routed_scaling_factor,

1933 + bias,

1934 + scoring_func,

1935 )

1936

1937

vllm/model_executor/layers/fused_moe/fused_moe.py

+27 -14

@@ -1330,24 +1330,37 @@ def fused_grouped_topk(

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

1331 assert hidden_states.size(0) == gating_output.size(0), "Number of tokens mismatch"

1332

1333 - if scoring_func == "softmax":

1333 + if scoring_func == "sigmoid":

1334 + # Fully fused kernel path for sigmoid

1335 + topk_values, topk_indices = ops.grouped_topk(

1336 + gating_output, # raw logits

1337 + num_expert_group,

1338 + topk_group,

1339 + topk,

1340 + renormalize,

1341 + routed_scaling_factor,

1342 + e_score_correction_bias.to(gating_output.dtype),

1343 + 1, # scoring_func=1 for sigmoid

1344 + )

1345 + elif scoring_func == "softmax":

1346 + # Apply softmax in Python, then use fused kernel

1347 + # TODO: Add support for softmax in kernel

1348 scores = torch.softmax(gating_output, dim=-1)

1335 - elif scoring_func == "sigmoid":

1336 - scores = gating_output.sigmoid()

1349 + topk_values, topk_indices = ops.grouped_topk(

1350 + scores, # pre-computed scores

1351 + num_expert_group,

1352 + topk_group,

1353 + topk,

1354 + renormalize,

1355 + routed_scaling_factor,

1356 + e_score_correction_bias.to(gating_output.dtype),

1357 + 0, # scoring_func=0 (no activation, scores already computed)

1358 + )

1359 else:

1360 raise ValueError(f"Unsupported scoring function: {scoring_func}")

1361

1340 - scores_with_bias = scores + e_score_correction_bias.unsqueeze(0)

1341 - topk_values, topk_indices = ops.grouped_topk(

1342 - scores,

1343 - scores_with_bias.to(scores.dtype),

1344 - num_expert_group,

1345 - topk_group,

1346 - topk,

1347 - renormalize,

1348 - routed_scaling_factor,

1349 - )

1350 - return topk_values.to(torch.float32), topk_indices.to(torch.int32)

1362 + # Fused kernel outputs float32 values and int32 indices directly

1363 + return topk_values, topk_indices

1364

1365

1366 def inplace_fused_experts(