f28125d87b · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -31,8 +31,6 @@ namespace moe {

31

32 constexpr unsigned FULL_WARP_MASK = 0xffffffff;

33 constexpr int32_t WARP_SIZE = 32;

34 -constexpr int32_t BLOCK_SIZE = 512;

35 -constexpr int32_t NUM_WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;

34

35 namespace warp_topk {

36

@@ -65,14 +63,6 @@ __forceinline__ __device__ bool is_better_than(T val, T baseline, idxT index,

63 return res;

64 }

65

68 -template <typename T, typename idxT>

69 -int calc_smem_size_for_block_wide(int num_of_warp, int64_t k) {

70 - int64_t cache_topk = (sizeof(T) + sizeof(idxT)) * num_of_warp * k;

71 - int64_t n = std::max<int>(num_of_warp / 2 * k, num_of_warp * WARP_SIZE);

72 - return max(cache_topk,

73 - round_up_to_multiple_of<256>(n * sizeof(T)) + n * sizeof(idxT));

74 -}

75 -

66 template <int size, bool ascending, bool reverse, typename T, typename idxT,

67 bool is_stable>

68 struct BitonicMerge {

@@ -267,6 +257,15 @@ class WarpSort {

257 }

258 }

259

260 + // Accessors for per-lane selected value/index.

261 + // NOTE: For the common case `capacity == WARP_SIZE`, `max_arr_len_ == 1`

262 + // and callers should use `i == 0`.

263 + __device__ __forceinline__ idxT get_idx(int i = 0) const {

264 + return idx_arr_[i];

265 + }

266 +

267 + __device__ __forceinline__ T get_val(int i = 0) const { return val_arr_[i]; }

268 +

269 protected:

270 static constexpr int max_arr_len_ = capacity / WARP_SIZE;

271

@@ -285,6 +284,7 @@ class WarpSelect : public WarpSort<capacity, greater, T, idxT, is_stable> {

284 __device__ WarpSelect(idxT k, T dummy)

285 : WarpSort<capacity, greater, T, idxT, is_stable>(k, dummy),

286 k_th_(dummy),

287 + k_th_idx_(0),

288 k_th_lane_((k - 1) % WARP_SIZE) {

289 extern __shared__ char smem_buf[]; // extern __shared__ T smem_buf[];

290

@@ -346,9 +346,6 @@ class WarpSelect : public WarpSort<capacity, greater, T, idxT, is_stable> {

346 idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0;

347 merge_buf_(val, idx);

348 }

349 -

350 - // after done(), smem is used for merging results among warps

351 - __syncthreads();

349 }

350

351 private:

@@ -503,255 +500,186 @@ __device__ void topk_with_k2(T* output, T const* input, BiasT const* bias,

500 }

501 }

502

506 -template <typename T, typename BiasT, ScoringFunc SF>

507 -__global__ void topk_with_k2_kernel(T* output, T* input, BiasT const* bias,

508 - int64_t const num_tokens,

509 - int64_t const num_cases,

510 - int64_t const n_group,

511 - int64_t const num_experts_per_group) {

512 - int32_t warp_id = threadIdx.x / WARP_SIZE;

513 - int32_t lane_id = threadIdx.x % WARP_SIZE;

514 -

515 - int32_t case_id = blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;

516 - if (case_id < num_cases) {

517 - input += case_id * num_experts_per_group;

518 - // bias is per expert group, offset to current group

519 - int32_t group_id = case_id % n_group;

520 - BiasT const* group_bias = bias + group_id * num_experts_per_group;

521 - output += case_id;

522 -

523 - cg::thread_block block = cg::this_thread_block();

524 - cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);

503 +template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>

504 +__global__ void grouped_topk_fused_kernel(

505 + T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias,

506 + int64_t const num_tokens, int64_t const num_experts, int64_t const n_group,

507 + int64_t const topk_group, int64_t const topk, bool renormalize,

508 + double routed_scaling_factor) {

509 + int32_t const token_id = static_cast<int32_t>(blockIdx.x);

510 + if (token_id >= num_tokens) {

511 + return;

512 + }

513

526 -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))

527 - asm volatile("griddepcontrol.wait;");

528 -#endif

529 - topk_with_k2<T, BiasT, SF>(output, input, group_bias, tile, lane_id,

530 - num_experts_per_group);

514 + int32_t const warp_id = threadIdx.x / WARP_SIZE;

515 + int32_t const lane_id = threadIdx.x % WARP_SIZE;

516 +

517 + int32_t const n_group_i32 = static_cast<int32_t>(n_group);

518 + int32_t const topk_group_i32 = static_cast<int32_t>(topk_group);

519 + int32_t const topk_i32 = static_cast<int32_t>(topk);

520 + int32_t const num_experts_i32 = static_cast<int32_t>(num_experts);

521 +

522 + int32_t const num_warps = blockDim.x / WARP_SIZE;

523 + if (warp_id >= n_group_i32 || num_warps < n_group_i32) {

524 + return;

525 }

532 -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))

