| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the fused MoE topk weight-and-reduce kernel. |
| 4 | |
| 5 | Run `pytest tests/kernels/moe/test_moe_fused_mul_sum.py`. |
| 6 | """ |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import ( |
| 12 | moe_fused_mul_sum, |
| 13 | ) |
| 14 | from vllm.utils.torch_utils import set_random_seed |
| 15 | |
| 16 | NUM_TOKENS = [1, 17, 256] |
| 17 | TOP_KS = [2, 8] |
| 18 | HIDDEN_SIZE = 512 |
| 19 | # One BLOCK_K tile, a partial trailing tile, and Kimi-K3's real dims |
| 20 | # (moe_intermediate_size 3072, hidden_size 7168) -- exercises the per-tile loop |
| 21 | # and its tail mask (BLOCK_K caps at 512 for <=2-byte dtypes). |
| 22 | HIDDEN_SIZES = [512, 1000, 3072, 7168] |
| 23 | DTYPES = [torch.bfloat16, torch.float32] |
| 24 | NUM_EXPERTS = 16 |
| 25 | |
| 26 | |
| 27 | def _reference( |
| 28 | inputs: torch.Tensor, |
| 29 | topk_weights: torch.Tensor, |
| 30 | topk_ids: torch.Tensor, |
| 31 | expert_map: torch.Tensor | None, |
| 32 | ) -> torch.Tensor: |
| 33 | valid = topk_ids >= 0 |
| 34 | if expert_map is not None: |
| 35 | valid &= expert_map[topk_ids.clamp(min=0)] >= 0 |
| 36 | weights = torch.where(valid, topk_weights, torch.zeros_like(topk_weights)) |
| 37 | return (inputs.float() * weights.float().unsqueeze(-1)).sum(dim=1) |
| 38 | |
| 39 | |
| 40 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 41 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 42 | @pytest.mark.parametrize("use_expert_map", [True, False]) |
| 43 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 44 | def test_invalid_slots_excluded(num_tokens: int, top_k: int, use_expert_map: bool): |
| 45 | """Slots marked -1 must not contribute, and must not be dereferenced. |
| 46 | |
| 47 | The expert GEMM never writes rows for topk slots whose expert id is -1 |
| 48 | (a non-local expert under EP, or an all2all padding row), so those rows |
| 49 | hold stale workspace data. Filling them with NaN here means any failure to |
| 50 | mask them shows up as NaN in the output rather than a small numeric drift. |
| 51 | Indexing `expert_map` with -1 is also an out-of-bounds read. |
| 52 | """ |
| 53 | set_random_seed(0) |
| 54 | device = "cuda" |
| 55 | |
| 56 | inputs = torch.randn( |
| 57 | num_tokens, top_k, HIDDEN_SIZE, dtype=torch.bfloat16, device=device |
| 58 | ) |
| 59 | topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device) |
| 60 | topk_ids = torch.randint( |
| 61 | 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device |
| 62 | ) |
| 63 | |
| 64 | # Mark roughly half the slots invalid and poison the rows behind them. |
| 65 | invalid = torch.rand(num_tokens, top_k, device=device) < 0.5 |
| 66 | topk_ids[invalid] = -1 |
| 67 | inputs[invalid] = float("nan") |
| 68 | |
| 69 | expert_map = None |
| 70 | if use_expert_map: |
| 71 | # Second half of the experts is non-local, so valid ids can still map |
| 72 | # to -1 -- exercising the second masking condition. |
| 73 | expert_map = torch.full((NUM_EXPERTS,), -1, dtype=torch.int32, device=device) |
| 74 | local = NUM_EXPERTS // 2 |
| 75 | expert_map[:local] = torch.arange(local, dtype=torch.int32, device=device) |
| 76 | inputs[(topk_ids >= 0) & (expert_map[topk_ids.clamp(min=0)] < 0)] = float("nan") |
| 77 | |
| 78 | # The _reference implementation zeros out outputs. We need to match that here, |
| 79 | # otherwise the test fails. |
| 80 | outputs = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) |
| 81 | out = moe_fused_mul_sum( |
| 82 | inputs=inputs, |
| 83 | topk_weights=topk_weights, |
| 84 | outputs=outputs, |
| 85 | topk_ids=topk_ids, |
| 86 | expert_map=expert_map, |
| 87 | ) |
| 88 | |
| 89 | assert not out.isnan().any(), "invalid slots leaked into the reduction" |
| 90 | ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map) |
| 91 | torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2) |
| 92 | |
| 93 | |
| 94 | def _local_expert_map(device: str) -> torch.Tensor: |
| 95 | """First half of the experts is local, second half maps to -1 (non-local).""" |
| 96 | expert_map = torch.full((NUM_EXPERTS,), -1, dtype=torch.int32, device=device) |
| 97 | local = NUM_EXPERTS // 2 |
| 98 | expert_map[:local] = torch.arange(local, dtype=torch.int32, device=device) |
| 99 | return expert_map |
| 100 | |
| 101 | |
| 102 | SENTINEL = 42.0 # exact in bf16; marks output the kernel must leave untouched |
| 103 | |
| 104 | |
| 105 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 106 | def test_nonlocal_row_is_zeroed(): |
| 107 | """A real row whose positive ids all map to non-local experts is written 0. |
| 108 | |
| 109 | Such a row is not padding (its ids are >= 0, so row_valid > 0 and the block |
| 110 | is not early-returned). Every slot is then masked out by expert_map, so the |
| 111 | accumulator stays 0 and 0 is stored -- the row must not be left holding stale |
| 112 | buffer data, because downstream reduction treats it as a real token. |
| 113 | """ |
| 114 | set_random_seed(0) |
| 115 | device = "cuda" |
| 116 | num_tokens, top_k = 17, 8 |
| 117 | local = NUM_EXPERTS // 2 |
| 118 | expert_map = _local_expert_map(device) |
| 119 | |
| 120 | # Poison inputs so any failure to mask the non-local slots surfaces as NaN. |
| 121 | inputs = torch.full( |
| 122 | (num_tokens, top_k, HIDDEN_SIZE), |
| 123 | float("nan"), |
| 124 | dtype=torch.bfloat16, |
| 125 | device=device, |
| 126 | ) |
| 127 | topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device) |
| 128 | # Every id lands in the non-local half [local, NUM_EXPERTS). |
| 129 | topk_ids = torch.randint( |
| 130 | local, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device |
| 131 | ) |
| 132 | |
| 133 | outputs = torch.full( |
| 134 | (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device |
| 135 | ) |
| 136 | out = moe_fused_mul_sum( |
| 137 | inputs=inputs, |
| 138 | topk_weights=topk_weights, |
| 139 | outputs=outputs, |
| 140 | topk_ids=topk_ids, |
| 141 | expert_map=expert_map, |
| 142 | ) |
| 143 | |
| 144 | assert not out.isnan().any(), "non-local row leaked NaN into the output" |
| 145 | torch.testing.assert_close(out.float(), torch.zeros_like(out.float())) |
| 146 | |
| 147 | |
| 148 | @pytest.mark.parametrize("use_expert_map", [True, False]) |
| 149 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 150 | def test_padding_row_left_untouched(use_expert_map: bool): |
| 151 | """An all -1 padding row is skipped: its output row is left untouched. |
| 152 | |
| 153 | This is the CUDA-graph decode contract -- padding rows past num_recv are |
| 154 | never read downstream, so the kernel elides them (whole-padding blocks |
| 155 | early-return) rather than zeroing. The skip is gated on topk_ids alone, so |
| 156 | it must hold whether or not expert_map is passed (the humming path passes |
| 157 | None); we pin both. |
| 158 | """ |
| 159 | set_random_seed(0) |
| 160 | device = "cuda" |
| 161 | num_tokens, top_k = 17, 8 |
| 162 | expert_map = _local_expert_map(device) if use_expert_map else None |
| 163 | |
| 164 | inputs = torch.full( |
| 165 | (num_tokens, top_k, HIDDEN_SIZE), |
| 166 | float("nan"), |
| 167 | dtype=torch.bfloat16, |
| 168 | device=device, |
| 169 | ) |
| 170 | topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device) |
| 171 | topk_ids = torch.full((num_tokens, top_k), -1, dtype=torch.int32, device=device) |
| 172 | |
| 173 | outputs = torch.full( |
| 174 | (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device |
| 175 | ) |
| 176 | out = moe_fused_mul_sum( |
| 177 | inputs=inputs, |
| 178 | topk_weights=topk_weights, |
| 179 | outputs=outputs, |
| 180 | topk_ids=topk_ids, |
| 181 | expert_map=expert_map, |
| 182 | ) |
| 183 | |
| 184 | assert torch.all(out == SENTINEL), "padding row output was modified" |
| 185 | |
| 186 | |
| 187 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 188 | def test_nonlocal_and_padding_adjacent(): |
| 189 | """A non-local real row next to an all -1 padding row. |
| 190 | |
| 191 | Each CTA owns one token row: the non-local row must be zeroed (its slots are |
| 192 | all masked by expert_map) while the padding row beside it is left untouched. |
| 193 | Guards against a change that lets one row's fate leak into its neighbour. |
| 194 | """ |
| 195 | set_random_seed(0) |
| 196 | device = "cuda" |
| 197 | num_tokens, top_k = 17, 8 |
| 198 | local = NUM_EXPERTS // 2 |
| 199 | expert_map = _local_expert_map(device) |
| 200 | |
| 201 | inputs = torch.full( |
| 202 | (num_tokens, top_k, HIDDEN_SIZE), |
| 203 | float("nan"), |
| 204 | dtype=torch.bfloat16, |
| 205 | device=device, |
| 206 | ) |
| 207 | topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device) |
| 208 | # All rows are non-local real rows except row 1, an all -1 padding row. |
| 209 | topk_ids = torch.randint( |
| 210 | local, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device |
| 211 | ) |
| 212 | topk_ids[1] = -1 |
| 213 | |
| 214 | outputs = torch.full( |
| 215 | (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device |
| 216 | ) |
| 217 | out = moe_fused_mul_sum( |
| 218 | inputs=inputs, |
| 219 | topk_weights=topk_weights, |
| 220 | outputs=outputs, |
| 221 | topk_ids=topk_ids, |
| 222 | expert_map=expert_map, |
| 223 | ) |
| 224 | |
| 225 | padding = torch.zeros(num_tokens, dtype=torch.bool, device=device) |
| 226 | padding[1] = True |
| 227 | assert torch.all(out[1] == SENTINEL), "padding row beside a real row was modified" |
| 228 | written = out[~padding].float() |
| 229 | assert not written.isnan().any(), "non-local row leaked NaN into the output" |
| 230 | torch.testing.assert_close(written, torch.zeros_like(written)) |
| 231 | |
| 232 | |
| 233 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 234 | @pytest.mark.parametrize("top_k", TOP_KS) |
| 235 | @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) |
| 236 | @pytest.mark.parametrize("dtype", DTYPES) |
| 237 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 238 | def test_weighted_sum_matches_reference( |
| 239 | num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype |
| 240 | ): |
| 241 | """Weighted sum over live slots matches the reference across hidden tiles. |
| 242 | |
| 243 | Covers the deployed humming path (globalized ids, expert_map=None): a random |
| 244 | -1 mix with NaN-poisoned dead slots, over hidden sizes that span a single |
| 245 | tile, a partial trailing tile, and many tiles, in both bf16 and fp32. |
| 246 | """ |
| 247 | set_random_seed(0) |
| 248 | device = "cuda" |
| 249 | |
| 250 | inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) |
| 251 | topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) |
| 252 | topk_ids = torch.randint( |
| 253 | 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device |
| 254 | ) |
| 255 | invalid = torch.rand(num_tokens, top_k, device=device) < 0.5 |
| 256 | topk_ids[invalid] = -1 |
| 257 | inputs[invalid] = float("nan") |
| 258 | |
| 259 | # All-padding rows early-return untouched, so pre-zero to match the reference. |
| 260 | outputs = torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device) |
| 261 | out = moe_fused_mul_sum( |
| 262 | inputs=inputs, |
| 263 | topk_weights=topk_weights, |
| 264 | outputs=outputs, |
| 265 | topk_ids=topk_ids, |
| 266 | expert_map=None, |
| 267 | ) |
| 268 | |
| 269 | assert not out.isnan().any(), "invalid slots leaked into the reduction" |
| 270 | ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None) |
| 271 | atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) |
| 272 | torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol) |
| 273 | |
| 274 | |
| 275 | @pytest.mark.parametrize("dtype", DTYPES) |
| 276 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 277 | def test_no_topk_ids_sums_all_slots(dtype: torch.dtype): |
| 278 | """With topk_ids=None every slot is live: plain weighted sum over top_k.""" |
| 279 | set_random_seed(0) |
| 280 | device = "cuda" |
| 281 | num_tokens, top_k, hidden_size = 17, 8, 1000 |
| 282 | |
| 283 | inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) |
| 284 | topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) |
| 285 | out = moe_fused_mul_sum(inputs=inputs, topk_weights=topk_weights) |
| 286 | |
| 287 | ref = (inputs.float() * topk_weights.float().unsqueeze(-1)).sum(dim=1) |
| 288 | atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) |
| 289 | torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol) |
| 290 | |
| 291 | |
| 292 | @pytest.mark.parametrize("hidden_size", [512, 7168]) |
| 293 | @pytest.mark.parametrize("dtype", DTYPES) |
| 294 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 295 | def test_expert_map_is_required_for_nonlocal_ids(hidden_size: int, dtype: torch.dtype): |
| 296 | """expert_map is load-bearing, not redundant: dropping it leaks non-local ids. |
| 297 | |
| 298 | A positive id whose expert is non-local marks a slot the GEMM never wrote, so |
| 299 | its `inputs` row is stale. `topk_ids >= 0` alone cannot tell it apart from a |
| 300 | real slot -- only expert_map can. This pins the contract the humming path |
| 301 | relies on (globalized ids are still routed through expert_map as a safety |
| 302 | net): with the map the stale slot is masked; without it, it corrupts the sum. |
| 303 | Regression guard against re-dropping expert_map as an optimization. |
| 304 | """ |
| 305 | set_random_seed(0) |
| 306 | device = "cuda" |
| 307 | num_tokens, top_k = 8, 8 |
| 308 | local = NUM_EXPERTS // 2 |
| 309 | expert_map = _local_expert_map(device) |
| 310 | |
| 311 | inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) |
| 312 | topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) |
| 313 | # Each row keeps at least one local slot (so the mapped output is finite) and |
| 314 | # at least one non-local slot whose stale row is poisoned with NaN. |
| 315 | half = top_k // 2 |
| 316 | topk_ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=device) |
| 317 | topk_ids[:, :half] = torch.randint( |
| 318 | 0, local, (num_tokens, half), dtype=torch.int32, device=device |
| 319 | ) |
| 320 | topk_ids[:, half:] = torch.randint( |
| 321 | local, NUM_EXPERTS, (num_tokens, top_k - half), dtype=torch.int32, device=device |
| 322 | ) |
| 323 | inputs[:, half:] = float("nan") |
| 324 | |
| 325 | atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) |
| 326 | |
| 327 | mapped = moe_fused_mul_sum( |
| 328 | inputs=inputs, |
| 329 | topk_weights=topk_weights, |
| 330 | outputs=torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device), |
| 331 | topk_ids=topk_ids, |
| 332 | expert_map=expert_map, |
| 333 | ) |
| 334 | assert not mapped.isnan().any(), "expert_map failed to mask stale non-local slot" |
| 335 | ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map) |
| 336 | torch.testing.assert_close(mapped.float(), ref, atol=atol, rtol=rtol) |
| 337 | |
| 338 | # Same call without the map: the non-local slot is summed and the NaN leaks. |
| 339 | unmapped = moe_fused_mul_sum( |
| 340 | inputs=inputs, |
| 341 | topk_weights=topk_weights, |
| 342 | outputs=torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device), |
| 343 | topk_ids=topk_ids, |
| 344 | expert_map=None, |
| 345 | ) |
| 346 | assert unmapped.isnan().any(), "expert_map is not redundant: dropping it must leak" |
| 347 | |
| 348 | |
| 349 | @pytest.mark.parametrize("dtype", DTYPES) |
| 350 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 351 | def test_num_valid_tokens_bounds_stale_padding(dtype: torch.dtype): |
| 352 | """Rows past num_valid_tokens are left untouched, even with stale ids/NaN. |
| 353 | |
| 354 | Under cudagraph decode the grid is the static padded token count and the |
| 355 | tail rows past num_recv hold stale (non -1) ids and stale NaN in `inputs` |
| 356 | from a prior replay -- the -1 mask alone cannot spot them. Bounding on |
| 357 | num_valid_tokens must skip the tail before any load so its NaN never reaches |
| 358 | the output; without the bound it is summed and a masked finalize can then |
| 359 | propagate the NaN. Regression guard for the decode accuracy fix. |
| 360 | """ |
| 361 | set_random_seed(0) |
| 362 | device = "cuda" |
| 363 | num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 512 |
| 364 | |
| 365 | inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) |
| 366 | topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) |
| 367 | # All ids valid (>= 0): the stale tail is indistinguishable by the -1 mask. |
| 368 | topk_ids = torch.randint( |
| 369 | 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device |
| 370 | ) |
| 371 | inputs[num_recv:] = float("nan") # stale down_output behind the padding tail |
| 372 | num_valid = torch.tensor([num_recv], dtype=torch.int32, device=device) |
| 373 | |
| 374 | atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) |
| 375 | |
| 376 | bounded = moe_fused_mul_sum( |
| 377 | inputs=inputs, |
| 378 | topk_weights=topk_weights, |
| 379 | outputs=torch.full( |
| 380 | (num_tokens, hidden_size), SENTINEL, dtype=dtype, device=device |
| 381 | ), |
| 382 | topk_ids=topk_ids, |
| 383 | num_valid_tokens=num_valid, |
| 384 | ) |
| 385 | assert torch.all(bounded[num_recv:] == SENTINEL), ( |
| 386 | "padding tail past num_recv written" |
| 387 | ) |
| 388 | assert not bounded[:num_recv].isnan().any(), "real rows leaked NaN" |
| 389 | ref = _reference( |
| 390 | inputs[:num_recv].nan_to_num(), |
| 391 | topk_weights[:num_recv], |
| 392 | topk_ids[:num_recv], |
| 393 | expert_map=None, |
| 394 | ) |
| 395 | torch.testing.assert_close(bounded[:num_recv].float(), ref, atol=atol, rtol=rtol) |
| 396 | |
| 397 | # Without the bound the stale tail is summed and leaks NaN into its rows. |
| 398 | unbounded = moe_fused_mul_sum( |
| 399 | inputs=inputs, |
| 400 | topk_weights=topk_weights, |
| 401 | outputs=torch.full( |
| 402 | (num_tokens, hidden_size), SENTINEL, dtype=dtype, device=device |
| 403 | ), |
| 404 | topk_ids=topk_ids, |
| 405 | ) |
| 406 | assert unbounded[num_recv:].isnan().any(), "num_valid_tokens bound is load-bearing" |