| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the MOE layers. |
| 4 | |
| 5 | Run `pytest tests/kernels/test_moe.py`. |
| 6 | """ |
| 7 | |
| 8 | import functools |
| 9 | from collections.abc import Callable |
| 10 | from dataclasses import dataclass |
| 11 | from typing import Any |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | from torch.nn import Parameter |
| 16 | from torch.nn import functional as F |
| 17 | |
| 18 | import vllm.model_executor.layers.fused_moe # noqa |
| 19 | import vllm.model_executor.layers.fused_moe.fused_moe as fused_moe_module |
| 20 | from tests.kernels.moe.utils import ( |
| 21 | fused_moe, |
| 22 | make_dummy_moe_config, |
| 23 | modular_triton_fused_moe, |
| 24 | ) |
| 25 | from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe |
| 26 | from vllm.config import VllmConfig, set_current_vllm_config |
| 27 | from vllm.model_executor.layers.fused_moe import ( |
| 28 | MoEActivation, |
| 29 | fused_topk, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 32 | ApplyMoEActivationConfig, |
| 33 | apply_moe_activation, |
| 34 | apply_moe_activation_supported, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.fused_moe.config import ( |
| 37 | FUSED_MOE_UNQUANTIZED_CONFIG, |
| 38 | int4_w4a16_moe_quant_config, |
| 39 | int8_w8a16_moe_quant_config, |
| 40 | ) |
| 41 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 42 | batched_fused_marlin_moe, |
| 43 | fused_marlin_moe, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 46 | moe_use_td_hw_supported, |
| 47 | ) |
| 48 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 49 | marlin_permute_bias, |
| 50 | ) |
| 51 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 52 | rand_marlin_weight_mxfp4_like, |
| 53 | rand_marlin_weight_nvfp4_like, |
| 54 | ) |
| 55 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 56 | marlin_quant_fp8_torch, |
| 57 | ) |
| 58 | from vllm.model_executor.layers.quantization.utils.marlin_utils_test import ( |
| 59 | awq_marlin_quantize, |
| 60 | marlin_quantize, |
| 61 | ) |
| 62 | from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights |
| 63 | from vllm.platforms import current_platform |
| 64 | from vllm.scalar_type import ScalarType, scalar_types |
| 65 | from vllm.triton_utils import tl |
| 66 | from vllm.utils.math_utils import next_power_of_2 |
| 67 | from vllm.utils.torch_utils import set_random_seed |
| 68 | |
| 69 | DEVICE_TYPE = current_platform.device_type |
| 70 | |
| 71 | |
| 72 | def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None: |
| 73 | captured: dict[str, torch.Tensor] = {} |
| 74 | |
| 75 | class FakeKernel: |
| 76 | def __getitem__(self, grid): |
| 77 | def launch(*args, **kwargs) -> None: |
| 78 | captured["a_scale"] = args[4] |
| 79 | |
| 80 | return launch |
| 81 | |
| 82 | monkeypatch.setattr(fused_moe_module, "fused_moe_kernel", FakeKernel()) |
| 83 | |
| 84 | a_scale = torch.tensor(0.5) |
| 85 | fused_moe_module.invoke_fused_moe_triton_kernel( |
| 86 | A=torch.ones((1, 1)), |
| 87 | B=torch.ones((1, 1, 1)), |
| 88 | C=torch.empty((1, 1, 1)), |
| 89 | A_scale=a_scale, |
| 90 | B_scale=torch.ones(1), |
| 91 | topk_weights=torch.ones((1, 1)), |
| 92 | sorted_token_ids=None, |
| 93 | expert_ids=torch.zeros(1, dtype=torch.int32), |
| 94 | num_tokens_post_padded=torch.ones(1, dtype=torch.int32), |
| 95 | mul_routed_weight=True, |
| 96 | top_k=1, |
| 97 | config={"BLOCK_SIZE_M": 1, "BLOCK_SIZE_N": 1, "BLOCK_SIZE_K": 1}, |
| 98 | compute_type=tl.float32, |
| 99 | use_fp8_w8a8=True, |
| 100 | use_int8_w8a8=False, |
| 101 | use_int8_w8a16=False, |
| 102 | use_int4_w4a16=False, |
| 103 | per_channel_quant=False, |
| 104 | ) |
| 105 | |
| 106 | captured_scale = captured["a_scale"] |
| 107 | assert a_scale.ndim == 0 |
| 108 | assert captured_scale.shape == (1,) |
| 109 | assert captured_scale.data_ptr() == a_scale.data_ptr() |
| 110 | |
| 111 | |
| 112 | def iterative_moe( |
| 113 | hidden_states: torch.Tensor, |
| 114 | w1: torch.Tensor, |
| 115 | w2: torch.Tensor, |
| 116 | gating_output: torch.Tensor, |
| 117 | topk: int, |
| 118 | global_num_experts: int, |
| 119 | expert_map: torch.Tensor = None, |
| 120 | renormalize: bool = False, |
| 121 | ) -> torch.Tensor: |
| 122 | """Baseline implementation of fused moe. |
| 123 | |
| 124 | Args: |
| 125 | hidden_states: [*, hidden_size] |
| 126 | w1: [num_experts, intermediate_size * 2, hidden_size] |
| 127 | w2: [num_experts, hidden_size, intermediate_size] |
| 128 | gating_output: [*, num_experts] |
| 129 | expert_map: [num_experts] |
| 130 | |
| 131 | """ |
| 132 | orig_shape = hidden_states.shape |
| 133 | hidden_size = hidden_states.shape[-1] |
| 134 | num_tokens = hidden_states.shape[:-1].numel() |
| 135 | num_experts = w1.shape[0] |
| 136 | intermediate_size = w2.shape[-1] |
| 137 | dtype = hidden_states.dtype |
| 138 | |
| 139 | hidden_states = hidden_states.view(num_tokens, hidden_size) |
| 140 | gating_output = gating_output.view(num_tokens, global_num_experts) |
| 141 | topk_weights = gating_output.softmax(dim=-1, dtype=torch.float) |
| 142 | topk_weights, selected_experts = topk_weights.topk(topk, dim=-1) |
| 143 | if renormalize: |
| 144 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 145 | topk_weights = topk_weights.to(dtype) |
| 146 | |
| 147 | if expert_map is not None: |
| 148 | selected_experts = expert_map[selected_experts] |
| 149 | |
| 150 | final_hidden_states = None |
| 151 | for expert_idx in range(num_experts): |
| 152 | expert_w1 = w1[expert_idx] |
| 153 | expert_w2 = w2[expert_idx] |
| 154 | expert_mask = selected_experts == expert_idx |
| 155 | expert_weights = (topk_weights * expert_mask).sum(dim=-1, keepdim=True) |
| 156 | x = F.linear(hidden_states, expert_w1) |
| 157 | gate = F.silu(x[:, :intermediate_size]) |
| 158 | x = x[:, intermediate_size:] * gate |
| 159 | x = F.linear(x, expert_w2) |
| 160 | current_hidden_states = x * expert_weights |
| 161 | if final_hidden_states is None: |
| 162 | final_hidden_states = current_hidden_states |
| 163 | else: |
| 164 | final_hidden_states = final_hidden_states + current_hidden_states |
| 165 | |
| 166 | return final_hidden_states.view(orig_shape) # type: ignore |
| 167 | |
| 168 | |
| 169 | NUM_EXPERTS = [8, 64, 192] |
| 170 | NUM_EXPERTS_LARGE = [128, 256] |
| 171 | EP_SIZE = [1, 4] |
| 172 | TOP_KS = [2, 6] |
| 173 | TOP_KS_SMALL = [1, 2] |
| 174 | |
| 175 | MOE_MARLIN_QUANT_TEST_CONFIGS = [ |
| 176 | # AWQ-INT4 |
| 177 | {"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]}, |
| 178 | # GPTQ-INT4 |
| 179 | { |
| 180 | "b_type": scalar_types.uint4b8, |
| 181 | "group_blocks": [-1, 2, 4, 8], |
| 182 | }, |
| 183 | # GPTQ-INT8 |
| 184 | { |
| 185 | "b_type": scalar_types.uint8b128, |
| 186 | "group_blocks": [-1, 2, 4, 8], |
| 187 | }, |
| 188 | # FP8 |
| 189 | {"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]}, |
| 190 | # NVFP4 |
| 191 | {"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]}, |
| 192 | # MXFP4 |
| 193 | { |
| 194 | "a_type": [scalar_types.bfloat16], |
| 195 | "b_type": scalar_types.float4_e2m1f, |
| 196 | "c_type": [scalar_types.bfloat16], |
| 197 | "group_blocks": [2], |
| 198 | }, |
| 199 | # MXFP8 |
| 200 | { |
| 201 | "a_type": [scalar_types.bfloat16], |
| 202 | "b_type": scalar_types.float8_e4m3fn, |
| 203 | "c_type": [scalar_types.bfloat16], |
| 204 | "group_blocks": [2], |
| 205 | }, |
| 206 | # AWQ-INT4 with INT8 activation |
| 207 | { |
| 208 | "a_type": [scalar_types.int8], |
| 209 | "b_type": scalar_types.uint4, |
| 210 | "group_blocks": [-1, 2, 4, 8], |
| 211 | }, |
| 212 | # GPTQ-INT4 with INT8 activation |
| 213 | { |
| 214 | "a_type": [scalar_types.int8], |
| 215 | "b_type": scalar_types.uint4b8, |
| 216 | "group_blocks": [-1, 2, 4, 8], |
| 217 | }, |
| 218 | # GPTQ-INT4 with FP8 activation |
| 219 | { |
| 220 | "a_type": [scalar_types.float8_e4m3fn], |
| 221 | "b_type": scalar_types.uint4b8, |
| 222 | "group_blocks": [-1, 2, 4, 8], |
| 223 | }, |
| 224 | # AWQ-INT4 with FP8 activation |
| 225 | { |
| 226 | "a_type": [scalar_types.float8_e4m3fn], |
| 227 | "b_type": scalar_types.uint4, |
| 228 | "group_blocks": [-1, 2, 4, 8], |
| 229 | }, |
| 230 | # MXFP4 with FP8 activation |
| 231 | { |
| 232 | "a_type": [scalar_types.float8_e4m3fn], |
| 233 | "b_type": scalar_types.float4_e2m1f, |
| 234 | "c_type": [scalar_types.bfloat16], |
| 235 | "group_blocks": [2], |
| 236 | }, |
| 237 | ] |
| 238 | |
| 239 | FUSED_MOE_MNK_FACTORS = [ |
| 240 | (1, 128, 128), |
| 241 | (1, 2048, 128), |
| 242 | (33, 2048, 128), |
| 243 | (32768, 2048, 511), |
| 244 | (40000, 1024, 1024), |
| 245 | ] |
| 246 | |
| 247 | FUSED_MOE_MNK_FACTORS_SMALL_M = [ |
| 248 | (1, 128, 128), |
| 249 | (1, 2048, 128), |
| 250 | (2, 2048, 128), |
| 251 | (2, 2048, 511), |
| 252 | ] |
| 253 | |
| 254 | FUSED_MOE_WN16_MNK_FACTORS = [ |
| 255 | (1, 128, 128), |
| 256 | (1, 1024, 1024), |
| 257 | (32, 2048, 128), |
| 258 | (222, 2048, 1024), |
| 259 | ] |
| 260 | |
| 261 | vllm_config = VllmConfig() |
| 262 | |
| 263 | |
| 264 | def run_moe_test( |
| 265 | baseline: Callable | torch.Tensor, |
| 266 | moe_fn: Callable, |
| 267 | a: torch.Tensor, |
| 268 | w1: torch.Tensor, |
| 269 | w2: torch.Tensor, |
| 270 | score: torch.Tensor, |
| 271 | topk: int, |
| 272 | global_num_experts: int = -1, |
| 273 | expert_map: torch.Tensor | None = None, |
| 274 | padding: bool = False, |
| 275 | use_compile: bool = False, |
| 276 | use_cudagraph: bool = False, |
| 277 | atol: float = 2e-2, |
| 278 | rtol: float = 0, |
| 279 | ) -> torch.Tensor: |
| 280 | if isinstance(baseline, torch.Tensor): |
| 281 | baseline_output = baseline |
| 282 | else: |
| 283 | baseline_output = baseline( |
| 284 | a, |
| 285 | w1, |
| 286 | w2, |
| 287 | score, |
| 288 | topk, |
| 289 | global_num_experts=global_num_experts, |
| 290 | expert_map=expert_map, |
| 291 | ) |
| 292 | |
| 293 | # Pad the weight if moe padding is enabled |
| 294 | if padding: |
| 295 | w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128] |
| 296 | w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128] |
| 297 | |
| 298 | if use_compile: |
| 299 | moe_fn = torch.compile(moe_fn, backend="inductor", fullgraph=True) |
| 300 | torch._dynamo.mark_dynamic(a, 0) |
| 301 | torch._dynamo.mark_dynamic(score, 0) |
| 302 | |
| 303 | test_output = moe_fn( |
| 304 | a, |
| 305 | w1, |
| 306 | w2, |
| 307 | score, |
| 308 | topk, |
| 309 | global_num_experts=global_num_experts, |
| 310 | expert_map=expert_map, |
| 311 | ) |
| 312 | |
| 313 | if use_cudagraph: |
| 314 | test_output.fill_(0) |
| 315 | stream = torch.cuda.Stream() |
| 316 | graph = torch.cuda.CUDAGraph() |
| 317 | with torch.cuda.graph(graph, stream=stream): |
| 318 | test_output = moe_fn( |
| 319 | a, |
| 320 | w1, |
| 321 | w2, |
| 322 | score, |
| 323 | topk, |
| 324 | global_num_experts=global_num_experts, |
| 325 | expert_map=expert_map, |
| 326 | ) |
| 327 | torch.accelerator.synchronize() |
| 328 | graph.replay() |
| 329 | torch.accelerator.synchronize() |
| 330 | |
| 331 | torch.testing.assert_close(test_output, baseline_output, atol=atol, rtol=rtol) |
| 332 | |
| 333 | return baseline_output |
| 334 | |
| 335 | |
| 336 | @pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS) |
| 337 | @pytest.mark.parametrize("e", NUM_EXPERTS) |
| 338 | @pytest.mark.parametrize("topk", TOP_KS) |
| 339 | @pytest.mark.parametrize("ep_size", EP_SIZE) |
| 340 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 341 | @pytest.mark.parametrize("padding", [True, False]) |
| 342 | @pytest.mark.parametrize("use_td", [False, True]) |
| 343 | def test_fused_moe( |
| 344 | m: int, |
| 345 | n: int, |
| 346 | k: int, |
| 347 | e: int, |
| 348 | topk: int, |
| 349 | ep_size: int, |
| 350 | dtype: torch.dtype, |
| 351 | padding: bool, |
| 352 | use_td: bool, |
| 353 | monkeypatch, |
| 354 | workspace_init, |
| 355 | ): |
| 356 | if use_td and not hasattr(tl, "make_tensor_descriptor"): |
| 357 | pytest.skip("Triton < 3.6 lacks tl.make_tensor_descriptor") |
| 358 | if use_td and not moe_use_td_hw_supported(): |
| 359 | pytest.skip( |
| 360 | "tensor_descriptor.gather requires XPU or NVIDIA Blackwell " |
| 361 | "(sm100+); lowers to tile::gather4 (tcgen05/TMEM), which ptxas " |
| 362 | "rejects on Hopper (sm90) and earlier" |
| 363 | ) |
| 364 | monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") |
| 365 | set_random_seed(7) |
| 366 | |
| 367 | # |
| 368 | # Setup test data |
| 369 | # |
| 370 | |
| 371 | # |
| 372 | # Setup test data |
| 373 | # |
| 374 | |
| 375 | a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 376 | w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 377 | w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 378 | |
| 379 | score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype) |
| 380 | |
| 381 | if ep_size > 1: |
| 382 | local_e = e // ep_size |
| 383 | e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32) |
| 384 | e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32) |
| 385 | e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32) |
| 386 | w1 = w1[e_ids] |
| 387 | w2 = w2[e_ids] |
| 388 | else: |
| 389 | e_map = None |
| 390 | |
| 391 | # |
| 392 | # Setup test functions |
| 393 | # |
| 394 | quant_config = FUSED_MOE_UNQUANTIZED_CONFIG |
| 395 | |
| 396 | m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) |
| 397 | |
| 398 | def m_fused_moe( |
| 399 | a: torch.Tensor, |
| 400 | w1: torch.Tensor, |
| 401 | w2: torch.Tensor, |
| 402 | score: torch.Tensor, |
| 403 | topk: int, |
| 404 | global_num_experts: int = -1, |
| 405 | expert_map: torch.Tensor | None = None, |
| 406 | ) -> torch.Tensor: |
| 407 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 408 | return m_fused_moe_fn.apply( |
| 409 | a, |
| 410 | w1, |
| 411 | w2, |
| 412 | topk_weights, |
| 413 | topk_ids, |
| 414 | activation=MoEActivation.SILU, |
| 415 | global_num_experts=global_num_experts, |
| 416 | expert_map=expert_map, |
| 417 | apply_router_weight_on_input=False, |
| 418 | ) |
| 419 | |
| 420 | fused_moe_fn = functools.partial(fused_moe, renormalize=False) |
| 421 | |
| 422 | # |
| 423 | # Run tests |
| 424 | # |
| 425 | runner = functools.partial( |
| 426 | run_moe_test, |
| 427 | a=a, |
| 428 | w1=w1, |
| 429 | w2=w2, |
| 430 | score=score, |
| 431 | topk=topk, |
| 432 | global_num_experts=e, |
| 433 | expert_map=e_map, |
| 434 | padding=padding, |
| 435 | ) |
| 436 | |
| 437 | # Note: for now use_compile will error out if the problem size is |
| 438 | # large enough to trigger chunking. I'm leaving the flag and |
| 439 | # setup code in case we are able to revisit this later. |
| 440 | use_compile = False |
| 441 | |
| 442 | use_cudagraph = n >= 1024 and k >= 1024 and current_platform.is_cuda_alike() |
| 443 | |
| 444 | with set_current_vllm_config(vllm_config): |
| 445 | baseline_output = runner(torch_moe, iterative_moe) |
| 446 | runner( |
| 447 | baseline_output, |
| 448 | fused_moe_fn, |
| 449 | use_compile=use_compile, |
| 450 | use_cudagraph=use_cudagraph, |
| 451 | ) |
| 452 | runner( |
| 453 | baseline_output, |
| 454 | m_fused_moe, |
| 455 | use_compile=use_compile, |
| 456 | use_cudagraph=use_cudagraph, |
| 457 | ) |
| 458 | |
| 459 | |
| 460 | def test_fused_shared_expert_alignment(workspace_init): |
| 461 | set_random_seed(7) |
| 462 | m, n, k = 4, 64, 128 |
| 463 | routed_experts = 8 |
| 464 | physical_experts = routed_experts + 1 |
| 465 | dtype = torch.bfloat16 |
| 466 | |
| 467 | a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 468 | w1 = torch.randn((physical_experts, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 469 | w2 = torch.randn((physical_experts, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 470 | topk_ids = torch.tensor( |
| 471 | [[0, 8], [1, 8], [2, 8], [3, 8]], device=DEVICE_TYPE, dtype=torch.int32 |
| 472 | ) |
| 473 | topk_weights = torch.tensor( |
| 474 | [[0.5, 1.0]] * m, device=DEVICE_TYPE, dtype=torch.float32 |
| 475 | ) |
| 476 | |
| 477 | moe_config = make_dummy_moe_config( |
| 478 | num_experts=physical_experts, |
| 479 | experts_per_token=2, |
| 480 | hidden_dim=k, |
| 481 | intermediate_size=n, |
| 482 | in_dtype=dtype, |
| 483 | max_num_tokens=m, |
| 484 | ) |
| 485 | modular_moe = modular_triton_fused_moe(moe_config, FUSED_MOE_UNQUANTIZED_CONFIG) |
| 486 | |
| 487 | with set_current_vllm_config(vllm_config): |
| 488 | expected = torch_experts(a, w1, w2, topk_weights, topk_ids) |
| 489 | actual = modular_moe.apply( |
| 490 | hidden_states=a, |
| 491 | w1=w1, |
| 492 | w2=w2, |
| 493 | topk_weights=topk_weights, |
| 494 | topk_ids=topk_ids, |
| 495 | activation=MoEActivation.SILU, |
| 496 | global_num_experts=routed_experts, |
| 497 | expert_map=None, |
| 498 | apply_router_weight_on_input=False, |
| 499 | ) |
| 500 | |
| 501 | torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0) |
| 502 | |
| 503 | |
| 504 | def test_fused_moe_int64_overflow(workspace_init): |
| 505 | """Regression test for int32 overflow in stride*offset products. |
| 506 | |
| 507 | With large M, stride_cm * offs_token can exceed int32 max. Verifies |
| 508 | the offs_token int64 cast (fix for #34413) prevents overflow and |
| 509 | produces correct results. |
| 510 | |
| 511 | Reproduces the scenario from PR #34279. |
| 512 | """ |
| 513 | # ~12 GB GPU memory needed for intermediate caches |
| 514 | free_mem = torch.accelerator.get_memory_info()[0] |
| 515 | if free_mem < 12 * 1024**3: |
| 516 | pytest.skip("Insufficient GPU memory for overflow test") |
| 517 | |
| 518 | set_random_seed(7) |
| 519 | |
| 520 | m, n, k, e, topk = 100000, 2048, 1024, 8, 6 |
| 521 | dtype = torch.bfloat16 |
| 522 | |
| 523 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 524 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 525 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 526 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 527 | |
| 528 | # Verify the test exercises the overflow condition: |
| 529 | # C has shape (M, topk, N) where N = w1.size(1) = 2*n |
| 530 | # stride_cm = C.stride(1) = N, max offs_token = M * topk |
| 531 | # Product must exceed int32 max for this test to be meaningful |
| 532 | N = w1.size(1) |
| 533 | assert N * m * topk > 2**31 - 1, "Test params don't trigger int32 overflow" |
| 534 | |
| 535 | fused_moe_fn = functools.partial(fused_moe, renormalize=False) |
| 536 | |
| 537 | with set_current_vllm_config(vllm_config): |
| 538 | run_moe_test( |
| 539 | torch_moe, |
| 540 | fused_moe_fn, |
| 541 | a=a, |
| 542 | w1=w1, |
| 543 | w2=w2, |
| 544 | score=score, |
| 545 | topk=topk, |
| 546 | global_num_experts=e, |
| 547 | ) |
| 548 | |
| 549 | |
| 550 | @pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS_SMALL_M) |
| 551 | @pytest.mark.parametrize("e", NUM_EXPERTS_LARGE) |
| 552 | @pytest.mark.parametrize("topk", TOP_KS_SMALL) |
| 553 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 554 | @pytest.mark.parametrize("padding", [True, False]) |
| 555 | def test_naive_block_assignment_moe( |
| 556 | m: int, |
| 557 | n: int, |
| 558 | k: int, |
| 559 | e: int, |
| 560 | topk: int, |
| 561 | dtype: torch.dtype, |
| 562 | padding: bool, |
| 563 | monkeypatch, |
| 564 | workspace_init, |
| 565 | ): |
| 566 | set_random_seed(7) |
| 567 | |
| 568 | # |
| 569 | # Setup test data |
| 570 | # |
| 571 | |
| 572 | # |
| 573 | # Setup test data |
| 574 | # |
| 575 | |
| 576 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 577 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 578 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 579 | |
| 580 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 581 | |
| 582 | e_map = None |
| 583 | |
| 584 | # |
| 585 | # Setup test functions |
| 586 | # |
| 587 | quant_config = FUSED_MOE_UNQUANTIZED_CONFIG |
| 588 | |
| 589 | m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) |
| 590 | |
| 591 | def m_fused_moe( |
| 592 | a: torch.Tensor, |
| 593 | w1: torch.Tensor, |
| 594 | w2: torch.Tensor, |
| 595 | score: torch.Tensor, |
| 596 | topk: int, |
| 597 | global_num_experts: int = -1, |
| 598 | expert_map: torch.Tensor | None = None, |
| 599 | ) -> torch.Tensor: |
| 600 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 601 | return m_fused_moe_fn.apply( |
| 602 | a, |
| 603 | w1, |
| 604 | w2, |
| 605 | topk_weights, |
| 606 | topk_ids, |
| 607 | activation=MoEActivation.SILU, |
| 608 | global_num_experts=global_num_experts, |
| 609 | expert_map=expert_map, |
| 610 | apply_router_weight_on_input=False, |
| 611 | ) |
| 612 | |
| 613 | fused_moe_fn = functools.partial(fused_moe, renormalize=False) |
| 614 | |
| 615 | # |
| 616 | # Run tests |
| 617 | # |
| 618 | runner = functools.partial( |
| 619 | run_moe_test, |
| 620 | a=a, |
| 621 | w1=w1, |
| 622 | w2=w2, |
| 623 | score=score, |
| 624 | topk=topk, |
| 625 | global_num_experts=e, |
| 626 | expert_map=e_map, |
| 627 | padding=padding, |
| 628 | ) |
| 629 | |
| 630 | # Note: for now use_compile will error out if the problem size is |
| 631 | # large enough to trigger chunking. I'm leaving the flag and |
| 632 | # setup code in case we are able to revisit this later. |
| 633 | use_compile = False |
| 634 | |
| 635 | use_cudagraph = n >= 1024 and k >= 1024 and current_platform.is_cuda_alike() |
| 636 | |
| 637 | with set_current_vllm_config(vllm_config): |
| 638 | baseline_output = runner(torch_moe, iterative_moe) |
| 639 | runner( |
| 640 | baseline_output, |
| 641 | fused_moe_fn, |
| 642 | use_compile=use_compile, |
| 643 | use_cudagraph=use_cudagraph, |
| 644 | ) |
| 645 | runner( |
| 646 | baseline_output, |
| 647 | m_fused_moe, |
| 648 | use_compile=use_compile, |
| 649 | use_cudagraph=use_cudagraph, |
| 650 | ) |
| 651 | |
| 652 | |
| 653 | @pytest.mark.parametrize("m,n,k", FUSED_MOE_WN16_MNK_FACTORS) |
| 654 | @pytest.mark.parametrize("e", NUM_EXPERTS) |
| 655 | @pytest.mark.parametrize("topk", TOP_KS) |
| 656 | @pytest.mark.parametrize("ep_size", EP_SIZE) |
| 657 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 658 | @pytest.mark.parametrize("group_size", [64, 128]) |
| 659 | @pytest.mark.parametrize("has_zp", [True, False]) |
| 660 | @pytest.mark.parametrize("weight_bits", [4, 8]) |
| 661 | @pytest.mark.skipif( |
| 662 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 663 | reason="MoE WNA16 Triton kernels require CUDA/ROCm or XPU.", |
| 664 | ) |
| 665 | def test_fused_moe_wn16( |
| 666 | m: int, |
| 667 | n: int, |
| 668 | k: int, |
| 669 | e: int, |
| 670 | topk: int, |
| 671 | ep_size: int, |
| 672 | dtype: torch.dtype, |
| 673 | group_size: int, |
| 674 | has_zp: bool, |
| 675 | weight_bits: int, |
| 676 | ): |
| 677 | a = torch.randn((m, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 678 | w1 = torch.randn((e, 2 * n, k), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 679 | w2 = torch.randn((e, k, n), device=DEVICE_TYPE, dtype=dtype) / 10 |
| 680 | score = torch.randn((m, e), device=DEVICE_TYPE, dtype=dtype) |
| 681 | |
| 682 | if weight_bits == 4: |
| 683 | pack_factor = 2 |
| 684 | quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8 |
| 685 | elif weight_bits == 8: |
| 686 | pack_factor = 1 |
| 687 | quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128 |
| 688 | |
| 689 | w1_ref = w1.clone() |
| 690 | w2_ref = w2.clone() |
| 691 | w1_qweight = torch.empty( |
| 692 | (e, 2 * n, k // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8 |
| 693 | ) |
| 694 | w2_qweight = torch.empty( |
| 695 | (e, k, n // pack_factor), device=DEVICE_TYPE, dtype=torch.uint8 |
| 696 | ) |
| 697 | w1_scales = torch.empty( |
| 698 | (e, 2 * n, k // group_size), device=DEVICE_TYPE, dtype=dtype |
| 699 | ) |
| 700 | w2_scales = torch.empty((e, k, n // group_size), device=DEVICE_TYPE, dtype=dtype) |
| 701 | w1_qzeros = torch.empty( |
| 702 | (e, 2 * n // pack_factor, k // group_size), |
| 703 | device=DEVICE_TYPE, |
| 704 | dtype=torch.uint8, |
| 705 | ) |
| 706 | w2_qzeros = torch.empty( |
| 707 | (e, k // pack_factor, n // group_size), device=DEVICE_TYPE, dtype=torch.uint8 |
| 708 | ) |
| 709 | |
| 710 | for i in range(e * 2): |
| 711 | expert_id = i % e |
| 712 | if i // e == 0: |
| 713 | w, w_ref, w_qweight, w_scales, w_qzeros = ( |
| 714 | w1, |
| 715 | w1_ref, |
| 716 | w1_qweight, |
| 717 | w1_scales, |
| 718 | w1_qzeros, |
| 719 | ) |
| 720 | else: |
| 721 | w, w_ref, w_qweight, w_scales, w_qzeros = ( |
| 722 | w2, |
| 723 | w2_ref, |
| 724 | w2_qweight, |
| 725 | w2_scales, |
| 726 | w2_qzeros, |
| 727 | ) |
| 728 | weight, qweight, scales, qzeros = quantize_weights( |
| 729 | w[expert_id].T, quant_type, group_size, has_zp, False |
| 730 | ) |
| 731 | weight = weight.T |
| 732 | qweight = qweight.T.contiguous().to(torch.uint8) |
| 733 | scales = scales.T |
| 734 | if has_zp: |
| 735 | qzeros = qzeros.T.contiguous().to(torch.uint8) |
| 736 | if weight_bits == 4: |
| 737 | qweight = qweight[:, 1::2] * 16 + qweight[:, ::2] |
| 738 | if has_zp: |
| 739 | qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :] |
| 740 | |
| 741 | w_ref[expert_id] = weight |
| 742 | w_qweight[expert_id] = qweight |
| 743 | w_scales[expert_id] = scales |
| 744 | if has_zp: |
| 745 | w_qzeros[expert_id] = qzeros |
| 746 | |
| 747 | if ep_size > 1: |
| 748 | local_e = e // ep_size |
| 749 | e_ids = torch.randint(0, e, (local_e,), device=DEVICE_TYPE, dtype=torch.int32) |
| 750 | e_map = torch.full((e,), -1, device=DEVICE_TYPE, dtype=torch.int32) |
| 751 | e_map[e_ids] = torch.arange(local_e, device=DEVICE_TYPE, dtype=torch.int32) |
| 752 | w1_ref = w1_ref[e_ids] |
| 753 | w2_ref = w2_ref[e_ids] |
| 754 | w1_qweight = w1_qweight[e_ids] |
| 755 | w2_qweight = w2_qweight[e_ids] |
| 756 | w1_scales = w1_scales[e_ids] |
| 757 | w2_scales = w2_scales[e_ids] |
| 758 | w1_qzeros = w1_qzeros[e_ids] |
| 759 | w2_qzeros = w2_qzeros[e_ids] |
| 760 | else: |
| 761 | e_map = None |
| 762 | |
| 763 | if weight_bits == 4: |
| 764 | quant_config_builder = int4_w4a16_moe_quant_config |
| 765 | else: |
| 766 | assert weight_bits == 8 |
| 767 | quant_config_builder = int8_w8a16_moe_quant_config |
| 768 | |
| 769 | quant_config = quant_config_builder( |
| 770 | w1_scale=w1_scales, |
| 771 | w2_scale=w2_scales, |
| 772 | w1_zp=w1_qzeros if has_zp else None, |
| 773 | w2_zp=w2_qzeros if has_zp else None, |
| 774 | block_shape=[0, group_size], |
| 775 | ) |
| 776 | |
| 777 | with set_current_vllm_config(vllm_config): |
| 778 | triton_output = fused_moe( |
| 779 | a, |
| 780 | w1_qweight, |
| 781 | w2_qweight, |
| 782 | score, |
| 783 | topk, |
| 784 | renormalize=False, |
| 785 | global_num_experts=e, |
| 786 | expert_map=e_map, |
| 787 | quant_config=quant_config, |
| 788 | ) |
| 789 | torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, expert_map=e_map) |
| 790 | |
| 791 | torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) |
| 792 | |
| 793 | |
| 794 | MARLIN_MOE_SCENARIOS = [ |
| 795 | # (m, n, k, e, topk, ep_size) |
| 796 | # N>=256 required for Marlin kernel thread config for MXFP8. |
| 797 | # Single token, small matrices |
| 798 | (1, 128, 256, 5, 2, 1), |
| 799 | # Single token, large matrices |
| 800 | (1, 1024, 2048, 5, 2, 1), |
| 801 | # Unaligned m, small matrices |
| 802 | (133, 256, 256, 5, 2, 1), |
| 803 | # Unaligned m, large matrices |
| 804 | (133, 1024, 2048, 12, 3, 1), |
| 805 | # Aligned batch, small matrices |
| 806 | (128, 256, 256, 5, 2, 1), |
| 807 | # Aligned batch, large matrices |
| 808 | (128, 1024, 2048, 12, 3, 1), |
| 809 | # Expert parallelism |
| 810 | (64, 1024, 2048, 12, 3, 4), |
| 811 | ] |
| 812 | |
| 813 | |
| 814 | def marlin_moe_generate_valid_test_cases(): |
| 815 | import itertools |
| 816 | |
| 817 | def is_valid( |
| 818 | a_type, |
| 819 | b_type, |
| 820 | c_type, |
| 821 | group_blocks, |
| 822 | m, |
| 823 | n, |
| 824 | k, |
| 825 | e, |
| 826 | topk, |
| 827 | ep_size, |
| 828 | ): |
| 829 | group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 |
| 830 | if group_size > 0 and k % group_size != 0: |
| 831 | return False |
| 832 | if group_size in [k, n]: |
| 833 | return False |
| 834 | if b_type == scalar_types.float8_e4m3fn and group_size == 32: |
| 835 | return False |
| 836 | return a_type.size_bits < 16 or a_type is c_type |
| 837 | |
| 838 | cases = [] |
| 839 | for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS: |
| 840 | f16_types = [scalar_types.float16] |
| 841 | inner_combinations = list( |
| 842 | itertools.product( |
| 843 | quant_test_config.get("a_type", f16_types), |
| 844 | [quant_test_config["b_type"]], |
| 845 | quant_test_config.get("c_type", f16_types), |
| 846 | quant_test_config["group_blocks"], |
| 847 | ) |
| 848 | ) |
| 849 | |
| 850 | for sub_case in inner_combinations: |
| 851 | if ( |
| 852 | sub_case[0] == scalar_types.float8_e4m3fn |
| 853 | and not current_platform.is_device_capability(89) |
| 854 | and not current_platform.is_device_capability_family(120) |
| 855 | ): |
| 856 | continue |
| 857 | |
| 858 | for scenario in MARLIN_MOE_SCENARIOS: |
| 859 | m, n, k, e, topk, ep_size = scenario |
| 860 | args = sub_case + (m, n, k, e, topk, ep_size) |
| 861 | if is_valid(*args): |
| 862 | cases.append(args) |
| 863 | return cases |
| 864 | |
| 865 | |
| 866 | @dataclass |
| 867 | class MarlinMoEWeightData: |
| 868 | w_ref: torch.Tensor |
| 869 | qweight: torch.Tensor |
| 870 | scales: torch.Tensor |
| 871 | global_scale: torch.Tensor | None |
| 872 | a_scales_factor: torch.Tensor | None |
| 873 | zeros: torch.Tensor | None |
| 874 | marlin_bias: torch.Tensor | None |
| 875 | |
| 876 | @staticmethod |
| 877 | def make( |
| 878 | w: torch.Tensor, |
| 879 | quant_type: ScalarType, |
| 880 | group_size: int, |
| 881 | bias: torch.Tensor | None = None, |
| 882 | input_type: ScalarType = None, |
| 883 | ) -> "MarlinMoEWeightData": |
| 884 | assert w.ndim == 3 |
| 885 | |
| 886 | has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] |
| 887 | if input_type == scalar_types.int8: |
| 888 | input_dtype = torch.int8 |
| 889 | elif input_type == scalar_types.float8_e4m3fn: |
| 890 | input_dtype = torch.float8_e4m3fn |
| 891 | else: |
| 892 | input_dtype = w.dtype |
| 893 | |
| 894 | w_ref_l: list[torch.Tensor] = [] |
| 895 | qweight_l: list[torch.Tensor] = [] |
| 896 | scales_l: list[torch.Tensor] = [] |
| 897 | global_scale_l: list[torch.Tensor] = [] |
| 898 | zeros_l: list[torch.Tensor] = [] |
| 899 | bias_l: list[torch.Tensor] = [] |
| 900 | |
| 901 | for i in range(w.shape[0]): |
| 902 | if quant_type == scalar_types.float4_e2m1f: |
| 903 | if group_size == 16: |
| 904 | w_ref, qweight, scales, global_scale = ( |
| 905 | rand_marlin_weight_nvfp4_like( |
| 906 | w[i], group_size, input_dtype=input_dtype |
| 907 | ) |
| 908 | ) |
| 909 | else: |
| 910 | w_ref, qweight, scales = rand_marlin_weight_mxfp4_like( |
| 911 | w[i], group_size, input_dtype=input_dtype |
| 912 | ) |
| 913 | global_scale = None |
| 914 | |
| 915 | w_ref_l.append(w_ref.T) |
| 916 | qweight_l.append(qweight) |
| 917 | scales_l.append(scales) |
| 918 | if global_scale is not None: |
| 919 | global_scale_l.append(global_scale) |
| 920 | elif quant_type == scalar_types.float8_e4m3fn: |
| 921 | w_ref, qweight, scales = marlin_quant_fp8_torch( |
| 922 | w[i], group_size, input_dtype=input_dtype |
| 923 | ) |
| 924 | w_ref_l.append(w_ref.T) |
| 925 | qweight_l.append(qweight) |
| 926 | scales_l.append(scales) |
| 927 | elif has_zp: |
| 928 | w_ref, qweight, scales, zeros = awq_marlin_quantize( |
| 929 | w[i].transpose(1, 0), |
| 930 | quant_type, |
| 931 | group_size, |
| 932 | input_dtype=input_dtype, |
| 933 | ) |
| 934 | |
| 935 | w_ref_l.append(w_ref.T) |
| 936 | qweight_l.append(qweight) |
| 937 | scales_l.append(scales) |
| 938 | zeros_l.append(zeros) |
| 939 | else: |
| 940 | w_ref, qweight, scales = marlin_quantize( |
| 941 | w[i].transpose(1, 0), |
| 942 | quant_type, |
| 943 | group_size, |
| 944 | input_dtype=input_dtype, |
| 945 | ) |
| 946 | |
| 947 | w_ref_l.append(w_ref.T) |
| 948 | qweight_l.append(qweight) |
| 949 | scales_l.append(scales) |
| 950 | |
| 951 | if bias is not None: |
| 952 | bias_l.append(marlin_permute_bias(bias[i])) |
| 953 | |
| 954 | w_ref = stack_and_dev(w_ref_l) |
| 955 | qweight = stack_and_dev(qweight_l).contiguous() |
| 956 | scales = stack_and_dev(scales_l) |
| 957 | global_scale = stack_and_dev(global_scale_l) if global_scale_l else None |
| 958 | zeros = stack_and_dev(zeros_l) if zeros_l else None |
| 959 | marlin_bias = stack_and_dev(bias_l) if bias_l else None |
| 960 | |
| 961 | a_scales_factor = None |
| 962 | if input_type == scalar_types.int8 and group_size != -1: |
| 963 | a_scales_factor = 1 / 4096 * scales.max().float() |
| 964 | scales = scales / scales.max() * 4096 |
| 965 | scales = scales.round().to(torch.int16).view(w.dtype) |
| 966 | |
| 967 | return MarlinMoEWeightData( |
| 968 | w_ref=w_ref, |
| 969 | qweight=qweight, |
| 970 | scales=scales, |
| 971 | global_scale=global_scale, |
| 972 | a_scales_factor=a_scales_factor, |
| 973 | zeros=zeros, |
| 974 | marlin_bias=marlin_bias, |
| 975 | ) |
| 976 | |
| 977 | |
| 978 | @pytest.mark.flaky(reruns=2) |
| 979 | @pytest.mark.parametrize( |
| 980 | ("a_type, b_type, c_type, group_blocks,m, n, k, e, topk, ep_size"), |
| 981 | marlin_moe_generate_valid_test_cases(), |
| 982 | ) |
| 983 | @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
| 984 | @pytest.mark.usefixtures("default_vllm_config") |
| 985 | def test_fused_marlin_moe( |
| 986 | a_type: ScalarType, |
| 987 | b_type: ScalarType, |
| 988 | c_type: ScalarType, |
| 989 | group_blocks: int, |
| 990 | m: int, |
| 991 | n: int, |
| 992 | k: int, |
| 993 | e: int, |
| 994 | topk: int, |
| 995 | ep_size: int, |
| 996 | ): |
| 997 | set_random_seed(1) |
| 998 | group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 |
| 999 | |
| 1000 | if c_type == scalar_types.float16: |
| 1001 | dtype = torch.float16 |
| 1002 | elif c_type == scalar_types.bfloat16: |
| 1003 | dtype = torch.bfloat16 |
| 1004 | else: |
| 1005 | raise RuntimeError("unsupported c_type") |
| 1006 | |
| 1007 | if a_type == scalar_types.int8: |
| 1008 | a_dtype = torch.int8 |
| 1009 | elif a_type == scalar_types.float8_e4m3fn: |
| 1010 | a_dtype = torch.float8_e4m3fn |
| 1011 | else: |
| 1012 | a_dtype = dtype |
| 1013 | |
| 1014 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 1015 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 1016 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 1017 | |
| 1018 | if ep_size > 1: |
| 1019 | local_e = e // ep_size |
| 1020 | e_ids = torch.randperm(e, device="cuda", dtype=torch.int32)[:local_e] |
| 1021 | e_map = torch.full((e,), -1, device="cuda", dtype=torch.int32) |
| 1022 | e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32) |
| 1023 | w1 = w1[e_ids] |
| 1024 | w2 = w2[e_ids] |
| 1025 | else: |
| 1026 | e_map = None |
| 1027 | |
| 1028 | w1_data = MarlinMoEWeightData.make( |
| 1029 | w=w1, |
| 1030 | quant_type=b_type, |
| 1031 | group_size=group_size, |
| 1032 | input_type=a_type, |
| 1033 | ) |
| 1034 | |
| 1035 | w2_data = MarlinMoEWeightData.make( |
| 1036 | w=w2, |
| 1037 | quant_type=b_type, |
| 1038 | group_size=group_size, |
| 1039 | input_type=a_type, |
| 1040 | ) |
| 1041 | |
| 1042 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 1043 | |
| 1044 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 1045 | |
| 1046 | with set_current_vllm_config(vllm_config): |
| 1047 | score = torch.softmax(score, dim=-1, dtype=torch.float32) |
| 1048 | topk_weight, topk_ids = torch.topk(score, topk) |
| 1049 | torch_output = torch_experts( |
| 1050 | a, |
| 1051 | w1_data.w_ref, |
| 1052 | w2_data.w_ref, |
| 1053 | topk_weight=topk_weight, |
| 1054 | topk_ids=topk_ids, |
| 1055 | global_num_experts=e, |
| 1056 | expert_map=e_map, |
| 1057 | quant_dtype=a_dtype, |
| 1058 | per_act_token_quant=True, |
| 1059 | ) |
| 1060 | |
| 1061 | def instance_activation( |
| 1062 | activation: MoEActivation, |
| 1063 | output: torch.Tensor, |
| 1064 | input: torch.Tensor, |
| 1065 | *, |
| 1066 | topk_ids: torch.Tensor | None = None, |
| 1067 | expert_map: torch.Tensor | None = None, |
| 1068 | ) -> None: |
| 1069 | apply_moe_activation( |
| 1070 | activation, |
| 1071 | output, |
| 1072 | input, |
| 1073 | activation_config=ApplyMoEActivationConfig(), |
| 1074 | topk_ids=topk_ids, |
| 1075 | expert_map=expert_map, |
| 1076 | ) |
| 1077 | |
| 1078 | marlin_output = fused_marlin_moe( |
| 1079 | a, |
| 1080 | w1_data.qweight, |
| 1081 | w2_data.qweight, |
| 1082 | None, |
| 1083 | None, |
| 1084 | w1_data.scales, |
| 1085 | w2_data.scales, |
| 1086 | topk_weights, |
| 1087 | topk_ids, |
| 1088 | global_num_experts=e, |
| 1089 | expert_map=e_map, |
| 1090 | global_scale1=w1_data.global_scale, |
| 1091 | global_scale2=w2_data.global_scale, |
| 1092 | input_global_scale1=w1_data.a_scales_factor, |
| 1093 | input_global_scale2=w2_data.a_scales_factor, |
| 1094 | w1_zeros=w1_data.zeros, |
| 1095 | w2_zeros=w2_data.zeros, |
| 1096 | input_dtype=a_dtype, |
| 1097 | quant_type_id=b_type.id, |
| 1098 | activation_func=instance_activation, |
| 1099 | ) |
| 1100 | |
| 1101 | torch.testing.assert_close(marlin_output, torch_output, atol=4e-2, rtol=0) |
| 1102 | |
| 1103 | |
| 1104 | @pytest.mark.flaky(reruns=2) |
| 1105 | @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
| 1106 | @pytest.mark.usefixtures("default_vllm_config") |
| 1107 | @pytest.mark.parametrize("m", [1, 256]) |
| 1108 | def test_fused_marlin_moe_with_bias(m): |
| 1109 | set_random_seed(0) |
| 1110 | |
| 1111 | e, topk = 32, 4 |
| 1112 | n, k = 2048, 2048 |
| 1113 | group_size = 128 |
| 1114 | quant_type = scalar_types.uint4b8 |
| 1115 | dtype = torch.half |
| 1116 | |
| 1117 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 1118 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 1119 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 1120 | b_bias1 = torch.randn((e, 2 * n), device="cuda", dtype=dtype) / 10 |
| 1121 | b_bias2 = torch.randn((e, k), device="cuda", dtype=dtype) / 10 |
| 1122 | |
| 1123 | w1_data = MarlinMoEWeightData.make( |
| 1124 | w=w1, |
| 1125 | quant_type=quant_type, |
| 1126 | group_size=group_size, |
| 1127 | bias=b_bias1, |
| 1128 | ) |
| 1129 | |
| 1130 | w2_data = MarlinMoEWeightData.make( |
| 1131 | w=w2, |
| 1132 | quant_type=quant_type, |
| 1133 | group_size=group_size, |
| 1134 | bias=b_bias2, |
| 1135 | ) |
| 1136 | |
| 1137 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 1138 | |
| 1139 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 1140 | |
| 1141 | with set_current_vllm_config(vllm_config): |
| 1142 | torch_output = torch_moe( |
| 1143 | a, w1_data.w_ref, w2_data.w_ref, score, topk, b_bias1, b_bias2 |
| 1144 | ) |
| 1145 | |
| 1146 | marlin_output = fused_marlin_moe( |
| 1147 | a, |
| 1148 | w1_data.qweight, |
| 1149 | w2_data.qweight, |
| 1150 | w1_data.marlin_bias, |
| 1151 | w2_data.marlin_bias, |
| 1152 | w1_data.scales, |
| 1153 | w2_data.scales, |
| 1154 | topk_weights, |
| 1155 | topk_ids, |
| 1156 | global_num_experts=e, |
| 1157 | expert_map=None, |
| 1158 | global_scale1=w1_data.global_scale, |
| 1159 | global_scale2=w2_data.global_scale, |
| 1160 | w1_zeros=w1_data.zeros, |
| 1161 | w2_zeros=w2_data.zeros, |
| 1162 | quant_type_id=quant_type.id, |
| 1163 | ) |
| 1164 | |
| 1165 | torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0) |
| 1166 | |
| 1167 | |
| 1168 | @pytest.mark.flaky(reruns=2) |
| 1169 | @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
| 1170 | @pytest.mark.usefixtures("default_vllm_config") |
| 1171 | @pytest.mark.parametrize("m", [1, 64, 256]) |
| 1172 | @pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)]) |
| 1173 | @pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)]) |
| 1174 | @pytest.mark.parametrize("activation", [MoEActivation.RELU2_NO_MUL]) |
| 1175 | def test_fused_marlin_moe_non_gated( |
| 1176 | m: int, n: int, k: int, e: int, topk: int, activation: MoEActivation |
| 1177 | ): |
| 1178 | """Test Marlin MoE with non-gated activation (relu2_no_mul). |
| 1179 | |
| 1180 | Non-gated activations like relu2 don't have the gate-up projection pattern, |
| 1181 | so w1 has shape (e, n, k) instead of (e, 2*n, k). |
| 1182 | """ |
| 1183 | set_random_seed(42) |
| 1184 | |
| 1185 | group_size = 16 # NVFP4 group size |
| 1186 | quant_type = scalar_types.float4_e2m1f |
| 1187 | dtype = torch.bfloat16 |
| 1188 | |
| 1189 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 1190 | # Non-gated: w1 shape is (e, n, k) not (e, 2*n, k) |
| 1191 | w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10 |
| 1192 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 1193 | |
| 1194 | w1_data = MarlinMoEWeightData.make( |
| 1195 | w=w1, |
| 1196 | quant_type=quant_type, |
| 1197 | group_size=group_size, |
| 1198 | ) |
| 1199 | |
| 1200 | w2_data = MarlinMoEWeightData.make( |
| 1201 | w=w2, |
| 1202 | quant_type=quant_type, |
| 1203 | group_size=group_size, |
| 1204 | ) |
| 1205 | |
| 1206 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 1207 | |
| 1208 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 1209 | |
| 1210 | with set_current_vllm_config(vllm_config): |
| 1211 | torch_output = torch_moe( |
| 1212 | a, |
| 1213 | w1_data.w_ref, |
| 1214 | w2_data.w_ref, |
| 1215 | score, |
| 1216 | topk, |
| 1217 | activation=activation, |
| 1218 | ) |
| 1219 | |
| 1220 | marlin_output = fused_marlin_moe( |
| 1221 | a, |
| 1222 | w1_data.qweight, |
| 1223 | w2_data.qweight, |
| 1224 | None, # bias1 |
| 1225 | None, # bias2 |
| 1226 | w1_data.scales, |
| 1227 | w2_data.scales, |
| 1228 | topk_weights, |
| 1229 | topk_ids, |
| 1230 | global_num_experts=e, |
| 1231 | expert_map=None, |
| 1232 | global_scale1=w1_data.global_scale, |
| 1233 | global_scale2=w2_data.global_scale, |
| 1234 | w1_zeros=w1_data.zeros, |
| 1235 | w2_zeros=w2_data.zeros, |
| 1236 | quant_type_id=quant_type.id, |
| 1237 | activation=activation, |
| 1238 | ) |
| 1239 | |
| 1240 | torch.testing.assert_close(marlin_output, torch_output, atol=1e-1, rtol=0) |
| 1241 | |
| 1242 | |
| 1243 | def _make_humming_indexed_experts(activation: MoEActivation): |
| 1244 | pytest.importorskip("humming") |
| 1245 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1246 | HummingIndexedExperts, |
| 1247 | ) |
| 1248 | from vllm.model_executor.layers.quantization.utils import humming as humming_utils |
| 1249 | from vllm.utils import humming |
| 1250 | |
| 1251 | top_k, num_experts = 6, 12 |
| 1252 | hidden_size, intermediate_size = 2688, 1856 |
| 1253 | gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size |
| 1254 | num_w13_stacks = 2 if activation.is_gated else 1 |
| 1255 | moe_config = make_dummy_moe_config( |
| 1256 | num_experts=num_experts, |
| 1257 | experts_per_token=top_k, |
| 1258 | hidden_dim=hidden_size, |
| 1259 | intermediate_size=intermediate_size, |
| 1260 | activation=activation, |
| 1261 | ) |
| 1262 | |
| 1263 | layer = torch.nn.Module() |
| 1264 | layer.moe_config = moe_config |
| 1265 | layer.params_dtype = torch.bfloat16 |
| 1266 | |
| 1267 | weight_schema = humming.ModeloptNvfp4WeightSchema() |
| 1268 | for sublayer_name, shape_n, shape_k, stack_size in ( |
| 1269 | ("w13", gate_up_size, hidden_size, num_w13_stacks), |
| 1270 | ("w2", hidden_size, intermediate_size, 1), |
| 1271 | ): |
| 1272 | tensor_attrs = weight_schema.get_tensors_attrs( |
| 1273 | shape_n=shape_n, |
| 1274 | shape_k=shape_k, |
| 1275 | param_dtype=layer.params_dtype, |
| 1276 | num_experts=num_experts, |
| 1277 | stack_size=stack_size, |
| 1278 | ) |
| 1279 | for tensor_name, attrs in tensor_attrs.items(): |
| 1280 | layer.register_parameter( |
| 1281 | f"{sublayer_name}_{tensor_name}", |
| 1282 | Parameter( |
| 1283 | torch.ones( |
| 1284 | attrs["shape"], |
| 1285 | dtype=attrs["dtype"], |
| 1286 | device="cuda", |
| 1287 | ), |
| 1288 | requires_grad=False, |
| 1289 | ), |
| 1290 | ) |
| 1291 | |
| 1292 | humming_utils.convert_to_humming_moe_kernel_format( |
| 1293 | layer, |
| 1294 | weight_schema=weight_schema, |
| 1295 | input_schema=humming.HummingInputSchema(a_dtype=humming.dtypes.bfloat16), |
| 1296 | ) |
| 1297 | |
| 1298 | layer.local_num_experts = layer.global_num_experts = num_experts |
| 1299 | layer.hidden_size = hidden_size |
| 1300 | layer.intermediate_size_per_partition = intermediate_size |
| 1301 | quant_config = humming_utils.get_humming_moe_quant_config(layer) |
| 1302 | experts = HummingIndexedExperts( |
| 1303 | moe_config=moe_config, |
| 1304 | quant_config=quant_config, |
| 1305 | ) |
| 1306 | return experts, layer |
| 1307 | |
| 1308 | |
| 1309 | @pytest.mark.parametrize("valid_shape_m", [8192, 8193]) |
| 1310 | @pytest.mark.parametrize("w2_block_size", [64, 96, 128]) |
| 1311 | def test_humming_indexed_routing_matches_each_projection_tile( |
| 1312 | valid_shape_m: int, w2_block_size: int |
| 1313 | ): |
| 1314 | """Independently tuned projections preserve routed rows and expert identity.""" |
| 1315 | from types import SimpleNamespace |
| 1316 | |
| 1317 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1318 | HummingIndexedExperts, |
| 1319 | ) |
| 1320 | |
| 1321 | topk_ids = torch.tensor([[0, 1], [1, 2], [2, 0]], device="cuda") |
| 1322 | expert_map = torch.tensor([1, -1, 0], dtype=torch.int32, device="cuda") |
| 1323 | experts = SimpleNamespace( |
| 1324 | global_num_experts=3, |
| 1325 | estimate_local_valid_shape_m=lambda _: valid_shape_m, |
| 1326 | w13_tuning_config=[ |
| 1327 | [0, 8192, {"block_shape": [64, 128, 64]}], |
| 1328 | [8192, 16384, {"block_shape": [128, 128, 64]}], |
| 1329 | ], |
| 1330 | w2_tuning_config=[ |
| 1331 | [0, 8192, {"block_shape": [64, 128, 64]}], |
| 1332 | [8192, 16384, {"block_shape": [w2_block_size, 128, 64]}], |
| 1333 | ], |
| 1334 | compute_config_str="compute", |
| 1335 | w13_tuning_config_str="w13", |
| 1336 | w2_tuning_config_str="w2", |
| 1337 | ) |
| 1338 | first, second, _ = HummingIndexedExperts.prepare_humming_moe_kwargs( |
| 1339 | experts, topk_ids, expert_map, None |
| 1340 | ) |
| 1341 | block_sizes = (64, 64) if valid_shape_m == 8192 else (128, w2_block_size) |
| 1342 | expected_rows = torch.where(expert_map[topk_ids.flatten()] >= 0)[0] |
| 1343 | for kwargs, block_size in zip((first, second), block_sizes): |
| 1344 | padded = kwargs["num_tokens_padded"].item() |
| 1345 | assert padded % block_size == 0 |
| 1346 | rows = kwargs["sorted_ids"][:padded].long() |
| 1347 | valid = rows < topk_ids.numel() |
| 1348 | torch.testing.assert_close(rows[valid].sort().values, expected_rows) |
| 1349 | routed_experts = kwargs["expert_ids"][: padded // block_size] |
| 1350 | torch.testing.assert_close( |
| 1351 | routed_experts.repeat_interleave(block_size)[valid], |
| 1352 | expert_map[topk_ids.flatten()[rows[valid]]], |
| 1353 | ) |
| 1354 | assert (first["top_k"], second["top_k"]) == (2, 1) |
| 1355 | for key in ("sorted_ids", "expert_ids", "num_tokens_padded"): |
| 1356 | assert (first[key] is second[key]) == (block_sizes[0] == block_sizes[1]) |
| 1357 | |
| 1358 | |
| 1359 | @pytest.mark.parametrize("activation", list(MoEActivation)) |
| 1360 | def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivation): |
| 1361 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1362 | HummingExpertsBase, |
| 1363 | ) |
| 1364 | |
| 1365 | assert HummingExpertsBase._supports_activation( |
| 1366 | activation |
| 1367 | ) == apply_moe_activation_supported(activation) |
| 1368 | |
| 1369 | |
| 1370 | @pytest.mark.parametrize( |
| 1371 | ("env_override", "use_ep", "expected"), |
| 1372 | [ |
| 1373 | (None, False, "indexed"), |
| 1374 | (None, True, "grouped_contiguous"), |
| 1375 | ("auto", True, "grouped_contiguous"), |
| 1376 | ("indexed", True, "indexed"), |
| 1377 | ("grouped", False, "grouped_contiguous"), |
| 1378 | ], |
| 1379 | ) |
| 1380 | def test_humming_selects_gemm_from_parallelism_and_override( |
| 1381 | monkeypatch: pytest.MonkeyPatch, |
| 1382 | env_override: str | None, |
| 1383 | use_ep: bool, |
| 1384 | expected: str, |
| 1385 | ): |
| 1386 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1387 | get_humming_moe_gemm_type, |
| 1388 | ) |
| 1389 | |
| 1390 | if env_override is None: |
| 1391 | monkeypatch.delenv("VLLM_HUMMING_MOE_GEMM_TYPE", raising=False) |
| 1392 | else: |
| 1393 | monkeypatch.setenv("VLLM_HUMMING_MOE_GEMM_TYPE", env_override) |
| 1394 | moe_config = make_dummy_moe_config() |
| 1395 | moe_config.moe_parallel_config.use_ep = use_ep |
| 1396 | |
| 1397 | assert get_humming_moe_gemm_type(moe_config) == expected |
| 1398 | |
| 1399 | |
| 1400 | @pytest.mark.parametrize( |
| 1401 | ("dp_token_counts", "topk_shape", "expected"), |
| 1402 | [ |
| 1403 | pytest.param(None, (7, 3), 21, id="runtime-layout"), |
| 1404 | pytest.param([3, 5], (48, 1), 48, id="expanded-dp-layout"), |
| 1405 | ], |
| 1406 | ) |
| 1407 | def test_humming_global_valid_shape_m( |
| 1408 | monkeypatch: pytest.MonkeyPatch, |
| 1409 | dp_token_counts: list[int] | None, |
| 1410 | topk_shape: tuple[int, int], |
| 1411 | expected: int, |
| 1412 | ): |
| 1413 | from types import SimpleNamespace |
| 1414 | |
| 1415 | import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming |
| 1416 | |
| 1417 | dp_metadata = ( |
| 1418 | None |
| 1419 | if dp_token_counts is None |
| 1420 | else SimpleNamespace( |
| 1421 | num_tokens_across_dp_cpu=torch.tensor(dp_token_counts, dtype=torch.int32) |
| 1422 | ) |
| 1423 | ) |
| 1424 | monkeypatch.setattr( |
| 1425 | humming, |
| 1426 | "get_forward_context", |
| 1427 | lambda: SimpleNamespace(dp_metadata=dp_metadata), |
| 1428 | ) |
| 1429 | experts = SimpleNamespace(moe_config=make_dummy_moe_config(experts_per_token=6)) |
| 1430 | topk_ids = torch.empty(topk_shape, dtype=torch.int64) |
| 1431 | |
| 1432 | result = humming.HummingExpertsBase.get_global_valid_shape_m(experts, topk_ids) |
| 1433 | |
| 1434 | assert result == expected |
| 1435 | |
| 1436 | |
| 1437 | def test_humming_permute_scratch_is_shared_by_config( |
| 1438 | monkeypatch: pytest.MonkeyPatch, |
| 1439 | ): |
| 1440 | from types import SimpleNamespace |
| 1441 | from unittest.mock import Mock |
| 1442 | |
| 1443 | import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming |
| 1444 | import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute |
| 1445 | import vllm.v1.worker.workspace as workspace |
| 1446 | |
| 1447 | manager = workspace.WorkspaceManager(torch.device("cpu")) |
| 1448 | monkeypatch.setattr(workspace, "_manager", manager) |
| 1449 | |
| 1450 | scratch_topk6 = Mock() |
| 1451 | scratch_topk1 = Mock() |
| 1452 | scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1]) |
| 1453 | monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True) |
| 1454 | monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type) |
| 1455 | moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6) |
| 1456 | moe_config.moe_parallel_config.dp_size = 2 |
| 1457 | experts = SimpleNamespace(moe_config=moe_config) |
| 1458 | other_layer = SimpleNamespace(moe_config=moe_config) |
| 1459 | |
| 1460 | assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 |
| 1461 | assert ( |
| 1462 | humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6 |
| 1463 | ) |
| 1464 | assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1 |
| 1465 | assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 |
| 1466 | |
| 1467 | assert scratch_type.call_count == 2 |
| 1468 | first_call, second_call = scratch_type.call_args_list |
| 1469 | assert first_call.kwargs["max_num_tokens"] == 1024 |
| 1470 | assert first_call.kwargs["topk"] == 6 |
| 1471 | assert second_call.kwargs["max_num_tokens"] == 6144 |
| 1472 | assert second_call.kwargs["topk"] == 1 |
| 1473 | |
| 1474 | manager.lock() |
| 1475 | assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6 |
| 1476 | with pytest.raises(AssertionError, match="was not allocated during warmup"): |
| 1477 | humming.HummingExpertsBase._get_permute_scratch(experts, 2) |
| 1478 | |
| 1479 | |
| 1480 | def test_humming_delegates_to_instance_activation(): |
| 1481 | from types import SimpleNamespace |
| 1482 | from unittest.mock import Mock |
| 1483 | |
| 1484 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1485 | HummingExpertsBase, |
| 1486 | ) |
| 1487 | |
| 1488 | activation_func = Mock() |
| 1489 | activation_config = ApplyMoEActivationConfig( |
| 1490 | clamp_limit=7.0, |
| 1491 | alpha=1.5, |
| 1492 | beta=0.25, |
| 1493 | activation_situ_beta=2.0, |
| 1494 | activation_situ_linear_beta=3.0, |
| 1495 | ) |
| 1496 | experts = SimpleNamespace( |
| 1497 | activation=activation_func, |
| 1498 | activation_config=activation_config, |
| 1499 | ) |
| 1500 | input = torch.empty(1, 2) |
| 1501 | output = torch.empty(1, 1) |
| 1502 | |
| 1503 | HummingExpertsBase.apply_activation( |
| 1504 | experts, MoEActivation.SWIGLUOAI_UNINTERLEAVE, output, input |
| 1505 | ) |
| 1506 | |
| 1507 | activation_func.assert_called_once_with( |
| 1508 | activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, |
| 1509 | input=input, |
| 1510 | output=output, |
| 1511 | ) |
| 1512 | |
| 1513 | |
| 1514 | def test_humming_grouped_apply_forwards_valid_prefix( |
| 1515 | monkeypatch: pytest.MonkeyPatch, |
| 1516 | ): |
| 1517 | from types import SimpleNamespace |
| 1518 | from unittest.mock import Mock |
| 1519 | |
| 1520 | import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming |
| 1521 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 1522 | HummingGroupedExperts, |
| 1523 | ) |
| 1524 | |
| 1525 | buffers = { |
| 1526 | "gate_up_output": torch.empty(6, 4), |
| 1527 | "quanted_down_input": torch.empty(6, 2), |
| 1528 | "down_output": torch.empty(6, 4), |
| 1529 | } |
| 1530 | experts = SimpleNamespace( |
| 1531 | num_experts=2, |
| 1532 | estimate_local_valid_shape_m=lambda _: 6, |
| 1533 | prepare_buffers=lambda *_: buffers, |
| 1534 | _get_permute_scratch=lambda _, *, indices_only=False: object(), |
| 1535 | process_input=Mock( |
| 1536 | side_effect=lambda _, **kwargs: ( |
| 1537 | kwargs["inputs"], |
| 1538 | kwargs.get("input_scale"), |
| 1539 | None, |
| 1540 | ) |
| 1541 | ), |
| 1542 | humming_forward=Mock(), |
| 1543 | compute_config_str="", |
| 1544 | w13_tuning_config_str="", |
| 1545 | w2_tuning_config_str="", |
| 1546 | ) |
| 1547 | |
| 1548 | hidden_states = torch.empty(3, 4) |
| 1549 | output = torch.empty_like(hidden_states) |
| 1550 | topk_ids = torch.zeros((3, 2), dtype=torch.int64) |
| 1551 | topk_weights = torch.ones_like(topk_ids, dtype=torch.float32) |
| 1552 | expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64) |
| 1553 | monkeypatch.setattr( |
| 1554 | humming, |
| 1555 | "moe_prepare_scatter", |
| 1556 | lambda *_: (expert_offsets, torch.arange(6).view(3, 2)), |
| 1557 | ) |
| 1558 | monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None) |
| 1559 | |
| 1560 | HummingGroupedExperts.apply( |
| 1561 | experts, |
| 1562 | output=output, |
| 1563 | hidden_states=hidden_states, |
| 1564 | w1=torch.empty(0), |
| 1565 | w2=torch.empty(0), |
| 1566 | topk_weights=topk_weights, |
| 1567 | topk_ids=topk_ids, |
| 1568 | activation=MoEActivation.SITU, |
| 1569 | global_num_experts=2, |
| 1570 | expert_map=None, |
| 1571 | a1q_scale=None, |
| 1572 | a2_scale=None, |
| 1573 | workspace13=torch.empty(0), |
| 1574 | workspace2=torch.empty(0), |
| 1575 | expert_tokens_meta=None, |
| 1576 | apply_router_weight_on_input=False, |
| 1577 | ) |
| 1578 | |
| 1579 | call_kwargs = experts.process_input.call_args.kwargs |
| 1580 | assert call_kwargs["activation"] == MoEActivation.SITU |
| 1581 | assert call_kwargs["inputs"].shape == (6, 4) |
| 1582 | assert call_kwargs["quanted_input"].shape == (6, 2) |
| 1583 | torch.testing.assert_close(call_kwargs["num_valid_tokens"], expert_offsets[-1:]) |
| 1584 | |
| 1585 | |
| 1586 | def test_batched_marlin_activation_uses_expert_token_counts( |
| 1587 | monkeypatch: pytest.MonkeyPatch, |
| 1588 | ): |
| 1589 | from types import SimpleNamespace |
| 1590 | from unittest.mock import Mock |
| 1591 | |
| 1592 | import vllm.model_executor.layers.fused_moe.experts.marlin_moe as marlin |
| 1593 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 1594 | BatchedMarlinExperts, |
| 1595 | ) |
| 1596 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 1597 | ExpertTokensMetadata, |
| 1598 | ) |
| 1599 | |
| 1600 | activation_func = Mock() |
| 1601 | experts = SimpleNamespace( |
| 1602 | activation=activation_func, |
| 1603 | w1_bias=None, |
| 1604 | w2_bias=None, |
| 1605 | w1_scale=None, |
| 1606 | w2_scale=None, |
| 1607 | quant_type_id=0, |
| 1608 | a1_gscale=None, |
| 1609 | a2_gscale=None, |
| 1610 | w1_zp=None, |
| 1611 | w2_zp=None, |
| 1612 | input_dtype=None, |
| 1613 | activation_config=ApplyMoEActivationConfig( |
| 1614 | clamp_limit=3.0, alpha=1.3, beta=0.5 |
| 1615 | ), |
| 1616 | ) |
| 1617 | hidden_states = torch.empty(2, 3, 4) |
| 1618 | expert_num_tokens = torch.tensor([1, 2], dtype=torch.int32) |
| 1619 | topk_ids = torch.zeros((2, 3), dtype=torch.int64) |
| 1620 | expert_map = torch.arange(2, dtype=torch.int32) |
| 1621 | |
| 1622 | def fake_batched_fused_marlin_moe(**kwargs): |
| 1623 | kwargs["activation_func"]( |
| 1624 | kwargs["activation"], |
| 1625 | torch.empty(6, 4), |
| 1626 | torch.empty(6, 8), |
| 1627 | topk_ids=topk_ids, |
| 1628 | expert_map=expert_map, |
| 1629 | ) |
| 1630 | return kwargs["output"] |
| 1631 | |
| 1632 | monkeypatch.setattr( |
| 1633 | marlin, "batched_fused_marlin_moe", fake_batched_fused_marlin_moe |
| 1634 | ) |
| 1635 | output = torch.empty_like(hidden_states) |
| 1636 | |
| 1637 | BatchedMarlinExperts.apply( |
| 1638 | experts, |
| 1639 | output=output, |
| 1640 | hidden_states=hidden_states, |
| 1641 | w1=torch.empty(0), |
| 1642 | w2=torch.empty(0), |
| 1643 | topk_weights=torch.ones_like(topk_ids, dtype=torch.float32), |
| 1644 | topk_ids=topk_ids, |
| 1645 | activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE, |
| 1646 | global_num_experts=2, |
| 1647 | expert_map=expert_map, |
| 1648 | a1q_scale=None, |
| 1649 | a2_scale=None, |
| 1650 | workspace13=torch.empty(0), |
| 1651 | workspace2=torch.empty(0), |
| 1652 | expert_tokens_meta=ExpertTokensMetadata(expert_num_tokens, None), |
| 1653 | apply_router_weight_on_input=False, |
| 1654 | ) |
| 1655 | |
| 1656 | activation_func.assert_called_once() |
| 1657 | call = activation_func.call_args |
| 1658 | assert call.args[0] == MoEActivation.SWIGLUOAI_UNINTERLEAVE |
| 1659 | assert call.args[1].shape == (2, 3, 4) |
| 1660 | assert call.args[2].shape == (2, 3, 8) |
| 1661 | torch.testing.assert_close(call.kwargs["valid_token_counts"], expert_num_tokens) |
| 1662 | |
| 1663 | |
| 1664 | @pytest.mark.parametrize( |
| 1665 | "activation", |
| 1666 | [ |
| 1667 | MoEActivation.SILU, |
| 1668 | MoEActivation.RELU2_NO_MUL, |
| 1669 | ], |
| 1670 | ids=["gated", "non_gated"], |
| 1671 | ) |
| 1672 | def test_humming_gated_non_gated_shape_contract(activation: MoEActivation): |
| 1673 | from vllm.utils import humming |
| 1674 | |
| 1675 | experts, layer = _make_humming_indexed_experts(activation) |
| 1676 | moe_config = experts.moe_config |
| 1677 | top_k = moe_config.experts_per_token |
| 1678 | num_experts = moe_config.num_experts |
| 1679 | hidden_size = moe_config.hidden_dim |
| 1680 | intermediate_size = moe_config.intermediate_size |
| 1681 | gate_up_size = intermediate_size * 2 if activation.is_gated else intermediate_size |
| 1682 | |
| 1683 | w13_meta, w2_meta = (experts.humming_configs[name] for name in ("w13", "w2")) |
| 1684 | for meta in (w13_meta, w2_meta): |
| 1685 | assert meta.a_dtype == humming.dtypes.bfloat16 |
| 1686 | assert meta.b_dtype == humming.dtypes.float4e2m1 |
| 1687 | |
| 1688 | assert w13_meta.shape_n - w13_meta.pad_shape_n == gate_up_size |
| 1689 | assert w2_meta.shape_k - w2_meta.pad_shape_k == intermediate_size |
| 1690 | |
| 1691 | buffer_metas, _ = experts.get_buffer_metas( |
| 1692 | M=1, |
| 1693 | topk=top_k, |
| 1694 | activation=moe_config.activation, |
| 1695 | ) |
| 1696 | assert buffer_metas["gate_up_output"]["shape"][-1] == gate_up_size |
| 1697 | assert "activation_output" not in buffer_metas |
| 1698 | assert buffer_metas["quanted_down_input"]["shape"][-1] == intermediate_size |
| 1699 | assert experts.moe_problem_size( |
| 1700 | a1=torch.empty(1, hidden_size), |
| 1701 | w1=torch.empty(num_experts, 1), |
| 1702 | w2=torch.empty(num_experts, 1), |
| 1703 | topk_ids=torch.empty(1, top_k, dtype=torch.long), |
| 1704 | ) == (num_experts, 1, intermediate_size, hidden_size, top_k) |
| 1705 | |
| 1706 | # Quantized batched input storage must cover every dispatcher's token slots. |
| 1707 | experts.is_batched = lambda: True |
| 1708 | experts.max_num_tokens = 17 |
| 1709 | experts.num_dispatchers = 2 |
| 1710 | batched_metas, _ = experts.get_buffer_metas(17, top_k, activation) |
| 1711 | assert batched_metas["quanted_gate_up_input"]["shape"][0] == num_experts * 17 * 2 |
| 1712 | |
| 1713 | |
| 1714 | def test_humming_indexed_writes_supplied_output_buffer(): |
| 1715 | from vllm.forward_context import set_forward_context |
| 1716 | |
| 1717 | activation = MoEActivation.SILU |
| 1718 | experts, layer = _make_humming_indexed_experts(activation) |
| 1719 | moe_config = experts.moe_config |
| 1720 | num_tokens = 1 |
| 1721 | top_k = moe_config.experts_per_token |
| 1722 | hidden_size = moe_config.hidden_dim |
| 1723 | num_experts = moe_config.num_experts |
| 1724 | workspace13_shape, workspace2_shape, _ = experts.workspace_shapes( |
| 1725 | M=num_tokens, |
| 1726 | N=moe_config.intermediate_size, |
| 1727 | K=hidden_size, |
| 1728 | topk=top_k, |
| 1729 | global_num_experts=num_experts, |
| 1730 | local_num_experts=num_experts, |
| 1731 | expert_tokens_meta=None, |
| 1732 | activation=activation, |
| 1733 | ) |
| 1734 | |
| 1735 | device = torch.device("cuda") |
| 1736 | dtype = layer.params_dtype |
| 1737 | workspace13 = torch.empty(workspace13_shape, dtype=dtype, device=device) |
| 1738 | workspace2 = torch.empty(workspace2_shape, dtype=dtype, device=device) |
| 1739 | hidden_states = torch.ones((num_tokens, hidden_size), dtype=dtype, device=device) |
| 1740 | output = torch.full_like(hidden_states, torch.nan) |
| 1741 | topk_weights = torch.full( |
| 1742 | (num_tokens, top_k), |
| 1743 | 1 / top_k, |
| 1744 | dtype=dtype, |
| 1745 | device=device, |
| 1746 | ) |
| 1747 | topk_ids = torch.arange(top_k, dtype=torch.int32, device=device).unsqueeze(0) |
| 1748 | with set_forward_context(None, vllm_config, num_tokens=num_tokens): |
| 1749 | experts.apply( |
| 1750 | output=output, |
| 1751 | hidden_states=hidden_states, |
| 1752 | w1=layer.w13_weight, |
| 1753 | w2=layer.w2_weight, |
| 1754 | topk_weights=topk_weights, |
| 1755 | topk_ids=topk_ids, |
| 1756 | activation=activation, |
| 1757 | global_num_experts=num_experts, |
| 1758 | expert_map=None, |
| 1759 | a1q_scale=None, |
| 1760 | a2_scale=None, |
| 1761 | workspace13=workspace13, |
| 1762 | workspace2=workspace2, |
| 1763 | expert_tokens_meta=None, |
| 1764 | apply_router_weight_on_input=False, |
| 1765 | ) |
| 1766 | |
| 1767 | assert torch.isfinite(output).all() |
| 1768 | |
| 1769 | |
| 1770 | @pytest.mark.parametrize("ep_size", [1, 2]) |
| 1771 | def test_moe_align_block_size_opcheck(ep_size): |
| 1772 | num_experts = 4 |
| 1773 | block_size = 4 |
| 1774 | |
| 1775 | expert_map = None |
| 1776 | if ep_size != 1: |
| 1777 | local_num_experts = num_experts // ep_size |
| 1778 | expert_ids = torch.randint( |
| 1779 | 0, num_experts, (local_num_experts,), device="cuda", dtype=torch.int32 |
| 1780 | ) |
| 1781 | expert_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32) |
| 1782 | expert_map[expert_ids] = torch.arange( |
| 1783 | local_num_experts, device="cuda", dtype=torch.int32 |
| 1784 | ) |
| 1785 | |
| 1786 | topk_ids = torch.randint(0, num_experts, (3, 4), dtype=torch.int32, device="cuda") |
| 1787 | |
| 1788 | max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) |
| 1789 | sorted_ids = torch.empty( |
| 1790 | (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device |
| 1791 | ) |
| 1792 | sorted_ids.fill_(topk_ids.numel()) |
| 1793 | max_num_m_blocks = max_num_tokens_padded // block_size |
| 1794 | expert_ids = torch.empty( |
| 1795 | (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device |
| 1796 | ) |
| 1797 | num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device) |
| 1798 | |
| 1799 | opcheck( |
| 1800 | torch.ops._moe_C.moe_align_block_size, |
| 1801 | ( |
| 1802 | topk_ids, |
| 1803 | num_experts, |
| 1804 | block_size, |
| 1805 | sorted_ids, |
| 1806 | expert_ids, |
| 1807 | num_tokens_post_pad, |
| 1808 | expert_map, |
| 1809 | ), |
| 1810 | ) |
| 1811 | |
| 1812 | |
| 1813 | def test_batched_moe_align_block_size_opcheck(): |
| 1814 | max_tokens_per_batch = 512 |
| 1815 | num_experts = 4 |
| 1816 | block_size = 16 |
| 1817 | |
| 1818 | expert_num_tokens = torch.randint( |
| 1819 | low=0, |
| 1820 | high=max_tokens_per_batch, |
| 1821 | size=(num_experts,), |
| 1822 | dtype=torch.int32, |
| 1823 | device="cuda", |
| 1824 | ) |
| 1825 | |
| 1826 | max_num_tokens_padded = num_experts * max(max_tokens_per_batch, block_size) |
| 1827 | sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda") |
| 1828 | |
| 1829 | assert max_num_tokens_padded % block_size == 0 |
| 1830 | max_num_m_blocks = max_num_tokens_padded // block_size |
| 1831 | expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda") |
| 1832 | |
| 1833 | num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device="cuda") |
| 1834 | |
| 1835 | opcheck( |
| 1836 | torch.ops._moe_C.batched_moe_align_block_size, |
| 1837 | ( |
| 1838 | max_tokens_per_batch, |
| 1839 | block_size, |
| 1840 | expert_num_tokens, |
| 1841 | sorted_ids, |
| 1842 | expert_ids, |
| 1843 | num_tokens_post_pad, |
| 1844 | ), |
| 1845 | ) |
| 1846 | |
| 1847 | |
| 1848 | # topk=8 covers topk > 4; k=511 covers the non-vectorized scalar path. The |
| 1849 | # layouts exercise contiguous input plus the two non-contiguous cases: a |
| 1850 | # transpose (strided hidden -> scalar gather) and a topk-slice (hidden still |
| 1851 | # contiguous -> vectorized). |
| 1852 | @pytest.mark.parametrize("m", [1, 33, 222]) |
| 1853 | @pytest.mark.parametrize("topk", [*TOP_KS, 8]) |
| 1854 | @pytest.mark.parametrize("k", [128, 511, 1024]) |
| 1855 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 1856 | @pytest.mark.parametrize("layout", ["contig", "transpose", "slice"]) |
| 1857 | def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype, layout: str): |
| 1858 | if layout == "transpose": |
| 1859 | input = torch.randn((m, k, topk), device="cuda", dtype=dtype).transpose(1, 2) |
| 1860 | elif layout == "slice": |
| 1861 | input = torch.randn((m, 2 * topk, k), device="cuda", dtype=dtype)[:, ::2, :] |
| 1862 | else: |
| 1863 | input = torch.randn((m, topk, k), device="cuda", dtype=dtype) |
| 1864 | assert input.is_contiguous() == (layout == "contig") |
| 1865 | actual = torch.empty((m, k), device="cuda", dtype=dtype) |
| 1866 | |
| 1867 | # Reduction accumulates in fp32. |
| 1868 | expected = input.float().sum(dim=1).to(dtype) |
| 1869 | torch.ops._moe_C.moe_sum(input, actual) |
| 1870 | |
| 1871 | torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0) |
| 1872 | |
| 1873 | opcheck(torch.ops._moe_C.moe_sum, (input, actual)) |
| 1874 | |
| 1875 | |
| 1876 | @pytest.mark.parametrize("topk", [2, 8]) |
| 1877 | @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| 1878 | @pytest.mark.parametrize("topk_ids_dtype", [torch.int32, torch.int64]) |
| 1879 | def test_moe_sum_pad_aware(topk: int, dtype: torch.dtype, topk_ids_dtype: torch.dtype): |
| 1880 | m, k = 17, 128 |
| 1881 | input = torch.randn((m, topk, k), device="cuda", dtype=dtype) |
| 1882 | |
| 1883 | topk_ids = torch.randint(0, 4, (m, topk), device="cuda", dtype=topk_ids_dtype) |
| 1884 | # Force some slots unrouted, independent of expert parallelism. |
| 1885 | topk_ids[0, :] = -1 |
| 1886 | topk_ids[1, 0] = -1 |
| 1887 | |
| 1888 | expert_map = torch.tensor([0, -1, 1, -1], device="cuda", dtype=torch.int32) |
| 1889 | |
| 1890 | actual = torch.empty((m, k), device="cuda", dtype=dtype) |
| 1891 | torch.ops._moe_C.moe_sum(input, actual, topk_ids, expert_map) |
| 1892 | |
| 1893 | routed = expert_map[topk_ids.clamp(min=0).long()] >= 0 |
| 1894 | routed &= topk_ids >= 0 |
| 1895 | expected = (input.float() * routed.unsqueeze(-1)).sum(dim=1).to(dtype) |
| 1896 | |
| 1897 | torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0) |
| 1898 | |
| 1899 | opcheck(torch.ops._moe_C.moe_sum, (input, actual, topk_ids, expert_map)) |
| 1900 | |
| 1901 | |
| 1902 | def _batched_fused_marlin_moe_cases() -> list[Any]: |
| 1903 | cases = [ |
| 1904 | pytest.param( |
| 1905 | m, |
| 1906 | 128, |
| 1907 | 128, |
| 1908 | e, |
| 1909 | topk, |
| 1910 | max_tokens_per_batch, |
| 1911 | torch.bfloat16, |
| 1912 | scalar_types.float4_e2m1f, |
| 1913 | None, |
| 1914 | 1e-3, |
| 1915 | id=( |
| 1916 | f"m{m}-n128-k128-e{e}-topk{topk}-max_tokens{max_tokens_per_batch}-mxfp4" |
| 1917 | ), |
| 1918 | ) |
| 1919 | for m in [16, 32, 64] |
| 1920 | for e in [8, 12, 16, 32] |
| 1921 | for topk in [2, 4] |
| 1922 | for max_tokens_per_batch in [16, 32, 64] |
| 1923 | ] |
| 1924 | cases.append( |
| 1925 | pytest.param( |
| 1926 | 32, |
| 1927 | 128, |
| 1928 | 128, |
| 1929 | 8, |
| 1930 | 2, |
| 1931 | 64, |
| 1932 | torch.float16, |
| 1933 | scalar_types.uint4, |
| 1934 | scalar_types.int8, |
| 1935 | 4e-2, |
| 1936 | id="awq-int8-activation-metadata", |
| 1937 | ) |
| 1938 | ) |
| 1939 | return cases |
| 1940 | |
| 1941 | |
| 1942 | @pytest.mark.parametrize( |
| 1943 | ("m,n,k,e,topk,max_tokens_per_batch,dtype,quant_dtype,input_type,atol"), |
| 1944 | _batched_fused_marlin_moe_cases(), |
| 1945 | ) |
| 1946 | @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") |
| 1947 | def test_batched_fused_marlin_moe( |
| 1948 | m: int, |
| 1949 | n: int, |
| 1950 | k: int, |
| 1951 | e: int, |
| 1952 | topk: int, |
| 1953 | max_tokens_per_batch: int, |
| 1954 | dtype: torch.dtype, |
| 1955 | quant_dtype: ScalarType, |
| 1956 | input_type: ScalarType | None, |
| 1957 | atol: float, |
| 1958 | ): |
| 1959 | print( |
| 1960 | f"testing m={m}, n={n}, k={k}, e={e}, " |
| 1961 | f"topk={topk}, " |
| 1962 | f"max_tokens_per_batch={max_tokens_per_batch}, " |
| 1963 | f"dtype={dtype}, quant_dtype={quant_dtype}, input_type={input_type}" |
| 1964 | ) |
| 1965 | set_random_seed(0) |
| 1966 | |
| 1967 | group_size = 32 |
| 1968 | if input_type == scalar_types.int8: |
| 1969 | input_dtype = torch.int8 |
| 1970 | elif input_type == scalar_types.float8_e4m3fn: |
| 1971 | input_dtype = torch.float8_e4m3fn |
| 1972 | else: |
| 1973 | input_dtype = None |
| 1974 | |
| 1975 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 1976 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20 |
| 1977 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20 |
| 1978 | |
| 1979 | w1_data = MarlinMoEWeightData.make( |
| 1980 | w=w1, |
| 1981 | quant_type=quant_dtype, |
| 1982 | group_size=group_size, |
| 1983 | input_type=input_type, |
| 1984 | ) |
| 1985 | w2_data = MarlinMoEWeightData.make( |
| 1986 | w=w2, |
| 1987 | quant_type=quant_dtype, |
| 1988 | group_size=group_size, |
| 1989 | input_type=input_type, |
| 1990 | ) |
| 1991 | |
| 1992 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 1993 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, False) |
| 1994 | |
| 1995 | class BatchedRun: |
| 1996 | @staticmethod |
| 1997 | def _make_expert_num_tokens_cpu( |
| 1998 | e: int, # num_experts |
| 1999 | topk_ids_cpu: torch.Tensor, |
| 2000 | ) -> torch.Tensor: |
| 2001 | expert_num_tokens_cpu = torch.zeros((e,), dtype=torch.int32, device="cpu") |
| 2002 | for topk_id in torch.flatten(topk_ids_cpu): |
| 2003 | expert_num_tokens_cpu[topk_id] += 1 |
| 2004 | return expert_num_tokens_cpu |
| 2005 | |
| 2006 | def __init__( |
| 2007 | self, |
| 2008 | max_tokens_per_batch: int, |
| 2009 | num_experts: int, |
| 2010 | _topk_ids: torch.Tensor, |
| 2011 | _topk_weights: torch.Tensor, |
| 2012 | ): |
| 2013 | self.max_tokens_per_batch = max_tokens_per_batch |
| 2014 | self.e = num_experts |
| 2015 | self.topk_ids_cpu = _topk_ids.to("cpu") |
| 2016 | self.topk_weights_cpu = _topk_weights.to("cpu") |
| 2017 | self.expert_num_tokens_cpu = self._make_expert_num_tokens_cpu( |
| 2018 | self.e, self.topk_ids_cpu |
| 2019 | ) |
| 2020 | |
| 2021 | def is_valid(self): |
| 2022 | """Return True only if the input can be represented in a Batched |
| 2023 | format. |
| 2024 | """ |
| 2025 | return torch.all(self.expert_num_tokens_cpu <= self.max_tokens_per_batch) |
| 2026 | |
| 2027 | def _scatter(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 2028 | hidden_states_cpu = hidden_states.to("cpu") |
| 2029 | K = hidden_states_cpu.size(1) |
| 2030 | batched_hidden_states_cpu = torch.empty( |
| 2031 | (e, max_tokens_per_batch, K), |
| 2032 | dtype=hidden_states_cpu.dtype, |
| 2033 | device="cpu", |
| 2034 | ) |
| 2035 | |
| 2036 | counter_cpu = torch.zeros_like(self.expert_num_tokens_cpu) |
| 2037 | for t_idx, token in enumerate(hidden_states_cpu): |
| 2038 | for topk_id in self.topk_ids_cpu[t_idx]: |
| 2039 | pos_in_batch = counter_cpu[topk_id] |
| 2040 | batched_hidden_states_cpu[topk_id, pos_in_batch] = token |
| 2041 | counter_cpu[topk_id] += 1 |
| 2042 | assert torch.allclose(counter_cpu, self.expert_num_tokens_cpu) |
| 2043 | return batched_hidden_states_cpu.to("cuda") |
| 2044 | |
| 2045 | def _gather( |
| 2046 | self, batched_outputs: torch.Tensor, gather_outputs: torch.Tensor |
| 2047 | ) -> torch.Tensor: |
| 2048 | batched_outputs_cpu = batched_outputs.to("cpu") |
| 2049 | gather_outputs_cpu = torch.zeros_like(gather_outputs) |
| 2050 | |
| 2051 | counter_cpu = torch.zeros((e,), device="cpu", dtype=torch.int32) |
| 2052 | md = gather_outputs_cpu.size(0) |
| 2053 | for t_idx in range(md): |
| 2054 | token = None |
| 2055 | for topk_id, topk_weight in zip( |
| 2056 | self.topk_ids_cpu[t_idx], self.topk_weights_cpu[t_idx] |
| 2057 | ): |
| 2058 | pos_in_batch = counter_cpu[topk_id] |
| 2059 | t = batched_outputs_cpu[topk_id, pos_in_batch] * topk_weight |
| 2060 | if token is None: |
| 2061 | token = t |
| 2062 | else: |
| 2063 | token += t |
| 2064 | counter_cpu[topk_id] += 1 |
| 2065 | assert token is not None |
| 2066 | gather_outputs_cpu[t_idx] = token |
| 2067 | gather_outputs.copy_(gather_outputs_cpu) |
| 2068 | return gather_outputs |
| 2069 | |
| 2070 | def run( |
| 2071 | self, hidden_states: torch.Tensor, fused_marlin_moe_kwargs: dict[Any, Any] |
| 2072 | ) -> torch.Tensor: |
| 2073 | assert hidden_states.ndim == 2 |
| 2074 | assert self.is_valid() |
| 2075 | |
| 2076 | batched_hidden_states = self._scatter(hidden_states) |
| 2077 | |
| 2078 | kwargs = fused_marlin_moe_kwargs | { |
| 2079 | "hidden_states": batched_hidden_states, |
| 2080 | "expert_num_tokens": self.expert_num_tokens_cpu.to("cuda"), |
| 2081 | } |
| 2082 | batched_outputs = batched_fused_marlin_moe(**kwargs) |
| 2083 | |
| 2084 | output = torch.zeros_like(hidden_states) |
| 2085 | output = self._gather(batched_outputs, output) |
| 2086 | return output |
| 2087 | |
| 2088 | kwargs = { |
| 2089 | "w1": w1_data.qweight, |
| 2090 | "w2": w2_data.qweight, |
| 2091 | "bias1": None, |
| 2092 | "bias2": None, |
| 2093 | "w1_scale": w1_data.scales, |
| 2094 | "w2_scale": w2_data.scales, |
| 2095 | "global_num_experts": e, |
| 2096 | "expert_map": None, |
| 2097 | "global_scale1": w1_data.global_scale, |
| 2098 | "global_scale2": w2_data.global_scale, |
| 2099 | "w1_zeros": w1_data.zeros, |
| 2100 | "w2_zeros": w2_data.zeros, |
| 2101 | "quant_type_id": quant_dtype.id, |
| 2102 | } |
| 2103 | if input_dtype is not None: |
| 2104 | kwargs["input_dtype"] = input_dtype |
| 2105 | if w1_data.a_scales_factor is not None: |
| 2106 | kwargs["input_global_scale1"] = w1_data.a_scales_factor |
| 2107 | if w2_data.a_scales_factor is not None: |
| 2108 | kwargs["input_global_scale2"] = w2_data.a_scales_factor |
| 2109 | |
| 2110 | # Reference |
| 2111 | fused_marlin_moe_kwargs = kwargs | { |
| 2112 | "hidden_states": a, |
| 2113 | "topk_ids": topk_ids, |
| 2114 | "topk_weights": topk_weights, |
| 2115 | } |
| 2116 | ref_marlin_output = fused_marlin_moe(**fused_marlin_moe_kwargs) |
| 2117 | |
| 2118 | # Batched |
| 2119 | br = BatchedRun(max_tokens_per_batch, e, topk_ids, topk_weights) |
| 2120 | if not br.is_valid(): |
| 2121 | pytest.skip("Cannot represent data in Batched Format.") |
| 2122 | marlin_output = br.run(a, kwargs) |
| 2123 | |
| 2124 | torch.testing.assert_close(marlin_output, ref_marlin_output, atol=atol, rtol=0) |
| 2125 | |
| 2126 | |
| 2127 | @pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)]) |
| 2128 | @pytest.mark.parametrize("e,topk", [(8, 2)]) |
| 2129 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 2130 | @pytest.mark.skipif( |
| 2131 | not current_platform.is_device_capability_family(100), |
| 2132 | reason="TRTLLM backend test only runs on Blackwell GPUs (SM10x).", |
| 2133 | ) |
| 2134 | def test_unquantized_bf16_flashinfer_trtllm_backend( |
| 2135 | m: int, |
| 2136 | n: int, |
| 2137 | k: int, |
| 2138 | e: int, |
| 2139 | topk: int, |
| 2140 | dtype: torch.dtype, |
| 2141 | workspace_init, |
| 2142 | ): |
| 2143 | """Test BF16 unquantized MoE with FlashInfer TRTLLM backend.""" |
| 2144 | set_random_seed(7) |
| 2145 | |
| 2146 | from vllm.model_executor.layers.fused_moe.config import ( |
| 2147 | FusedMoEConfig, |
| 2148 | FusedMoEParallelConfig, |
| 2149 | RoutingMethodType, |
| 2150 | ) |
| 2151 | from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( |
| 2152 | UnquantizedMoeBackend, |
| 2153 | ) |
| 2154 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 2155 | UnquantizedFusedMoEMethod, |
| 2156 | ) |
| 2157 | |
| 2158 | # Setup test data |
| 2159 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 2160 | w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 2161 | w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10 |
| 2162 | router_logits = torch.randn((m, e), device="cuda", dtype=dtype) |
| 2163 | |
| 2164 | moe_config = FusedMoEConfig( |
| 2165 | num_experts=e, |
| 2166 | experts_per_token=topk, |
| 2167 | hidden_dim=k, |
| 2168 | intermediate_size=n, |
| 2169 | num_local_experts=e, |
| 2170 | num_logical_experts=e, |
| 2171 | activation=MoEActivation.SILU, |
| 2172 | device="cuda", |
| 2173 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 2174 | in_dtype=dtype, |
| 2175 | routing_method=RoutingMethodType.Renormalize, |
| 2176 | max_num_tokens=next_power_of_2(m), |
| 2177 | moe_backend="flashinfer_trtllm", |
| 2178 | ) |
| 2179 | |
| 2180 | with set_current_vllm_config(vllm_config): |
| 2181 | quant_method = UnquantizedFusedMoEMethod(moe_config) |
| 2182 | |
| 2183 | # Verify TRTLLM backend was selected |
| 2184 | assert ( |
| 2185 | quant_method.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM |
| 2186 | ), f"Expected FLASHINFER_TRTLLM backend, got {quant_method.unquantized_backend}" |
| 2187 | |
| 2188 | # Verify it's using monolithic path |
| 2189 | assert quant_method.is_monolithic, ( |
| 2190 | "FLASHINFER_TRTLLM backend should use monolithic forward" |
| 2191 | ) |
| 2192 | layer = torch.nn.Module() |
| 2193 | layer.w13_weight = Parameter(w1.clone(), requires_grad=False) |
| 2194 | layer.w2_weight = Parameter(w2.clone(), requires_grad=False) |
| 2195 | layer.global_num_experts = e |
| 2196 | layer.local_num_experts = e |
| 2197 | layer.top_k = topk |
| 2198 | layer.num_expert_group = 1 |
| 2199 | layer.topk_group = 1 |
| 2200 | layer.intermediate_size_per_partition = n |
| 2201 | layer.ep_rank = 0 |
| 2202 | layer.activation = MoEActivation.SILU |
| 2203 | layer.e_score_correction_bias = None |
| 2204 | layer.routing_method_type = RoutingMethodType.Renormalize |
| 2205 | layer.expert_map = None |
| 2206 | layer.apply_router_weight_on_input = False |
| 2207 | layer.routed_scaling_factor = 2.446 |
| 2208 | layer.shared_experts = None |
| 2209 | layer._expert_routing_tables = lambda: None |
| 2210 | |
| 2211 | quant_method.process_weights_after_loading(layer) |
| 2212 | |
| 2213 | assert quant_method.moe_kernel is not None, ( |
| 2214 | "moe_kernel should be set after process_weights_after_loading" |
| 2215 | ) |
| 2216 | assert quant_method.supports_internal_mk, ( |
| 2217 | "supports_internal_mk should be True after setup" |
| 2218 | ) |
| 2219 | |
| 2220 | trtllm_output = quant_method.apply_monolithic( |
| 2221 | layer=layer, |
| 2222 | x=a, |
| 2223 | router_logits=router_logits, |
| 2224 | ) |
| 2225 | |
| 2226 | # Compute torch baseline |
| 2227 | w1_original = w1.clone() |
| 2228 | w2_original = w2.clone() |
| 2229 | baseline_output = ( |
| 2230 | torch_moe(a, w1_original, w2_original, router_logits, topk) |
| 2231 | * layer.routed_scaling_factor |
| 2232 | ) |
| 2233 | |
| 2234 | close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85) |
| 2235 | assert close.float().mean() > 0.925 |