@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 160f7f0840

[ROCm][CI] Extended Fused MoE and FP8 MoE test support (#41100)

Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>

Andreas Karatzas committed Aug 19, 2026 at 10:19 UTC 160f7f0840e1773b2418199800c7ac66f1b4f074

9 files changed +331 -31

.buildkite/test-amd.yaml

+58 -1

@@ -3500,9 +3500,44 @@ steps:

3500 - vllm/_aiter_ops.py

3501 - vllm/platforms/rocm.py

3502 commands:

3503 - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

3503 + - pytest -v -s kernels/moe

3504 + --ignore=kernels/moe/test_modular_oai_triton_moe.py

3505 + --ignore=kernels/moe/test_gpt_oss_triton_kernels.py

3506 + --ignore=kernels/moe/test_moe.py

3507 + --ignore=kernels/moe/test_block_int8.py

3508 + --ignore=kernels/moe/test_triton_moe_no_act_mul.py

3509 + --ignore=kernels/moe/test_triton_moe_ptpc_fp8.py

3510 + --ignore=kernels/moe/test_deepep_moe.py

3511 + --ignore=kernels/moe/test_moe_layer.py

3512 + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

3513 - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

3514

3515 +- label: Kernels FusedMoE Layer Test (2xB200-2xMI355) # TBD

3516 + timeout_in_minutes: 180

3517 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

3518 + dind: false

3519 + agent_pool: mi355_2

3520 + num_gpus: 2

3521 + optional: true

3522 + working_dir: "/vllm-workspace/tests"

3523 + source_file_dependencies:

3524 + - csrc/moe/

3525 + - csrc/rocm/

3526 + - tests/kernels/moe

3527 + - vllm/model_executor/layers/fused_moe/

3528 + - vllm/model_executor/layers/quantization/

3529 + - vllm/distributed/

3530 + - vllm/config/

3531 + - vllm/forward_context.py

3532 + - vllm/v1/worker/workspace.py

3533 + - vllm/utils/import_utils.py

3534 + - vllm/utils/math_utils.py

3535 + - vllm/utils/torch_utils.py

3536 + - vllm/platforms/

3537 + - vllm/_aiter_ops.py

3538 + commands:

3539 + - pytest -v -s kernels/moe/test_moe_layer.py

3540 +

3541 - label: Kernels Quantization Test %N # TBD

3542 timeout_in_minutes: 180

3543 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

@@ -3526,12 +3561,34 @@ steps:

3561 commands:

3562 - pytest -v -s kernels/quantization

3563

3564 +- label: Kernels FP8 MoE Test (2xH100-1xMI355) # TBD

3565 + timeout_in_minutes: 180

3566 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

3567 + dind: false

3568 + agent_pool: mi355_1

3569 + working_dir: "/vllm-workspace/tests"

3570 + source_file_dependencies:

3571 + - csrc/moe/

3572 + - vllm/model_executor/layers/fused_moe/

3573 + - tests/kernels/moe/test_deepep_moe.py

3574 + - vllm/_aiter_ops.py

3575 + - vllm/platforms/rocm.py

3576 + - vllm/envs.py

3577 + commands:

3578 + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py

3579 + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py

3580 + - pytest -v -s kernels/moe/test_moe.py

3581 + - pytest -v -s kernels/moe/test_block_int8.py

3582 + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py

3583 + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

3584 +

3585 - label: Kernels FP8 MoE Test (2xH100-2xMI355) # TBD

3586 timeout_in_minutes: 180

3587 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

3588 dind: false

3589 agent_pool: mi355_2

3590 num_gpus: 2

3591 + optional: true

3592 working_dir: "/vllm-workspace/tests"

3593 source_file_dependencies:

3594 - csrc/moe/

.buildkite/test_areas/kernels.yaml

+53

@@ -451,3 +451,56 @@ steps:

451 commands:

452 - pytest -v -s kernels/moe/test_moe_layer.py

453 - pytest -v -s kernels/moe/test_deepep_v2_moe.py

454 +

455 +- label: Kernels FusedMoE Layer Test (2xMI355)

456 + key: kernels-fusedmoe-layer-test-2-mi355

457 + depends_on:

458 + - image-build-amd

459 + timeout_in_minutes: 180

460 + dind: false

461 + device: mi355_2

462 + soft_fail: true

463 + optional: true

464 + working_dir: "/vllm-workspace/tests"

465 + source_file_dependencies:

466 + - csrc/moe/

467 + - csrc/rocm/

468 + - tests/kernels/moe

469 + - vllm/model_executor/layers/fused_moe/

470 + - vllm/model_executor/layers/quantization/

471 + - vllm/distributed/

472 + - vllm/config/

473 + - vllm/forward_context.py

474 + - vllm/v1/worker/workspace.py

475 + - vllm/utils/import_utils.py

476 + - vllm/utils/math_utils.py

477 + - vllm/utils/torch_utils.py

478 + - vllm/platforms/

479 + - vllm/_aiter_ops.py

480 + commands:

481 + - pytest -v -s kernels/moe/test_moe_layer.py

482 +

483 +- label: Kernels FP8 MoE Test (MI355)

484 + key: kernels-fp8-moe-test-mi355

485 + depends_on:

486 + - image-build-amd

487 + timeout_in_minutes: 180

488 + dind: false

489 + device: mi355_1

490 + soft_fail: true

491 + optional: true

492 + working_dir: "/vllm-workspace/tests"

493 + source_file_dependencies:

494 + - csrc/moe/

495 + - vllm/model_executor/layers/fused_moe/

496 + - tests/kernels/moe/test_deepep_moe.py

497 + - vllm/_aiter_ops.py

498 + - vllm/platforms/rocm.py

499 + - vllm/envs.py

500 + commands:

501 + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py

502 + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py

503 + - pytest -v -s kernels/moe/test_moe.py

504 + - pytest -v -s kernels/moe/test_block_int8.py

505 + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py

506 + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

tests/kernels/moe/test_modular_oai_triton_moe.py

+27 -2

@@ -8,6 +8,7 @@ from __future__ import annotations

9 import pytest

10 import torch

11 +import torch.nn.functional as F

13 from tests.utils import wait_for_gpu_memory_to_clear

14 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

@@ -37,6 +38,7 @@ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe imp

38 )

