| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Test modular OAI Triton MoE.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | |
| 11 | from tests.utils import wait_for_gpu_memory_to_clear |
| 12 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 13 | from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels |
| 14 | |
| 15 | if not has_triton_kernels(): |
| 16 | pytest.skip( |
| 17 | "triton_kernels not found, skipping all related tests", |
| 18 | allow_module_level=True, |
| 19 | ) |
| 20 | |
| 21 | if get_triton_kernels_version() == "3.8": |
| 22 | # 3.8: matmul_ogs -> matmul |
| 23 | from triton_kernels.matmul import FlexCtx, PrecisionConfig |
| 24 | else: |
| 25 | from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig |
| 26 | from triton_kernels.numerics import InFlexData |
| 27 | from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp |
| 28 | from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor |
| 29 | from triton_kernels.testing import assert_close |
| 30 | |
| 31 | from vllm.config import VllmConfig, set_current_vllm_config |
| 32 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 33 | maybe_make_prepare_finalize, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config |
| 36 | from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( |
| 37 | OAITritonExperts, |
| 38 | UnfusedOAITritonExperts, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 41 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs |
| 42 | from vllm.platforms import current_platform |
| 43 | from vllm.utils.math_utils import round_up |
| 44 | from vllm.utils.torch_utils import set_random_seed |
| 45 | |
| 46 | from .utils import make_dummy_moe_config, mxfp4_w_layouts, shuffle_weight |
| 47 | |
| 48 | MNK = [ |
| 49 | (1, 512, 384), |
| 50 | (1, 2880, 2880), |
| 51 | (2, 512, 384), |
| 52 | (2, 2880, 2880), |
| 53 | (16, 2880, 2880), |
| 54 | ] |
| 55 | |
| 56 | |
| 57 | def deepseek_v4_flash_moe_topology(): |
| 58 | """MoE sizes for the DeepSeek-V4-Flash`. |
| 59 | Default weights from: ``deepseek-ai/DeepSeek-V4-Flash`. |
| 60 | ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. |
| 61 | """ |
| 62 | defaults = { |
| 63 | "hidden_size": 4096, |
| 64 | "moe_intermediate_size": 2048, |
| 65 | "n_routed_experts": 256, |
| 66 | "num_experts_per_tok": 6, |
| 67 | } |
| 68 | |
| 69 | return defaults |
| 70 | |
| 71 | |
| 72 | def scaled_deepseek_v4_flash_problem( |
| 73 | *, |
| 74 | dim_scale: int = 8, |
| 75 | expert_scale: int = 8, |
| 76 | ): |
| 77 | """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" |
| 78 | t = deepseek_v4_flash_moe_topology() |
| 79 | k = max(128, t["hidden_size"] // dim_scale) |
| 80 | n = max(64, t["moe_intermediate_size"] // dim_scale) |
| 81 | num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) |
| 82 | topk = t["num_experts_per_tok"] |
| 83 | return k, n, num_experts, topk |
| 84 | |
| 85 | |
| 86 | def unshuffle_weight(w: torch.Tensor): |
| 87 | first = w[..., ::2] |
| 88 | second = w[..., 1::2] |
| 89 | return torch.concat((first, second), dim=-1) |
| 90 | |
| 91 | |
| 92 | def make_weights(dtype, k, n, e): |
| 93 | w1 = torch.randn((e, k, 2 * n), dtype=dtype, device="cuda") |
| 94 | w1_bias = torch.randn((e, 2 * n), dtype=dtype, device="cuda") |
| 95 | |
| 96 | w2 = torch.randn((e, n, k), dtype=dtype, device="cuda") |
| 97 | w2_bias = torch.randn((e, k), dtype=dtype, device="cuda") |
| 98 | |
| 99 | w1_tri = w1.clone() |
| 100 | w2_tri = w2.clone() |
| 101 | |
| 102 | w1_bias_tri = w1_bias.clone() |
| 103 | w2_bias_tri = w2_bias.clone() |
| 104 | w1_bias_tri = w1_bias_tri.to(torch.float32) |
| 105 | w2_bias_tri = w2_bias_tri.to(torch.float32) |
| 106 | |
| 107 | # shuffle weights |
| 108 | w1_tri = shuffle_weight(w1_tri) |
| 109 | w1_bias_tri = shuffle_weight(w1_bias_tri) |
| 110 | |
| 111 | if current_platform.is_rocm(): |
| 112 | k_align, n2_align = 256, 512 |
| 113 | else: |
| 114 | k_align, n2_align = 64, 128 |
| 115 | |
| 116 | w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1] |
| 117 | w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2] |
| 118 | w2_bottom_pad = w1_right_pad // 2 |
| 119 | w2_right_pad = w1_bottom_pad |
| 120 | |
| 121 | w1_tri = F.pad(w1_tri, (0, w1_right_pad, 0, w1_bottom_pad, 0, 0)) |
| 122 | w2_tri = F.pad(w2_tri, (0, w2_right_pad, 0, w2_bottom_pad, 0, 0)) |
| 123 | w1_bias_tri = F.pad(w1_bias_tri, (0, w1_right_pad, 0, 0)) |
| 124 | w2_bias_tri = F.pad(w2_bias_tri, (0, w2_right_pad, 0, 0)) |
| 125 | |
| 126 | # quant triton_weights |
| 127 | w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1) |
| 128 | w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, dtype, axis=1) |
| 129 | w1 = w1[..., :k, : 2 * n] |
| 130 | w1 = unshuffle_weight(w1) |
| 131 | |
| 132 | w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1) |
| 133 | w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, dtype, axis=1) |
| 134 | w2 = w2[..., :n, :k] |
| 135 | |
| 136 | num_warps = 8 |
| 137 | ( |
| 138 | w_layout, |
| 139 | w_layout_opts, |
| 140 | w_scale_layout, |
| 141 | w_scale_layout_opts, |
| 142 | ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps) |
| 143 | |
| 144 | w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts) |
| 145 | w1_scale_tri = convert_layout( |
| 146 | wrap_torch_tensor(w1_scale_tri), |
| 147 | w_scale_layout, |
| 148 | **w_scale_layout_opts, |
| 149 | ) |
| 150 | |
| 151 | w2_tri = convert_layout(wrap_torch_tensor(w2_tri, FP4), w_layout, **w_layout_opts) |
| 152 | w2_scale_tri = convert_layout( |
| 153 | wrap_torch_tensor(w2_scale_tri), |
| 154 | w_scale_layout, |
| 155 | **w_scale_layout_opts, |
| 156 | ) |
| 157 | |
| 158 | w1_precision_config = PrecisionConfig( |
| 159 | **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) |
| 160 | ) |
| 161 | w2_precision_config = PrecisionConfig( |
| 162 | **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) |
| 163 | ) |
| 164 | |
| 165 | return ( |
| 166 | w1, |
| 167 | w2, |
| 168 | w1_bias, |
| 169 | w2_bias, |
| 170 | w1_tri, |
| 171 | w2_tri, |
| 172 | w1_bias_tri, |
| 173 | w2_bias_tri, |
| 174 | w1_precision_config, |
| 175 | w2_precision_config, |
| 176 | w1_bottom_pad, |
| 177 | ) |
| 178 | |
| 179 | |
| 180 | def swiglu(x, alpha: float = 1.702, limit: float = 1.0): |
| 181 | # Note we add an extra bias of 1 to the linear layer |
| 182 | x_glu, x_linear = torch.chunk(x, 2, dim=-1) |
| 183 | if limit is not None: |
| 184 | x_glu = x_glu.clamp(max=limit) |
| 185 | out_glu = x_glu * torch.sigmoid(alpha * x_glu) |
| 186 | if limit is not None: |
| 187 | x_linear = x_linear.clamp(min=-limit, max=limit) |
| 188 | return out_glu * (x_linear + 1) |
| 189 | |
| 190 | |
| 191 | def torch_moe_impl( |
| 192 | hidden_states: torch.Tensor, # (M, K) |
| 193 | w1: torch.Tensor, # (E, K, 2N) |
| 194 | w2: torch.Tensor, # (E, N, K) |
| 195 | w1_bias: torch.Tensor, # (E, 2N) |
| 196 | w2_bias: torch.Tensor, # (E, K) |
| 197 | topk_weights: torch.Tensor, # (M, topk) |
| 198 | topk_ids: torch.Tensor, # (M, topk) |
| 199 | ): |
| 200 | w1 = w1[topk_ids, ...] |
| 201 | w1_bias = w1_bias[topk_ids, ...] |
| 202 | hidden_states = torch.einsum("bekc,bk->bec", w1, hidden_states) + w1_bias |
| 203 | hidden_states = swiglu(hidden_states, limit=7) |
| 204 | |
| 205 | w2 = w2[topk_ids, ...] |
| 206 | w2_bias = w2_bias[topk_ids, ...] |
| 207 | hidden_states = torch.einsum("bekc,bek->bec", w2, hidden_states) + w2_bias |
| 208 | |
| 209 | # Weighted sum of experts |
| 210 | hidden_states = torch.einsum("bec,be->bc", hidden_states, topk_weights) |
| 211 | return hidden_states |
| 212 | |
| 213 | |
| 214 | def oai_triton_moe_impl( |
| 215 | x: torch.Tensor, |
| 216 | w1: torch.Tensor, |
| 217 | w2: torch.Tensor, |
| 218 | w1_scale: PrecisionConfig, |
| 219 | w2_scale: PrecisionConfig, |
| 220 | w1_bias: torch.Tensor | None, |
| 221 | w2_bias: torch.Tensor | None, |
| 222 | num_experts: int, |
| 223 | topk_weights: torch.Tensor, |
| 224 | topk_ids: torch.Tensor, |
| 225 | unfused: bool = False, |
| 226 | ) -> torch.Tensor: |
| 227 | quant_config = mxfp4_w4a16_moe_quant_config( |
| 228 | w1_bias=w1_bias, |
| 229 | w2_bias=w2_bias, |
| 230 | w1_scale=w1_scale, |
| 231 | w2_scale=w2_scale, |
| 232 | ) |
| 233 | moe_config = make_dummy_moe_config() |
| 234 | |
| 235 | if unfused: |
| 236 | fused_experts = UnfusedOAITritonExperts(moe_config, quant_config) |
| 237 | else: |
| 238 | fused_experts = OAITritonExperts(moe_config, quant_config) |
| 239 | |
| 240 | mk = FusedMoEKernel( |
| 241 | maybe_make_prepare_finalize( |
| 242 | moe=moe_config, |
| 243 | quant_config=quant_config, |
| 244 | allow_new_interface=True, |
| 245 | use_monolithic=False, |
| 246 | ), |
| 247 | fused_experts, |
| 248 | ) |
| 249 | |
| 250 | return mk.apply( |
| 251 | hidden_states=x, |
| 252 | w1=w1, |
| 253 | w2=w2, |
| 254 | topk_weights=topk_weights, |
| 255 | topk_ids=topk_ids, |
| 256 | activation=MoEActivation.SWIGLUOAI, |
| 257 | global_num_experts=num_experts, |
| 258 | expert_map=None, |
| 259 | apply_router_weight_on_input=False, |
| 260 | ) |
| 261 | |
| 262 | |
| 263 | @pytest.mark.skipif( |
| 264 | not OAITritonExperts._supports_current_device(), |
| 265 | reason="OAI Triton MoE is not supported on this device.", |
| 266 | ) |
| 267 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 268 | @pytest.mark.parametrize("m,n,k", MNK) |
| 269 | @pytest.mark.parametrize("num_experts", [32, 128]) |
| 270 | @pytest.mark.parametrize("topk", [4]) |
| 271 | @pytest.mark.parametrize("unfused", [True, False]) |
| 272 | def test_oai_triton_moe( |
| 273 | dtype: torch.dtype, |
| 274 | m: int, |
| 275 | n: int, |
| 276 | k: int, |
| 277 | num_experts: int, |
| 278 | topk: int, |
| 279 | unfused: bool, |
| 280 | workspace_init, |
| 281 | ): |
| 282 | wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) |
| 283 | set_random_seed(0) |
| 284 | |
| 285 | ( |
| 286 | w1, |
| 287 | w2, |
| 288 | w1_bias, |
| 289 | w2_bias, |
| 290 | w1_tri, |
| 291 | w2_tri, |
| 292 | w1_bias_tri, |
| 293 | w2_bias_tri, |
| 294 | w1_precision_config, |
| 295 | w2_precision_config, |
| 296 | x_pad, |
| 297 | ) = make_weights(dtype, k, n, num_experts) |
| 298 | |
| 299 | x = torch.randn((m, k), dtype=dtype, device="cuda") |
| 300 | x_tri = F.pad(x, (0, x_pad, 0, 0)) |
| 301 | router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) |
| 302 | topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) |
| 303 | topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) |
| 304 | |
| 305 | with set_current_vllm_config(VllmConfig()): |
| 306 | out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) |
| 307 | |
| 308 | out = oai_triton_moe_impl( |
| 309 | x_tri, |
| 310 | w1_tri, |
| 311 | w2_tri, |
| 312 | w1_precision_config, |
| 313 | w2_precision_config, |
| 314 | w1_bias_tri, |
| 315 | w2_bias_tri, |
| 316 | num_experts, |
| 317 | topk_weights, |
| 318 | topk_ids, |
| 319 | unfused, |
| 320 | ) |
| 321 | out = out[..., :k] |
| 322 | |
| 323 | assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) |
| 324 | |
| 325 | |
| 326 | @pytest.mark.skipif( |
| 327 | not UnfusedOAITritonExperts._supports_current_device(), |
| 328 | reason="Unfused OAI Triton MoE is not supported on this device.", |
| 329 | ) |
| 330 | def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): |
| 331 | """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. |
| 332 | |
| 333 | Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, |
| 334 | with linear dimensions and expert count scaled down for test GPU memory. |
| 335 | """ |
| 336 | wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) |
| 337 | set_random_seed(0) |
| 338 | |
| 339 | k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() |
| 340 | m = 7 |
| 341 | dtype = torch.bfloat16 |
| 342 | |
| 343 | ( |
| 344 | w1, |
| 345 | w2, |
| 346 | w1_bias, |
| 347 | w2_bias, |
| 348 | w1_tri, |
| 349 | w2_tri, |
| 350 | w1_bias_tri, |
| 351 | w2_bias_tri, |
| 352 | w1_precision_config, |
| 353 | w2_precision_config, |
| 354 | x_pad, |
| 355 | ) = make_weights(dtype, k, n, num_experts) |
| 356 | |
| 357 | x = torch.randn((m, k), dtype=dtype, device="cuda") |
| 358 | x_tri = F.pad(x, (0, x_pad, 0, 0)) |
| 359 | router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) |
| 360 | topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) |
| 361 | topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) |
| 362 | |
| 363 | quant_config = mxfp4_w4a16_moe_quant_config( |
| 364 | w1_bias=w1_bias_tri, |
| 365 | w2_bias=w2_bias_tri, |
| 366 | w1_scale=w1_precision_config, |
| 367 | w2_scale=w2_precision_config, |
| 368 | ) |
| 369 | moe_config = make_dummy_moe_config( |
| 370 | num_experts=num_experts, |
| 371 | experts_per_token=topk, |
| 372 | hidden_dim=k, |
| 373 | intermediate_size=n, |
| 374 | ) |
| 375 | experts = UnfusedOAITritonExperts(moe_config, quant_config) |
| 376 | |
| 377 | _, _, N, K, top_k = experts.moe_problem_size(x_tri, w1_tri, w2_tri, topk_ids) |
| 378 | assert top_k == topk |
| 379 | ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( |
| 380 | m, |
| 381 | N, |
| 382 | K, |
| 383 | topk, |
| 384 | num_experts, |
| 385 | num_experts, |
| 386 | None, |
| 387 | MoEActivation.SWIGLUOAI, |
| 388 | ) |
| 389 | workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") |
| 390 | workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") |
| 391 | output = torch.empty(out_shape, dtype=dtype, device="cuda") |
| 392 | |
| 393 | with set_current_vllm_config(VllmConfig()): |
| 394 | out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) |
| 395 | experts.apply( |
| 396 | output=output, |
| 397 | hidden_states=x_tri, |
| 398 | w1=w1_tri, |
| 399 | w2=w2_tri, |
| 400 | topk_weights=topk_weights, |
| 401 | topk_ids=topk_ids, |
| 402 | activation=MoEActivation.SWIGLUOAI, |
| 403 | global_num_experts=num_experts, |
| 404 | expert_map=None, |
| 405 | a1q_scale=None, |
| 406 | a2_scale=None, |
| 407 | workspace13=workspace13, |
| 408 | workspace2=workspace2, |
| 409 | expert_tokens_meta=None, |
| 410 | apply_router_weight_on_input=False, |
| 411 | ) |
| 412 | output = output[..., :k] |
| 413 | |
| 414 | assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) |