533 - asm volatile("griddepcontrol.launch_dependents;");

534 -#endif

535 -}

526

537 -template <typename T, typename BiasT, typename IdxT, ScoringFunc SF,

538 - int NGroup = -1>

539 -__global__ void group_idx_and_topk_idx_kernel(

540 - T* scores, T const* group_scores, float* topk_values, IdxT* topk_indices,

541 - BiasT const* bias, int64_t const num_tokens, int64_t const n_group,

542 - int64_t const topk_group, int64_t const topk, int64_t const num_experts,

543 - int64_t const num_experts_per_group, bool renormalize,

544 - double routed_scaling_factor) {

545 - int32_t warp_id = threadIdx.x / WARP_SIZE;

546 - int32_t lane_id = threadIdx.x % WARP_SIZE;

547 - int32_t case_id =

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

549 - scores += case_id * num_experts;

550 - group_scores += case_id * n_group;

551 - topk_values += case_id * topk;

552 - topk_indices += case_id * topk;

553 -

554 - constexpr bool kUseStaticNGroup = (NGroup > 0);

555 - // use int32 to avoid implicit conversion

556 - int32_t const n_group_i32 =

557 - kUseStaticNGroup ? NGroup : static_cast<int32_t>(n_group);

558 -

559 - int32_t align_num_experts_per_group =

560 - warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);

527 + int32_t const num_experts_per_group = num_experts_i32 / n_group_i32;

528 +

529 + T* scores_token = scores + static_cast<int64_t>(token_id) * num_experts;

530

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

532 cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);

533

565 - extern __shared__ char smem_buf[]; // NOTE: reuse the shared memory here to

566 - // store the target topk idx

567 - int32_t* s_topk_idx = reinterpret_cast<int32_t*>(smem_buf);

568 - T* s_topk_value =

569 - reinterpret_cast<T*>(s_topk_idx + NUM_WARPS_PER_BLOCK * topk) +

570 - warp_id * topk;

571 - s_topk_idx += warp_id * topk;

572 -

573 - T value = neg_inf<T>();

574 - T topk_group_value = neg_inf<T>();

575 - int32_t num_equalto_topkth_group;

534 + extern __shared__ char smem_buf[];

535 + // warpSelect internal staging buffer layout

536 + size_t const val_bytes =

537 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);

538 + size_t const val_bytes_aligned =

539 + warp_topk::round_up_to_multiple_of<256>(val_bytes);

540 + size_t const idx_bytes =

541 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);

542 + size_t const internal_bytes = val_bytes_aligned + idx_bytes;

543 +

544 + // user-managed shared memory starts after warpSelect internal staging.

545 + uintptr_t ptr_u = reinterpret_cast<uintptr_t>(smem_buf + internal_bytes);

546 + ptr_u = (ptr_u + 15) & ~static_cast<uintptr_t>(15); // align to 16B

547 + T* s_group_scores = reinterpret_cast<T*>(ptr_u);

548

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

550 asm volatile("griddepcontrol.wait;"); // I think all prolog can be put before

551 // acqbulk because it's ptr arithmetic

552 #endif

553

582 - if (case_id < num_tokens) {

583 - // calculate group_idx

584 - int32_t target_num_min =

585 - WARP_SIZE - n_group_i32 + static_cast<int32_t>(topk_group);

586 - // The check is necessary to avoid abnormal input

587 - if (lane_id < n_group_i32 && is_finite(group_scores[lane_id])) {

588 - value = group_scores[lane_id];

589 - }

554 + // phase 1: per-group scan

555 + int32_t const group_offset = warp_id * num_experts_per_group;

556 + topk_with_k2<T, BiasT, SF>(s_group_scores + warp_id,

557 + scores_token + group_offset, bias + group_offset,

558 + tile, lane_id, num_experts_per_group);

559

591 - int count_equal_to_top_value = WARP_SIZE - n_group_i32;

592 - int pre_count_equal_to_top_value = 0;

593 - // Use loop to find the largset top_group

594 - while (count_equal_to_top_value < target_num_min) {

595 - topk_group_value = cg::reduce(tile, value, cg::greater<T>());

596 - if (value == topk_group_value) {

597 - value = neg_inf<T>();

598 - }

599 - pre_count_equal_to_top_value = count_equal_to_top_value;

600 - count_equal_to_top_value =

601 - __popc(__ballot_sync(FULL_WARP_MASK, (value == neg_inf<T>())));

602 - }

603 - num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;

604 - }