39 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

40 from vllm.platforms import current_platform

41 +from vllm.utils.math_utils import round_up

42 from vllm.utils.torch_utils import set_random_seed

44 from .utils import make_dummy_moe_config, shuffle_weight

@@ -104,13 +106,30 @@ def make_weights(dtype, k, n, e):

106 w1_tri = shuffle_weight(w1_tri)

107 w1_bias_tri = shuffle_weight(w1_bias_tri)

108

109 + if current_platform.is_rocm():

110 + k_align, n2_align = 256, 512

111 + else:

112 + k_align, n2_align = 64, 128

113 +

114 + w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1]

115 + w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2]

116 + w2_bottom_pad = w1_right_pad // 2

117 + w2_right_pad = w1_bottom_pad

118 +

119 + w1_tri = F.pad(w1_tri, (0, w1_right_pad, 0, w1_bottom_pad, 0, 0))

120 + w2_tri = F.pad(w2_tri, (0, w2_right_pad, 0, w2_bottom_pad, 0, 0))

121 + w1_bias_tri = F.pad(w1_bias_tri, (0, w1_right_pad, 0, 0))

122 + w2_bias_tri = F.pad(w2_bias_tri, (0, w2_right_pad, 0, 0))

123 +

124 # quant triton_weights

125 w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)

126 w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, dtype, axis=1)

127 + w1 = w1[..., :k, : 2 * n]

128 w1 = unshuffle_weight(w1)

129

130 w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1)

131 w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, dtype, axis=1)

