[ROCm][MoE] Pad the AITER MoE intermediate size at allocation time, and round the expert-group count to a kernel that exists (#55368)
Signed-off-by: Sergei Shliapnikov <sergei.shliapnikov@amd.com>
Sergey Shlyapnikov committed
Sep 30, 2026 at 19:08 UTC
4e1182a3a609a5165afd6de99b1269099b5d5999
4 files changed
+452
-15
tests/kernels/moe/test_rocm_aiter_moe.py
+250
-1
@@ -30,7 +30,7 @@ import torch.nn.functional as F
30
31
from tests.kernels.utils import _assert_deterministic
32
from vllm.platforms import current_platform
33
-from vllm.platforms.rocm import on_gfx942, on_gfx950
33
+from vllm.platforms.rocm import on_cdna, on_gfx942, on_gfx950
34
from vllm.utils.torch_utils import set_random_seed
35
36
pytestmark = pytest.mark.skipif(
@@ -1317,3 +1317,252 @@ def test_aiter_fused_moe_mi3xx_fp8_accuracy():
1317
pass_rate=1.0,
1318
max_violation_factor=1.5,
1319
)
1320
+
1321
+
1322
+# Weight alignment tests --------------------------------------------------
1323
+#
1324
+# AITER's CK 2stages MoE kernel rejects an intermediate size not divisible by
1325
+# its tile width (64 at or below inter_dim 192, 128 above). Some model + TP
1326
+# splits land on an unaligned size (e.g. 1792 / TP=8 = 224), so the AITER path
1327
+# rounds the intermediate dim up in ``maybe_roundup_sizes`` and allocates the
1328
+# weights at the padded size.
1329
+
1330
+ALIGNMENT_HIDDEN = 64
1331
+ALIGNMENT_NUM_EXPERTS = 2
1332
+
1333
+
1334
+def _make_alignment_moe_config(intermediate: int):
1335
+ from tests.kernels.moe.utils import make_dummy_moe_config
1336
+
1337
+ return make_dummy_moe_config(
1338
+ num_experts=ALIGNMENT_NUM_EXPERTS,
1339
+ hidden_dim=ALIGNMENT_HIDDEN,
1340
+ intermediate_size=intermediate,
1341
+ )
1342
+
1343
+
1344
+def _make_aiter_method(moe_config):
1345
+ """Build the unquantized method with the backend pinned to AITER.
1346
+
1347
+ ``select_unquantized_moe_backend`` needs a real ROCm + AITER runtime, so
1348
+ stub it out and set the backend directly.
1349
+ """
1350
+ from unittest.mock import patch
1351
+
1352
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1353
+ UnquantizedMoeBackend,
1354
+ )
1355
+ from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
1356
+ UnquantizedFusedMoEMethod,
1357
+ )
1358
+
1359
+ with patch(
1360
+ "vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method"
1361
+ ".select_unquantized_moe_backend",
1362
+ return_value=(UnquantizedMoeBackend.AITER, None),
1363
+ ):
1364
+ return UnquantizedFusedMoEMethod(moe_config)
1365
+
1366
+
1367
+def _roundup(method, moe_config, intermediate):
1368
+ from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
1369
+
1370
+ return method.maybe_roundup_sizes(
1371
+ hidden_size=ALIGNMENT_HIDDEN,
1372
+ intermediate_size_per_partition=intermediate,
1373
+ act_dtype=moe_config.in_dtype,
1374
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1375
+ )
1376
+
1377
+
1378
+@pytest.mark.parametrize(
1379
+ ("intermediate", "expected_padded"),
1380
+ [
1381
+ (64, 64), # already aligned, untouched
1382
+ (96, 128),
1383
+ (160, 192),
1384
+ (192, 192), # must stay 192: it is valid and has tuned configs
1385
+ (224, 256), # K2-Horizon-375B at TP=8
1386
+ (256, 256),
1387
+ (448, 512), # K2-Horizon-375B at TP=4
1388
+ (4096, 4096),
1389
+ ],
1390
+)
1391
+def test_aiter_moe_roundup_pads_intermediate(
1392
+ intermediate,
1393
+ expected_padded,
1394
+ default_vllm_config,
1395
+):
1396
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1397
+ aiter_moe_intermediate_alignment,
1398
+ )
1399
+
1400
+ moe_config = _make_alignment_moe_config(intermediate)
1401
+ method = _make_aiter_method(moe_config)
1402
+
1403
+ hidden, padded = _roundup(method, moe_config, intermediate)
1404
+
1405
+ assert padded == expected_padded
1406
+ assert hidden == ALIGNMENT_HIDDEN
1407
+ # Padding must not move a size across the <= 192 threshold, or the
1408
+ # alignment we picked is not the one the kernel dispatches the padded
1409
+ # shape to.
1410
+ assert aiter_moe_intermediate_alignment(padded) == (
1411
+ aiter_moe_intermediate_alignment(intermediate)
1412
+ )
1413
+
1414
+
1415
+@pytest.mark.parametrize("backend_name", ["TRITON", "FLASHINFER_CUTLASS"])
1416
+def test_aiter_moe_roundup_is_not_applied_to_other_backends(
1417
+ backend_name,
1418
+ default_vllm_config,
1419
+):
1420
+ """Backends with no alignment requirement keep the unaligned size."""
1421
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1422
+ UnquantizedMoeBackend,
1423
+ )
1424
+
1425
+ intermediate = 224
1426
+ moe_config = _make_alignment_moe_config(intermediate)
1427
+ method = _make_aiter_method(moe_config)
1428
+ method.unquantized_backend = UnquantizedMoeBackend[backend_name]
1429
+
1430
+ _, padded = _roundup(method, moe_config, intermediate)
1431
+
1432
+ assert padded == intermediate
1433
+
1434
+
1435
+@pytest.mark.parametrize("intermediate", [224, 448])
1436
+def test_aiter_moe_reload_zeroes_intermediate_padding(
1437
+ intermediate,
1438
+ monkeypatch,
1439
+ default_vllm_config,
1440
+):
1441
+ """A reload must re-zero pad lanes left dirty in the weight storage.
1442
+
1443
+ The reload writes back only the logical slices, so the pad lanes hold
1444
+ garbage (NaN here) that the AITER conversion must zero. Two passes also
1445
+ cover reloading into storage that already holds a previous conversion.
1446
+ """
1447
+ from unittest.mock import MagicMock
1448
+
1449
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1450
+ aiter_moe_intermediate_alignment,
1451
+ )
1452
+
1453
+ alignment = aiter_moe_intermediate_alignment(intermediate)
1454
+ padded = -(-intermediate // alignment) * alignment
1455
+ assert padded != intermediate
1456
+
1457
+ moe_config = _make_alignment_moe_config(intermediate)
1458
+ assert moe_config.is_act_and_mul
1459
+ method = _make_aiter_method(moe_config)
1460
+ # RoutedExperts records the rounded-up size after maybe_roundup_sizes.
1461
+ moe_config.intermediate_size_per_partition = padded
1462
+ layer = torch.nn.Module()
1463
+ layer.moe_config = moe_config
1464
+
1465
+ method.create_weights(
1466
+ layer=layer,
1467
+ num_experts=ALIGNMENT_NUM_EXPERTS,
1468
+ hidden_size=ALIGNMENT_HIDDEN,
1469
+ intermediate_size_per_partition=padded,
1470
+ params_dtype=torch.float32,
1471
+ )
1472
+
1473
+ # Identity shuffle keeps pad lanes sliceable; both stubs need real ROCm.
1474
+ monkeypatch.setattr(
1475
+ "vllm.model_executor.layers.fused_moe.oracle.unquantized"
1476
+ ".rocm_aiter_ops.shuffle_weights",
1477
+ lambda w13, w2: (w13, w2),
1478
+ )
1479
+ monkeypatch.setattr(
1480
+ method,
1481
+ "_init_moe_kernel",
1482
+ lambda _layer: setattr(method, "moe_kernel", MagicMock()),
1483
+ )
1484
+
1485
+ up = padded # up projection starts one padded block into the fused rows
1486
+
1487
+ for _ in range(2):
1488
+ gate = torch.randn(ALIGNMENT_NUM_EXPERTS, intermediate, ALIGNMENT_HIDDEN)
1489
+ up_proj = torch.randn(ALIGNMENT_NUM_EXPERTS, intermediate, ALIGNMENT_HIDDEN)
1490
+ down = torch.randn(ALIGNMENT_NUM_EXPERTS, ALIGNMENT_HIDDEN, intermediate)
1491
+
1492
+ # NaN the pads, then write only the logical slices like the loader.
1493
+ layer.w13_weight.data.fill_(float("nan"))
1494
+ layer.w2_weight.data.fill_(float("nan"))
1495
+ layer.w13_weight.data[:, :intermediate].copy_(gate)
1496
+ layer.w13_weight.data[:, up : up + intermediate].copy_(up_proj)
1497
+ layer.w2_weight.data[:, :, :intermediate].copy_(down)
1498
+
1499
+ method.process_weights_after_loading(layer)
1500
+
1501
+ w13 = layer.w13_weight
1502
+ w2 = layer.w2_weight
1503
+ assert torch.equal(w13[:, :intermediate], gate)
1504
+ assert torch.equal(w13[:, up : up + intermediate], up_proj)
1505
+ assert torch.equal(w2[:, :, :intermediate], down)
1506
+ assert torch.all(w13[:, intermediate:up] == 0)
1507
+ assert torch.all(w13[:, up + intermediate :] == 0)
1508
+ assert torch.all(w2[:, :, intermediate:] == 0)
1509
+
1510
+
1511
+def _aiter_accepts_intermediate(intermediate: int, num_tokens: int) -> bool:
1512
+ """Run the real AITER CK MoE kernel and report whether it dispatched."""
1513
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1514
+ ActivationMethod,
1515
+ QuantMethod,
1516
+ )
1517
+
1518
+ case = _make_moe_case(
1519
+ num_tokens=num_tokens,
1520
+ hidden_dim=1024,
1521
+ intermediate_dim=intermediate,
1522
+ num_experts=8,
1523
+ topk=2,
1524
+ seed=0,
1525
+ )
1526
+ try:
1527
+ _run_fused_moe(
1528
+ case["hidden_states"],
1529
+ case["w1"],
1530
+ case["w2"],
1531
+ case["topk_weights"],
1532
+ case["topk_ids"],
1533
+ activation_method=int(ActivationMethod.SILU),
1534
+ quant_method=int(QuantMethod.NO),
1535
+ )
1536
+ except RuntimeError:
1537
+ return False
1538
+ return True
1539
+
1540
+
1541
+@pytest.mark.skipif(
1542
+ not on_cdna(),
1543
+ reason="CDNA ROCm only",
1544
+)
1545
+@pytest.mark.parametrize("intermediate", [192, 224, 256])
1546
+# With topk=2 and 8 experts, these token counts select block_m 32 / 64 / 128
1547
+# respectively (AITER's get_block_size_M, on both 256- and 304-CU parts).
1548
+@pytest.mark.parametrize("num_tokens", [1024, 2560, 5120])
1549
+def test_aiter_moe_alignment_rule_holds_across_block_m(intermediate, num_tokens):
1550
+ """The alignment rule must hold at every reachable ``block_m``.
1551
+
1552
+ Stage 1 checks ``NPerBlock`` and stage 2 ``KPerBlock``; both vary with
1553
+ ``block_m``, so a size validated at one tile size says nothing about the
1554
+ others.
1555
+ """
1556
+ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
1557
+ aiter_moe_intermediate_alignment,
1558
+ )
1559
+
1560
+ _assert_aiter_supported()
1561
+
1562
+ alignment = aiter_moe_intermediate_alignment(intermediate)
1563
+ padded = -(-intermediate // alignment) * alignment
1564
+
1565
+ assert _aiter_accepts_intermediate(intermediate, num_tokens) == (
1566
+ intermediate % alignment == 0
1567
+ )
1568
+ assert _aiter_accepts_intermediate(padded, num_tokens)
tests/kernels/moe/test_rocm_aiter_num_expert_group.py
new
+136
@@ -0,0 +1,136 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the AITER biased_grouped_topk expert-group count.
4
+
5
+The AITER kernel is only instantiated for ``NUM_GRP`` in {1, 2, 4, 8}, so the
6
+group count derived from ``num_experts`` has to be rounded to a supported
7
+value that still divides ``num_experts``. Grouping is a no-op on this path
8
+(``topk_group == num_expert_group``), so any such value routes identically;
9
+the constraint is purely about which kernel exists.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm._aiter_ops import rocm_aiter_ops
16
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
17
+ AITER_MAX_EXPERTS_PER_GROUP as MAX_EXPERTS_PER_GROUP,
18
+)
19
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
20
+ AITER_SUPPORTED_NUM_GRP as SUPPORTED_NUM_GRP,
21
+)
22
+from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23
+ _aiter_can_use_biased_grouped_topk,
24
+ _aiter_get_num_expert_group,
25
+)
26
+from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
27
+ grouped_topk,
28
+)
29
+from vllm.platforms import current_platform
30
+from vllm.platforms.rocm import on_cdna
31
+
32
+pytestmark = pytest.mark.skipif(
33
+ not current_platform.is_rocm(), reason="ROCm-specific tests"
34
+)
35
+
36
+
37
+@pytest.mark.parametrize(
38
+ ("num_experts", "expected"),
39
+ [
40
+ (8, 1), # fits in a single group
41
+ (32, 1), # exactly at the per-group limit
42
+ (64, 2),
43
+ (128, 4),
44
+ (256, 8),
45
+ (72, 8), # naive ceil gives 3 -> unsupported, rounds up to 8
46
+ (96, 8), # naive ceil gives 3 -> unsupported, rounds up to 8
47
+ (160, 8), # naive ceil gives 5 -> unsupported, rounds up to 8
48
+ (192, 8), # naive ceil gives 6 -> unsupported, rounds up to 8
49
+ # No supported NUM_GRP divides these within the group-size limit, so
50
+ # the naive value is kept and the call site declines the kernel.
51
+ (33, 3),
52
+ (129, 43),
53
+ (257, 257),
54
+ (320, 10),
55
+ (384, 12),
56
+ ],
57
+)
58
+def test_group_count_and_guard(num_experts, expected):
59
+ g = _aiter_get_num_expert_group(num_experts)
60
+ assert g == expected
61
+ # The router asserts both unconditionally; rounding must never break them.
62
+ assert num_experts % g == 0
63
+ assert num_experts // g <= MAX_EXPERTS_PER_GROUP
64
+
65
+ supported = g in SUPPORTED_NUM_GRP
66
+ # Ample topk, so only the membership check can decline.
67
+ assert _aiter_can_use_biased_grouped_topk(num_experts, topk=64) == supported
68
+ if supported:
69
+ assert _aiter_can_use_biased_grouped_topk(num_experts, topk=g)
70
+ assert not _aiter_can_use_biased_grouped_topk(num_experts, topk=g - 1)
71
+
72
+
73
+@pytest.mark.skipif(not on_cdna(), reason="CDNA ROCm only")
74
+@pytest.mark.parametrize(
75
+ "num_experts",
76
+ [
77
+ 64, # g == 2, not rounded -- covers the no-op claim itself
78
+ 128, # g == 4, not rounded
79
+ 256, # g == 8, not rounded
80
+ 96, # naive 3 -> rounded to 8
81
+ 192, # naive 6 -> rounded to 8
82
+ ],
83
+)
84
+def test_rounded_group_count_routes_like_the_reference(num_experts):
85
+ """The rounded group count must route identically to ``grouped_topk``.
86
+
87
+ Gating is continuous random, so exact score ties -- the one place the two
88
+ may legitimately differ -- have measure zero.
89
+ """
90
+ torch.manual_seed(num_experts)
91
+ device = "cuda"
92
+ topk = 8
93
+ num_tokens = 83 # not a multiple of any warp/tile size
94
+ g = _aiter_get_num_expert_group(num_experts)
95
+ assert _aiter_can_use_biased_grouped_topk(num_experts, topk)
96
+
97
+ gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device)
98
+ bias = torch.randn(num_experts, dtype=torch.float32, device=device)
99
+
100
+ ref_weights, ref_ids = grouped_topk(
101
+ hidden_states=torch.empty(num_tokens, 1, device=device),
102
+ gating_output=gating,
103
+ topk=topk,
104
+ renormalize=True,
105
+ num_expert_group=g,
106
+ topk_group=g,
107
+ scoring_func="sigmoid",
108
+ e_score_correction_bias=bias,
109
+ )
110
+
111
+ weights = torch.empty(num_tokens, topk, dtype=torch.float32, device=device)
112
+ ids = torch.empty(num_tokens, topk, dtype=torch.int32, device=device)
113
+ rocm_aiter_ops.biased_grouped_topk(
114
+ gating,
115
+ bias,
116
+ weights,
117
+ ids,
118
+ num_expert_group=g,
119
+ topk_group=g,
120
+ need_renorm=True,
121
+ )
122
+
123
+ # Neither side promises an order within a token's top-k, so sort by expert
124
+ # id and carry the weights through the same permutation -- sorting the two
125
+ # independently would not catch a weight attached to the wrong expert.
126
+ order = ids.argsort(dim=-1)
127
+ ref_order = ref_ids.argsort(dim=-1)
128
+ torch.testing.assert_close(
129
+ ids.gather(1, order), ref_ids.to(torch.int32).gather(1, ref_order)
130
+ )
131
+ torch.testing.assert_close(
132
+ weights.gather(1, order),
133
+ ref_weights.gather(1, ref_order),
134
+ atol=1e-4,
135
+ rtol=1e-4,
136
+ )
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+33
-9
@@ -371,6 +371,16 @@ def select_unquantized_moe_backend(
371
)
372
373
374
+def aiter_moe_intermediate_alignment(intermediate: int) -> int:
375
+ """Intermediate-size alignment required by AITER's CK 2stages MoE kernel.
376
+
377
+ AITER dispatches on ``inter_dim <= 192``: below the threshold both stages
378
+ use 64-wide tiles, above it at least one stage uses a 128-wide tile, and
379
+ CK's ``IsSupportedArgument`` rejects a size not divisible by that width.
380
+ """
381
+ return 64 if intermediate <= 192 else 128
382
+
383
+
384
def unquantized_round_up_hidden_size_and_intermediate_size(
385
backend: UnquantizedMoeBackend,
386
hidden_size: int,
@@ -379,9 +389,30 @@ def unquantized_round_up_hidden_size_and_intermediate_size(
389
"""Round up dimensions before allocation to satisfy the selected kernel."""
390
if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
391
intermediate_size = round_up(intermediate_size, 128)
392
+ elif backend == UnquantizedMoeBackend.AITER:
393
+ intermediate_size = round_up(
394
+ intermediate_size, aiter_moe_intermediate_alignment(intermediate_size)
395
+ )
396
return hidden_size, intermediate_size
397
398
399
+def _zero_intermediate_padding(
400
+ moe_config: FusedMoEConfig,
401
+ w13_weight: torch.Tensor,
402
+ w2_weight: torch.Tensor,
403
+) -> None:
404
+ """Zero intermediate padding; reloads rewrite only the checkpoint slices."""
405
+ unpadded = moe_config.intermediate_size_per_partition_unpadded
406
+ assert unpadded is not None
407
+ padded = moe_config.intermediate_size_per_partition
408
+ if padded <= unpadded:
409
+ return
410
+ w13_weight[:, unpadded:padded].zero_()
411
+ if moe_config.is_act_and_mul:
412
+ w13_weight[:, padded + unpadded :].zero_()
413
+ w2_weight[:, :, unpadded:].zero_()
414
+
415
+
416
def convert_to_unquantized_kernel_format(
417
unquantized_backend: UnquantizedMoeBackend,
418
moe_config: FusedMoEConfig,
@@ -424,6 +455,7 @@ def convert_to_unquantized_kernel_format(
455
return layout.full_gate_weight, layout.full_down_weight
456
457
if unquantized_backend == UnquantizedMoeBackend.AITER:
458
+ _zero_intermediate_padding(moe_config, w13_weight, w2_weight)
459
w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
460
w13_weight.is_shuffled = True
461
w2_weight.is_shuffled = True
@@ -445,15 +477,7 @@ def convert_to_unquantized_kernel_format(
477
)
478
moe_config.intermediate_size_per_partition = padded_intermediate
479
448
- # Reloads only overwrite checkpoint slices. An earlier in-place
449
- # permutation can leave nonzero values in the raw padding slots.
450
- unpadded = moe_config.intermediate_size_per_partition_unpadded
451
- assert unpadded is not None
452
- if padded_intermediate > unpadded:
453
- w13_weight[:, unpadded:padded_intermediate].zero_()
454
- if is_act_and_mul:
455
- w13_weight[:, padded_intermediate + unpadded :].zero_()
456
- w2_weight[:, :, unpadded:].zero_()
480
+ _zero_intermediate_padding(moe_config, w13_weight, w2_weight)
481
482
_cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
483
convert_moe_weights_to_flashinfer_trtllm_block_layout(
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+33
-5
@@ -19,6 +19,11 @@ from vllm.model_executor.layers.fused_moe.router.dsv4_topk import (
19
dsv4_topk,
20
)
21
22
+# AITER instantiates biased_grouped_topk only for these NUM_GRP values, with at
23
+# most this many experts per group.
24
+AITER_SUPPORTED_NUM_GRP = (1, 2, 4, 8)
25
+AITER_MAX_EXPERTS_PER_GROUP = 32
26
+
27
28
def _get_padding_mask(num_tokens: int) -> torch.Tensor | None:
29
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
@@ -104,17 +109,40 @@ def vllm_topk_softplus_sqrt(
109
110
@functools.lru_cache(maxsize=8)
111
def _aiter_get_num_expert_group(num_experts: int) -> int:
107
- _AITER_MAX_EXPERTS_PER_GROUP = 32
108
- g = max(1, -(-num_experts // _AITER_MAX_EXPERTS_PER_GROUP))
112
+ g = max(1, -(-num_experts // AITER_MAX_EXPERTS_PER_GROUP))
113
while num_experts % g != 0:
114
g += 1
115
assert num_experts % g == 0, f"{num_experts=} not divisible by {g=}"
112
- assert num_experts // g <= _AITER_MAX_EXPERTS_PER_GROUP, (
113
- f"group size {num_experts // g} exceeds limit {_AITER_MAX_EXPERTS_PER_GROUP}"
116
+ # Grouping is a no-op here (topk_group == num_expert_group), so any divisor
117
+ # within the limit routes identically and we are free to pick one AITER has
118
+ # a kernel for. If none fits, the unsupported value is returned as-is --
119
+ # callers must gate on _aiter_can_use_biased_grouped_topk.
120
+ if g not in AITER_SUPPORTED_NUM_GRP:
121
+ g = next(
122
+ (
123
+ c
124
+ for c in sorted(AITER_SUPPORTED_NUM_GRP, reverse=True)
125
+ if num_experts % c == 0
126
+ and num_experts // c <= AITER_MAX_EXPERTS_PER_GROUP
127
+ ),
128
+ g,
129
+ )
130
+ assert num_experts // g <= AITER_MAX_EXPERTS_PER_GROUP, (
131
+ f"group size {num_experts // g} exceeds limit {AITER_MAX_EXPERTS_PER_GROUP}"
132
)
133
return g
134
135
136
+def _aiter_can_use_biased_grouped_topk(num_experts: int, topk: int) -> bool:
137
+ """Whether AITER has an instantiated kernel for this routing shape.
138
+
139
+ Both checks are needed: an unsupported group count can still be small
140
+ enough to satisfy ``topk >= g`` (num_experts=33 gives g=3).
141
+ """
142
+ g = _aiter_get_num_expert_group(num_experts)
143
+ return topk >= g and g in AITER_SUPPORTED_NUM_GRP
144
+
145
+
146
def fused_topk_bias(
147
hidden_states: torch.Tensor,
148
gating_output: torch.Tensor,
@@ -229,7 +257,7 @@ def fused_topk_bias(
257
M = hidden_states.size(0)
258
num_experts = gating_output.shape[-1]
259
num_expert_group = _aiter_get_num_expert_group(num_experts)
232
- if topk >= num_expert_group:
260
+ if _aiter_can_use_biased_grouped_topk(num_experts, topk):
261
topk_weights = torch.empty(
262
M, topk, dtype=torch.float32, device=hidden_states.device
263
)