560 __syncthreads();

561

562 + // phase 2: warp0 selects groups + merges candidates to final topk

563 + if (warp_id != 0) {

564 + return;

565 + }

566 +

567 + topk_values += static_cast<int64_t>(token_id) * topk;

568 + topk_indices += static_cast<int64_t>(token_id) * topk;

569 +

570 + // select topk_group groups by group score

571 warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,

572 /* is_stable */ true>

609 - queue((int32_t)topk, neg_inf<T>());

610 -

611 - int count_equalto_topkth_group = 0;

612 - bool if_proceed_next_topk = topk_group_value != neg_inf<T>();

613 - if (case_id < num_tokens && if_proceed_next_topk) {

614 - auto process_group = [&](int i_group) {

615 - if ((group_scores[i_group] > topk_group_value) ||

616 - ((group_scores[i_group] == topk_group_value) &&

617 - (count_equalto_topkth_group < num_equalto_topkth_group))) {

618 - int32_t offset = i_group * num_experts_per_group;

619 - for (int32_t i = lane_id; i < align_num_experts_per_group;

620 - i += WARP_SIZE) {

621 - T candidates = neg_inf<T>();

622 - if (i < num_experts_per_group) {

623 - // apply scoring function (if any) and add bias

624 - T input = scores[offset + i];

625 - if (is_finite(input)) {

626 - T score = apply_scoring<SF>(input);

627 - candidates = score + static_cast<T>(bias[offset + i]);

628 - }

629 - }

630 - queue.add(candidates, offset + i);

631 - }

632 - if (group_scores[i_group] == topk_group_value) {

633 - count_equalto_topkth_group++;

634 - }

635 - }

636 - };

573 + group_sel(static_cast<int32_t>(topk_group_i32), neg_inf<T>());

574 +

575 + // all lanes must participate in WarpSelect::add().

576 + T gscore = (lane_id < n_group_i32) ? s_group_scores[lane_id] : neg_inf<T>();

577 + group_sel.add(gscore, lane_id);

578 + group_sel.done();

579 +

580 + // proceed only if the k-th selected group score is not -inf

581 + bool proceed = false;

582 + if (topk_group_i32 > 0) {

583 + int const kth_lane = topk_group_i32 - 1;

584 + // broadcast the k-th selected group score to all lanes

585 + T kth_val = __shfl_sync(FULL_WARP_MASK, group_sel.get_val(0), kth_lane);

586 + proceed = (kth_val != neg_inf<T>());

587 + }

588

638 - if constexpr (kUseStaticNGroup) {

639 -#pragma unroll

640 - for (int i_group = 0; i_group < NGroup; ++i_group) {

641 - process_group(i_group);

642 - }

643 - } else {

644 - for (int i_group = 0; i_group < n_group_i32; ++i_group) {

645 - process_group(i_group);

646 - }

589 + if (!proceed) {

590 + for (int i = lane_id; i < topk_i32; i += WARP_SIZE) {

591 + topk_indices[i] = static_cast<IdxT>(i);

592 + topk_values[i] = 1.0f / static_cast<float>(topk_i32);

593 }

648 - queue.done();

649 - // Get the topk_idx

650 - queue.dumpIdx(s_topk_idx);

651 - }

652 -

653 - // Load the valid score value

654 - // Calculate the summation

655 - float topk_sum = 1e-20;

656 - if (case_id < num_tokens && if_proceed_next_topk) {

657 - for (int i = lane_id;

658 - i < warp_topk::round_up_to_multiple_of<WARP_SIZE>(topk);

659 - i += WARP_SIZE) {

660 - T value = cuda_cast<T, float>(0.0f);

661 - if (i < topk) {

662 - // Load the score value (without bias) for normalization

663 - T input = scores[s_topk_idx[i]];

664 - value = apply_scoring<SF>(input);

665 - s_topk_value[i] = value;

666 - }

667 - if (renormalize) {

668 - topk_sum +=

669 - cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());

594 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))

595 + asm volatile("griddepcontrol.launch_dependents;");

596 +#endif

597 + return;

598 + }

599 +

600 + // merge per-group topk candidates for selected groups, then select topk

601 + warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,

602 + /* is_stable */ true>

603 + expert_sel(static_cast<int32_t>(topk_i32), neg_inf<T>());

604 +

605 + // selected group ids reside in lanes [0, topk_group)

606 + int32_t sel_gid_lane = (lane_id < topk_group_i32) ? group_sel.get_idx(0) : 0;

607 +

608 + // add candidates from selected groups to expert_sel

609 + for (int32_t g = 0; g < topk_group_i32; ++g) {

610 + int32_t gid = __shfl_sync(FULL_WARP_MASK, sel_gid_lane, g);

611 + int32_t const offset = gid * num_experts_per_group;

612 + int32_t const align_num_experts_per_group =

613 + warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);

614 + for (int32_t i = lane_id; i < align_num_experts_per_group; i += WARP_SIZE) {

615 + // all lanes must call `add()` the same number of times.

616 + T cand = neg_inf<T>();

617 + int32_t idx = 0;

618 + if (i < num_experts_per_group) {

619 + idx = offset + i;

620 + T input = scores_token[idx];

621 + if (is_finite(input)) {

622 + T score = apply_scoring<SF>(input);

623 + cand = score + static_cast<T>(bias[idx]);

624 + }

625 }

626 + expert_sel.add(cand, idx);

627 }

628 }