132 + w2 = w2[..., :n, :k]

133

134 num_warps = 8

135 w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)

@@ -150,6 +169,7 @@ def make_weights(dtype, k, n, e):

169 w2_bias_tri,

170 w1_precision_config,

171 w2_precision_config,

172 + w1_bottom_pad,

173 )

174

175

@@ -237,7 +257,7 @@ def oai_triton_moe_impl(

257

258

259 @pytest.mark.skipif(

240 - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."

260 + not current_platform.is_cuda_alike(), reason="Requires CUDA-alike platform."

261 )

262 @pytest.mark.parametrize("dtype", [torch.bfloat16])

263 @pytest.mark.parametrize("m,n,k", MNK)

@@ -256,6 +276,7 @@ def test_oai_triton_moe(

276 ):

277 wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1)

278 set_random_seed(0)

279 +

280 (

281 w1,

282 w2,

@@ -267,9 +288,11 @@ def test_oai_triton_moe(

288 w2_bias_tri,

289 w1_precision_config,

290 w2_precision_config,

291 + x_pad,

292 ) = make_weights(dtype, k, n, num_experts)

293

294 x = torch.randn((m, k), dtype=dtype, device="cuda")

295 + x_tri = F.pad(x, (0, x_pad, 0, 0))

296 router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype)

297 topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)

298 topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)

@@ -278,7 +301,7 @@ def test_oai_triton_moe(

301 out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids)

302

303 out = oai_triton_moe_impl(

281 - x,

304 + x_tri,

305 w1_tri,

306 w2_tri,

307 w1_precision_config,

@@ -290,6 +313,7 @@ def test_oai_triton_moe(

313 topk_ids,

314 unfused,

315 )

316 + out = out[..., :k]

317

318 assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005)

319

@@ -321,6 +345,7 @@ def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_

345 w2_bias_tri,

346 w1_precision_config,

347 w2_precision_config,

348 + _x_pad,

349 ) = make_weights(dtype, k, n, num_experts)

350

351 x = torch.randn((m, k), dtype=dtype, device="cuda")

tests/kernels/moe/test_moe.py

+41

@@ -16,6 +16,7 @@ from torch.nn import Parameter

16 from torch.nn import functional as F

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,

@@ -68,6 +69,46 @@ from vllm.utils.torch_utils import set_random_seed

69 DEVICE_TYPE = current_platform.device_type

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,

tests/kernels/moe/test_moe_layer.py

+106 -26

@@ -7,9 +7,11 @@ Run `pytest tests/kernels/test_moe_layer.py`.

8 import functools

9 import os

10 +import tempfile

11 import traceback

12 import types

13 from collections.abc import Callable

14 +from contextlib import suppress

15 from dataclasses import astuple, dataclass, fields

16 from itertools import product

17 from typing import get_args

@@ -71,6 +73,15 @@ from vllm.v1.worker.workspace import (

73 is_workspace_manager_initialized,

74 )

76 +

77 +def on_gfx950() -> bool:

78 + if current_platform.is_rocm():

79 + from vllm.platforms.rocm import on_gfx950 as rocm_on_gfx950

80 +

81 + return rocm_on_gfx950()

82 + return False

83 +

84 +

85 fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype

