| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.platforms import current_platform |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | |
| 9 | _TOPK = 6 |
| 10 | |
| 11 | # Adapted from: |
| 12 | # https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/moe_fused_gate.py |
| 13 | |
| 14 | |
| 15 | def can_use_dsv4_topk( |
| 16 | gating_output: torch.Tensor, |
| 17 | correction_bias: torch.Tensor | None, |
| 18 | topk: int, |
| 19 | renormalize: bool, |
| 20 | indices_dtype: torch.dtype, |
| 21 | ) -> bool: |
| 22 | return ( |
| 23 | current_platform.is_cuda() |
| 24 | and gating_output.dtype == torch.float32 |
| 25 | and gating_output.ndim == 2 |
| 26 | and gating_output.shape[1] in (256, 384) |
| 27 | and gating_output.is_contiguous() |
| 28 | and correction_bias is not None |
| 29 | and correction_bias.dtype == torch.float32 |
| 30 | and correction_bias.shape == (gating_output.shape[1],) |
| 31 | and correction_bias.is_contiguous() |
| 32 | and topk == _TOPK |
| 33 | and renormalize |
| 34 | and indices_dtype in (torch.int32, torch.uint32, torch.int64) |
| 35 | ) |
| 36 | |
| 37 | |
| 38 | if current_platform.is_cuda(): |
| 39 | |
| 40 | @triton.jit |
| 41 | def _dsv4_topk_kernel( |
| 42 | gating_output_ptr, |
| 43 | correction_bias_ptr, |
| 44 | topk_weights_ptr, |
| 45 | topk_ids_ptr, |
| 46 | routed_scaling_factor, |
| 47 | input_ids_ptr, |
| 48 | bias_vl_ptr, |
| 49 | image_sentinel_lo, |
| 50 | NUM_EXPERTS: tl.constexpr, |
| 51 | BLOCK_N: tl.constexpr, |
| 52 | HAS_VL: tl.constexpr, |
| 53 | launch_pdl: tl.constexpr, |
| 54 | ): |
| 55 | row = tl.program_id(0) |
| 56 | expert_offsets = tl.arange(0, BLOCK_N) |
| 57 | expert_mask = expert_offsets < NUM_EXPERTS |
| 58 | bias = tl.load( |
| 59 | correction_bias_ptr + expert_offsets, mask=expert_mask, other=0.0 |
| 60 | ).to(tl.float32) |
| 61 | if HAS_VL: |
| 62 | # Image tokens carry five consecutive in-vocab sentinel ids |
| 63 | # starting at image_sentinel_lo and use bias_vl for expert |
| 64 | # selection instead of the regular correction bias. Ids above the |
| 65 | # sentinel block are regular special tokens and must not match. |
| 66 | token_id = tl.load(input_ids_ptr + row).to(tl.int64) |
| 67 | bias_vl = tl.load( |
| 68 | bias_vl_ptr + expert_offsets, mask=expert_mask, other=0.0 |
| 69 | ).to(tl.float32) |
| 70 | is_image = (token_id >= image_sentinel_lo) & ( |
| 71 | token_id < image_sentinel_lo + 5 |
| 72 | ) |
| 73 | bias = tl.where(is_image, bias_vl, bias) |
| 74 | |
| 75 | if launch_pdl: |
| 76 | tl.extra.cuda.gdc_wait() |
| 77 | |
| 78 | logits = tl.load( |
| 79 | gating_output_ptr + row * NUM_EXPERTS + expert_offsets, |
| 80 | mask=expert_mask, |
| 81 | other=0.0, |
| 82 | ).to(tl.float32) |
| 83 | weights = tl.sqrt(tl.where(logits > 20.0, logits, tl.log(1.0 + tl.exp(logits)))) |
| 84 | current = tl.where(expert_mask, weights + bias, -float("inf")) |
| 85 | current = tl.where(current == current, current, -1e30) |
| 86 | |
| 87 | topk_offsets = tl.arange(0, 8) |
| 88 | selected_weights = tl.zeros([8], dtype=tl.float32) |
| 89 | selected_ids = tl.zeros([8], dtype=tl.int32) |
| 90 | for slot in tl.static_range(6): |
| 91 | max_value = tl.max(current, axis=0) |
| 92 | candidate = tl.where(current == max_value, expert_offsets, NUM_EXPERTS) |
| 93 | expert_id = tl.min(candidate, axis=0).to(tl.int32) |
| 94 | selected_weight = tl.sum( |
| 95 | tl.where(expert_offsets == expert_id, weights, 0.0), axis=0 |
| 96 | ) |
| 97 | is_slot = topk_offsets == slot |
| 98 | selected_weights = tl.where(is_slot, selected_weight, selected_weights) |
| 99 | selected_ids = tl.where(is_slot, expert_id, selected_ids) |
| 100 | current = tl.where(expert_offsets == expert_id, -float("inf"), current) |
| 101 | |
| 102 | weight_sum = tl.sum(selected_weights, axis=0) |
| 103 | selected_weights *= routed_scaling_factor / tl.where( |
| 104 | weight_sum > 0.0, weight_sum, 1.0 |
| 105 | ) |
| 106 | output_mask = topk_offsets < 6 |
| 107 | output_offsets = row * 6 + topk_offsets |
| 108 | |
| 109 | if launch_pdl: |
| 110 | tl.extra.cuda.gdc_launch_dependents() |
| 111 | |
| 112 | tl.store(topk_weights_ptr + output_offsets, selected_weights, mask=output_mask) |
| 113 | tl.store(topk_ids_ptr + output_offsets, selected_ids, mask=output_mask) |
| 114 | |
| 115 | |
| 116 | def dsv4_topk( |
| 117 | gating_output: torch.Tensor, |
| 118 | correction_bias: torch.Tensor, |
| 119 | indices_dtype: torch.dtype, |
| 120 | routed_scaling_factor: float, |
| 121 | input_ids: torch.Tensor | None = None, |
| 122 | bias_vl: torch.Tensor | None = None, |
| 123 | image_sentinel_lo: int = 0, |
| 124 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 125 | num_tokens, num_experts = gating_output.shape |
| 126 | has_vl = bias_vl is not None and image_sentinel_lo > 0 |
| 127 | if bias_vl is not None: |
| 128 | assert input_ids is not None, "bias_vl routing requires input_ids" |
| 129 | assert bias_vl.dtype == torch.float32 and bias_vl.is_contiguous() |
| 130 | assert bias_vl.shape == (num_experts,) |
| 131 | assert input_ids.is_contiguous() |
| 132 | shape = (num_tokens, _TOPK) |
| 133 | topk_weights = gating_output.new_empty(shape, dtype=torch.float32) |
| 134 | topk_ids = gating_output.new_empty(shape, dtype=indices_dtype) |
| 135 | if num_tokens > 0: |
| 136 | _dsv4_topk_kernel[(num_tokens,)]( |
| 137 | gating_output, |
| 138 | correction_bias, |
| 139 | topk_weights, |
| 140 | topk_ids, |
| 141 | routed_scaling_factor, |
| 142 | input_ids, |
| 143 | bias_vl, |
| 144 | image_sentinel_lo, |
| 145 | NUM_EXPERTS=num_experts, |
| 146 | BLOCK_N=triton.next_power_of_2(num_experts), |
| 147 | HAS_VL=has_vl, |
| 148 | num_warps=1, |
| 149 | launch_pdl=current_platform.is_arch_support_pdl(), |
| 150 | ) |
| 151 | return topk_weights, topk_ids |