629 + expert_sel.done();

630 +

631 + // compute unbiased routing weights + optional renorm.

632 + float lane_unbiased = 0.0f;

633 + IdxT lane_idx = 0;

634 + if (lane_id < topk_i32) {

635 + lane_idx = static_cast<IdxT>(expert_sel.get_idx(0));

636 + T in = scores_token[static_cast<int32_t>(lane_idx)];

637 + lane_unbiased = cuda_cast<float, T>(apply_scoring<SF>(in));

638 + }

639

674 - __syncthreads();

640 + float topk_sum = 1e-20f;

641 + if (renormalize) {

642 + topk_sum += cg::reduce(tile, lane_unbiased, cg::plus<float>());

643 + }

644

676 - if (case_id < num_tokens) {

677 - if (if_proceed_next_topk) {

678 - float scale = routed_scaling_factor;

679 - if (renormalize) {

680 - scale /= topk_sum;

681 - }

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

683 - float base = cuda_cast<float, T>(s_topk_value[i]);

684 - float value = base * scale;

685 - topk_indices[i] = s_topk_idx[i];

686 - topk_values[i] = value;

687 - }

688 - } else {

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

690 - topk_indices[i] = i;

691 - topk_values[i] = 1.0f / topk;

692 - }

693 - }

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

695 - // default result.

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

646 + if (renormalize) {

647 + scale /= topk_sum;

648 + }

649 +

650 + if (lane_id < topk_i32) {

651 + topk_indices[lane_id] = lane_idx;

652 + topk_values[lane_id] = lane_unbiased * scale;

653 }

654 +

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

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

657 #endif

658 }

659

702 -template <typename T, typename BiasT, typename IdxT, ScoringFunc SF>

703 -inline void launch_group_idx_and_topk_kernel(

704 - cudaLaunchConfig_t const& config, T* scores, T* group_scores,

705 - float* topk_values, IdxT* topk_indices, BiasT const* bias,

706 - int64_t const num_tokens, int64_t const n_group, int64_t const topk_group,

707 - int64_t const topk, int64_t const num_experts,

708 - int64_t const num_experts_per_group, bool const renormalize,

709 - double const routed_scaling_factor) {

710 - auto launch = [&](auto* kernel_instance2) {

711 - cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,

712 - topk_values, topk_indices, bias, num_tokens, n_group,

713 - topk_group, topk, num_experts, num_experts_per_group,

714 - renormalize, routed_scaling_factor);

715 - };

716 -

717 - switch (n_group) {

718 - case 4: {

719 - launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 4>);

720 - break;

721 - }

722 - case 8: {

723 - launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 8>);

724 - break;

725 - }

726 - case 16: {

727 - launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 16>);

