| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright (c) 2025 FlyDSL Project Contributors |
| 4 | |
| 5 | |
| 6 | import importlib.util |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.fused_moe import fused_experts |
| 12 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 13 | from vllm.model_executor.layers.fused_moe.config import ( |
| 14 | int4_w4a16_moe_quant_config, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.platforms.rocm import on_gfx950 |
| 18 | from vllm.utils.torch_utils import set_random_seed |
| 19 | |
| 20 | if not (current_platform.is_rocm() and on_gfx950()): |
| 21 | pytest.skip("This test can only run on ROCm and gfx950.", allow_module_level=True) |
| 22 | |
| 23 | aiter_available = importlib.util.find_spec("aiter") is not None |
| 24 | |
| 25 | if not aiter_available: |
| 26 | pytest.skip("These tests require AITER to run.", allow_module_level=True) |
| 27 | |
| 28 | from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import ( # noqa: E402 |
| 29 | fused_flydsl_moe, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E402, E501 |
| 32 | compressed_tensors_moe_w4a16_flydsl, |
| 33 | ) |
| 34 | |
| 35 | RoutingBuffers = tuple[ |
| 36 | torch.Tensor, # sorted_token_ids |
| 37 | torch.Tensor, # sorted_weights |
| 38 | torch.Tensor, # sorted_expert_ids |
| 39 | torch.Tensor, # num_valid_ids (shape [1], i32) |
| 40 | int, # sorted_size |
| 41 | int, # blocks |
| 42 | ] |
| 43 | |
| 44 | |
| 45 | def _assert_flydsl_matches_reference( |
| 46 | actual: torch.Tensor, |
| 47 | expected: torch.Tensor, |
| 48 | *, |
| 49 | atol: float = 0.5, |
| 50 | rtol: float = 0.1, |
| 51 | max_mismatch_fraction: float = 1e-5, |
| 52 | max_error_ratio: float = 3.0, |
| 53 | ) -> None: |
| 54 | """Compare kernels without making accuracy depend on output size. |
| 55 | |
| 56 | The FlyDSL and reference kernels accumulate BF16 values in different |
| 57 | orders. A strict allclose makes one expected rounding outlier fail an |
| 58 | otherwise accurate output, which becomes increasingly likely for the |
| 59 | largest token counts (up to 117 million output elements here). |
| 60 | |
| 61 | Keep the original elementwise tolerance for 99.999% of values, while also |
| 62 | bounding every tolerated outlier to three times its elementwise tolerance. |
| 63 | """ |
| 64 | assert actual.shape == expected.shape |
| 65 | |
| 66 | mismatch = ~torch.isclose(actual, expected, atol=atol, rtol=rtol) |
| 67 | mismatch_count = int(mismatch.sum().item()) |
| 68 | total = actual.numel() |
| 69 | allowed_mismatches = int(total * max_mismatch_fraction) |
| 70 | mismatch_msg = ( |
| 71 | f"FlyDSL/reference mismatch: {mismatch_count}/{total} values " |
| 72 | f"({mismatch_count / total:.6%}) exceed atol={atol}, rtol={rtol}; " |
| 73 | f"allowed <= {allowed_mismatches}/{total} " |
| 74 | f"({max_mismatch_fraction:.6%})" |
| 75 | ) |
| 76 | assert mismatch_count <= allowed_mismatches, mismatch_msg |
| 77 | |
| 78 | if mismatch_count: |
| 79 | actual_mismatch = actual[mismatch].float() |
| 80 | expected_mismatch = expected[mismatch].float() |
| 81 | abs_error = (actual_mismatch - expected_mismatch).abs() |
| 82 | tolerance = atol + rtol * expected_mismatch.abs() |
| 83 | worst_error_ratio = (abs_error / tolerance).max().item() |
| 84 | assert worst_error_ratio <= max_error_ratio, ( |
| 85 | f"{mismatch_msg}; worst error is {worst_error_ratio:.4f}x its " |
| 86 | f"elementwise tolerance (allowed <= {max_error_ratio:.4f}x)" |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | @pytest.mark.parametrize( |
| 91 | "num_tokens", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384] |
| 92 | ) |
| 93 | @pytest.mark.parametrize("inter_dim", [256, 512]) |
| 94 | def test_flydsl_moe(num_tokens: int, inter_dim: int): |
| 95 | set_random_seed(0) |
| 96 | |
| 97 | device = "cuda" |
| 98 | topk = 8 |
| 99 | num_experts = 384 |
| 100 | hidden_size = 7168 |
| 101 | packed_factor = 8 |
| 102 | w13_num_shards = 2 |
| 103 | params_dtype = torch.bfloat16 |
| 104 | group_size = 32 |
| 105 | w2_scales_size = inter_dim |
| 106 | scale_factor = 0.01 |
| 107 | |
| 108 | num_groups_w2 = w2_scales_size // group_size |
| 109 | num_groups_w13 = hidden_size // group_size |
| 110 | |
| 111 | w13_weight = torch.randint( |
| 112 | 0, |
| 113 | 255, |
| 114 | (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim), |
| 115 | dtype=torch.int32, |
| 116 | device=device, |
| 117 | ) |
| 118 | |
| 119 | w2_weight = torch.randint( |
| 120 | 0, |
| 121 | 255, |
| 122 | (num_experts, inter_dim // packed_factor, hidden_size), |
| 123 | dtype=torch.int32, |
| 124 | device=device, |
| 125 | ) |
| 126 | w13_scale = scale_factor * torch.randn( |
| 127 | num_experts, |
| 128 | num_groups_w13, |
| 129 | w13_num_shards * inter_dim, |
| 130 | dtype=params_dtype, |
| 131 | device=device, |
| 132 | ) |
| 133 | w2_scale = scale_factor * torch.randn( |
| 134 | num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device |
| 135 | ) |
| 136 | |
| 137 | w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8) |
| 138 | w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8) |
| 139 | w13_weight_scale = w13_scale.transpose(1, 2).contiguous() |
| 140 | w2_weight_scale = w2_scale.transpose(1, 2).contiguous() |
| 141 | |
| 142 | moe_quant_config = int4_w4a16_moe_quant_config( |
| 143 | w1_scale=w13_weight_scale, |
| 144 | w2_scale=w2_weight_scale, |
| 145 | w1_zp=None, |
| 146 | w2_zp=None, |
| 147 | block_shape=[0, group_size], |
| 148 | ) |
| 149 | score = torch.rand((num_tokens, num_experts), device=device, dtype=torch.float32) |
| 150 | topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) |
| 151 | topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) |
| 152 | x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=device) |
| 153 | out_ref = fused_experts( |
| 154 | x, |
| 155 | w13_weight_packed, |
| 156 | w2_weight_packed, |
| 157 | topk_weights=topk_weights, |
| 158 | topk_ids=topk_ids, |
| 159 | activation=MoEActivation.SILU, |
| 160 | apply_router_weight_on_input=False, |
| 161 | global_num_experts=num_experts, |
| 162 | expert_map=None, |
| 163 | quant_config=moe_quant_config, |
| 164 | ) |
| 165 | |
| 166 | w13 = w13_weight |
| 167 | w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13) |
| 168 | w13 = w13.view(-1).contiguous() |
| 169 | |
| 170 | w2 = w2_weight |
| 171 | w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2) |
| 172 | w2 = w2.view(-1).contiguous() |
| 173 | |
| 174 | w13_scale_flydsl = w13_scale |
| 175 | w2_scale_flydsl = w2_scale |
| 176 | |
| 177 | if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: |
| 178 | E, G, N = w13_scale.shape |
| 179 | w13_scale_flydsl = ( |
| 180 | w13_scale_flydsl.view(E, G // 2, 2, N) |
| 181 | .permute(0, 1, 3, 2) |
| 182 | .contiguous() |
| 183 | .view(-1) |
| 184 | .contiguous() |
| 185 | ) |
| 186 | elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: |
| 187 | w13_scale_flydsl = w13_scale_flydsl.squeeze(1) |
| 188 | |
| 189 | if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: |
| 190 | E, G, N = w2_scale.shape |
| 191 | w2_scale_flydsl = ( |
| 192 | w2_scale_flydsl.view(E, G // 2, 2, N) |
| 193 | .permute(0, 1, 3, 2) |
| 194 | .contiguous() |
| 195 | .view(-1) |
| 196 | .contiguous() |
| 197 | ) |
| 198 | elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: |
| 199 | w2_scale_flydsl = w2_scale_flydsl.squeeze(1) |
| 200 | |
| 201 | w13_scale_flydsl = w13_scale_flydsl.contiguous() |
| 202 | w2_scale_flydsl = w2_scale_flydsl.contiguous() |
| 203 | |
| 204 | w13.is_shuffled = True |
| 205 | w2.is_shuffled = True |
| 206 | |
| 207 | out = fused_flydsl_moe( |
| 208 | x, |
| 209 | w13, |
| 210 | w2, |
| 211 | num_experts, |
| 212 | inter_dim, |
| 213 | topk_weights, |
| 214 | topk_ids, |
| 215 | w1_scale=w13_scale_flydsl, |
| 216 | w2_scale=w2_scale_flydsl, |
| 217 | topk=topk_weights.shape[-1], |
| 218 | group_size=group_size, |
| 219 | doweight_stage1=False, |
| 220 | scale_is_bf16=True, |
| 221 | ) |
| 222 | |
| 223 | _assert_flydsl_matches_reference(out, out_ref) |
| 224 | |
| 225 | |
| 226 | if __name__ == "__main__": |
| 227 | test_flydsl_moe(512, 256) |