@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 15ba07ef25
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 15, 16:00 UTC

[Minor] Fused experts refactor (#15914)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Apr 3, 2025 at 13:19 UTC 15ba07ef25d079c8569959e3ede637f432c51d98
8 files changed +790 -737
tests/kernels/test_block_fp8.py
+6 -3
@@ -9,8 +9,11 @@ import torch
9 from vllm.config import VllmConfig, set_current_vllm_config
10 from vllm.model_executor.layers.activation import SiluAndMul
11 from vllm.model_executor.layers.fused_moe import fused_moe
12 -from vllm.model_executor.layers.fused_moe.fused_moe import (
13 - deep_gemm_moe_fp8, fused_topk, moe_align_block_size)
12 +from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
13 + deep_gemm_moe_fp8)
14 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
15 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
16 + moe_align_block_size)
17 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 per_token_group_quant_fp8, w8a8_block_fp8_matmul)
19 from vllm.platforms import current_platform
@@ -437,7 +440,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
440 pytest.skip(
441 f"Skipping test; bad size m={M}, n={N}, k={K}, topk={topk}, E={E}")
442
440 - if (N <= 512):
443 + if N <= 512:
444 pytest.skip("Skipping N <= 512 until performance issues solved.")
445
446 vllm_config = VllmConfig()
tests/kernels/test_cutlass_moe.py
+8 -8
@@ -4,8 +4,8 @@ import torch
4
5 from vllm import _custom_ops as ops
6 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
7 -from vllm.model_executor.layers.fused_moe.fused_moe import (cutlass_moe_fp8,
8 - fused_experts,
7 +from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
8 +from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts,
9 fused_topk)
10 from vllm.platforms import current_platform
11
@@ -131,9 +131,9 @@ def test_cutlass_moe_no_graph(
131 c_strides2,
132 a1_scale=a_scale1)
133
134 - print(triton_output)
135 - print(cutlass_output)
136 - print("*")
134 + #print(triton_output)
135 + #print(cutlass_output)
136 + #print("*")
137
138 torch.testing.assert_close(triton_output,
139 cutlass_output,
@@ -234,9 +234,9 @@ def test_cutlass_moe_cuda_graph(
234 graph.replay()
235 torch.cuda.synchronize()
236
237 - print(triton_output)
238 - print(cutlass_output)
239 - print("*")
237 + #print(triton_output)
238 + #print(cutlass_output)
239 + #print("*")
240
241 torch.testing.assert_close(triton_output,
242 cutlass_output,
vllm/model_executor/layers/fused_moe/__init__.py
+4 -2
@@ -35,9 +35,11 @@ if HAS_TRITON:
35 # import to register the custom ops
36 import vllm.model_executor.layers.fused_moe.fused_marlin_moe # noqa
37 import vllm.model_executor.layers.fused_moe.fused_moe # noqa
38 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
39 + cutlass_moe_fp8)
40 from vllm.model_executor.layers.fused_moe.fused_moe import (
39 - cutlass_moe_fp8, fused_experts, fused_moe, fused_topk,
40 - get_config_file_name, grouped_topk)
41 + fused_experts, fused_moe, fused_topk, get_config_file_name,
42 + grouped_topk)
43
44 __all__ += [
45 "fused_moe",
vllm/model_executor/layers/fused_moe/cutlass_moe.py new
+144
@@ -0,0 +1,144 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Fused MoE kernel."""
3 +from typing import Optional
4 +
5 +import torch
6 +
7 +from vllm import _custom_ops as ops
8 +
9 +
10 +#TODO make the grouped gemm kernel consistent with scaled gemm kernel
11 +def cutlass_moe_fp8(
12 + a: torch.Tensor,
13 + w1_q: torch.Tensor,
14 + w2_q: torch.Tensor,
15 + w1_scale: torch.Tensor,
16 + w2_scale: torch.Tensor,
17 + topk_weights: torch.Tensor,
18 + topk_ids: torch.Tensor,
19 + ab_strides1: torch.Tensor,
20 + c_strides1: torch.Tensor,
21 + ab_strides2: torch.Tensor,
22 + c_strides2: torch.Tensor,
23 + a1_scale: Optional[torch.Tensor] = None,
24 + a2_scale: Optional[torch.Tensor] = None,
25 + out_dtype: torch.dtype = torch.half,
26 +) -> torch.Tensor:
27 + """
28 + This function computes a a8w8-quantized Mixture of Experts (MoE) layer
29 + using two sets of quantized weights, w1_q and w2_q, and top-k gating
30 + mechanism. The matrix multiplications are implemented with CUTLASS
31 + grouped gemm.
32 +
33 + Parameters:
34 + - a (torch.Tensor): The input tensor to the MoE layer.
35 + Shape: [M, K]
36 + - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
37 + Shape: [num_experts, K, 2N] (the weights are passed transposed)
38 + - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
39 + Shape: [num_experts, N, K] (the weights are passed transposed)
40 + - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
41 + Shape: [num_experts] or [num_experts, 2N]
42 + - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
43 + Shape: [num_experts] or [num_experts, K]
44 + - gating_output (torch.Tensor): The output of the gating operation
45 + (before softmax).
46 + - topk_weights (torch.Tensor): The weights of each token->expert mapping.
47 + - ab_strides1 (torch.Tensor): The input and weights strides of the first
48 + grouped gemm.
49 + - c_strides1 (torch.Tensor): The output strides of the first grouped gemm.
50 + - ab_strides2 (torch.Tensor): The input and weights strides of the second
51 + grouped gemm.
52 + - c_strides2 (torch.Tensor): The output strides of the second grouped gemm.
53 + - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
54 + Shape: scalar or [M]
55 + - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
56 + quantize the intermediate result between the gemms.
57 + Shape: scalar or [M]
58 + - out_dtype (torch.Tensor): The output tensor type.
59 +
60 + Returns:
61 + - torch.Tensor: The fp16 output tensor after applying the MoE layer.
62 + """
63 +
64 + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
65 + assert w1_q.dtype == torch.float8_e4m3fn
66 + assert w2_q.dtype == torch.float8_e4m3fn
67 + assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
68 + assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
69 + assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
70 + assert a1_scale is None or a1_scale.dim(
71 + ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[0] == a.shape[
72 + 0], "Input scale shape mismatch"
73 + assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
74 + 1] == w1_q.shape[2], "W1 scale shape mismatch"
75 + assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
76 + 1] == w2_q.shape[2], "W2 scale shape mismatch"
77 + assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch"
78 + assert w1_q.shape[0] == w1_scale.shape[
79 + 0], "w1 scales expert number mismatch"
80 + assert w1_q.shape[0] == w2_scale.shape[
81 + 0], "w2 scales expert number mismatch"
82 + assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
83 + assert ab_strides1.shape[0] == w1_q.shape[
84 + 0], "AB Strides 1 expert number mismatch"
85 + assert c_strides1.shape[0] == w1_q.shape[
86 + 0], "C Strides 1 expert number mismatch"
87 + assert ab_strides2.shape[0] == w2_q.shape[
88 + 0], "AB Strides 2 expert number mismatch"
89 + assert c_strides2.shape[0] == w2_q.shape[
90 + 0], "C Strides 2 expert number mismatch"
91 + assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
92 +
93 + num_experts = w1_q.size(0)
94 + m = a.size(0)
95 + k = w1_q.size(1)
96 + n = w2_q.size(1)
97 +
98 + topk = topk_ids.size(1)
99 + per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
100 + a2_scale.numel() != 1 if a2_scale is not None else False)
101 +
102 + a_q, a1_scale = ops.scaled_fp8_quant(
103 + a, a1_scale, use_per_token_if_dynamic=per_act_token)
104 + device = a_q.device
105 +
106 + expert_offsets = torch.empty((num_experts + 1),
107 + dtype=torch.int32,
108 + device=device)
109 + problem_sizes1 = torch.empty((num_experts, 3),
110 + dtype=torch.int32,
111 + device=device)
112 + problem_sizes2 = torch.empty((num_experts, 3),
113 + dtype=torch.int32,
114 + device=device)
115 +
116 + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
117 + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
118 +
119 + ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
120 + problem_sizes2, a_map, c_map, num_experts, n,
121 + k)
122 +
123 + rep_a_q = a_q.view(dtype=torch.uint8)[a_map].view(dtype=a_q.dtype)
124 + rep_a1_scales = a1_scale[a_map] if per_act_token else a1_scale
125 +
126 + c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
127 + c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
128 +
129 + ops.cutlass_moe_mm(c1, rep_a_q, w1_q, rep_a1_scales, w1_scale,
130 + expert_offsets[:-1], problem_sizes1, ab_strides1,
131 + ab_strides1, c_strides1)
132 +
133 + intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
134 + torch.ops._C.silu_and_mul(intermediate, c1)
135 +
136 + intemediate_q, a2_scale = ops.scaled_fp8_quant(
137 + intermediate, a2_scale, use_per_token_if_dynamic=per_act_token)
138 +
139 + ops.cutlass_moe_mm(c2, intemediate_q, w2_q, a2_scale, w2_scale,
140 + expert_offsets[:-1], problem_sizes2, ab_strides2,
141 + ab_strides2, c_strides2)
142 +
143 + return (c2[c_map].view(m, topk, k) *
144 + topk_weights.view(m, topk, 1).to(out_dtype)).sum(dim=1)
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py new
+294
@@ -0,0 +1,294 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import importlib.util
3 +from typing import Optional, Tuple
4 +
5 +import torch
6 +
7 +import vllm.envs as envs
8 +from vllm import _custom_ops as ops
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
11 + moe_align_block_size)
12 +from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,
13 + _fp8_quantize,
14 + _resize_cache)
15 +from vllm.utils import round_up
16 +
17 +logger = init_logger(__name__)
18 +
19 +has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
20 +
21 +
22 +def _valid_deep_gemm(hidden_states: torch.Tensor,
23 + w1: torch.Tensor,
24 + w2: torch.Tensor,
25 + expert_map: Optional[torch.Tensor] = None) -> bool:
26 + """
27 + Check if the given problem size is supported by the DeepGemm grouped
28 + gemm kernel. All of M, N, K and the quantization block_shape must be
29 + aligned by `dg.get_m_alignment_for_contiguous_layout()`.
30 + """
31 + if not has_deep_gemm:
32 + return False
33 +
34 + # Lazy import to avoid CUDA initialization problems.
35 + import deep_gemm as dg
36 +
37 + # Expert maps not supported yet.
38 + if expert_map is not None:
39 + return False
40 +
41 + align = dg.get_m_alignment_for_contiguous_layout()
42 + M = hidden_states.shape[0]
43 + _, K, N = w2.shape
44 +
45 + # For now, disable DeepGemm for small N until better permute/unpermute
46 + # ops are available.
47 + if N <= 512:
48 + return False
49 +
50 + if align > M or N % align != 0 or K % align != 0:
51 + return False
52 +
53 + return (hidden_states.is_contiguous() and w1.is_contiguous()
54 + and w2.is_contiguous())
55 +
56 +
57 +def _moe_permute(
58 + curr_hidden_states: torch.Tensor,
59 + a1q_scale: Optional[torch.Tensor],
60 + curr_topk_ids: torch.Tensor,
61 + global_num_experts: int,
62 + expert_map: Optional[torch.Tensor],
63 + block_m: int,
64 +) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
65 + Optional[torch.Tensor]]:
66 + """
67 + Determine the sorted_token_ids, expert_ids for the given problem size.
68 + Permute the hidden states and scales according to `sorted_token_ids`.
69 + """
70 + top_k_num = curr_topk_ids.shape[1]
71 +
72 + tokens_in_chunk, _ = curr_hidden_states.shape
73 +
74 + sorted_token_ids, expert_ids, num_tokens_post_padded = (
75 + moe_align_block_size(curr_topk_ids,
76 + block_m,
77 + global_num_experts,
78 + expert_map,
79 + pad_sorted_ids=True))
80 +
81 + inv_perm: Optional[torch.Tensor] = None
82 +
83 + num_tokens = top_k_num * tokens_in_chunk
84 + sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
85 + expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
86 + inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
87 +
88 + # Permute according to sorted token ids.
89 + curr_hidden_states = _fp8_perm(curr_hidden_states,
90 + sorted_token_ids // top_k_num)
91 +
92 + if a1q_scale is not None:
93 + a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
94 +
95 + return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
96 + inv_perm)
97 +
98 +
99 +def _moe_unpermute_and_reduce(
100 + out: torch.Tensor,
101 + curr_hidden: torch.Tensor,
102 + inv_perm: Optional[torch.Tensor],
103 + topk_weight: torch.Tensor,
104 +) -> None:
105 + """
106 + Unpermute the final result and apply topk_weights, then perform the final
107 + reduction on the hidden states.
108 + """
109 + M, topk = topk_weight.shape
110 + K = curr_hidden.shape[1]
111 + curr_hidden = curr_hidden[inv_perm, ...]
112 + curr_hidden = curr_hidden.view(-1, topk, K)
113 + curr_hidden.mul_(topk_weight.view(M, -1, 1))
114 + ops.moe_sum(curr_hidden, out)
115 +
116 +
117 +def deep_gemm_moe_fp8(
118 + hidden_states: torch.Tensor,
119 + w1: torch.Tensor,
120 + w2: torch.Tensor,
121 + w1_scale: torch.Tensor,
122 + w2_scale: torch.Tensor,
123 + topk_weights: torch.Tensor,
124 + topk_ids: torch.Tensor,
125 + inplace: bool = False,
126 + activation: str = "silu",
127 + global_num_experts: int = -1,
128 + expert_map: Optional[torch.Tensor] = None,
129 + a1_scale: Optional[torch.Tensor] = None,
130 + a2_scale: Optional[torch.Tensor] = None,
131 +) -> torch.Tensor:
132 + """
133 + This function computes a a8w8-quantized Mixture of Experts (MoE) layer
134 + using two sets of quantized weights, w1_q and w2_q, and top-k gating
135 + mechanism. The matrix multiplications are implemented with DeepGemm
136 + grouped gemm.
137 +
138 + Parameters:
139 + - hidden_states (torch.Tensor): The input tensor to the MoE layer.
140 + Shape: [M, K]
141 + - w1 (torch.Tensor): The first set of fp8 quantized expert weights.
142 + Shape: [num_experts, K, 2N] (the weights are passed transposed)
143 + - w2 (torch.Tensor): The second set of fp8 quantized expert weights.
144 + Shape: [num_experts, N, K] (the weights are passed transposed)
145 + - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
146 + Shape: [num_experts] or [num_experts, 2N]
147 + - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
148 + Shape: [num_experts] or [num_experts, K]
149 + - topk_weights (torch.Tensor): The weights of each token->expert mapping.
150 + - topk_ids (torch.Tensor): The token->expert mapping for topk_weights.
151 + - inplace (bool): If True, perform the operation in-place.
152 + Defaults to False.
153 + - activation (str): The activation function to apply after the first
154 + MoE layer.
155 + - global_num_experts (int): The total number of experts in the global
156 + expert space.
157 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
158 + from the global expert space to the local expert space of the expert
159 + parallel shard.
160 + - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
161 + Shape: scalar or [M]
162 + - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
163 + quantize the intermediate result between the gemms.
164 + Shape: scalar or [M]
165 +
166 + Returns:
167 + - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
168 + """
169 + # Lazy import to avoid CUDA initialization problems.
170 + import deep_gemm as dg
171 +
172 + assert expert_map is None, "Expert maps not supported yet"
173 +
174 + assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
175 +
176 + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
177 + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
178 + assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
179 + assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
180 + assert hidden_states.dtype in [
181 + torch.float32, torch.float16, torch.bfloat16
182 + ]
183 + assert w1.dtype == torch.float8_e4m3fn
184 + assert w2.dtype == torch.float8_e4m3fn
185 + assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
186 + assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
187 + assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
188 + assert a1_scale is None or a1_scale.dim(
189 + ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[
190 + 0] == hidden_states.shape[0], "Input scale shape mismatch"
191 + assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
192 +
193 + num_tokens, _ = hidden_states.shape
194 + E, N, _ = w1.shape
195 + K = w2.shape[1]
196 + if global_num_experts == -1:
197 + global_num_experts = E
198 +
199 + # We execute the fused_moe kernel in chunks to circumvent this issue:
200 + # https://github.com/vllm-project/vllm/issues/5938
201 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
202 +
203 + assert _valid_deep_gemm(hidden_states, w1, w2, expert_map)
204 +
205 + if inplace:
206 + out_hidden_states = hidden_states
207 + else:
208 + out_hidden_states = torch.empty_like(hidden_states)
209 +
210 + block_m = dg.get_m_alignment_for_contiguous_layout()
211 + block_shape = [block_m, block_m]
212 +
213 + assert w1_scale is not None
214 + assert w2_scale is not None
215 +
216 + # We attempt to transpose and align offline in Fp8MoEMethod, in which
217 + # case these calls will be nops. Otherwise, they'll be performed every
218 + # time the layer is executed.
219 + w1_scale = dg.get_col_major_tma_aligned_tensor(w1_scale).contiguous()
220 + w2_scale = dg.get_col_major_tma_aligned_tensor(w2_scale).contiguous()
221 +
222 + M_sum = topk_ids.numel() + global_num_experts * (block_m - 1)
223 + M_sum = round_up(M_sum, block_m)
224 +
225 + num_chunks = (num_tokens // CHUNK_SIZE) + 1
226 +
227 + # We can reuse the memory between cache1 and cache3 because by the time
228 + # we need cache3, we're done with cache1
229 + workspace13 = torch.empty(M_sum * max(N, K),
230 + device=hidden_states.device,
231 + dtype=hidden_states.dtype)
232 +
233 + workspace1 = workspace13[:M_sum * N].view(M_sum, N)
234 + workspace2 = torch.empty((M_sum, N // 2),
235 + device=hidden_states.device,
236 + dtype=hidden_states.dtype)
237 + workspace3 = workspace13[:M_sum * K].view(M_sum, K)
238 +
239 + for chunk in range(num_chunks):
240 + begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
241 + min((chunk + 1) * CHUNK_SIZE,
242 + num_tokens))
243 + curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
244 + tokens_in_chunk, _ = curr_hidden_states.shape
245 +
246 + if tokens_in_chunk == 0:
247 + break
248 +
249 + curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
250 + curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
251 +
252 + a1q_scale: Optional[torch.Tensor] = None
253 +
254 + qcurr_hidden_states, a1q_scale = _fp8_quantize(curr_hidden_states,
255 + a1_scale, block_shape)
256 +
257 + (qcurr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
258 + inv_perm) = _moe_permute(qcurr_hidden_states, a1q_scale,
259 + curr_topk_ids, global_num_experts,
260 + expert_map, block_m)
261 +
262 + # Adjust the intermediate cache size and config for the last chunk.
263 + # Note that in most cases we only have one chunk so the cache size
264 + # and config are already set correctly and do not need to be adjusted.
265 + if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
266 + curr_M = sorted_token_ids.numel()
267 + workspace1 = _resize_cache(workspace1, (curr_M, N))
268 + workspace2 = _resize_cache(workspace2, (curr_M, N // 2))
269 + workspace3 = _resize_cache(workspace3, (curr_M, K))
270 +
271 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
272 + (qcurr_hidden_states, a1q_scale), (w1, w1_scale), workspace1,
273 + expert_ids)
274 +
275 + if activation == "silu":
276 + torch.ops._C.silu_and_mul(workspace2, workspace1.view(-1, N))
277 + elif activation == "gelu":
278 + torch.ops._C.gelu_and_mul(workspace2, workspace1.view(-1, N))
279 + else:
280 + raise ValueError(f"Unsupported FusedMoe activation: {activation}")
281 +
282 + a2q_scale: Optional[torch.Tensor] = None
283 +
284 + qworkspace2, a2q_scale = _fp8_quantize(workspace2, a2_scale,
285 + block_shape)
286 +
287 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
288 + (qworkspace2, a2q_scale), (w2, w2_scale), workspace3, expert_ids)
289 +
290 + _moe_unpermute_and_reduce(
291 + out_hidden_states[begin_chunk_idx:end_chunk_idx],
292 + workspace3.view(*workspace3.shape), inv_perm, curr_topk_weights)
293 +
294 + return out_hidden_states
vllm/model_executor/layers/fused_moe/fused_moe.py
+43 -724
@@ -1,10 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """Fused MoE kernel."""
3 import functools
4 -import importlib.util
4 import json
5 import os
7 -from math import prod
6 from typing import Any, Callable, Dict, List, Optional, Tuple
7
8 import torch
@@ -14,10 +12,13 @@ import triton.language as tl
12 import vllm.envs as envs
13 from vllm import _custom_ops as ops
14 from vllm.logger import init_logger
17 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 - per_token_group_quant_fp8)
15 +from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
16 + _valid_deep_gemm, deep_gemm_moe_fp8)
17 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
18 + moe_align_block_size)
19 +from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize
20 from vllm.platforms import current_platform
20 -from vllm.utils import direct_register_custom_op, round_up
21 +from vllm.utils import direct_register_custom_op
22
23 from .rocm_aiter_fused_moe import (is_rocm_aiter_moe_enabled,
24 rocm_aiter_fused_experts,
@@ -25,8 +26,6 @@ from .rocm_aiter_fused_moe import (is_rocm_aiter_moe_enabled,
26
27 logger = init_logger(__name__)
28
28 -has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
29 -
29
30 @triton.jit
31 def write_zeros_to_output(c_ptr, stride_cm, stride_cn, pid_n, N, offs_token,
@@ -443,300 +442,13 @@ def fused_moe_kernel(
442 tl.store(c_ptrs, accumulator, mask=c_mask)
443
444
446 -def ceil_div(a, b):
447 - return (a + b - 1) // b
448 -
449 -
450 -@triton.jit
451 -def moe_align_block_size_stage1(
452 - topk_ids_ptr,
453 - tokens_cnts_ptr,
454 - num_experts: tl.constexpr,
455 - numel: tl.constexpr,
456 - tokens_per_thread: tl.constexpr,
457 -):
458 - pid = tl.program_id(0)
459 -
460 - start_idx = pid * tokens_per_thread
461 -
462 - off_c = (pid + 1) * num_experts
463 -
464 - for i in range(tokens_per_thread):
465 - if start_idx + i < numel:
466 - idx = tl.load(topk_ids_ptr + start_idx + i)
467 - token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
468 - tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
469 -
470 -
471 -@triton.jit
472 -def moe_align_block_size_stage2(
473 - tokens_cnts_ptr,
474 - num_experts: tl.constexpr,
475 -):
476 - pid = tl.program_id(0)
477 -
478 - last_cnt = 0
479 - for i in range(1, num_experts + 1):
480 - token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
481 - last_cnt = last_cnt + token_cnt
482 - tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
483 -
484 -
485 -@triton.jit
486 -def moe_align_block_size_stage3(
487 - total_tokens_post_pad_ptr,
488 - tokens_cnts_ptr,
489 - cumsum_ptr,
490 - num_experts: tl.constexpr,
491 - block_size: tl.constexpr,
492 -):
493 - last_cumsum = 0
494 - off_cnt = num_experts * num_experts
495 - for i in range(1, num_experts + 1):
496 - token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
497 - last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
498 - tl.store(cumsum_ptr + i, last_cumsum)
499 - tl.store(total_tokens_post_pad_ptr, last_cumsum)
500 -
501 -
502 -@triton.jit
503 -def moe_align_block_size_stage4(
504 - topk_ids_ptr,
505 - sorted_token_ids_ptr,
506 - expert_ids_ptr,
507 - tokens_cnts_ptr,
508 - cumsum_ptr,
509 - num_experts: tl.constexpr,
510 - block_size: tl.constexpr,
511 - numel: tl.constexpr,
512 - tokens_per_thread: tl.constexpr,
513 -):
514 - pid = tl.program_id(0)
515 - start_idx = tl.load(cumsum_ptr + pid)
516 - end_idx = tl.load(cumsum_ptr + pid + 1)
517 -
518 - for i in range(start_idx, end_idx, block_size):
519 - tl.store(expert_ids_ptr + i // block_size, pid)
520 -
521 - start_idx = pid * tokens_per_thread
522 - off_t = pid * num_experts
523 -
524 - for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread,
525 - numel)):
526 - expert_id = tl.load(topk_ids_ptr + i)
527 - token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
528 - rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
529 - tl.store(sorted_token_ids_ptr + rank_post_pad, i)
530 - tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
531 -
532 -
533 -# Triton implementation based on:
534 -# https://github.com/sgl-project/sglang/commit/ba5112ff691d791a9e38c6c71f59324a5fcb49d0
535 -def moe_align_block_size_triton(
536 - topk_ids: torch.Tensor,
537 - num_experts: int,
538 - block_size: int,
539 - sorted_token_ids: torch.Tensor,
540 - expert_ids: torch.Tensor,
541 - num_tokens_post_pad: torch.Tensor,
542 -) -> None:
543 - numel = topk_ids.numel()
544 - grid = (num_experts, )
545 - tokens_cnts = torch.zeros((num_experts + 1, num_experts),
546 - dtype=torch.int32,
547 - device=topk_ids.device)
548 - cumsum = torch.zeros((num_experts + 1, ),
549 - dtype=torch.int32,
550 - device=topk_ids.device)
551 - tokens_per_thread = ceil_div(numel, num_experts)
552 -
553 - moe_align_block_size_stage1[grid](
554 - topk_ids,
555 - tokens_cnts,
556 - num_experts,
557 - numel,
558 - tokens_per_thread,
559 - )
560 - moe_align_block_size_stage2[grid](
561 - tokens_cnts,
562 - num_experts,
563 - )
564 - moe_align_block_size_stage3[(1, )](
565 - num_tokens_post_pad,
566 - tokens_cnts,
567 - cumsum,
568 - num_experts,
569 - block_size,
570 - )
571 - moe_align_block_size_stage4[grid](
572 - topk_ids,
573 - sorted_token_ids,
574 - expert_ids,
575 - tokens_cnts,
576 - cumsum,
577 - num_experts,
578 - block_size,
579 - numel,
580 - tokens_per_thread,
581 - )
582 -
583 -
584 -def moe_align_block_size(
585 - topk_ids: torch.Tensor,
586 - block_size: int,
587 - num_experts: int,
588 - expert_map: Optional[torch.Tensor] = None,
589 - pad_sorted_ids: bool = False
590 -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
591 - """
592 - Aligns the token distribution across experts to be compatible with block
593 - size for matrix multiplication.
594 -
595 - Parameters:
596 - - topk_ids: A tensor of shape [total_tokens, top_k] representing the
597 - top-k expert indices for each token.
598 - - block_size: The block size used in block matrix multiplication.
599 - - num_experts: The total number of experts.
600 - - expert_map: A tensor of shape [num_experts] that maps the expert index
601 - from the global space to the local index space of the current
602 - expert parallel shard. If the expert is not in the current expert
603 - parallel shard, the mapping is set to -1.
604 - - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
605 - should be padded to a multiple of block_size,
606 -
607 - Returns:
608 - - sorted_token_ids: A tensor containing the sorted token indices according
609 - to their allocated expert.
610 - - expert_ids: A tensor indicating the assigned expert index for each block.
611 - - num_tokens_post_padded: The total number of tokens after padding,
612 - ensuring divisibility by block_size.
613 -
614 - This function pads the number of tokens that each expert needs to process
615 - so that it is divisible by block_size.
616 - Padding ensures that during block matrix multiplication, the dimensions
617 - align correctly.
618 -
619 - Example:
620 - Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
621 - block_size = 4, and num_experts = 4:
622 - - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
623 - with each expert needing to process 3 tokens.
624 - - As block_size is 4, we pad 1 token for each expert.
625 - - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
626 - - Then append padding tokens [12, 12, 12, 12] for each block.
627 - - After sorting by expert index, we obtain token_ids
628 - [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
629 - Tokens 12 are non-existent (padding) and are ignored in
630 - the subsequent matrix multiplication.
631 - - The padding ensures that the total number of tokens is now divisible
632 - by block_size for proper block matrix operations.
633 - """
634 - max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
635 - if pad_sorted_ids:
636 - max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
637 - sorted_ids = torch.empty((max_num_tokens_padded, ),
638 - dtype=torch.int32,
639 - device=topk_ids.device)
640 - sorted_ids.fill_(topk_ids.numel())
641 - max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
642 - # Expert ids must be zeroed out to prevent index out of bounds error while
643 - # mapping global expert ids to local expert ids in expert parallelism.
644 - expert_ids = torch.zeros((max_num_m_blocks, ),
645 - dtype=torch.int32,
646 - device=topk_ids.device)
647 - num_tokens_post_pad = torch.empty((1),
648 - dtype=torch.int32,
649 - device=topk_ids.device)
650 - if num_experts >= 224:
651 - if envs.VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON or num_experts != 256:
652 - moe_align_block_size_triton(
653 - topk_ids,
654 - num_experts,
655 - block_size,
656 - sorted_ids,
657 - expert_ids,
658 - num_tokens_post_pad,
659 - )
660 - else:
661 - # Currently requires num_experts=256
662 - ops.sgl_moe_align_block_size(
663 - topk_ids,
664 - num_experts,
665 - block_size,
666 - sorted_ids,
667 - expert_ids,
668 - num_tokens_post_pad,
669 - )
670 - else:
671 - ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
672 - expert_ids, num_tokens_post_pad)
673 - if expert_map is not None:
674 - expert_ids = expert_map[expert_ids]
675 -
676 - return sorted_ids, expert_ids, num_tokens_post_pad
677 -
678 -
679 -def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
680 - w2: torch.Tensor,
681 - expert_map: Optional[torch.Tensor]) -> bool:
682 - """
683 - Check if the given problem size is supported by the DeepGemm grouped
684 - gemm kernel. All of M, N, K and the quantization block_shape must be
685 - aligned by `dg.get_m_alignment_for_contiguous_layout()`.
686 - """
687 - if not has_deep_gemm:
688 - return False
689 -
690 - # Lazy import to avoid CUDA initialization problems.
691 - import deep_gemm as dg
692 -
693 - # Expert maps not supported yet.
694 - if expert_map is not None:
695 - return False
696 -
697 - align = dg.get_m_alignment_for_contiguous_layout()
698 - M = hidden_states.shape[0]
699 - _, K, N = w2.shape
700 -
701 - # For now, disable DeepGemm for small N until better permute/unpermute
702 - # ops are available.
703 - if N <= 512:
704 - return False
705 -
706 - if align > M or N % align != 0 or K % align != 0:
707 - return False
708 -
709 - return (hidden_states.is_contiguous() and w1.is_contiguous()
710 - and w2.is_contiguous())
711 -
712 -
713 -def _fp8_quantize(
714 - A: torch.Tensor,
715 - A_scale: Optional[torch.Tensor],
716 - block_shape: Optional[List[int]],
717 -) -> Tuple[torch.Tensor, torch.Tensor]:
718 - """
719 - Perform fp8 quantization on the inputs. If a block_shape
720 - is provided, the output will be blocked.
721 - """
722 - if block_shape is None:
723 - A, A_scale = ops.scaled_fp8_quant(A, A_scale)
724 - else:
725 - assert len(block_shape) == 2
726 - _, block_k = block_shape[0], block_shape[1]
727 - A, A_scale = per_token_group_quant_fp8(A, block_k)
728 - assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
729 - return A, A_scale
730 -
731 -
445 def invoke_fused_moe_kernel(A: torch.Tensor,
446 B: torch.Tensor,
447 C: torch.Tensor,
448 A_scale: Optional[torch.Tensor],
449 B_scale: Optional[torch.Tensor],
450 B_zp: Optional[torch.Tensor],
738 - topk_weights: torch.Tensor,
739 - topk_ids: torch.Tensor,
451 + topk_weights: Optional[torch.Tensor],
452 sorted_token_ids: torch.Tensor,
453 expert_ids: torch.Tensor,
454 num_tokens_post_padded: torch.Tensor,
@@ -748,7 +460,8 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
460 use_int8_w8a16: bool,
461 use_int4_w4a16: bool,
462 block_shape: Optional[List[int]] = None) -> None:
751 - assert topk_weights.stride(1) == 1
463 + assert topk_weights is not None or not mul_routed_weight
464 + assert topk_weights is None or topk_weights.stride(1) == 1
465 assert sorted_token_ids.stride(0) == 1
466
467 if use_fp8_w8a8:
@@ -765,6 +478,9 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
478 assert A_scale is None
479 assert B_scale is None
480
481 + M = A.shape[0]
482 + num_tokens = M * top_k
483 +
484 EM = sorted_token_ids.shape[0]
485 if A.shape[0] < config["BLOCK_SIZE_M"]:
486 # optimize for small batch_size.
@@ -782,7 +498,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
498 assert B_zp is None or B_zp.ndim == 3
499
500 use_moe_wna16_cuda = should_moe_wna16_use_cuda(
785 - num_valid_tokens=topk_ids.numel(),
501 + num_valid_tokens=num_tokens,
502 group_size=block_shape[1],
503 num_experts=B.shape[0],
504 bit=4 if use_int4_w4a16 else 8)
@@ -790,12 +506,12 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
506 config.update(
507 get_moe_wna16_block_config(config=config,
508 use_moe_wna16_cuda=use_moe_wna16_cuda,
793 - num_valid_tokens=topk_ids.numel(),
509 + num_valid_tokens=num_tokens,
510 size_k=A.shape[1],
511 size_n=B.shape[1],
512 num_experts=B.shape[1],
513 group_size=block_shape[1],
798 - real_top_k=topk_ids.shape[1],
514 + real_top_k=top_k,
515 block_size_m=config["BLOCK_SIZE_M"]))
516
517 if use_moe_wna16_cuda:
@@ -821,7 +537,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
537 B.shape[1],
538 A.shape[1],
539 EM,
824 - topk_ids.numel(),
540 + num_tokens,
541 A.stride(0),
542 A.stride(1),
543 B.stride(0),
@@ -864,7 +580,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
580 B.shape[1],
581 B.shape[2],
582 EM,
867 - topk_ids.numel(),
583 + num_tokens,
584 A.stride(0),
585 A.stride(1),
586 B.stride(0),
@@ -1389,6 +1105,7 @@ def fused_experts(hidden_states: torch.Tensor,
1105 w2=w2,
1106 topk_weights=topk_weights,
1107 topk_ids=topk_ids,
1108 + inplace=inplace,
1109 activation=activation,
1110 global_num_experts=global_num_experts,
1111 expert_map=expert_map,
@@ -1419,85 +1136,6 @@ def fused_experts(hidden_states: torch.Tensor,
1136 block_shape=block_shape)
1137
1138
1422 -def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
1423 - """
1424 - A permutation routine that works on fp8 types.
1425 - """
1426 - if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
1427 - return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
1428 - else:
1429 - return m[idx, ...]
1430 -
1431 -
1432 -def _moe_permute(
1433 - curr_hidden_states: torch.Tensor,
1434 - a1q_scale: Optional[torch.Tensor],
1435 - curr_topk_ids: torch.Tensor,
1436 - global_num_experts: int,
1437 - expert_map: Optional[torch.Tensor],
1438 - top_k_num: int,
1439 - block_m: int,
1440 -) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
1441 - torch.Tensor]:
1442 - """
1443 - Determine the sorted_token_ids, expert_ids for the given problem size.
1444 - Permute the hidden states and scales according to `sorted_token_ids`.
1445 - """
1446 - tokens_in_chunk, _ = curr_hidden_states.shape
1447 -
1448 - sorted_token_ids, expert_ids, num_tokens_post_padded = (
1449 - moe_align_block_size(curr_topk_ids,
1450 - block_m,
1451 - global_num_experts,
1452 - expert_map,
1453 - pad_sorted_ids=True))
1454 -
1455 - inv_perm: Optional[torch.Tensor] = None
1456 -
1457 - num_tokens = top_k_num * tokens_in_chunk
1458 - sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
1459 - expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
1460 - inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
1461 -
1462 - # Permute according to sorted token ids.
1463 - curr_hidden_states = _fp8_perm(curr_hidden_states,
1464 - sorted_token_ids // top_k_num)
1465 -
1466 - if a1q_scale is not None:
1467 - a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
1468 -
1469 - return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
1470 - inv_perm)
1471 -
1472 -
1473 -def _moe_unpermute_and_reduce(
1474 - out: torch.Tensor,
1475 - curr_hidden: torch.Tensor,
1476 - inv_perm: Optional[torch.Tensor],
1477 - topk: int,
1478 - K: int,
1479 - topk_weight: torch.Tensor,
1480 -) -> None:
1481 - """
1482 - Unpermute the final result and apply topk_weights, then perform the final
1483 - reduction on the hidden states.
1484 - """
1485 - M = topk_weight.shape[0]
1486 - curr_hidden = curr_hidden[inv_perm, ...]
1487 - curr_hidden = curr_hidden.view(-1, topk, K)
1488 - curr_hidden.mul_(topk_weight.view(M, -1, 1))
1489 - ops.moe_sum(curr_hidden, out)
1490 -
1491 -
1492 -def _resize_cache(x: torch.Tensor, v: Tuple[int, ...]) -> torch.Tensor:
1493 - """
1494 - Shrink the given tensor and apply the given view to it. This is
1495 - used to resize the intermediate fused_moe caches.
1496 - """
1497 - assert prod(v) <= x.numel()
1498 - return x.flatten()[:prod(v)].view(*v)
1499 -
1500 -
1139 def fused_experts_impl(hidden_states: torch.Tensor,
1140 w1: torch.Tensor,
1141 w2: torch.Tensor,
@@ -1629,7 +1267,6 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1267 w1_scale,
1268 w1_zp,
1269 curr_topk_weights,
1632 - curr_topk_ids,
1270 sorted_token_ids,
1271 expert_ids,
1272 num_tokens_post_padded,
@@ -1660,28 +1297,34 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1297 qintermediate_cache2 = intermediate_cache2
1298 a2q_scale = a2_scale
1299
1663 - invoke_fused_moe_kernel(qintermediate_cache2,
1664 - w2,
1665 - intermediate_cache3,
1666 - a2q_scale,
1667 - w2_scale,
1668 - w2_zp,
1669 - curr_topk_weights,
1670 - curr_topk_ids,
1671 - sorted_token_ids,
1672 - expert_ids,
1673 - num_tokens_post_padded,
1674 - True,
1675 - 1,
1676 - config,
1677 - compute_type=compute_type,
1678 - use_fp8_w8a8=use_fp8_w8a8,
1679 - use_int8_w8a16=use_int8_w8a16,
1680 - use_int4_w4a16=use_int4_w4a16,
1681 - block_shape=block_shape)
1300 + invoke_fused_moe_kernel(
1301 + qintermediate_cache2,
1302 + w2,
1303 + intermediate_cache3,
1304 + a2q_scale,
1305 + w2_scale,
1306 + w2_zp,
1307 + curr_topk_weights,
1308 + sorted_token_ids,
1309 + expert_ids,
1310 + num_tokens_post_padded,
1311 + False, #True,
1312 + 1,
1313 + config,
1314 + compute_type=compute_type,
1315 + use_fp8_w8a8=use_fp8_w8a8,
1316 + use_int8_w8a16=use_int8_w8a16,
1317 + use_int4_w4a16=use_int4_w4a16,
1318 + block_shape=block_shape)
1319 +
1320 + if True:
1321 + intermediate_cache3 = intermediate_cache3.view(-1, top_k_num, K)
1322 + intermediate_cache3.mul_(
1323 + curr_topk_weights.view(tokens_in_chunk, -1, 1))
1324
1325 ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
1326 out_hidden_states[begin_chunk_idx:end_chunk_idx])
1327 +
1328 return out_hidden_states
1329
1330
@@ -1790,327 +1433,3 @@ def fused_moe(
1433 a1_scale=a1_scale,
1434 a2_scale=a2_scale,
1435 block_shape=block_shape)
1793 -
1794 -
1795 -def deep_gemm_moe_fp8(
1796 - hidden_states: torch.Tensor,
1797 - w1: torch.Tensor,
1798 - w2: torch.Tensor,
1799 - w1_scale: torch.Tensor,
1800 - w2_scale: torch.Tensor,
1801 - topk_weights: torch.Tensor,
1802 - topk_ids: torch.Tensor,
1803 - inplace: bool = False,
1804 - activation: str = "silu",
1805 - global_num_experts: int = -1,
1806 - expert_map: Optional[torch.Tensor] = None,
1807 - a1_scale: Optional[torch.Tensor] = None,
1808 - a2_scale: Optional[torch.Tensor] = None,
1809 -) -> torch.Tensor:
1810 - """
1811 - This function computes a a8w8-quantized Mixture of Experts (MoE) layer
1812 - using two sets of quantized weights, w1_q and w2_q, and top-k gating
1813 - mechanism. The matrix multiplications are implemented with DeepGemm
1814 - grouped gemm.
1815 -
1816 - Parameters:
1817 - - hidden_states (torch.Tensor): The input tensor to the MoE layer.
1818 - Shape: [M, K]
1819 - - w1 (torch.Tensor): The first set of fp8 quantized expert weights.
1820 - Shape: [num_experts, K, 2N] (the weights are passed transposed)
1821 - - w2 (torch.Tensor): The second set of fp8 quantized expert weights.
1822 - Shape: [num_experts, N, K] (the weights are passed transposed)
1823 - - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
1824 - Shape: [num_experts] or [num_experts, 2N]
1825 - - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
1826 - Shape: [num_experts] or [num_experts, K]
1827 - - topk_weights (torch.Tensor): The weights of each token->expert mapping.
1828 - - topk_ids (torch.Tensor): The token->expert mapping for topk_weights.
1829 - - inplace (bool): If True, perform the operation in-place.
1830 - Defaults to False.
1831 - - activation (str): The activation function to apply after the first
1832 - MoE layer.
1833 - - global_num_experts (int): The total number of experts in the global
1834 - expert space.
1835 - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
1836 - from the global expert space to the local expert space of the expert
1837 - parallel shard.
1838 - - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
1839 - Shape: scalar or [M]
1840 - - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
1841 - quantize the intermediate result between the gemms.
1842 - Shape: scalar or [M]
1843 -
1844 - Returns:
1845 - - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
1846 - """
1847 - # Lazy import to avoid CUDA initialization problems.
1848 - import deep_gemm as dg
1849 -
1850 - assert expert_map is None, "Expert maps not supported yet"
1851 -
1852 - assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
1853 -
1854 - assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1855 - assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
1856 - assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
1857 - assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
1858 - assert hidden_states.dtype in [
1859 - torch.float32, torch.float16, torch.bfloat16
1860 - ]
1861 - assert w1.dtype == torch.float8_e4m3fn
1862 - assert w2.dtype == torch.float8_e4m3fn
1863 - assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
1864 - assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
1865 - assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
1866 - assert a1_scale is None or a1_scale.dim(
1867 - ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[
1868 - 0] == hidden_states.shape[0], "Input scale shape mismatch"
1869 - assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
1870 -
1871 - num_tokens, _ = hidden_states.shape
1872 - E, N, _ = w1.shape
1873 - K = w2.shape[1]
1874 - if global_num_experts == -1:
1875 - global_num_experts = E
1876 - top_k_num = topk_ids.shape[1]
1877 - # We execute the fused_moe kernel in chunks to circumvent this issue:
1878 - # https://github.com/vllm-project/vllm/issues/5938
1879 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
1880 -
1881 - assert _valid_deep_gemm(hidden_states, w1, w2, expert_map)
1882 -
1883 - if inplace:
1884 - out_hidden_states = hidden_states
1885 - else:
1886 - out_hidden_states = torch.empty_like(hidden_states)
1887 -
1888 - block_m = dg.get_m_alignment_for_contiguous_layout()
1889 - block_shape = [block_m, block_m]
1890 -
1891 - assert w1_scale is not None
1892 - assert w2_scale is not None
1893 -
1894 - # We attempt to transpose and align offline in Fp8MoEMethod, in which
1895 - # case these calls will be nops. Otherwise, they'll be performed every
1896 - # time the layer is executed.
1897 - w1_scale = dg.get_col_major_tma_aligned_tensor(w1_scale).contiguous()
1898 - w2_scale = dg.get_col_major_tma_aligned_tensor(w2_scale).contiguous()
1899 -
1900 - M_sum = topk_ids.numel() + global_num_experts * (block_m - 1)
1901 - M_sum = round_up(M_sum, block_m)
1902 -
1903 - num_chunks = (num_tokens // CHUNK_SIZE) + 1
1904 -
1905 - # We can reuse the memory between cache1 and cache3 because by the time
1906 - # we need cache3, we're done with cache1
1907 - cache13 = torch.empty(M_sum * max(N, K),
1908 - device=hidden_states.device,
1909 - dtype=hidden_states.dtype)
1910 -
1911 - intermediate_cache1 = cache13[:M_sum * N].view(M_sum, N)
1912 - intermediate_cache2 = torch.empty((M_sum, N // 2),
1913 - device=hidden_states.device,
1914 - dtype=hidden_states.dtype)
1915 - intermediate_cache3 = cache13[:M_sum * K].view(M_sum, K)
1916 -
1917 - for chunk in range(num_chunks):
1918 - begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
1919 - min((chunk + 1) * CHUNK_SIZE,
1920 - num_tokens))
1921 - curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
1922 - tokens_in_chunk, _ = curr_hidden_states.shape
1923 -
1924 - if tokens_in_chunk == 0:
1925 - break
1926 -
1927 - curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1928 - curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1929 -
1930 - a1q_scale: Optional[torch.Tensor] = None
1931 -
1932 - qcurr_hidden_states, a1q_scale = _fp8_quantize(curr_hidden_states,
1933 - a1_scale, block_shape)
1934 -
1935 - (qcurr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
1936 - inv_perm) = _moe_permute(qcurr_hidden_states, a1q_scale,
1937 - curr_topk_ids, global_num_experts,
1938 - expert_map, top_k_num, block_m)
1939 -
1940 - # Adjust the intermediate cache size and config for the last chunk.
1941 - # Note that in most cases we only have one chunk so the cache size
1942 - # and config are already set correctly and do not need to be adjusted.
1943 - if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
1944 - curr_M = sorted_token_ids.numel()
1945 - intermediate_cache1 = _resize_cache(intermediate_cache1,
1946 - (curr_M, N))
1947 - intermediate_cache2 = _resize_cache(intermediate_cache2,
1948 - (curr_M, N // 2))
1949 - intermediate_cache3 = _resize_cache(intermediate_cache3,
1950 - (curr_M, K))
1951 -
1952 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
1953 - (qcurr_hidden_states, a1q_scale), (w1, w1_scale),
1954 - intermediate_cache1, expert_ids)
1955 -
1956 - if activation == "silu":
1957 - torch.ops._C.silu_and_mul(intermediate_cache2,
1958 - intermediate_cache1.view(-1, N))
1959 - elif activation == "gelu":
1960 - torch.ops._C.gelu_and_mul(intermediate_cache2,
1961 - intermediate_cache1.view(-1, N))
1962 - else:
1963 - raise ValueError(f"Unsupported FusedMoe activation: {activation}")
1964 -
1965 - a2q_scale: Optional[torch.Tensor] = None
1966 -
1967 - qintermediate_cache2, a2q_scale = _fp8_quantize(
1968 - intermediate_cache2, a2_scale, block_shape)
1969 -
1970 - dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
1971 - (qintermediate_cache2, a2q_scale), (w2, w2_scale),
1972 - intermediate_cache3, expert_ids)
1973 -
1974 - _moe_unpermute_and_reduce(
1975 - out_hidden_states[begin_chunk_idx:end_chunk_idx],
1976 - intermediate_cache3.view(*intermediate_cache3.shape), inv_perm,
1977 - top_k_num, K, curr_topk_weights)
1978 -
1979 - return out_hidden_states
1980 -
1981 -
1982 -#TODO make the grouped gemm kernel consistent with scaled gemm kernel
1983 -def cutlass_moe_fp8(
1984 - a: torch.Tensor,
1985 - w1_q: torch.Tensor,
1986 - w2_q: torch.Tensor,
1987 - w1_scale: torch.Tensor,
1988 - w2_scale: torch.Tensor,
1989 - topk_weights: torch.Tensor,
1990 - topk_ids: torch.Tensor,
1991 - ab_strides1: torch.Tensor,
1992 - c_strides1: torch.Tensor,
1993 - ab_strides2: torch.Tensor,
1994 - c_strides2: torch.Tensor,
1995 - a1_scale: Optional[torch.Tensor] = None,
1996 - a2_scale: Optional[torch.Tensor] = None,
1997 - out_dtype: torch.dtype = torch.half,
1998 -) -> torch.Tensor:
1999 - """
2000 - This function computes a a8w8-quantized Mixture of Experts (MoE) layer
2001 - using two sets of quantized weights, w1_q and w2_q, and top-k gating
2002 - mechanism. The matrix multiplications are implemented with CUTLASS
2003 - grouped gemm.
2004 -
2005 - Parameters:
2006 - - a (torch.Tensor): The input tensor to the MoE layer.
2007 - Shape: [M, K]
2008 - - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
2009 - Shape: [num_experts, K, 2N] (the weights are passed transposed)
2010 - - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
2011 - Shape: [num_experts, N, K] (the weights are passed transposed)
2012 - - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
2013 - Shape: [num_experts] or [num_experts, 2N]
2014 - - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
2015 - Shape: [num_experts] or [num_experts, K]
2016 - - gating_output (torch.Tensor): The output of the gating operation
2017 - (before softmax).
2018 - - topk_weights (torch.Tensor): The weights of each token->expert mapping.
2019 - - ab_strides1 (torch.Tensor): The input and weights strides of the first
2020 - grouped gemm.
2021 - - c_strides1 (torch.Tensor): The output strides of the first grouped gemm.
2022 - - ab_strides2 (torch.Tensor): The input and weights strides of the second
2023 - grouped gemm.
2024 - - c_strides2 (torch.Tensor): The output strides of the second grouped gemm.
2025 - - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
2026 - Shape: scalar or [M]
2027 - - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
2028 - quantize the intermediate result between the gemms.
2029 - Shape: scalar or [M]
2030 - - out_dtype (torch.Tensor): The output tensor type.
2031 -
2032 - Returns:
2033 - - torch.Tensor: The fp16 output tensor after applying the MoE layer.
2034 - """
2035 -
2036 - assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
2037 - assert w1_q.dtype == torch.float8_e4m3fn
2038 - assert w2_q.dtype == torch.float8_e4m3fn
2039 - assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
2040 - assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
2041 - assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
2042 - assert a1_scale is None or a1_scale.dim(
2043 - ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[0] == a.shape[
2044 - 0], "Input scale shape mismatch"
2045 - assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
2046 - 1] == w1_q.shape[2], "W1 scale shape mismatch"
2047 - assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
2048 - 1] == w2_q.shape[2], "W2 scale shape mismatch"
2049 - assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch"
2050 - assert w1_q.shape[0] == w1_scale.shape[
2051 - 0], "w1 scales expert number mismatch"
2052 - assert w1_q.shape[0] == w2_scale.shape[
2053 - 0], "w2 scales expert number mismatch"
2054 - assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
2055 - assert ab_strides1.shape[0] == w1_q.shape[
2056 - 0], "AB Strides 1 expert number mismatch"
2057 - assert c_strides1.shape[0] == w1_q.shape[
2058 - 0], "C Strides 1 expert number mismatch"
2059 - assert ab_strides2.shape[0] == w2_q.shape[
2060 - 0], "AB Strides 2 expert number mismatch"
2061 - assert c_strides2.shape[0] == w2_q.shape[
2062 - 0], "C Strides 2 expert number mismatch"
2063 - assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
2064 -
2065 - num_experts = w1_q.size(0)
2066 - m = a.size(0)
2067 - k = w1_q.size(1)
2068 - n = w2_q.size(1)
2069 -
2070 - topk = topk_ids.size(1)
2071 - per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
2072 - a2_scale.numel() != 1 if a2_scale is not None else False)
2073 -
2074 - a_q, a1_scale = ops.scaled_fp8_quant(
2075 - a, a1_scale, use_per_token_if_dynamic=per_act_token)
2076 - device = a_q.device
2077 -
2078 - expert_offsets = torch.empty((num_experts + 1),
2079 - dtype=torch.int32,
2080 - device=device)
2081 - problem_sizes1 = torch.empty((num_experts, 3),
2082 - dtype=torch.int32,
2083 - device=device)
2084 - problem_sizes2 = torch.empty((num_experts, 3),
2085 - dtype=torch.int32,
2086 - device=device)
2087 -
2088 - a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
2089 - c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
2090 -
2091 - ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
2092 - problem_sizes2, a_map, c_map, num_experts, n,
2093 - k)
2094 -
2095 - rep_a_q = a_q.view(dtype=torch.uint8)[a_map].view(dtype=a_q.dtype)
2096 - rep_a1_scales = a1_scale[a_map] if per_act_token else a1_scale
2097 -
2098 - c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
2099 - c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
2100 -
2101 - ops.cutlass_moe_mm(c1, rep_a_q, w1_q, rep_a1_scales, w1_scale,
2102 - expert_offsets[:-1], problem_sizes1, ab_strides1,
2103 - ab_strides1, c_strides1)
2104 -
2105 - intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
2106 - torch.ops._C.silu_and_mul(intermediate, c1)
2107 -
2108 - intemediate_q, a2_scale = ops.scaled_fp8_quant(
2109 - intermediate, a2_scale, use_per_token_if_dynamic=per_act_token)
2110 -
2111 - ops.cutlass_moe_mm(c2, intemediate_q, w2_q, a2_scale, w2_scale,
2112 - expert_offsets[:-1], problem_sizes2, ab_strides2,
2113 - ab_strides2, c_strides2)
2114 -
2115 - return (c2[c_map].view(m, topk, k) *
2116 - topk_weights.view(m, topk, 1).to(out_dtype)).sum(dim=1)
vllm/model_executor/layers/fused_moe/moe_align_block_size.py new
+243
@@ -0,0 +1,243 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional, Tuple
3 +
4 +import torch
5 +import triton
6 +import triton.language as tl
7 +
8 +import vllm.envs as envs
9 +from vllm import _custom_ops as ops
10 +from vllm.utils import round_up
11 +
12 +
13 +def ceil_div(a, b):
14 + return (a + b - 1) // b
15 +
16 +
17 +@triton.jit
18 +def moe_align_block_size_stage1(
19 + topk_ids_ptr,
20 + tokens_cnts_ptr,
21 + num_experts: tl.constexpr,
22 + numel: tl.constexpr,
23 + tokens_per_thread: tl.constexpr,
24 +):
25 + pid = tl.program_id(0)
26 +
27 + start_idx = pid * tokens_per_thread
28 +
29 + off_c = (pid + 1) * num_experts
30 +
31 + for i in range(tokens_per_thread):
32 + if start_idx + i < numel:
33 + idx = tl.load(topk_ids_ptr + start_idx + i)
34 + token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
35 + tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
36 +
37 +
38 +@triton.jit
39 +def moe_align_block_size_stage2(
40 + tokens_cnts_ptr,
41 + num_experts: tl.constexpr,
42 +):
43 + pid = tl.program_id(0)
44 +
45 + last_cnt = 0
46 + for i in range(1, num_experts + 1):
47 + token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
48 + last_cnt = last_cnt + token_cnt
49 + tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
50 +
51 +
52 +@triton.jit
53 +def moe_align_block_size_stage3(
54 + total_tokens_post_pad_ptr,
55 + tokens_cnts_ptr,
56 + cumsum_ptr,
57 + num_experts: tl.constexpr,
58 + block_size: tl.constexpr,
59 +):
60 + last_cumsum = 0
61 + off_cnt = num_experts * num_experts
62 + for i in range(1, num_experts + 1):
63 + token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
64 + last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
65 + tl.store(cumsum_ptr + i, last_cumsum)
66 + tl.store(total_tokens_post_pad_ptr, last_cumsum)
67 +
68 +
69 +@triton.jit
70 +def moe_align_block_size_stage4(
71 + topk_ids_ptr,
72 + sorted_token_ids_ptr,
73 + expert_ids_ptr,
74 + tokens_cnts_ptr,
75 + cumsum_ptr,
76 + num_experts: tl.constexpr,
77 + block_size: tl.constexpr,
78 + numel: tl.constexpr,
79 + tokens_per_thread: tl.constexpr,
80 +):
81 + pid = tl.program_id(0)
82 + start_idx = tl.load(cumsum_ptr + pid)
83 + end_idx = tl.load(cumsum_ptr + pid + 1)
84 +
85 + for i in range(start_idx, end_idx, block_size):
86 + tl.store(expert_ids_ptr + i // block_size, pid)
87 +
88 + start_idx = pid * tokens_per_thread
89 + off_t = pid * num_experts
90 +
91 + for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread,
92 + numel)):
93 + expert_id = tl.load(topk_ids_ptr + i)
94 + token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
95 + rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
96 + tl.store(sorted_token_ids_ptr + rank_post_pad, i)
97 + tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
98 +
99 +
100 +# Triton implementation based on:
101 +# https://github.com/sgl-project/sglang/commit/ba5112ff691d791a9e38c6c71f59324a5fcb49d0
102 +def moe_align_block_size_triton(
103 + topk_ids: torch.Tensor,
104 + num_experts: int,
105 + block_size: int,
106 + sorted_token_ids: torch.Tensor,
107 + expert_ids: torch.Tensor,
108 + num_tokens_post_pad: torch.Tensor,
109 +) -> None:
110 + numel = topk_ids.numel()
111 + grid = (num_experts, )
112 + tokens_cnts = torch.zeros((num_experts + 1, num_experts),
113 + dtype=torch.int32,
114 + device=topk_ids.device)
115 + cumsum = torch.zeros((num_experts + 1, ),
116 + dtype=torch.int32,
117 + device=topk_ids.device)
118 + tokens_per_thread = ceil_div(numel, num_experts)
119 +
120 + moe_align_block_size_stage1[grid](
121 + topk_ids,
122 + tokens_cnts,
123 + num_experts,
124 + numel,
125 + tokens_per_thread,
126 + )
127 + moe_align_block_size_stage2[grid](
128 + tokens_cnts,
129 + num_experts,
130 + )
131 + moe_align_block_size_stage3[(1, )](
132 + num_tokens_post_pad,
133 + tokens_cnts,
134 + cumsum,
135 + num_experts,
136 + block_size,
137 + )
138 + moe_align_block_size_stage4[grid](
139 + topk_ids,
140 + sorted_token_ids,
141 + expert_ids,
142 + tokens_cnts,
143 + cumsum,
144 + num_experts,
145 + block_size,
146 + numel,
147 + tokens_per_thread,
148 + )
149 +
150 +
151 +def moe_align_block_size(
152 + topk_ids: torch.Tensor,
153 + block_size: int,
154 + num_experts: int,
155 + expert_map: Optional[torch.Tensor] = None,
156 + pad_sorted_ids: bool = False
157 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
158 + """
159 + Aligns the token distribution across experts to be compatible with block
160 + size for matrix multiplication.
161 +
162 + Parameters:
163 + - topk_ids: A tensor of shape [total_tokens, top_k] representing the
164 + top-k expert indices for each token.
165 + - block_size: The block size used in block matrix multiplication.
166 + - num_experts: The total number of experts.
167 + - expert_map: A tensor of shape [num_experts] that maps the expert index
168 + from the global space to the local index space of the current
169 + expert parallel shard. If the expert is not in the current expert
170 + parallel shard, the mapping is set to -1.
171 + - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
172 + should be padded to a multiple of block_size,
173 +
174 + Returns:
175 + - sorted_token_ids: A tensor containing the sorted token indices according
176 + to their allocated expert.
177 + - expert_ids: A tensor indicating the assigned expert index for each block.
178 + - num_tokens_post_padded: The total number of tokens after padding,
179 + ensuring divisibility by block_size.
180 +
181 + This function pads the number of tokens that each expert needs to process
182 + so that it is divisible by block_size.
183 + Padding ensures that during block matrix multiplication, the dimensions
184 + align correctly.
185 +
186 + Example:
187 + Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
188 + block_size = 4, and num_experts = 4:
189 + - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
190 + with each expert needing to process 3 tokens.
191 + - As block_size is 4, we pad 1 token for each expert.
192 + - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
193 + - Then append padding tokens [12, 12, 12, 12] for each block.
194 + - After sorting by expert index, we obtain token_ids
195 + [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
196 + Tokens 12 are non-existent (padding) and are ignored in
197 + the subsequent matrix multiplication.
198 + - The padding ensures that the total number of tokens is now divisible
199 + by block_size for proper block matrix operations.
200 + """
201 + max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
202 + if pad_sorted_ids:
203 + max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
204 + sorted_ids = torch.empty((max_num_tokens_padded, ),
205 + dtype=torch.int32,
206 + device=topk_ids.device)
207 + sorted_ids.fill_(topk_ids.numel())
208 + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
209 + # Expert ids must be zeroed out to prevent index out of bounds error while
210 + # mapping global expert ids to local expert ids in expert parallelism.
211 + expert_ids = torch.zeros((max_num_m_blocks, ),
212 + dtype=torch.int32,
213 + device=topk_ids.device)
214 + num_tokens_post_pad = torch.empty((1),
215 + dtype=torch.int32,
216 + device=topk_ids.device)
217 + if num_experts >= 224:
218 + if envs.VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON or num_experts != 256:
219 + moe_align_block_size_triton(
220 + topk_ids,
221 + num_experts,
222 + block_size,
223 + sorted_ids,
224 + expert_ids,
225 + num_tokens_post_pad,
226 + )
227 + else:
228 + # Currently requires num_experts=256
229 + ops.sgl_moe_align_block_size(
230 + topk_ids,
231 + num_experts,
232 + block_size,
233 + sorted_ids,
234 + expert_ids,
235 + num_tokens_post_pad,
236 + )
237 + else:
238 + ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
239 + expert_ids, num_tokens_post_pad)
240 + if expert_map is not None:
241 + expert_ids = expert_map[expert_ids]
242 +
243 + return sorted_ids, expert_ids, num_tokens_post_pad
vllm/model_executor/layers/fused_moe/utils.py new
+48
@@ -0,0 +1,48 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from math import prod
3 +from typing import List, Optional, Tuple
4 +
5 +import torch
6 +
7 +from vllm import _custom_ops as ops
8 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
9 + per_token_group_quant_fp8)
10 +from vllm.utils import cdiv
11 +
12 +
13 +def _resize_cache(x: torch.Tensor, v: Tuple[int, ...]) -> torch.Tensor:
14 + """
15 + Shrink the given tensor and apply the given view to it. This is
16 + used to resize the intermediate fused_moe caches.
17 + """
18 + assert prod(v) <= x.numel()
19 + return x.flatten()[:prod(v)].view(*v)
20 +
21 +
22 +def _fp8_quantize(
23 + A: torch.Tensor,
24 + A_scale: Optional[torch.Tensor],
25 + block_shape: Optional[List[int]],
26 +) -> Tuple[torch.Tensor, torch.Tensor]:
27 + """
28 + Perform fp8 quantization on the inputs. If a block_shape
29 + is provided, the output will be blocked.
30 + """
31 + if block_shape is None:
32 + A, A_scale = ops.scaled_fp8_quant(A, A_scale)
33 + else:
34 + assert len(block_shape) == 2
35 + _, block_k = block_shape[0], block_shape[1]
36 + A, A_scale = per_token_group_quant_fp8(A, block_k)
37 + assert cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
38 + return A, A_scale
39 +
40 +
41 +def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
42 + """
43 + A permutation routine that works on fp8 types.
44 + """
45 + if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
46 + return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
47 + else:
48 + return m[idx, ...]