@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e789cad6b8
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 14, 05:30 UTC

[gpt-oss] triton kernel mxfp4 (#22421)

Signed-off-by: <zyy1102000@gmail.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Aug 8, 2025 at 08:24 UTC e789cad6b8b5d2a01aa6521b9208bb8d6501ee5b
8 files changed +755 -9
.gitignore
+3
@@ -4,6 +4,9 @@
4 # vllm-flash-attn built from source
5 vllm/vllm_flash_attn/*
6
7 +# triton jit
8 +.triton
9 +
10 # Byte-compiled / optimized / DLL files
11 __pycache__/
12 *.py[cod]
tests/kernels/moe/test_gpt_oss_triton_kernels.py new
+375
@@ -0,0 +1,375 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass, fields
4 +
5 +import pytest
6 +import torch
7 +import torch.nn.functional as F
8 +import triton_kernels.swiglu
9 +from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
10 +from triton_kernels.numerics import InFlexData
11 +from triton_kernels.numerics_details.mxfp import (downcast_to_mxfp,
12 + upcast_from_mxfp)
13 +from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
14 +from triton_kernels.tensor_details import layout
15 +from triton_kernels.testing import assert_close
16 +
17 +from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
18 + BatchedPrepareAndFinalize)
19 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
20 +from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
21 + BatchedOAITritonExperts, triton_kernel_moe_forward)
22 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
23 + FusedMoEModularKernel)
24 +from vllm.model_executor.layers.utils import shuffle_weight
25 +from vllm.utils import round_up
26 +
27 +
28 +def deshuffle(w: torch.Tensor):
29 + first = w[..., ::2]
30 + second = w[..., 1::2]
31 +
32 + deshuffled = torch.concat((first, second), dim=-1)
33 + return deshuffled
34 +
35 +
36 +def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
37 + randbits = [torch.randperm(E) for _ in range(M)]
38 + x_list = [
39 + (-1)**i *
40 + ((16384 +
41 + ((i * 512) % 4096) + bits).to(torch.int16).view(torch.bfloat16))
42 + for i, bits in enumerate(randbits)
43 + ]
44 + exp_data = torch.stack(x_list).to(
45 + device="cuda") # simulating gate_output (M, E)
46 +
47 + # create input tensor
48 + x = torch.randn((M, K), dtype=torch.bfloat16, device="cuda")
49 + w1 = torch.randn((E, 2 * N, K), dtype=torch.bfloat16, device="cuda")
50 + w1_bias = torch.randn((E, 2 * N), dtype=torch.bfloat16, device="cuda")
51 +
52 + w2 = torch.randn((E, K, N), dtype=torch.bfloat16, device="cuda")
53 + w2_bias = torch.randn((E, K), dtype=torch.bfloat16, device="cuda")
54 +
55 + exp_data_tri = exp_data.clone()
56 + x_tri = x.clone()
57 + w1_tri = w1.clone()
58 + w2_tri = w2.clone()
59 +
60 + w1_bias_tri = w1_bias.clone()
61 + w2_bias_tri = w2_bias.clone()
62 + w1_bias_tri = w1_bias_tri.to(torch.float32)
63 + w2_bias_tri = w2_bias_tri.to(torch.float32)
64 +
65 + dtype_dict = {
66 + "bf16": torch.bfloat16,
67 + "fp8_e4m3": torch.float8_e4m3fn,
68 + "fp8_e5m2": torch.float8_e5m2
69 + }
70 +
71 + x = x.to(dtype_dict[a_dtype]).to(torch.bfloat16)
72 + if w_dtype != "mx4":
73 + # simulate quantization support on reference impl
74 + w1 = w1.to(dtype_dict[w_dtype]).to(torch.bfloat16)
75 + w2 = w2.to(dtype_dict[w_dtype]).to(torch.bfloat16)
76 +
77 + # triton moe kernel use transposed shape for matmul
78 + w1_tri = w1_tri.transpose(-2, -1)
79 + w2_tri = w2_tri.transpose(-2, -1)
80 +
81 + # shuffle weights
82 + w1_tri = shuffle_weight(w1_tri)
83 + w1_bias_tri = shuffle_weight(w1_bias_tri)
84 +
85 + # quant triton_weights
86 + x_tri = x.to(dtype_dict[a_dtype])
87 + if w_dtype != "mx4":
88 + pytest.skip("NYI")
89 + else: # quantize to mx4
90 + # careful on the padding here, the activation padding need to be
91 + # multiple of 64, the actual engine is not implemented
92 + w1_bottom_pad = round_up(w1_tri.shape[1], 64) - w1_tri.shape[1]
93 + w1_right_pad = round_up(w1_tri.shape[2], 128) - w1_tri.shape[2]
94 +
95 + w2_bottom_pad = w1_right_pad // 2
96 + w2_right_pad = w1_bottom_pad
97 +
98 + x_pad = w1_bottom_pad
99 +
100 + w1_tri = F.pad(w1_tri, (0, w1_right_pad, 0, w1_bottom_pad, 0, 0),
101 + mode="constant",
102 + value=0)
103 + w2_tri = F.pad(w2_tri, (0, w2_right_pad, 0, w2_bottom_pad, 0, 0),
104 + mode="constant",
105 + value=0)
106 +
107 + w1_bias_tri = F.pad(w1_bias_tri, (0, w1_right_pad, 0, 0),
108 + mode="constant",
109 + value=0)
110 + w2_bias_tri = F.pad(w2_bias_tri, (0, w2_right_pad, 0, 0),
111 + mode="constant",
112 + value=0)
113 +
114 + x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0)
115 +
116 + w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
117 + mx_axis=1)
118 + w_scale_layout, w_scale_layout_opts = (
119 + layout.make_default_matmul_mxfp4_w_scale_layout(
120 + mx_axis=1, num_warps=num_warps))
121 +
122 + w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)
123 + w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, torch.bfloat16, axis=1)
124 +
125 + w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1)
126 + w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, torch.bfloat16, axis=1)
127 +
128 + w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout,
129 + **w_layout_opts)
130 + w1_scale_tri = convert_layout(wrap_torch_tensor(w1_scale_tri),
131 + w_scale_layout, **w_scale_layout_opts)
132 +
133 + w2_tri = convert_layout(wrap_torch_tensor(w2_tri, FP4), w_layout,
134 + **w_layout_opts)
135 + w2_scale_tri = convert_layout(wrap_torch_tensor(w2_scale_tri),
136 + w_scale_layout, **w_scale_layout_opts)
137 +
138 + pc1 = PrecisionConfig(weight_scale=w1_scale_tri,
139 + flex_ctx=FlexCtx(rhs_data=InFlexData()))
140 + pc2 = PrecisionConfig(weight_scale=w2_scale_tri,
141 + flex_ctx=FlexCtx(rhs_data=InFlexData()))
142 +
143 + # tucuate so the rest can run properly
144 + w1 = w1[..., :K, :2 * N]
145 + w2 = w2[..., :N, :K]
146 +
147 + w1 = deshuffle(w1)
148 +
149 + w1 = w1.transpose(-1, -2).contiguous()
150 + w2 = w2.transpose(-1, -2).contiguous()
151 +
152 + return (x, w1, w1_bias, w2, w2_bias, exp_data, x_tri, w1_tri, w2_tri,
153 + exp_data_tri, w1_bias_tri, w2_bias_tri, pc1, pc2)
154 +
155 +
156 +@dataclass
157 +class ModelConfig:
158 + num_hidden_layers: int = 36
159 + num_experts: int = 128
160 + experts_per_token: int = 4
161 + vocab_size: int = 201088
162 + hidden_size: int = 2880
163 + intermediate_size: int = 2880
164 + head_dim: int = 64
165 + num_attention_heads: int = 64
166 + num_key_value_heads: int = 8
167 + sliding_window: int = 128
168 + initial_context_length: int = 4096
169 + rope_theta: float = 150000.0
170 + rope_scaling_factor: float = 32.0
171 + rope_ntk_alpha: float = 1.0
172 + rope_ntk_beta: float = 32.0
173 +
174 +
175 +def swiglu(x, alpha: float = 1.702, limit: float = 1.0):
176 + # Note we add an extra bias of 1 to the linear layer
177 + x_glu, x_linear = torch.chunk(x, 2, dim=-1)
178 + if limit is not None:
179 + x_glu = x_glu.clamp(max=limit)
180 + out_glu = x_glu * torch.sigmoid(alpha * x_glu)
181 + if limit is not None:
182 + x_linear = x_linear.clamp(min=-limit, max=limit)
183 + return out_glu * (x_linear + 1)
184 +
185 +
186 +def oai_moe_forward(
187 + hidden_states: torch.Tensor, # (M, K)
188 + w1: torch.Tensor, # (E, 2N)
189 + w1_bias: torch.Tensor, # (E, 2N, K)
190 + w2: torch.Tensor, # (E, K, N)
191 + w2_bias: torch.Tensor, # (E, N)
192 + gating_output: torch.Tensor, # (M, E)
193 + topk: int):
194 + # model.py 309:330, assuming gating and norm
195 + t = hidden_states
196 + experts = torch.topk(gating_output, k=topk, dim=-1, sorted=True)
197 + expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
198 + expert_indices = experts.indices
199 +
200 + # MLP #1
201 + mlp1_weight = w1[expert_indices, ...]
202 + mlp1_bias = w1_bias[expert_indices, ...]
203 + t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias
204 + t = swiglu(t, limit=7)
205 +
206 + # MLP #2
207 + mlp2_weight = w2[expert_indices, ...]
208 + mlp2_bias = w2_bias[expert_indices, ...]
209 + t = torch.einsum("beck,bek->bec", mlp2_weight, t)
210 + t += mlp2_bias
211 +
212 + # Weighted sum of experts
213 + t = torch.einsum("bec,be->bc", t, expert_weights)
214 +
215 + return t
216 +
217 +
218 +@dataclass
219 +class Case:
220 + a_dtype: str
221 + w_dtype: str
222 +
223 +
224 +@pytest.mark.parametrize(
225 + ", ".join(f.name for f in fields(Case)),
226 + [
227 + tuple(getattr(case, f.name) for f in fields(Case)) for case in [
228 + # Case(a_dtype="bf16", w_dtype="bf16"),
229 + # Case(a_dtype="fp8_e4m3", w_dtype="fp8_e5m2"),
230 + Case(a_dtype="bf16", w_dtype="mx4")
231 + ]
232 + ],
233 +)
234 +@pytest.mark.parametrize("num_token", [2])
235 +@pytest.mark.parametrize("tp", [1, 2, 4, 8])
236 +def test_equiv(num_token, a_dtype, w_dtype, tp):
237 + M = num_token
238 + E = ModelConfig.num_experts
239 + K = ModelConfig.hidden_size
240 + N = ModelConfig.intermediate_size // tp
241 + topk = ModelConfig.experts_per_token
242 +
243 + x, w1, w1_bias, w2, w2_bias, exp_data, \
244 + x_tri, w1_tri, w2_tri, exp_data_tri, w1_bias_tri,\
245 + w2_bias_tri, pc1, pc2 = init_compute_data(
246 + M, K, N, E, a_dtype, w_dtype, num_warps=8)
247 +
248 + out_triton_monolithic = triton_kernel_moe_forward(
249 + hidden_states=x_tri,
250 + w1=w1_tri,
251 + w2=w2_tri,
252 + gating_output=exp_data_tri,
253 + topk=topk,
254 + renormalize=True,
255 + w1_bias=w1_bias_tri,
256 + w2_bias=w2_bias_tri,
257 + w1_precision=pc1,
258 + w2_precision=pc2)
259 + out_triton_monolithic = out_triton_monolithic[..., :K]
260 +
261 + out_ref = oai_moe_forward(hidden_states=x,
262 + w1=w1,
263 + w1_bias=w1_bias,
264 + w2=w2,
265 + w2_bias=w2_bias,
266 + gating_output=exp_data,
267 + topk=topk)
268 + assert_close(ref=out_ref,
269 + tri=out_triton_monolithic,
270 + maxtol=0.025,
271 + rmstol=0.005)
272 +
273 +
274 +def batched_moe(a: torch.Tensor, w1, w2, gating_output: torch.Tensor,
275 + topk: int, renormalize: bool, w1_bias: torch.Tensor,
276 + w2_bias: torch.Tensor, w1_precision: PrecisionConfig,
277 + w2_precision: PrecisionConfig) -> torch.Tensor:
278 + max_num_tokens = round_up(a.shape[0], 64)
279 +
280 + fused_experts = FusedMoEModularKernel(
281 + BatchedPrepareAndFinalize(max_num_tokens,
282 + num_dispatchers=1,
283 + num_local_experts=w1.shape[0],
284 + rank=0),
285 + BatchedOAITritonExperts(
286 + None,
287 + max_num_tokens=max_num_tokens,
288 + num_dispatchers=1,
289 + w1_precision=w1_precision,
290 + w2_precision=w2_precision,
291 + ),
292 + )
293 +
294 + extra_expert_args = {
295 + "w1_bias": w1_bias,
296 + "w2_bias": w2_bias,
297 + }
298 +
299 + topk_weight, topk_ids, _ = fused_topk(a, gating_output, topk, renormalize)
300 +
301 + return fused_experts(
302 + a,
303 + w1,
304 + w2,
305 + topk_weight,
306 + topk_ids,
307 + extra_expert_args=extra_expert_args,
308 + )
309 +
310 +
311 +@pytest.mark.parametrize(
312 + ", ".join(f.name for f in fields(Case)),
313 + [
314 + tuple(getattr(case, f.name) for f in fields(Case)) for case in [
315 + # Case(a_dtype="bf16", w_dtype="bf16"),
316 + # Case(a_dtype="fp8_e4m3", w_dtype="fp8_e5m2"),
317 + Case(a_dtype="bf16", w_dtype="mx4")
318 + ]
319 + ],
320 +)
321 +@pytest.mark.parametrize("num_token", [64])
322 +@pytest.mark.parametrize("ep", [1, 2, 4, 8])
323 +def test_triton_kernel_batched_moe(num_token, a_dtype, w_dtype, ep):
324 + M = num_token
325 + E = ModelConfig.num_experts // ep
326 + K = ModelConfig.hidden_size
327 + N = ModelConfig.intermediate_size
328 + topk = ModelConfig.experts_per_token
329 +
330 + x, w1, w1_bias, w2, w2_bias, exp_data, \
331 + x_tri, w1_tri, w2_tri, exp_data_tri, w1_bias_tri, \
332 + w2_bias_tri, pc1, pc2 = init_compute_data(
333 + M, K, N, E, a_dtype, w_dtype, num_warps=4)
334 +
335 + out_tri = batched_moe(a=x_tri,
336 + w1=w1_tri,
337 + w2=w2_tri,
338 + gating_output=exp_data_tri,
339 + topk=topk,
340 + renormalize=True,
341 + w1_bias=w1_bias_tri,
342 + w2_bias=w2_bias_tri,
343 + w1_precision=pc1,
344 + w2_precision=pc2)
345 + out_tri = out_tri[..., :K]
346 +
347 + out_ref = oai_moe_forward(hidden_states=x,
348 + w1=w1,
349 + w1_bias=w1_bias,
350 + w2=w2,
351 + w2_bias=w2_bias,
352 + gating_output=exp_data,
353 + topk=topk)
354 + assert_close(ref=out_ref, tri=out_tri, maxtol=0.025, rmstol=0.005)
355 +
356 +
357 +def test_unit_shuffle():
358 + N = ModelConfig.intermediate_size
359 + K = ModelConfig.hidden_size
360 + m = torch.randn((K, 2 * N), dtype=torch.bfloat16, device="cuda")
361 +
362 + x = torch.randn(K, dtype=torch.bfloat16, device="cuda")
363 +
364 + m_shuffled = shuffle_weight(m)
365 +
366 + out_ref = x @ m
367 + out_ref = swiglu(out_ref, limit=1.0)
368 +
369 + out = x @ m_shuffled
370 + out = triton_kernels.swiglu.swiglu_torch(
371 + out,
372 + alpha=1.702,
373 + precision_config=triton_kernels.swiglu.PrecisionConfig(limit=1.0))
374 +
375 + assert_close(ref=out_ref, tri=out)
\ No newline at end of file
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py new
+230
@@ -0,0 +1,230 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any, Optional
4 +
5 +import torch
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
9 + TopKWeightAndReduceDelegate)
10 +from vllm.model_executor.layers.fused_moe.utils import extract_required_args
11 +
12 +if True:
13 + import triton_kernels.swiglu
14 + from triton_kernels.matmul_ogs import (FnSpecs, FusedActivation,
15 + PrecisionConfig, matmul_ogs)
16 + from triton_kernels.routing import routing
17 +
18 +
19 +def triton_kernel_moe_forward(
20 + hidden_states: torch.Tensor,
21 + w1, # Tensor or triton_kernels.Tensor
22 + w2, # Tensor or triton_kernels.Tensor
23 + gating_output: torch.Tensor,
24 + topk: int,
25 + renormalize: bool,
26 + activation: str = "silu",
27 + apply_router_weight_on_input: bool = False,
28 + use_fp8_w8a8: bool = False,
29 + per_channel_quant: bool = False,
30 + global_num_experts: int = -1,
31 + expert_map: Optional[torch.Tensor] = None,
32 + w1_scale: Optional[torch.Tensor] = None,
33 + w2_scale: Optional[torch.Tensor] = None,
34 + w1_bias: Optional[torch.Tensor] = None,
35 + w2_bias: Optional[torch.Tensor] = None,
36 + w1_precision=None, # PrecisionConfig or None
37 + w2_precision=None, # PrecisionConfig or None
38 + a1_scale: Optional[torch.Tensor] = None,
39 + a2_scale: Optional[torch.Tensor] = None,
40 + block_shape: Optional[list[int]] = None,
41 +) -> torch.Tensor:
42 +
43 + routing_data, gather_idx, scatter_idx = routing(gating_output,
44 + topk,
45 + sm_first=not renormalize)
46 +
47 + return triton_kernel_fused_experts(
48 + None,
49 + hidden_states,
50 + w1,
51 + w2,
52 + routing_data,
53 + gather_idx,
54 + scatter_idx,
55 + activation=activation,
56 + apply_router_weight_on_input=apply_router_weight_on_input,
57 + use_fp8_w8a8=use_fp8_w8a8,
58 + per_channel_quant=per_channel_quant,
59 + global_num_experts=global_num_experts,
60 + expert_map=expert_map,
61 + w1_scale=w1_scale,
62 + w2_scale=w2_scale,
63 + w1_bias=w1_bias,
64 + w2_bias=w2_bias,
65 + w1_precision=w1_precision,
66 + w2_precision=w2_precision,
67 + a1_scale=a1_scale,
68 + a2_scale=a2_scale,
69 + block_shape=block_shape)
70 +
71 +
72 +# This is a triton implementation of the fused_experts function
73 +def triton_kernel_fused_experts(
74 + output_tensor: torch.Tensor,
75 + hidden_states: torch.Tensor,
76 + w1, # Tensor or triton_kernels.Tensor
77 + w2, # Tensor or triton_kernels.Tensor
78 + routing_data, # RoutingData
79 + gather_indx, # GatherIndx
80 + scatter_indx, # ScatterIndx
81 + activation: str = "silu",
82 + swiglu_alpha: float = 1.702,
83 + swiglu_limit: float = 7.0,
84 + apply_router_weight_on_input: bool = False,
85 + use_fp8_w8a8: bool = False,
86 + per_channel_quant: bool = False,
87 + global_num_experts: int = -1,
88 + expert_map: Optional[torch.Tensor] = None,
89 + w1_scale: Optional[torch.Tensor] = None,
90 + w2_scale: Optional[torch.Tensor] = None,
91 + w1_bias: Optional[torch.Tensor] = None,
92 + w2_bias: Optional[torch.Tensor] = None,
93 + w1_precision=None, # PrecisionConfig or None
94 + w2_precision=None, # PrecisionConfig or None
95 + a1_scale: Optional[torch.Tensor] = None,
96 + a2_scale: Optional[torch.Tensor] = None,
97 + block_shape: Optional[list[int]] = None,
98 +) -> torch.Tensor:
99 +
100 + # type check, uint8 means mxfp4
101 + assert hidden_states.dtype == torch.bfloat16
102 + assert w1_bias is None or w1_bias.dtype == torch.float32
103 + assert w2_bias is None or w2_bias.dtype == torch.float32
104 +
105 + # Shape check, only check non-mxfp4
106 + assert hidden_states.shape[-1] == w1.shape[-2]
107 + assert w2.shape[-1] == w1.shape[1]
108 +
109 + E, _, N = w1.shape
110 +
111 + if global_num_experts == -1:
112 + global_num_experts = E
113 +
114 + act = FusedActivation(
115 + FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
116 + (swiglu_alpha, swiglu_limit), 2)
117 + gammas = routing_data.gate_scal if routing_data else None
118 +
119 + intermediate_cache1 = matmul_ogs(
120 + hidden_states,
121 + w1,
122 + w1_bias,
123 + routing_data,
124 + gather_indx=gather_indx,
125 + precision_config=w1_precision,
126 + gammas=gammas if apply_router_weight_on_input else None,
127 + fused_activation=act)
128 +
129 + intermediate_cache3 = matmul_ogs(
130 + intermediate_cache1,
131 + w2,
132 + w2_bias,
133 + routing_data,
134 + scatter_indx=scatter_indx,
135 + precision_config=w2_precision,
136 + gammas=None if apply_router_weight_on_input else gammas,
137 + y=output_tensor,
138 + )
139 + return intermediate_cache3
140 +
141 +
142 +class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
143 +
144 + def __init__(self, quant_config, max_num_tokens: int, num_dispatchers: int,
145 + w1_precision: PrecisionConfig, w2_precision: PrecisionConfig):
146 + super().__init__(quant_config)
147 + self.max_num_tokens = max_num_tokens
148 + self.num_dispatchers = num_dispatchers
149 + self.w1_precision = w1_precision
150 + self.w2_precision = w2_precision
151 +
152 + @property
153 + def activation_formats(
154 + self
155 + ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
156 + return (mk.FusedMoEActivationFormat.BatchedExperts,
157 + mk.FusedMoEActivationFormat.BatchedExperts)
158 +
159 + def supports_chunking(self) -> bool:
160 + return False
161 +
162 + def supports_expert_map(self) -> bool:
163 + return False
164 +
165 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
166 + # Let PrepareAndFinalize::finalize() decide the impl.
167 + return TopKWeightAndReduceDelegate()
168 +
169 + def workspace_shapes(
170 + self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
171 + topk: int, global_num_experts: int, local_num_experts: int,
172 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata]
173 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
174 + # workspace are allocated inside the kernel
175 + assert a.dim() == 2
176 + num_dp = self.num_dispatchers
177 + num_experts = local_num_experts
178 + max_num_tokens = self.max_num_tokens
179 + workspace2 = (0, 0, 0)
180 + output = (num_experts, max_num_tokens * num_dp, N)
181 + return (output, workspace2, output, a.dtype)
182 +
183 + def apply(
184 + self,
185 + output: torch.Tensor,
186 + hidden_states: torch.Tensor,
187 + w1: torch.Tensor,
188 + w2: torch.Tensor,
189 + topk_weights: torch.Tensor,
190 + topk_ids: torch.Tensor,
191 + activation: str,
192 + global_num_experts: int,
193 + expert_map: Optional[torch.Tensor],
194 + w1_scale: Optional[torch.Tensor],
195 + w2_scale: Optional[torch.Tensor],
196 + w1_zp: Optional[torch.Tensor],
197 + w2_zp: Optional[torch.Tensor],
198 + a1q_scale: Optional[torch.Tensor],
199 + a2_scale: Optional[torch.Tensor],
200 + workspace13: torch.Tensor,
201 + workspace2: torch.Tensor,
202 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
203 + apply_router_weight_on_input: bool,
204 + extra_expert_args: Optional[dict[str, Any]],
205 + ):
206 + w1_bias, w2_bias = (extract_required_args(extra_expert_args,
207 + ["w1_bias", "w2_bias"]))
208 +
209 + return triton_kernel_fused_experts(
210 + output,
211 + hidden_states,
212 + w1,
213 + w2,
214 + None,
215 + None,
216 + None,
217 + activation=activation,
218 + apply_router_weight_on_input=False,
219 + use_fp8_w8a8=False,
220 + per_channel_quant=False,
221 + global_num_experts=global_num_experts,
222 + expert_map=expert_map,
223 + w1_scale=w1_scale,
224 + w2_scale=w2_scale,
225 + w1_bias=w1_bias,
226 + w2_bias=w2_bias,
227 + w1_precision=self.w1_precision,
228 + w2_precision=self.w2_precision,
229 + a1_scale=a1q_scale,
230 + a2_scale=a2_scale)
vllm/model_executor/layers/fused_moe/layer.py
+12 -5
@@ -36,7 +36,7 @@ from vllm.model_executor.utils import set_weight_attrs
36 from vllm.platforms import current_platform
37 from vllm.platforms.interface import CpuArchEnum
38 from vllm.utils import (direct_register_custom_op, has_deep_ep, has_pplx,
39 - round_up)
39 + has_triton_kernels, is_torch_equal_or_newer, round_up)
40 from vllm.utils.flashinfer import has_flashinfer
41
42 if current_platform.is_cuda_alike():
@@ -723,10 +723,17 @@ class FusedMoE(torch.nn.Module):
723 self.global_num_experts = num_experts + num_redundant_experts
724
725 # we padding globally so EP buffer allocation works
726 - if quant_config and quant_config.get_name() == "mxfp4" and (
727 - envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
728 - or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
729 - hidden_size = round_up(hidden_size, 256)
726 + if quant_config and quant_config.get_name() == "mxfp4":
727 + if not is_torch_equal_or_newer("2.8.0"):
728 + raise RuntimeError("Mxfp4 on hopper requires torch >= 2.8.0")
729 + if current_platform.is_device_capability(
730 + 90) and not has_triton_kernels():
731 + raise NotImplementedError(
732 + "Triton kernels must be installed for mxfp4 on hopper")
733 + if (current_platform.is_rocm()
734 + or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
735 + or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
736 + hidden_size = round_up(hidden_size, 256)
737
738 # For smuggling this layer into the fused moe custom op
739 compilation_config = vllm_config.compilation_config
vllm/model_executor/layers/quantization/mxfp4.py
+63 -3
@@ -8,16 +8,19 @@ from torch.nn.parameter import Parameter
8 from vllm import envs
9 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
10 FusedMoEMethodBase)
11 +from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
12 + triton_kernel_moe_forward)
13 from vllm.model_executor.layers.linear import (LinearBase,
14 UnquantizedLinearMethod)
15 from vllm.model_executor.layers.quantization import QuantizationMethods
16 from vllm.model_executor.layers.quantization.base_config import (
17 QuantizationConfig, QuantizeMethodBase)
18 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
17 - _can_support_mxfp4)
19 + _can_support_mxfp4, _swizzle_mxfp4)
20 from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 is_layer_skipped)
22 from vllm.model_executor.utils import set_weight_attrs
23 +from vllm.platforms import current_platform
24 from vllm.utils import next_power_of_2, round_up
25
26 if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
@@ -39,7 +42,7 @@ class Mxfp4Config(QuantizationConfig):
42
43 @classmethod
44 def get_min_capability(cls) -> int:
42 - return 100
45 + return 90
46
47 @classmethod
48 def get_name(cls) -> QuantizationMethods:
@@ -100,11 +103,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
103 intermediate_size_per_partition
104 # pad the intermediate size to be a multiple of 2 * mxfp4_block
105 # for to hold non-uniform sharded tensor as well as swizzling
106 + # other padding to increase performance
107 if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
108 or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
109 intermediate_size_per_partition_after_pad = round_up(
110 intermediate_size_per_partition, 256)
111 hidden_size = round_up(hidden_size, 256)
112 + elif current_platform.is_rocm():
113 + intermediate_size_per_partition_after_pad = round_up(
114 + intermediate_size_per_partition, 128)
115 + else:
116 + intermediate_size_per_partition_after_pad = round_up(
117 + intermediate_size_per_partition, 64)
118
119 self.intermediate_size = intermediate_size_per_partition_after_pad
120 self.hidden_size = hidden_size
@@ -303,7 +313,41 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
313 layer.w2_bias = Parameter(torch.stack(gemm2_bias_shuffled).reshape(
314 self.num_experts, -1),
315 requires_grad=False)
306 - return
316 + else:
317 + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
318 +
319 + w13_bias = layer.w13_bias.to(torch.float32)
320 + w2_bias = layer.w2_bias.to(torch.float32)
321 +
322 + layer.w13_bias = Parameter(w13_bias, requires_grad=False)
323 + layer.w2_bias = Parameter(w2_bias, requires_grad=False)
324 +
325 + # FIXME warp need to be adjusted based on batch size
326 + # only apply to batched mode
327 + if self.moe.use_ep:
328 + num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
329 + else:
330 + num_warps = 8
331 +
332 + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
333 + layer.w13_weight, layer.w13_weight_scale, num_warps)
334 + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
335 + layer.w2_weight, layer.w2_weight_scale, num_warps)
336 +
337 + self.w13_precision_config = PrecisionConfig(
338 + weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex))
339 + self.w2_precision_config = PrecisionConfig(
340 + weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex))
341 +
342 + self.w13_weight_triton_tensor = w13_weight
343 + self.w2_weight_triton_tensor = w2_weight
344 +
345 + # need to delete the original weights to save memory on single GPU
346 + del layer.w13_weight
347 + del layer.w2_weight
348 + layer.w13_weight = None
349 + layer.w2_weight = None
350 + torch.cuda.empty_cache()
351
352 def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int):
353 # Number of tokens in the input tensor.
@@ -404,3 +448,19 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
448 True, # do finalize
449 )[0]
450 return trtllm_gen_output
451 + else:
452 + return triton_kernel_moe_forward(
453 + hidden_states=x,
454 + w1=self.w13_weight_triton_tensor,
455 + w2=self.w2_weight_triton_tensor,
456 + gating_output=router_logits,
457 + topk=top_k,
458 + renormalize=renormalize,
459 + global_num_experts=global_num_experts,
460 + expert_map=expert_map,
461 + w1_bias=layer.w13_bias,
462 + w2_bias=layer.w2_bias,
463 + w1_precision=self.w13_precision_config,
464 + w2_precision=self.w2_precision_config,
465 + apply_router_weight_on_input=apply_router_weight_on_input,
466 + )
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+45 -1
@@ -4,11 +4,55 @@ from typing import Callable, Optional
4
5 import torch
6
7 -from vllm.utils import direct_register_custom_op
7 +from vllm.logger import init_logger
8 +from vllm.platforms import current_platform
9 +from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
10 +
11 +logger = init_logger(__name__)
12
13 OCP_MX_BLOCK_SIZE = 32
14
15
16 +def _swizzle_mxfp4(quant_tensor, scale, num_warps):
17 + """ weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel
18 + """
19 + import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
20 + from triton_kernels.numerics import InFlexData
21 + from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
22 + from triton_kernels.tensor_details import layout
23 + from triton_kernels.tensor_details.layout import StridedLayout
24 + if (current_platform.is_cuda()
25 + and current_platform.is_device_capability(90)
26 + and not is_torch_equal_or_newer("2.8.1")):
27 + logger.warning_once(
28 + "Mxfp4 on hopper is running on torch < 2.8.1, "
29 + "this cause swizling to be disabled, which may "
30 + "cause performance degradation. Please upgrade to torch nightly")
31 + value_layout, value_layout_opts = StridedLayout, dict()
32 + scale_layout, scale_layout_opts = StridedLayout, dict()
33 + else:
34 + value_layout, value_layout_opts = \
35 + layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
36 + scale_layout, scale_layout_opts = (
37 + layout.make_default_matmul_mxfp4_w_scale_layout(
38 + mx_axis=1, num_warps=num_warps))
39 + if current_platform.is_cuda() and \
40 + current_platform.is_device_capability(100):
41 + constraints = {
42 + "is_persistent": True,
43 + "epilogue_subtile": 1,
44 + }
45 + opt_flags.update_opt_flags_constraints(constraints)
46 + # transpose the tensor so that the quantization axis is on dim1
47 + quant_tensor = quant_tensor.transpose(-2, -1)
48 + scale = scale.transpose(-2, -1)
49 + quant_tensor = convert_layout(wrap_torch_tensor(quant_tensor, dtype=FP4),
50 + value_layout, **value_layout_opts)
51 + scale = convert_layout(wrap_torch_tensor(scale), scale_layout,
52 + **scale_layout_opts)
53 + return quant_tensor, InFlexData(), scale
54 +
55 +
56 def _can_support_mxfp4(use_grouped_topk: bool = False,
57 topk_group: Optional[int] = None,
58 num_expert_group: Optional[int] = None,
vllm/model_executor/layers/utils.py
+21
@@ -11,6 +11,27 @@ from vllm.platforms import current_platform
11 from vllm.utils import direct_register_custom_op
12
13
14 +def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
15 + # Shuffle weight along the last dimension so that
16 + # we folded the weights to adjance location
17 + # Example:
18 + # input:
19 + # [[1, 2, 3, 4, 5, 6],
20 + # [7, 8, 9, 10, 11, 12]]
21 + # output:
22 + # [[1, 4, 2, 5, 3, 6],
23 + # [7, 10, 8, 11, 9, 12]]
24 + # This will be used together with triton swiglu kernel
25 + shape = w.shape
26 + N = shape[-1]
27 + first = w[..., :N // 2]
28 + second = w[..., N // 2:]
29 +
30 + stacked = torch.stack((first, second), dim=-1)
31 + w_shuffled = stacked.reshape(shape)
32 + return w_shuffled
33 +
34 +
35 def get_token_bin_counts_and_mask(
36 tokens: torch.Tensor,
37 vocab_size: int,
vllm/utils/__init__.py
+6
@@ -3254,6 +3254,12 @@ def has_deep_gemm() -> bool:
3254 return _has_module("deep_gemm")
3255
3256
3257 +def has_triton_kernels() -> bool:
3258 + """Whether the optional `triton_kernels` package is available."""
3259 +
3260 + return _has_module("triton_kernels")
3261 +
3262 +
3263 def set_process_title(name: str,
3264 suffix: str = "",
3265 append: bool = False) -> None: