| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for INT8 (W8A8) fused-MoE oracle backend selection. |
| 4 | |
| 5 | These exercise ``select_int8_moe_backend`` only (no MoE kernels are launched), |
| 6 | so they run on any platform where the Triton INT8 MoE kernel is available — |
| 7 | CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.config.kernel import MoEBackend |
| 14 | from vllm.model_executor.layers.fused_moe.config import ( |
| 15 | FusedMoEConfig, |
| 16 | FusedMoEParallelConfig, |
| 17 | RoutingMethodType, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.oracle.int8 import ( |
| 20 | Int8MoeBackend, |
| 21 | select_int8_moe_backend, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 24 | kInt8DynamicTensorSym, |
| 25 | kInt8DynamicTokenSym, |
| 26 | kInt8StaticChannelSym, |
| 27 | kInt8StaticTensorSym, |
| 28 | ) |
| 29 | from vllm.platforms import current_platform |
| 30 | |
| 31 | # The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+), on |
| 32 | # ROCm CDNA GPUs and on XPU. Gate on that rather than on a specific arch. |
| 33 | INT8_MOE_SUPPORTED = ( |
| 34 | (current_platform.is_cuda() and current_platform.has_device_capability((7, 5))) |
| 35 | or current_platform.is_rocm() |
| 36 | or current_platform.is_xpu() |
| 37 | ) |
| 38 | |
| 39 | requires_int8_moe = pytest.mark.skipif( |
| 40 | not INT8_MOE_SUPPORTED, |
| 41 | reason="Requires Triton INT8 MoE support (CUDA SM>=7.5, ROCm, or XPU)", |
| 42 | ) |
| 43 | |
| 44 | # So FusedMoEConfig.device names the device the test actually runs on. |
| 45 | DEVICE = current_platform.device_type |
| 46 | |
| 47 | |
| 48 | def _make_int8_moe_config(moe_backend: MoEBackend = "auto") -> FusedMoEConfig: |
| 49 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 50 | |
| 51 | return FusedMoEConfig( |
| 52 | num_experts=8, |
| 53 | experts_per_token=2, |
| 54 | hidden_dim=256, |
| 55 | intermediate_size=256, |
| 56 | num_local_experts=8, |
| 57 | num_logical_experts=8, |
| 58 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 59 | activation=MoEActivation.SILU, |
| 60 | in_dtype=torch.bfloat16, |
| 61 | device=DEVICE, |
| 62 | routing_method=RoutingMethodType.Renormalize, |
| 63 | moe_backend=moe_backend, |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | @requires_int8_moe |
| 68 | @pytest.mark.parametrize( |
| 69 | "weight_key,activation_key", |
| 70 | [ |
| 71 | pytest.param( |
| 72 | kInt8StaticChannelSym, |
| 73 | kInt8DynamicTokenSym, |
| 74 | id="per_channel_weight-per_token_act", |
| 75 | ), |
| 76 | pytest.param( |
| 77 | kInt8StaticTensorSym, |
| 78 | kInt8DynamicTensorSym, |
| 79 | id="per_tensor_weight-per_tensor_act", |
| 80 | ), |
| 81 | ], |
| 82 | ) |
| 83 | def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): |
| 84 | """Both dynamic-activation INT8 MoE schemes select the Triton backend.""" |
| 85 | config = _make_int8_moe_config() |
| 86 | backend, experts_cls = select_int8_moe_backend( |
| 87 | config, weight_key=weight_key, activation_key=activation_key |
| 88 | ) |
| 89 | assert backend == Int8MoeBackend.TRITON |
| 90 | assert experts_cls is not None |
| 91 | |
| 92 | |
| 93 | @pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU-only behaviour") |
| 94 | def test_scaled_int8_quant_is_available_on_xpu(): |
| 95 | """The per-tensor scheme's activation quantization works on XPU. |
| 96 | |
| 97 | ``_int8_quantize`` sends per-tensor activations through |
| 98 | ``ops.scaled_int8_quant``, which has an XPU branch. This is what lets the |
| 99 | per-tensor scheme be enabled here alongside CUDA and ROCm; pin it so a |
| 100 | future change cannot silently remove that branch. |
| 101 | """ |
| 102 | from vllm import _custom_ops as ops |
| 103 | |
| 104 | x = torch.randn(4, 16, device=DEVICE, dtype=torch.bfloat16) |
| 105 | scale = torch.full((1,), 0.05, device=DEVICE, dtype=torch.float32) |
| 106 | |
| 107 | q, returned_scale, azp = ops.scaled_int8_quant(x, scale=scale) |
| 108 | |
| 109 | assert q.dtype == torch.int8 |
| 110 | assert q.shape == x.shape |
| 111 | assert azp is None |
| 112 | torch.testing.assert_close(returned_scale, scale) |
| 113 | |
| 114 | |
| 115 | @requires_int8_moe |
| 116 | def test_int8_explicit_moe_backend_triton(): |
| 117 | """An explicit --moe-backend triton selects the Triton INT8 backend.""" |
| 118 | config = _make_int8_moe_config(moe_backend="triton") |
| 119 | backend, experts_cls = select_int8_moe_backend( |
| 120 | config, |
| 121 | weight_key=kInt8StaticChannelSym, |
| 122 | activation_key=kInt8DynamicTokenSym, |
| 123 | ) |
| 124 | assert backend == Int8MoeBackend.TRITON |
| 125 | assert experts_cls is not None |
| 126 | |
| 127 | |
| 128 | @requires_int8_moe |
| 129 | def test_int8_unsupported_moe_backend_raises(): |
| 130 | """An unsupported --moe-backend for INT8 MoE raises a clear error.""" |
| 131 | config = _make_int8_moe_config(moe_backend="cutlass") |
| 132 | with pytest.raises(ValueError, match="not supported for Int8 MoE"): |
| 133 | select_int8_moe_backend( |
| 134 | config, |
| 135 | weight_key=kInt8StaticChannelSym, |
| 136 | activation_key=kInt8DynamicTokenSym, |
| 137 | ) |