| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, fields, is_dataclass |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.import_utils import get_triton_kernels_version, has_triton_kernels |
| 11 | |
| 12 | if not has_triton_kernels(): |
| 13 | pytest.skip( |
| 14 | "triton_kernels not found, skipping all related tests", |
| 15 | allow_module_level=True, |
| 16 | ) |
| 17 | |
| 18 | import triton_kernels.swiglu |
| 19 | |
| 20 | if get_triton_kernels_version() == "3.8": |
| 21 | # 3.8: matmul_ogs -> matmul, matmul_ogs_details -> matmul_details |
| 22 | import triton_kernels.matmul_details.opt_flags as opt_flags |
| 23 | from triton_kernels.matmul import FlexCtx, PrecisionConfig |
| 24 | else: |
| 25 | import triton_kernels.matmul_ogs_details.opt_flags as opt_flags |
| 26 | from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig |
| 27 | from triton_kernels.numerics import InFlexData |
| 28 | from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp |
| 29 | from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor |
| 30 | from triton_kernels.tensor_details import layout |
| 31 | from triton_kernels.testing import assert_close |
| 32 | |
| 33 | from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config |
| 34 | from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( |
| 35 | triton_kernel_moe_forward, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mx_scale_kwargs |
| 38 | from vllm.utils.math_utils import round_up |
| 39 | |
| 40 | from .utils import mxfp4_w_layouts, shuffle_weight |
| 41 | |
| 42 | |
| 43 | def deshuffle(w: torch.Tensor): |
| 44 | first = w[..., ::2] |
| 45 | second = w[..., 1::2] |
| 46 | |
| 47 | deshuffled = torch.concat((first, second), dim=-1) |
| 48 | return deshuffled |
| 49 | |
| 50 | |
| 51 | def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int): |
| 52 | randbits = [torch.randperm(E) for _ in range(M)] |
| 53 | x_list = [ |
| 54 | (-1) ** i |
| 55 | * ((16384 + ((i * 512) % 4096) + bits).to(torch.int16).view(torch.bfloat16)) |
| 56 | for i, bits in enumerate(randbits) |
| 57 | ] |
| 58 | exp_data = torch.stack(x_list).to(device="cuda") # simulating gate_output (M, E) |
| 59 | |
| 60 | # create input tensor |
| 61 | x = torch.randn((M, K), dtype=torch.bfloat16, device="cuda") |
| 62 | w1 = torch.randn((E, 2 * N, K), dtype=torch.bfloat16, device="cuda") |
| 63 | w1_bias = torch.randn((E, 2 * N), dtype=torch.bfloat16, device="cuda") |
| 64 | |
| 65 | w2 = torch.randn((E, K, N), dtype=torch.bfloat16, device="cuda") |
| 66 | w2_bias = torch.randn((E, K), dtype=torch.bfloat16, device="cuda") |
| 67 | |
| 68 | exp_data_tri = exp_data.clone() |
| 69 | x_tri = x.clone() |
| 70 | w1_tri = w1.clone() |
| 71 | w2_tri = w2.clone() |
| 72 | |
| 73 | w1_bias_tri = w1_bias.clone() |
| 74 | w2_bias_tri = w2_bias.clone() |
| 75 | w1_bias_tri = w1_bias_tri.to(torch.float32) |
| 76 | w2_bias_tri = w2_bias_tri.to(torch.float32) |
| 77 | |
| 78 | dtype_dict = { |
| 79 | "bf16": torch.bfloat16, |
| 80 | "fp8_e4m3": torch.float8_e4m3fn, |
| 81 | "fp8_e5m2": torch.float8_e5m2, |
| 82 | } |
| 83 | |
| 84 | x = x.to(dtype_dict[a_dtype]).to(torch.bfloat16) |
| 85 | if w_dtype != "mx4": |
| 86 | # simulate quantization support on reference impl |
| 87 | w1 = w1.to(dtype_dict[w_dtype]).to(torch.bfloat16) |
| 88 | w2 = w2.to(dtype_dict[w_dtype]).to(torch.bfloat16) |
| 89 | |
| 90 | # triton moe kernel use transposed shape for matmul |
| 91 | w1_tri = w1_tri.transpose(-2, -1) |
| 92 | w2_tri = w2_tri.transpose(-2, -1) |
| 93 | |
| 94 | # shuffle weights |
| 95 | w1_tri = shuffle_weight(w1_tri) |
| 96 | w1_bias_tri = shuffle_weight(w1_bias_tri) |
| 97 | |
| 98 | # quant triton_weights |
| 99 | x_tri = x.to(dtype_dict[a_dtype]) |
| 100 | if w_dtype != "mx4": |
| 101 | pytest.skip("NYI") |
| 102 | else: # quantize to mx4 |
| 103 | # Padding alignment depends on the platform. On CDNA4 the scale |
| 104 | # swizzle requires SCALE_K % 8 == 0 (K % 256) and |
| 105 | # SCALE_N % 32 == 0 (2*N % 512), matching the production |
| 106 | # alignment in mxfp4_round_up_hidden_size_and_intermediate_size. |
| 107 | # On CUDA (Hopper) the scale layout pads internally, so the |
| 108 | # original 64/128 alignment is sufficient. |
| 109 | if current_platform.is_rocm(): |
| 110 | k_align, n2_align = 256, 512 |
| 111 | else: |
| 112 | k_align, n2_align = 64, 128 |
| 113 | w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1] |
| 114 | w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2] |
| 115 | |
| 116 | w2_bottom_pad = w1_right_pad // 2 |
| 117 | w2_right_pad = w1_bottom_pad |
| 118 | |
| 119 | x_pad = w1_bottom_pad |
| 120 | |
| 121 | w1_tri = F.pad( |
| 122 | w1_tri, |
| 123 | (0, w1_right_pad, 0, w1_bottom_pad, 0, 0), |
| 124 | mode="constant", |
| 125 | value=0, |
| 126 | ) |
| 127 | w2_tri = F.pad( |
| 128 | w2_tri, |
| 129 | (0, w2_right_pad, 0, w2_bottom_pad, 0, 0), |
| 130 | mode="constant", |
| 131 | value=0, |
| 132 | ) |
| 133 | |
| 134 | w1_bias_tri = F.pad( |
| 135 | w1_bias_tri, (0, w1_right_pad, 0, 0), mode="constant", value=0 |
| 136 | ) |
| 137 | w2_bias_tri = F.pad( |
| 138 | w2_bias_tri, (0, w2_right_pad, 0, 0), mode="constant", value=0 |
| 139 | ) |
| 140 | |
| 141 | x_tri = F.pad(x_tri, (0, x_pad, 0, 0), mode="constant", value=0) |
| 142 | |
| 143 | ( |
| 144 | w_layout, |
| 145 | w_layout_opts, |
| 146 | w_scale_layout, |
| 147 | w_scale_layout_opts, |
| 148 | ) = mxfp4_w_layouts(mx_axis=1, num_warps=num_warps) |
| 149 | |
| 150 | w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1) |
| 151 | w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, torch.bfloat16, axis=1) |
| 152 | |
| 153 | w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1) |
| 154 | w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, torch.bfloat16, axis=1) |
| 155 | |
| 156 | w1_tri = convert_layout( |
| 157 | wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts |
| 158 | ) |
| 159 | w1_scale_tri = convert_layout( |
| 160 | wrap_torch_tensor(w1_scale_tri), |
| 161 | w_scale_layout, |
| 162 | **w_scale_layout_opts, |
| 163 | ) |
| 164 | |
| 165 | w2_tri = convert_layout( |
| 166 | wrap_torch_tensor(w2_tri, FP4), w_layout, **w_layout_opts |
| 167 | ) |
| 168 | w2_scale_tri = convert_layout( |
| 169 | wrap_torch_tensor(w2_scale_tri), |
| 170 | w_scale_layout, |
| 171 | **w_scale_layout_opts, |
| 172 | ) |
| 173 | |
| 174 | pc1 = PrecisionConfig( |
| 175 | **mx_scale_kwargs(w1_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) |
| 176 | ) |
| 177 | pc2 = PrecisionConfig( |
| 178 | **mx_scale_kwargs(w2_scale_tri), flex_ctx=FlexCtx(rhs_data=InFlexData()) |
| 179 | ) |
| 180 | |
| 181 | # tucuate so the rest can run properly |
| 182 | w1 = w1[..., :K, : 2 * N] |
| 183 | w2 = w2[..., :N, :K] |
| 184 | |
| 185 | w1 = deshuffle(w1) |
| 186 | |
| 187 | w1 = w1.transpose(-1, -2).contiguous() |
| 188 | w2 = w2.transpose(-1, -2).contiguous() |
| 189 | |
| 190 | return ( |
| 191 | x, |
| 192 | w1, |
| 193 | w1_bias, |
| 194 | w2, |
| 195 | w2_bias, |
| 196 | exp_data, |
| 197 | x_tri, |
| 198 | w1_tri, |
| 199 | w2_tri, |
| 200 | exp_data_tri, |
| 201 | w1_bias_tri, |
| 202 | w2_bias_tri, |
| 203 | pc1, |
| 204 | pc2, |
| 205 | ) |
| 206 | |
| 207 | |
| 208 | @dataclass |
| 209 | class ModelConfig: |
| 210 | num_hidden_layers: int = 36 |
| 211 | num_experts: int = 128 |
| 212 | experts_per_token: int = 4 |
| 213 | vocab_size: int = 201088 |
| 214 | hidden_size: int = 2880 |
| 215 | intermediate_size: int = 2880 |
| 216 | head_dim: int = 64 |
| 217 | num_attention_heads: int = 64 |
| 218 | num_key_value_heads: int = 8 |
| 219 | sliding_window: int = 128 |
| 220 | initial_context_length: int = 4096 |
| 221 | rope_theta: float = 150000.0 |
| 222 | rope_parameters_factor: float = 32.0 |
| 223 | rope_ntk_alpha: float = 1.0 |
| 224 | rope_ntk_beta: float = 32.0 |
| 225 | |
| 226 | |
| 227 | def swiglu(x, alpha: float = 1.702, limit: float = 1.0): |
| 228 | # Note we add an extra bias of 1 to the linear layer |
| 229 | x_glu, x_linear = torch.chunk(x, 2, dim=-1) |
| 230 | if limit is not None: |
| 231 | x_glu = x_glu.clamp(max=limit) |
| 232 | out_glu = x_glu * torch.sigmoid(alpha * x_glu) |
| 233 | if limit is not None: |
| 234 | x_linear = x_linear.clamp(min=-limit, max=limit) |
| 235 | return out_glu * (x_linear + 1) |
| 236 | |
| 237 | |
| 238 | def oai_moe_forward( |
| 239 | hidden_states: torch.Tensor, # (M, K) |
| 240 | w1: torch.Tensor, # (E, 2N) |
| 241 | w1_bias: torch.Tensor, # (E, 2N, K) |
| 242 | w2: torch.Tensor, # (E, K, N) |
| 243 | w2_bias: torch.Tensor, # (E, N) |
| 244 | gating_output: torch.Tensor, # (M, E) |
| 245 | topk: int, |
| 246 | ): |
| 247 | # model.py 309:330, assuming gating and norm |
| 248 | t = hidden_states |
| 249 | experts = torch.topk(gating_output, k=topk, dim=-1, sorted=True) |
| 250 | expert_weights = torch.nn.functional.softmax(experts.values, dim=1) |
| 251 | expert_indices = experts.indices |
| 252 | |
| 253 | # MLP #1 |
| 254 | mlp1_weight = w1[expert_indices, ...] |
| 255 | mlp1_bias = w1_bias[expert_indices, ...] |
| 256 | t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias |
| 257 | t = swiglu(t, limit=7) |
| 258 | |
| 259 | # MLP #2 |
| 260 | mlp2_weight = w2[expert_indices, ...] |
| 261 | mlp2_bias = w2_bias[expert_indices, ...] |
| 262 | t = torch.einsum("beck,bek->bec", mlp2_weight, t) |
| 263 | t += mlp2_bias |
| 264 | |
| 265 | # Weighted sum of experts |
| 266 | t = torch.einsum("bec,be->bc", t, expert_weights) |
| 267 | |
| 268 | return t |
| 269 | |
| 270 | |
| 271 | @dataclass |
| 272 | class Case: |
| 273 | a_dtype: str |
| 274 | w_dtype: str |
| 275 | |
| 276 | |
| 277 | @pytest.mark.parametrize( |
| 278 | ", ".join(f.name for f in fields(Case)), |
| 279 | [ |
| 280 | tuple(getattr(case, f.name) for f in fields(Case)) |
| 281 | for case in [ |
| 282 | # Case(a_dtype="bf16", w_dtype="bf16"), |
| 283 | # Case(a_dtype="fp8_e4m3", w_dtype="fp8_e5m2"), |
| 284 | Case(a_dtype="bf16", w_dtype="mx4") |
| 285 | ] |
| 286 | ], |
| 287 | ) |
| 288 | @pytest.mark.parametrize("num_token", [2]) |
| 289 | @pytest.mark.parametrize("tp", [1, 2, 4, 8]) |
| 290 | def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init): |
| 291 | if not hasattr(layout, "make_default_matmul_mxfp4_w_layout"): |
| 292 | pytest.skip("make_default_matmul_mxfp4_w_layout not available") |
| 293 | |
| 294 | M = num_token |
| 295 | E = ModelConfig.num_experts |
| 296 | K = ModelConfig.hidden_size |
| 297 | N = ModelConfig.intermediate_size // tp |
| 298 | topk = ModelConfig.experts_per_token |
| 299 | |
| 300 | ( |
| 301 | x, |
| 302 | w1, |
| 303 | w1_bias, |
| 304 | w2, |
| 305 | w2_bias, |
| 306 | exp_data, |
| 307 | x_tri, |
| 308 | w1_tri, |
| 309 | w2_tri, |
| 310 | exp_data_tri, |
| 311 | w1_bias_tri, |
| 312 | w2_bias_tri, |
| 313 | pc1, |
| 314 | pc2, |
| 315 | ) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8) |
| 316 | |
| 317 | if current_platform.is_device_capability_family(100): |
| 318 | constraints = { |
| 319 | "is_persistent": True, |
| 320 | } |
| 321 | opt_flags.update_opt_flags_constraints(constraints) |
| 322 | |
| 323 | if a_dtype == "bf16" and w_dtype == "mx4": |
| 324 | quant_config = mxfp4_w4a16_moe_quant_config( |
| 325 | w1_scale=pc1, |
| 326 | w2_scale=pc2, |
| 327 | w1_bias=w1_bias_tri, |
| 328 | w2_bias=w2_bias_tri, |
| 329 | ) |
| 330 | else: |
| 331 | raise NotImplementedError( |
| 332 | f"Quantization configuration for activation={a_dtype} and weight={w_dtype} " |
| 333 | f"has not been implemented." |
| 334 | ) |
| 335 | |
| 336 | out_triton_monolithic = triton_kernel_moe_forward( |
| 337 | hidden_states=x_tri, |
| 338 | w1=w1_tri, |
| 339 | w2=w2_tri, |
| 340 | gating_output=exp_data_tri, |
| 341 | topk=topk, |
| 342 | renormalize=True, |
| 343 | quant_config=quant_config, |
| 344 | ) |
| 345 | out_triton_monolithic = out_triton_monolithic[..., :K] |
| 346 | |
| 347 | out_ref = oai_moe_forward( |
| 348 | hidden_states=x, |
| 349 | w1=w1, |
| 350 | w1_bias=w1_bias, |
| 351 | w2=w2, |
| 352 | w2_bias=w2_bias, |
| 353 | gating_output=exp_data, |
| 354 | topk=topk, |
| 355 | ) |
| 356 | assert_close(ref=out_ref, tri=out_triton_monolithic, maxtol=0.025, rmstol=0.005) |
| 357 | |
| 358 | |
| 359 | def test_unit_shuffle(): |
| 360 | N = ModelConfig.intermediate_size |
| 361 | K = ModelConfig.hidden_size |
| 362 | m = torch.randn((K, 2 * N), dtype=torch.bfloat16, device="cuda") |
| 363 | |
| 364 | x = torch.randn(K, dtype=torch.bfloat16, device="cuda") |
| 365 | |
| 366 | m_shuffled = shuffle_weight(m) |
| 367 | |
| 368 | out_ref = x @ m |
| 369 | out_ref = swiglu(out_ref, limit=1.0) |
| 370 | |
| 371 | out = x @ m_shuffled |
| 372 | out = triton_kernels.swiglu.swiglu_torch( |
| 373 | out, |
| 374 | alpha=1.702, |
| 375 | precision_config=triton_kernels.swiglu.PrecisionConfig(limit=1.0), |
| 376 | ) |
| 377 | |
| 378 | assert_close(ref=out_ref, tri=out) |
| 379 | |
| 380 | |
| 381 | @pytest.mark.parametrize("n_tokens", [1, 33, 512]) |
| 382 | @pytest.mark.parametrize("n_experts,topk", [(32, 4), (128, 4)]) |
| 383 | def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk): |
| 384 | """routing_data_from_sparse_topk must produce routing structures |
| 385 | identical to make_routing_data for the same topk result.""" |
| 386 | from vllm.model_executor.layers.fused_moe.experts import ( |
| 387 | gpt_oss_triton_kernels_moe as gptoss_moe, |
| 388 | ) |
| 389 | |
| 390 | make_routing_data = gptoss_moe.make_routing_data |
| 391 | routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk |
| 392 | |
| 393 | if gptoss_moe.triton_kernels_version == "3.5.1": |
| 394 | pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+") |
| 395 | if gptoss_moe.triton_kernels_version == "3.8": |
| 396 | pytest.skip( |
| 397 | "3.8 rebuilds routing from raw ids, so routing_data_from_sparse_topk " |
| 398 | "just forwards to make_routing_data and the comparison is vacuous" |
| 399 | ) |
| 400 | |
| 401 | from triton_kernels.topk import topk as topk_fn |
| 402 | |
| 403 | torch.manual_seed(0) |
| 404 | logits = torch.randn(n_tokens, n_experts, dtype=torch.bfloat16, device="cuda") |
| 405 | sparse_topk = topk_fn(logits, topk, apply_softmax=True) |
| 406 | assert not isinstance(sparse_topk, tuple) |
| 407 | |
| 408 | rd_new, gather_new, scatter_new = routing_data_from_sparse_topk( |
| 409 | sparse_topk, n_experts |
| 410 | ) |
| 411 | rd_ref, gather_ref, scatter_ref = make_routing_data( |
| 412 | sparse_topk.indx.to(torch.long), sparse_topk.vals, n_experts |
| 413 | ) |
| 414 | |
| 415 | def assert_equivalent(new, ref, path="RoutingData"): |
| 416 | """Recursively compare tensor-valued leaves of routing structures.""" |
| 417 | if isinstance(ref, torch.Tensor): |
| 418 | torch.testing.assert_close(new, ref, msg=lambda m: f"{path}: {m}") |
| 419 | elif isinstance(ref, (int, float, bool, str)) or ref is None: |
| 420 | assert new == ref, f"{path}: {new!r} != {ref!r}" |
| 421 | elif is_dataclass(ref): |
| 422 | for f in fields(ref): |
| 423 | assert_equivalent( |
| 424 | getattr(new, f.name), getattr(ref, f.name), f"{path}.{f.name}" |
| 425 | ) |
| 426 | elif callable(ref): |
| 427 | assert callable(new), f"{path}: callable vs {type(new)}" |
| 428 | elif hasattr(ref, "__dict__"): |
| 429 | for k, v in vars(ref).items(): |
| 430 | assert_equivalent(getattr(new, k), v, f"{path}.{k}") |
| 431 | else: |
| 432 | assert type(new) is type(ref), f"{path}: {type(new)} vs {type(ref)}" |
| 433 | |
| 434 | assert_equivalent(rd_new, rd_ref) |
| 435 | torch.testing.assert_close(gather_new.src_indx, gather_ref.src_indx) |
| 436 | torch.testing.assert_close(gather_new.dst_indx, gather_ref.dst_indx) |
| 437 | torch.testing.assert_close(scatter_new.src_indx, scatter_ref.src_indx) |
| 438 | torch.testing.assert_close(scatter_new.dst_indx, scatter_ref.dst_indx) |