| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the MoE grouped topk kernel. |
| 4 | |
| 5 | Run `pytest tests/kernels/moe/test_grouped_topk.py`. |
| 6 | """ |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.config import ( |
| 13 | CompilationConfig, |
| 14 | VllmConfig, |
| 15 | get_cached_compilation_config, |
| 16 | set_current_vllm_config, |
| 17 | ) |
| 18 | from vllm.forward_context import set_forward_context |
| 19 | from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( |
| 20 | GroupedTopk, |
| 21 | GroupedTopKRouter, |
| 22 | fused_grouped_topk, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.utils.torch_utils import set_random_seed |
| 26 | |
| 27 | |
| 28 | def _run_single_group_topk( |
| 29 | logits: torch.Tensor, |
| 30 | bias: torch.Tensor, |
| 31 | topk: int, |
| 32 | *, |
| 33 | scoring_func: str, |
| 34 | renormalize: bool, |
| 35 | routed_scaling_factor: float = 1.0, |
| 36 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 37 | return fused_grouped_topk( |
| 38 | hidden_states=torch.empty( |
| 39 | (logits.shape[0], 0), dtype=logits.dtype, device=logits.device |
| 40 | ), |
| 41 | gating_output=logits, |
| 42 | topk=topk, |
| 43 | renormalize=renormalize, |
| 44 | e_score_correction_bias=bias, |
| 45 | num_expert_group=1, |
| 46 | topk_group=1, |
| 47 | scoring_func=scoring_func, |
| 48 | routed_scaling_factor=routed_scaling_factor, |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def _single_group_reference( |
| 53 | logits: torch.Tensor, |
| 54 | bias: torch.Tensor, |
| 55 | topk: int, |
| 56 | *, |
| 57 | scoring_func: str, |
| 58 | renormalize: bool, |
| 59 | routed_scaling_factor: float = 1.0, |
| 60 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 61 | if scoring_func == "sigmoid": |
| 62 | scores = 0.5 * torch.tanh(0.5 * logits.float()) + 0.5 |
| 63 | else: |
| 64 | scores = torch.softmax(logits, dim=-1).float() |
| 65 | indices = torch.argsort( |
| 66 | scores + bias.float(), dim=-1, descending=True, stable=True |
| 67 | )[:, :topk] |
| 68 | values = scores.gather(1, indices) |
| 69 | if renormalize: |
| 70 | values /= values.sum(dim=-1, keepdim=True) + 1e-20 |
| 71 | values *= routed_scaling_factor |
| 72 | return values, indices.to(torch.int32) |
| 73 | |
| 74 | |
| 75 | @pytest.mark.skipif( |
| 76 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 77 | ) |
| 78 | @pytest.mark.parametrize("skip_padding", [False, True]) |
| 79 | def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool): |
| 80 | hidden_states = torch.randn((2, 16), device="cuda") |
| 81 | gating_output = torch.randn((2, 896), device="cuda") |
| 82 | is_padding = torch.tensor([False, True], device="cuda") |
| 83 | router = GroupedTopKRouter( |
| 84 | top_k=16, |
| 85 | global_num_experts=896, |
| 86 | num_expert_group=1, |
| 87 | topk_group=1, |
| 88 | skip_padding=skip_padding, |
| 89 | ) |
| 90 | |
| 91 | with set_forward_context(None, VllmConfig(), is_padding=is_padding): |
| 92 | topk_weights, topk_ids = router._compute_routing( |
| 93 | hidden_states, gating_output, torch.int32 |
| 94 | ) |
| 95 | |
| 96 | assert torch.all(topk_ids[0] >= 0) |
| 97 | if skip_padding: |
| 98 | assert torch.all(topk_ids[1] == -1) |
| 99 | assert torch.all(topk_weights[1] == 0) |
| 100 | else: |
| 101 | assert torch.all(topk_ids[1] >= 0) |
| 102 | assert torch.all(topk_weights[1] > 0) |
| 103 | |
| 104 | |
| 105 | @pytest.mark.skipif( |
| 106 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 107 | ) |
| 108 | @pytest.mark.parametrize("n_token", [1, 33, 64]) |
| 109 | @pytest.mark.parametrize("n_hidden", [1024, 2048]) |
| 110 | @pytest.mark.parametrize( |
| 111 | "n_expert,topk,num_expert_group,topk_group", |
| 112 | [ |
| 113 | (16, 2, 8, 2), |
| 114 | (128, 2, 8, 2), |
| 115 | (256, 8, 8, 4), |
| 116 | (384, 8, 1, 1), |
| 117 | (512, 22, 1, 1), |
| 118 | ], |
| 119 | ) |
| 120 | @pytest.mark.parametrize("renormalize", [True, False]) |
| 121 | @pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"]) |
| 122 | @pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) |
| 123 | @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float32]) |
| 124 | @pytest.mark.parametrize("bias_dtype", [torch.float32]) |
| 125 | def test_grouped_topk( |
| 126 | monkeypatch: pytest.MonkeyPatch, |
| 127 | n_token: int, |
| 128 | n_hidden: int, |
| 129 | n_expert: int, |
| 130 | topk: int, |
| 131 | num_expert_group: int, |
| 132 | topk_group: int, |
| 133 | renormalize: bool, |
| 134 | scoring_func: str, |
| 135 | routed_scaling_factor: float, |
| 136 | input_dtype: torch.dtype, |
| 137 | bias_dtype: torch.dtype, |
| 138 | ): |
| 139 | vllm_config = VllmConfig( |
| 140 | compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"]) |
| 141 | ) |
| 142 | get_cached_compilation_config.cache_clear() |
| 143 | |
| 144 | set_random_seed(0) |
| 145 | hidden_states = torch.randn((n_token, n_hidden), dtype=input_dtype, device="cuda") |
| 146 | gating_output = torch.randn((n_token, n_expert), dtype=input_dtype, device="cuda") |
| 147 | e_score_correction_bias = torch.randn((n_expert,), dtype=bias_dtype, device="cuda") |
| 148 | |
| 149 | with set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 150 | m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0") |
| 151 | m.setattr(envs, "VLLM_BATCH_INVARIANT", True) |
| 152 | grouped_topk = GroupedTopk( |
| 153 | topk=topk, |
| 154 | renormalize=renormalize, |
| 155 | num_expert_group=num_expert_group, |
| 156 | topk_group=topk_group, |
| 157 | scoring_func=scoring_func, |
| 158 | routed_scaling_factor=routed_scaling_factor, |
| 159 | ) |
| 160 | assert grouped_topk._forward_method.__name__ == "forward_cuda" |
| 161 | baseline_topk_weights, baseline_topk_ids = grouped_topk( |
| 162 | hidden_states=hidden_states, |
| 163 | gating_output=gating_output, |
| 164 | e_score_correction_bias=e_score_correction_bias, |
| 165 | ) |
| 166 | |
| 167 | test_topk_weights, test_topk_ids = fused_grouped_topk( |
| 168 | hidden_states=hidden_states, |
| 169 | gating_output=gating_output, |
| 170 | topk=topk, |
| 171 | renormalize=renormalize, |
| 172 | num_expert_group=num_expert_group, |
| 173 | topk_group=topk_group, |
| 174 | scoring_func=scoring_func, |
| 175 | routed_scaling_factor=routed_scaling_factor, |
| 176 | e_score_correction_bias=e_score_correction_bias, |
| 177 | ) |
| 178 | |
| 179 | torch.testing.assert_close( |
| 180 | baseline_topk_weights, test_topk_weights, atol=2e-2, rtol=0 |
| 181 | ) |
| 182 | torch.testing.assert_close(baseline_topk_ids, test_topk_ids, atol=0, rtol=0) |
| 183 | |
| 184 | |
| 185 | @pytest.mark.skipif( |
| 186 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 187 | ) |
| 188 | def test_grouped_topk_single_group_large_batch(): |
| 189 | set_random_seed(0) |
| 190 | logits = torch.randn((1536, 896), dtype=torch.bfloat16, device="cuda") |
| 191 | bias = torch.randn((896,), dtype=torch.float32, device="cuda") |
| 192 | |
| 193 | expected_values, expected_ids = _single_group_reference( |
| 194 | logits, bias, 16, scoring_func="sigmoid", renormalize=True |
| 195 | ) |
| 196 | actual_values, actual_ids = _run_single_group_topk( |
| 197 | logits, bias, 16, scoring_func="sigmoid", renormalize=True |
| 198 | ) |
| 199 | |
| 200 | torch.testing.assert_close(actual_ids, expected_ids) |
| 201 | torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) |
| 202 | |
| 203 | |
| 204 | @pytest.mark.skipif( |
| 205 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 206 | ) |
| 207 | @pytest.mark.parametrize( |
| 208 | "num_experts,topk,input_dtype,bias_dtype", |
| 209 | [ |
| 210 | (512, 9, torch.bfloat16, torch.float32), |
| 211 | (512, 16, torch.float16, torch.float16), |
| 212 | (513, 9, torch.float32, torch.bfloat16), |
| 213 | (513, 16, torch.bfloat16, torch.float32), |
| 214 | (895, 9, torch.float16, torch.bfloat16), |
| 215 | (896, 16, torch.float32, torch.float16), |
| 216 | (897, 9, torch.bfloat16, torch.bfloat16), |
| 217 | (897, 16, torch.float16, torch.float32), |
| 218 | (1024, 9, torch.float32, torch.bfloat16), |
| 219 | (1024, 16, torch.bfloat16, torch.float16), |
| 220 | ], |
| 221 | ) |
| 222 | @pytest.mark.parametrize( |
| 223 | "scoring_func,renormalize,routed_scaling_factor", |
| 224 | [ |
| 225 | ("sigmoid", True, 1.0), |
| 226 | ("sigmoid", False, 2.5), |
| 227 | ("softmax", True, 2.5), |
| 228 | ("softmax", False, 1.0), |
| 229 | ], |
| 230 | ) |
| 231 | def test_grouped_topk_single_group_tiers( |
| 232 | num_experts: int, |
| 233 | topk: int, |
| 234 | input_dtype: torch.dtype, |
| 235 | bias_dtype: torch.dtype, |
| 236 | scoring_func: str, |
| 237 | renormalize: bool, |
| 238 | routed_scaling_factor: float, |
| 239 | ): |
| 240 | set_random_seed(7) |
| 241 | logits = torch.randn((17, num_experts), dtype=input_dtype, device="cuda") |
| 242 | bias = torch.randn((num_experts,), dtype=bias_dtype, device="cuda") |
| 243 | |
| 244 | expected_values, expected_ids = _single_group_reference( |
| 245 | logits, |
| 246 | bias, |
| 247 | topk, |
| 248 | scoring_func=scoring_func, |
| 249 | renormalize=renormalize, |
| 250 | routed_scaling_factor=routed_scaling_factor, |
| 251 | ) |
| 252 | actual_values, actual_ids = _run_single_group_topk( |
| 253 | logits, |
| 254 | bias, |
| 255 | topk, |
| 256 | scoring_func=scoring_func, |
| 257 | renormalize=renormalize, |
| 258 | routed_scaling_factor=routed_scaling_factor, |
| 259 | ) |
| 260 | |
| 261 | torch.testing.assert_close(actual_ids, expected_ids) |
| 262 | torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) |
| 263 | |
| 264 | |
| 265 | @pytest.mark.skipif( |
| 266 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 267 | ) |
| 268 | @pytest.mark.parametrize( |
| 269 | "num_experts,topk,scoring_func", |
| 270 | [ |
| 271 | (128, 8, "sigmoid"), |
| 272 | (129, 8, "sigmoid"), |
| 273 | (257, 8, "sigmoid"), |
| 274 | (385, 8, "sigmoid"), |
| 275 | (512, 9, "sigmoid"), |
| 276 | (513, 9, "sigmoid"), |
| 277 | (769, 9, "sigmoid"), |
| 278 | (897, 16, "sigmoid"), |
| 279 | (1024, 16, "sigmoid"), |
| 280 | (128, 4, "softmax"), |
| 281 | (128, 5, "softmax"), |
| 282 | (129, 8, "softmax"), |
| 283 | (161, 8, "softmax"), |
| 284 | (256, 9, "softmax"), |
| 285 | (257, 8, "softmax"), |
| 286 | (512, 9, "softmax"), |
| 287 | (512, 17, "softmax"), |
| 288 | (512, 23, "softmax"), |
| 289 | (513, 8, "softmax"), |
| 290 | (577, 9, "softmax"), |
| 291 | (769, 9, "softmax"), |
| 292 | (897, 9, "softmax"), |
| 293 | (1024, 16, "softmax"), |
| 294 | ], |
| 295 | ) |
| 296 | def test_grouped_topk_single_group_capacity_tiers( |
| 297 | num_experts: int, |
| 298 | topk: int, |
| 299 | scoring_func: str, |
| 300 | ): |
| 301 | set_random_seed(11) |
| 302 | logits = torch.randn((3, num_experts), dtype=torch.bfloat16, device="cuda") |
| 303 | bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda") |
| 304 | expected_values, expected_ids = _single_group_reference( |
| 305 | logits, |
| 306 | bias, |
| 307 | topk, |
| 308 | scoring_func=scoring_func, |
| 309 | renormalize=True, |
| 310 | routed_scaling_factor=2.5, |
| 311 | ) |
| 312 | actual_values, actual_ids = _run_single_group_topk( |
| 313 | logits, |
| 314 | bias, |
| 315 | topk, |
| 316 | scoring_func=scoring_func, |
| 317 | renormalize=True, |
| 318 | routed_scaling_factor=2.5, |
| 319 | ) |
| 320 | |
| 321 | torch.testing.assert_close(actual_ids, expected_ids) |
| 322 | torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) |
| 323 | |
| 324 | |
| 325 | @pytest.mark.skipif( |
| 326 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 327 | ) |
| 328 | @pytest.mark.parametrize("num_experts", [512, 896, 1024]) |
| 329 | def test_grouped_topk_single_group_stable_ties(num_experts: int): |
| 330 | logits = torch.zeros((1, num_experts), dtype=torch.bfloat16, device="cuda") |
| 331 | bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") |
| 332 | |
| 333 | actual_values, actual_ids = _run_single_group_topk( |
| 334 | logits, |
| 335 | bias, |
| 336 | 16, |
| 337 | scoring_func="sigmoid", |
| 338 | renormalize=True, |
| 339 | routed_scaling_factor=2.5, |
| 340 | ) |
| 341 | |
| 342 | expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] |
| 343 | expected_values = torch.full((1, 16), 2.5 / 16, dtype=torch.float32, device="cuda") |
| 344 | torch.testing.assert_close(actual_ids, expected_ids) |
| 345 | torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) |
| 346 | |
| 347 | |
| 348 | @pytest.mark.skipif( |
| 349 | not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." |
| 350 | ) |
| 351 | @pytest.mark.parametrize("num_experts", [512, 896, 1024]) |
| 352 | @pytest.mark.parametrize("num_finite", [0, 15]) |
| 353 | @pytest.mark.parametrize("renormalize", [False, True]) |
| 354 | def test_grouped_topk_single_group_nonfinite_scores( |
| 355 | num_experts: int, num_finite: int, renormalize: bool |
| 356 | ): |
| 357 | logits = torch.full( |
| 358 | (1, num_experts), float("nan"), dtype=torch.bfloat16, device="cuda" |
| 359 | ) |
| 360 | if num_finite: |
| 361 | logits[0, :num_finite] = torch.arange( |
| 362 | num_finite, dtype=torch.bfloat16, device="cuda" |
| 363 | ) |
| 364 | logits[0, num_finite] = torch.inf |
| 365 | logits[0, num_finite + 1] = -torch.inf |
| 366 | bias = torch.zeros((num_experts,), dtype=torch.float32, device="cuda") |
| 367 | |
| 368 | actual_values, actual_ids = _run_single_group_topk( |
| 369 | logits, |
| 370 | bias, |
| 371 | 16, |
| 372 | scoring_func="sigmoid", |
| 373 | renormalize=renormalize, |
| 374 | routed_scaling_factor=2.5, |
| 375 | ) |
| 376 | |
| 377 | if num_finite == 0: |
| 378 | expected_ids = torch.arange(16, dtype=torch.int32, device="cuda")[None] |
| 379 | if renormalize: |
| 380 | expected_values = torch.full( |
| 381 | (1, 16), 1 / 16, dtype=torch.float32, device="cuda" |
| 382 | ) |
| 383 | else: |
| 384 | expected_values = torch.zeros((1, 16), dtype=torch.float32, device="cuda") |
| 385 | else: |
| 386 | expected_ids = torch.cat( |
| 387 | ( |
| 388 | torch.arange(num_finite - 1, -1, -1, dtype=torch.int32, device="cuda"), |
| 389 | torch.tensor([num_finite], dtype=torch.int32, device="cuda"), |
| 390 | ) |
| 391 | )[None] |
| 392 | finite_values = logits[0, :num_finite].float().sigmoid().flip(0) |
| 393 | if renormalize: |
| 394 | finite_values /= finite_values.sum() |
| 395 | finite_values *= 2.5 |
| 396 | expected_values = torch.cat( |
| 397 | (finite_values, torch.zeros(1, dtype=torch.float32, device="cuda")) |
| 398 | )[None] |
| 399 | |
| 400 | torch.testing.assert_close(actual_ids, expected_ids) |
| 401 | torch.testing.assert_close(actual_values, expected_values, atol=2e-5, rtol=0) |