87 SHAPE_COMBOS = [

@@ -114,6 +125,7 @@ if has_nixl_ep():

125 BACKENDS += ["nixl_ep"]

126

127 DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"}

128 +MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"}

129

130 QUANT_METHODS = [

131 None,

@@ -484,12 +496,12 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

496 "leads to large differences.",

497 )

498

487 - # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+)

488 - if (

489 - config.quantization == "modelopt_fp4"

490 - and not current_platform.has_device_capability(100)

499 + # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+) or gfx950.

500 + if config.quantization == "modelopt_fp4" and not (

501 + (current_platform.is_rocm() and on_gfx950())

502 + or current_platform.has_device_capability(100)

503 ):

492 - return False, "modelopt_fp4 not supported on H100+ GPUs"

504 + return False, "modelopt_fp4 requires native NVFP4 or emulation"

505

506 # Skip flashinfer_nvlink if not on H100+ (compute capability 10.0+)

507 if (

@@ -508,6 +520,17 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

520 f"{config.backend} does not support quantization={config.quantization}",

521 )

522

523 + if config.backend in MORI_BACKENDS:

524 + if os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") != "1":

525 + return False, "mori MoE layer matrix is opt-in"

526 +

527 + from vllm._aiter_ops import rocm_aiter_ops

528 +

529 + if not rocm_aiter_ops.is_fused_moe_enabled():

530 + return False, "mori requires AITER fused MoE"

531 + if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():

532 + return False, "mori does not support AITER shared expert fusion"

533 +

534 if config.backend == "deepep_low_latency":

535 from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( # noqa: E501

536 DeepEPLLPrepareAndFinalize,

@@ -1695,14 +1718,16 @@ def _parallel_worker(

1718 cpu_group,

1719 test_configs: list[MoETestConfig],

1720 verbosity: int,

1721 + failure_report_path: str | None = None,

1722 **kwargs,

1723 ) -> None:

1724 set_random_seed(7)

1725 + is_logging_rank = pgi.rank == 0

1726

1727 total = 0

1728 passed = 0

1729 failed = 0

1705 - fail_ids = []

1730 + failure_details = []

1731

1732 dp_rank = vllm_config.parallel_config.data_parallel_rank

1733

@@ -1717,9 +1742,11 @@ def _parallel_worker(

1742

1743 tp_rank = pgi.rank % test_config.tp_size

1744

1720 - if verbosity > 0:

1745 + if verbosity > 0 and is_logging_rank:

1746 print(f"subtest: {test_config.id()}", end="")

1747

1748 + local_failed = False

1749 + local_error: str | None = None

1750 try:

1751 _run_one_config(

1752 vllm_config,

@@ -1743,19 +1770,9 @@ def _parallel_worker(

1770 use_gate=test_config.use_gate,

1771 use_routed_input_transform=test_config.use_routed_input_transform,

1772 )

1746 - if verbosity > 0:

1747 - print(" PASSED")

1748 - else:

1749 - print(".", end="")

1750 - passed = passed + 1

1751 - except Exception as ex:

1752 - fail_ids.append(test_config.id())

1753 - failed = failed + 1

1754 - if verbosity > 0:

1755 - traceback.print_exc()

1756 - print(f"\n{str(ex)}\nFAILED")

1757 - else:

1758 - print("F", end="")

1773 + except Exception:

1774 + local_failed = True

1775 + local_error = traceback.format_exc()

1776 finally:

1777 # DeepEP managers are not reliably reusable across many subtests in

1778 # a single worker process. Tear them down after each DeepEP case so

@@ -1771,6 +1788,40 @@ def _parallel_worker(

1788 total = total + 1

1789 torch.distributed.barrier()

1790

1791 + any_failed_tensor = torch.tensor(

1792 + [int(local_failed)], device=pgi.device, dtype=torch.int32

1793 + )

1794 + torch.distributed.all_reduce(

1795 + any_failed_tensor, op=torch.distributed.ReduceOp.MAX

1796 + )

1797 + any_failed = bool(any_failed_tensor.item())

1798 +

1799 + if any_failed:

1800 + failed = failed + 1

1801 +

1802 + gathered_errors = [None] * pgi.world_size

1803 + torch.distributed.all_gather_object(

1804 + gathered_errors, local_error, group=cpu_group

1805 + )

1806 + first_error = next(

1807 + (error for error in gathered_errors if error is not None),

1808 + "unknown distributed failure",

1809 + )

1810 + assert first_error is not None

1811 + failure_details.append(f"{test_config.id()}\n{first_error.rstrip()}")

1812 +

1813 + if verbosity > 0 and is_logging_rank:

1814 + print(" FAILED")

1815 + print(first_error.rstrip())

1816 + elif is_logging_rank:

1817 + print("F", end="")

1818 + else:

1819 + passed = passed + 1

1820 + if verbosity > 0 and is_logging_rank:

1821 + print(" PASSED")

1822 + elif is_logging_rank:

1823 + print(".", end="")

1824 +

1825 skipped = total - (passed + failed)

1826

1827 fails = f"{failed} failed" if failed > 0 else ""

@@ -1780,17 +1831,24 @@ def _parallel_worker(

1831 passes = f"{sep}{passed} passed" if passed > 0 else ""

1832

1833 report = (

1783 - f"============= {fails}{skips}{passes} of {total} total tests ============="

1834 + f"============= {fails}{skips}{passes} of {total} total subcases ============="

1835 )

1836

1786 - sep = "\n" if verbosity == 0 else ""

1787 - print(f"{sep}{report}")

1837 + if is_logging_rank:

1838 + sep = "\n" if verbosity == 0 else ""

1839 + print(f"{sep}{report}")

1840

1841 if failed > 0:

1790 - fail_ids_str = "\n".join(fail_ids)

1791 - raise RuntimeError(

1792 - f"\n============= Failed subtests =============\n{fail_ids_str}\n{report}"

1842 + failure_details_str = "\n\n".join(failure_details)

1843 + failure_report = (

1844 + f"\n============= Failed subcases =============\n"

1845 + f"{failure_details_str}\n{report}"

1846 )

1847 + if is_logging_rank and failure_report_path is not None:

1848 + with open(failure_report_path, "w", encoding="utf-8") as report_file:

1849 + report_file.write(failure_report)

1850 + if is_logging_rank:

1851 + raise RuntimeError(failure_report)

1852

1853

1854 # TODO: add cudagraphs/torch.compile tests

@@ -1836,6 +1894,17 @@ def test_moe_layer(

1894 if os.environ.get("VLLM_LOGGING_LEVEL") is None:

1895 monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")

1896

1897 + if (

1898 + backend in MORI_BACKENDS

1899 + and os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") == "1"

1900 + ):

1901 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

1902 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")

1903 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")

1904 + from vllm._aiter_ops import rocm_aiter_ops

1905 +

1906 + rocm_aiter_ops.refresh_env_variables()

1907 +

1908 # TODO: cover FlashInfer MoE backends via moe_backend, e.g.

1909 # moe_backend=flashinfer_trtllm / flashinfer_cutlass / flashinfer_cutedsl

1910 # (BF16, FP8 and NVFP4 paths), and VLLM_USE_FLASHINFER_MOE_INT4=1.

@@ -1881,6 +1950,11 @@ def test_moe_layer(

1950 if len(test_configs) == 0:

1951 pytest.skip("No supported configs found for this testpoint.")

1952

1953 + with tempfile.NamedTemporaryFile(

1954 + prefix="moe-layer-failures-", delete=False

1955 + ) as failure_report_file:

1956 + failure_report_path = failure_report_file.name

1957 +

1958 try:

1959 parallel_launch_with_config(

1960 world_size,

@@ -1889,7 +1963,13 @@ def test_moe_layer(

1963 None,

1964 test_configs,

1965 verbosity,

1966 + failure_report_path=failure_report_path,

1967 )

1968 + if os.path.getsize(failure_report_path) > 0:

1969 + with open(failure_report_path, encoding="utf-8") as report_file:

1970 + pytest.fail(report_file.read())

1971 finally:

1972 + with suppress(FileNotFoundError):

1973 + os.remove(failure_report_path)

1974 torch.accelerator.synchronize() # TODO: Is this needed?

1975 torch.accelerator.empty_cache()

tests/kernels/moe/utils.py

+26 -1

@@ -33,6 +33,10 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (

33 )

34 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk

35 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

36 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (

37 + ref_nvfp4_quant,

38 +)

39 +from vllm.platforms import current_platform

40 from vllm.utils.deep_gemm import per_block_cast_to_fp8

41 from vllm.utils.math_utils import round_up

@@ -294,13 +298,34 @@ def moe_quantize_weights_2d(

298 assert not per_token_quant

299 w_amax = torch.abs(w).max().to(torch.float32)

300 w_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w_amax

297 - w, w_s = ops.scaled_fp4_quant(w, w_gs)

301 + if current_platform.is_rocm():

302 + w, w_s = _scaled_fp4_quant_emulated(w, w_gs)

303 + else:

304 + w, w_s = ops.scaled_fp4_quant(w, w_gs)

305 else:

306 raise RuntimeError(f"Unsupported quant type {quant_dtype}")

307

308 return w, w_s, w_gs

309

310

311 +def _pack_e2m1_fp4(fp4_values: torch.Tensor) -> torch.Tensor:

312 + assert fp4_values.shape[-1] % 2 == 0

313 +

314 + abs_values = fp4_values.abs()

315 + codes = torch.empty_like(abs_values, dtype=torch.uint8)

316 + for code, value in enumerate((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)):

317 + codes[abs_values == value] = code

318 + codes = codes | ((fp4_values < 0).to(torch.uint8) << 3)

319 + return codes[..., 0::2] | (codes[..., 1::2] << 4)

320 +

321 +

322 +def _scaled_fp4_quant_emulated(

323 + w: torch.Tensor, w_gs: torch.Tensor

324 +) -> tuple[torch.Tensor, torch.Tensor]:

325 + fp4_values, w_s = ref_nvfp4_quant(w, w_gs, block_size=16)

326 + return _pack_e2m1_fp4(fp4_values), w_s.to(torch.float8_e4m3fn)

327 +

328 +

329 def moe_quantize_weights(

330 w: torch.Tensor,

331 w_s: torch.Tensor | None,

tests/quantization/test_fp8.py

+13

@@ -31,6 +31,9 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

31 from vllm.model_executor.layers.quantization.online.fp8 import (

32 Fp8PerTensorOnlineLinearMethod,

33 )

34 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

35 + process_fp8_input_tensor_strategy_moe,

36 +)

37 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

38 from vllm.platforms import current_platform

@@ -47,6 +50,16 @@ MODELS = [

50 ]

53 +def test_static_fp8_moe_input_scales_remain_scalar() -> None:

54 + a1_scale, a2_scale = process_fp8_input_tensor_strategy_moe(

55 + torch.tensor([0.25, 0.5]),

56 + torch.tensor([0.75, 0.6]),

57 + enable_eplb=False,

58 + )

59 +

60 + assert a1_scale.ndim == a2_scale.ndim == 0

61 +

62 +

63 @pytest.mark.skipif(

64 not is_quant_method_supported("fp8"),

65 reason="FP8 is not supported on this GPU type.",

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -857,6 +857,12 @@ def invoke_fused_moe_triton_kernel(

857 BLOCK_SIZE_K,

858 )

859 use_td = False

860 +

861 + # Triton treats 0-D tensor arguments as scalar values, but the kernel

862 + # loads tensor-wise activation scales through a pointer.

863 + if A_scale is not None and A_scale.ndim == 0:

864 + A_scale = A_scale.reshape(1)

865 +

866 fused_moe_kernel[grid](

867 A,

868 B,

vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py

+1 -1

@@ -313,7 +313,7 @@ def break_fp4_bytes(a, dtype):

313 signs = (combined & 0x08).to(torch.bool) # Sign bits

314 abs_vals = (combined & 0x07).to(torch.long)

315

316 - kE2M1 = kE2M1ToFloat_handle.val

316 + kE2M1 = kE2M1ToFloat_handle.val.to(device=a.device)

317 # Device-aware lookup and sign application

318 values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)

319 # Reshape to final form