728 - break;

729 - }

730 - case 32: {

731 - launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF, 32>);

732 - break;

733 - }

734 - default: {

735 - launch(&group_idx_and_topk_idx_kernel<T, BiasT, IdxT, SF>);

736 - break;

737 - }

738 - }

739 -}

740 -

660 template <typename T, typename BiasT, typename IdxT>

742 -void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,

743 - IdxT* topk_indices, BiasT const* bias,

744 - int64_t const num_tokens, int64_t const num_experts,

745 - int64_t const n_group, int64_t const topk_group,

746 - int64_t const topk, bool const renormalize,

747 - double const routed_scaling_factor, int const scoring_func,

748 - bool enable_pdl = false, cudaStream_t const stream = 0) {

749 - int64_t num_cases = num_tokens * n_group;

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

661 +void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices,

662 + BiasT const* bias, int64_t const num_tokens,

663 + int64_t const num_experts, int64_t const n_group,

664 + int64_t const topk_group, int64_t const topk,

665 + bool const renormalize, double const routed_scaling_factor,

666 + int const scoring_func, bool enable_pdl = false,

667 + cudaStream_t const stream = 0) {

668 cudaLaunchConfig_t config;

752 - config.gridDim = topk_with_k2_num_blocks;

753 - config.blockDim = BLOCK_SIZE;

754 - config.dynamicSmemBytes = 0;

669 + // One block per token; one warp per group.

670 + config.gridDim = static_cast<uint32_t>(num_tokens);

671 + config.blockDim = static_cast<uint32_t>(n_group) * WARP_SIZE;

672 + // Dynamic shared memory: WarpSelect staging + per-group topk buffers.

673 + int32_t const num_warps = static_cast<int32_t>(n_group);

674 + size_t const val_bytes =

675 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(T);

676 + size_t const val_bytes_aligned =

677 + warp_topk::round_up_to_multiple_of<256>(val_bytes);

678 + size_t const idx_bytes =

679 + static_cast<size_t>(num_warps) * WARP_SIZE * sizeof(int32_t);

680 + size_t const internal_bytes = val_bytes_aligned + idx_bytes;

681 + size_t const extra_bytes = 16 + static_cast<size_t>(n_group) * sizeof(T);

682 + config.dynamicSmemBytes = internal_bytes + extra_bytes;

683 config.stream = stream;

684 cudaLaunchAttribute attrs[1];

685 attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;

@@ -759,66 +687,35 @@ void invokeNoAuxTc(T* scores, T* group_scores, float* topk_values,

687 config.numAttrs = 1;

688 config.attrs = attrs;

689 auto const sf = static_cast<ScoringFunc>(scoring_func);

762 - int64_t const num_experts_per_group = num_experts / n_group;

763 - auto launch_topk_with_k2 = [&](auto* kernel_instance1) {

764 - cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores, bias,

765 - num_tokens, num_cases, n_group, num_experts_per_group);

766 - };

690 switch (sf) {

691 case SCORING_NONE: {

769 - auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_NONE>;

770 - launch_topk_with_k2(kernel_instance1);

771 - break;

692 + auto* kernel_instance =

693 + &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_NONE>;

694 + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,

695 + topk_indices, bias, num_tokens, num_experts, n_group,

696 + topk_group, topk, renormalize, routed_scaling_factor);

697 + return;

698 }

699 case SCORING_SIGMOID: {

774 - auto* kernel_instance1 = &topk_with_k2_kernel<T, BiasT, SCORING_SIGMOID>;

775 - launch_topk_with_k2(kernel_instance1);

776 - break;

700 + auto* kernel_instance =

701 + &grouped_topk_fused_kernel<T, BiasT, IdxT, SCORING_SIGMOID>;

702 + cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values,

703 + topk_indices, bias, num_tokens, num_experts, n_group,

704 + topk_group, topk, renormalize, routed_scaling_factor);

705 + return;

706 }

707 default:

708 // should be guarded by higher level checks.

709 TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");

710 }

782 -

783 - int64_t topk_with_k_group_num_blocks =

784 - (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;

785 - size_t dynamic_smem_in_bytes =

786 - warp_topk::calc_smem_size_for_block_wide<T, int32_t>(NUM_WARPS_PER_BLOCK,

787 - topk);

788 - config.gridDim = topk_with_k_group_num_blocks;

789 - config.blockDim = BLOCK_SIZE;

790 - config.dynamicSmemBytes = dynamic_smem_in_bytes;

791 - config.stream = stream;

792 - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;

793 - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;

794 - config.numAttrs = 1;

795 - config.attrs = attrs;

796 - switch (sf) {

797 - case SCORING_NONE: {

798 - launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_NONE>(

799 - config, scores, group_scores, topk_values, topk_indices, bias,

800 - num_tokens, n_group, topk_group, topk, num_experts,

801 - num_experts_per_group, renormalize, routed_scaling_factor);

802 - break;

803 - }

804 - case SCORING_SIGMOID: {

805 - launch_group_idx_and_topk_kernel<T, BiasT, IdxT, SCORING_SIGMOID>(

806 - config, scores, group_scores, topk_values, topk_indices, bias,

807 - num_tokens, n_group, topk_group, topk, num_experts,

808 - num_experts_per_group, renormalize, routed_scaling_factor);

809 - break;

810 - }

811 - default:

812 - TORCH_CHECK(false, "Unsupported scoring_func in invokeNoAuxTc");

813 - }

711 }

712

816 -#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT) \

817 - template void invokeNoAuxTc<T, BiasT, IdxT>( \

818 - T * scores, T * group_scores, float* topk_values, IdxT* topk_indices, \

819 - BiasT const* bias, int64_t const num_tokens, int64_t const num_experts, \

820 - int64_t const n_group, int64_t const topk_group, int64_t const topk, \

821 - bool const renormalize, double const routed_scaling_factor, \

713 +#define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT) \

714 + template void invokeNoAuxTc<T, BiasT, IdxT>( \

715 + T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \

716 + int64_t const num_tokens, int64_t const num_experts, \

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

718 + bool const renormalize, double const routed_scaling_factor, \

719 int const scoring_func, bool enable_pdl, cudaStream_t const stream);

720

721 INSTANTIATE_NOAUX_TC(float, float, int32_t);

@@ -843,17 +740,21 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

740 int64_t num_tokens = input_size[0];

741 int64_t num_experts = input_size[1];

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

743 + TORCH_CHECK(n_group > 0, "n_group must be positive");

744 + TORCH_CHECK(topk > 0, "topk must be positive");

745 + TORCH_CHECK(topk_group > 0, "topk_group must be positive");

746 + TORCH_CHECK(topk_group <= n_group, "topk_group must be <= n_group");

747 TORCH_CHECK(num_experts % n_group == 0,

748 "num_experts should be divisible by n_group");

749 TORCH_CHECK(n_group <= 32,

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

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

752 + TORCH_CHECK(topk <= topk_group * (num_experts / n_group),

753 + "topk must be <= topk_group * (num_experts / n_group)");

754 TORCH_CHECK(scoring_func == vllm::moe::SCORING_NONE ||

755 scoring_func == vllm::moe::SCORING_SIGMOID,

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

757

855 - torch::Tensor group_scores = torch::empty(

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

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

759 torch::Tensor topk_values = torch::empty(

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

@@ -868,7 +769,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

769 case torch::kFloat16: \

770 vllm::moe::invokeNoAuxTc<T, half, IdxT>( \

771 reinterpret_cast<T*>(scores.mutable_data_ptr()), \

871 - reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \

772 reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \

773 reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \

774 reinterpret_cast<half const*>(bias.data_ptr()), num_tokens, \

@@ -879,7 +779,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

779 case torch::kFloat32: \

780 vllm::moe::invokeNoAuxTc<T, float, IdxT>( \

781 reinterpret_cast<T*>(scores.mutable_data_ptr()), \

882 - reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \

782 reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \

783 reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \

784 reinterpret_cast<float const*>(bias.data_ptr()), num_tokens, \

@@ -890,7 +789,6 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(

789 case torch::kBFloat16: \

790 vllm::moe::invokeNoAuxTc<T, __nv_bfloat16, IdxT>( \

791 reinterpret_cast<T*>(scores.mutable_data_ptr()), \

893 - reinterpret_cast<T*>(group_scores.mutable_data_ptr()), \

792 reinterpret_cast<float*>(topk_values.mutable_data_ptr()), \

793 reinterpret_cast<IdxT*>(topk_indices.mutable_data_ptr()), \

794 reinterpret_cast<__nv_bfloat16 const*>(bias.data_ptr()), \

[Perf] Optimize grouped topk kernel, 1.2%~2% E2E Throughput improvement (#32058)