A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 12:15 UTC
[Kernels][MoE] Fix legacy_routing to use bitmatrix-based routing path (#38504)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Apr 6, 2026 at 21:57 UTC
2df2c85be494d1b98f8a66d0acb7c65f0b119a37
7 files changed
+84
-216
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml
+1
-1
@@ -3,4 +3,4 @@
3
model_name: openai/gpt-oss-20b
4
metric_threshold: 0.568
5
reasoning_effort: low
6
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN"
\ No newline at end of file
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2"
\ No newline at end of file
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml
+2
-2
@@ -3,6 +3,6 @@
3
model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
4
metric_threshold: 0.568
5
reasoning_effort: low
6
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter"
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2"
7
env:
8
- VLLM_ROCM_USE_AITER: "1"
8
+ VLLM_ROCM_USE_AITER: "1"
\ No newline at end of file
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml
+1
-1
@@ -3,4 +3,4 @@
3
model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
4
metric_threshold: 0.568
5
reasoning_effort: low
6
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton"
\ No newline at end of file
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2"
\ No newline at end of file
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml
+1
-1
@@ -3,6 +3,6 @@
3
model_name: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8
4
metric_threshold: 0.568
5
reasoning_effort: low
6
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN"
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2"
7
env:
8
VLLM_ROCM_USE_AITER: "1"
\ No newline at end of file
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+12
-57
@@ -23,16 +23,12 @@ from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_m
23
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
24
from triton_kernels.tensor_details import layout
25
from triton_kernels.testing import assert_close
26
-from triton_kernels.topk import topk as topk_fn
26
27
from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
28
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
30
- legacy_routing,
31
- make_routing_data,
29
triton_kernel_moe_forward,
30
)
31
from vllm.utils.math_utils import round_up
35
-from vllm.utils.torch_utils import set_random_seed
32
33
from .utils import shuffle_weight
34
@@ -97,10 +93,18 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
93
if w_dtype != "mx4":
94
pytest.skip("NYI")
95
else: # quantize to mx4
100
- # careful on the padding here, the activation padding need to be
101
- # multiple of 64, the actual engine is not implemented
102
- w1_bottom_pad = round_up(w1_tri.shape[1], 64) - w1_tri.shape[1]
103
- w1_right_pad = round_up(w1_tri.shape[2], 128) - w1_tri.shape[2]
96
+ # Padding alignment depends on the platform. On CDNA4 the scale
97
+ # swizzle requires SCALE_K % 8 == 0 (K % 256) and
98
+ # SCALE_N % 32 == 0 (2*N % 512), matching the production
99
+ # alignment in mxfp4_round_up_hidden_size_and_intermediate_size.
100
+ # On CUDA (Hopper) the scale layout pads internally, so the
101
+ # original 64/128 alignment is sufficient.
102
+ if current_platform.is_rocm():
103
+ k_align, n2_align = 256, 512
104
+ else:
105
+ k_align, n2_align = 64, 128
106
+ w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1]
107
+ w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2]
108
109
w2_bottom_pad = w1_right_pad // 2
110
w2_right_pad = w1_bottom_pad
@@ -367,52 +371,3 @@ def test_unit_shuffle():
371
)
372
373
assert_close(ref=out_ref, tri=out)
370
-
371
-
372
-@pytest.mark.parametrize("num_tokens", [2, 8, 64])
373
-@pytest.mark.parametrize("num_experts", [32, 128])
374
-@pytest.mark.parametrize("topk", [1, 4])
375
-@pytest.mark.parametrize("renormalize", [True, False])
376
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
377
-def test_legacy_routing(
378
- num_tokens: int, num_experts: int, topk: int, renormalize: bool, dtype: torch.dtype
379
-):
380
- set_random_seed(0)
381
- gating_output = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
382
-
383
- sm_first = not renormalize
384
- logits = gating_output
385
- if sm_first:
386
- logits = torch.softmax(logits, dim=-1)
387
- topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
388
- # topk_fn returns SparseMatrix on NVIDIA, plain tuple on ROCm.
389
- if isinstance(topk_result, tuple):
390
- topk_weights, topk_ids_raw, bitmatrix = topk_result
391
- from triton_kernels.routing import routing_from_bitmatrix
392
-
393
- routing_data_ref, gather_indx_ref, scatter_indx_ref = routing_from_bitmatrix(
394
- bitmatrix, topk_weights, topk_ids_raw, num_experts, topk
395
- )
396
- else:
397
- topk_ids = topk_result.indx.to(torch.long)
398
- topk_weights = topk_result.vals
399
- routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data(
400
- topk_ids, topk_weights, num_experts
401
- )
402
-
403
- routing_data, gather_indx, scatter_indx = legacy_routing(
404
- gating_output, topk, sm_first=sm_first
405
- )
406
-
407
- assert_close(
408
- ref=gather_indx_ref.src_indx, tri=gather_indx.src_indx, maxtol=0, rmstol=0
409
- )
410
- assert_close(
411
- ref=gather_indx_ref.dst_indx, tri=gather_indx.dst_indx, maxtol=0, rmstol=0
412
- )
413
- assert_close(
414
- ref=scatter_indx_ref.src_indx, tri=scatter_indx.src_indx, maxtol=0, rmstol=0
415
- )
416
- assert_close(
417
- ref=scatter_indx_ref.dst_indx, tri=scatter_indx.dst_indx, maxtol=0, rmstol=0
418
- )
tests/kernels/quantization/test_mxfp4_triton_ep.py
+18
-41
@@ -4,12 +4,9 @@
4
Tests that triton_kernel_moe_forward correctly applies expert_map
5
remapping when expert parallelism (EP) is enabled.
6
7
-Previously, legacy_routing was always used and it produced routing data
8
-with global expert IDs that didn't correspond to local weight indices,
9
-causing illegal memory access with EP. The fix splits routing: when
10
-expert_map is provided, topk selection is performed first, expert_map is
11
-applied to remap global→local IDs, and make_routing_data builds routing
12
-structures from the local IDs.
7
+Both EP and non-EP paths use topk + make_routing_data. When expert_map
8
+is provided, global expert IDs are remapped to local IDs before building
9
+routing structures.
10
"""
11
12
from unittest.mock import MagicMock, patch
@@ -24,21 +21,15 @@ class TestTritonMoeForwardExpertMap:
21
22
@pytest.mark.parametrize("expert_map_present", [False, True])
23
def test_routing_path_selection(self, expert_map_present):
27
- """Verify that the EP-aware routing path is taken when expert_map
28
- is present, and the legacy_routing path is taken otherwise."""
24
+ """Verify that both EP and non-EP paths use topk + make_routing_data,
25
+ and that expert_map remapping is applied when present."""
26
27
device = "cuda" if torch.cuda.is_available() else "cpu"
31
- # This is a structural test: we mock the routing functions to
32
- # verify the correct path is exercised.
28
mock_expert_map = (
29
torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
30
)
31
32
with (
38
- patch(
39
- "vllm.model_executor.layers.fused_moe."
40
- "gpt_oss_triton_kernels_moe.legacy_routing"
41
- ) as mock_legacy,
33
patch("triton_kernels.topk.topk") as mock_topk,
34
patch(
35
"vllm.model_executor.layers.fused_moe."
@@ -53,27 +44,19 @@ class TestTritonMoeForwardExpertMap:
44
triton_kernel_moe_forward,
45
)
46
56
- # Set up return values
47
mock_routing_data = MagicMock()
48
mock_gather = MagicMock()
49
mock_scatter = MagicMock()
50
61
- if expert_map_present:
62
- sparse_result = MagicMock()
63
- sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
64
- sparse_result.vals = torch.tensor([[0.6, 0.4]])
65
- mock_topk.return_value = sparse_result
66
- mock_make_routing.return_value = (
67
- mock_routing_data,
68
- mock_gather,
69
- mock_scatter,
70
- )
71
- else:
72
- mock_legacy.return_value = (
73
- mock_routing_data,
74
- mock_gather,
75
- mock_scatter,
76
- )
51
+ sparse_result = MagicMock()
52
+ sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
53
+ sparse_result.vals = torch.tensor([[0.6, 0.4]])
54
+ mock_topk.return_value = sparse_result
55
+ mock_make_routing.return_value = (
56
+ mock_routing_data,
57
+ mock_gather,
58
+ mock_scatter,
59
+ )
60
61
mock_fused_experts.return_value = torch.zeros((1, 8), device=device)
62
@@ -92,20 +75,14 @@ class TestTritonMoeForwardExpertMap:
75
expert_map=mock_expert_map,
76
)
77
78
+ # Both paths use topk + make_routing_data
79
+ mock_topk.assert_called_once()
80
+ mock_make_routing.assert_called_once()
81
+
82
if expert_map_present:
96
- # EP path: should use topk + make_routing_data, NOT
97
- # legacy_routing
98
- mock_topk.assert_called_once()
99
- mock_make_routing.assert_called_once()
100
- mock_legacy.assert_not_called()
83
# expert_map should be None in the fused_experts call
84
# (already applied)
85
call_kwargs = mock_fused_experts.call_args
86
assert call_kwargs[1].get("expert_map") is None or (
87
len(call_kwargs[0]) > 0
88
)
107
- else:
108
- # Non-EP path: should use legacy_routing
109
- mock_legacy.assert_called_once()
110
- mock_topk.assert_not_called()
111
- mock_make_routing.assert_not_called()
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+49
-113
@@ -47,7 +47,6 @@ if has_triton_kernels():
47
BIT,
48
Bitmatrix,
49
)
50
- from triton_kernels.topk import topk
50
51
try:
52
from triton_kernels.tensor import (
@@ -89,6 +88,7 @@ def pack_bitmatrix(
88
offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :]
89
mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :]
90
indices = tl.load(topk_ids + offsets, mask=mask, other=-1)
91
+ valid = indices >= 0
92
div = indices // 32
93
rem = indices % 32
94
one = tl.cast(1, tl.uint32)
@@ -99,8 +99,13 @@ def pack_bitmatrix(
99
offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32)
100
# All topks that need to go into this column has the correct bit set.
101
# Other bits are 0. x is a 2D tensor.
102
+ # Guard with `valid` to prevent negative indices from producing
103
+ # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets
104
+ # bit 31).
105
x = tl.where(
103
- div[:, :, None] == offs[None, None, :], (one << rem)[:, :, None], 0
106
+ valid[:, :, None] & (div[:, :, None] == offs[None, None, :]),
107
+ (one << rem)[:, :, None],
108
+ 0,
109
)
110
# Reduce x to get a single int32_t bitpack.
111
y = tl.reduce_or(x, axis=1)
@@ -108,93 +113,6 @@ def pack_bitmatrix(
113
tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows)
114
115
111
-def legacy_routing_from_bitmatrix(
112
- bitmatrix: "Bitmatrix",
113
- expt_scal: torch.Tensor,
114
- expt_indx: torch.Tensor,
115
- n_expts_tot: int,
116
- n_expts_act: int,
117
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
118
- """
119
- Replacement for the removed triton_kernels.routing.routing_from_bitmatrix.
120
- Creates routing data from a bitmatrix representation.
121
- """
122
- if use_legacy_triton_kernels:
123
- from triton_kernels.routing import routing_from_bitmatrix
124
-
125
- return routing_from_bitmatrix(
126
- bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act
127
- )
128
- sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix)
129
- dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
130
- combine_indx = sparse_logits.mask_metadata.col_sorted_indx
131
- ragged_batch_metadata = make_ragged_tensor_metadata(
132
- sparse_logits.mask_metadata.col_sum,
133
- dispatch_indx.shape[0],
134
- )
135
- gate_scal = sparse_logits.vals.flatten()[combine_indx]
136
- routing_data = RoutingData(
137
- gate_scal,
138
- ragged_batch_metadata.block_sizes,
139
- n_expts_tot,
140
- n_expts_act,
141
- ragged_batch_metadata,
142
- )
143
- gather_idx = GatherIndx(combine_indx, dispatch_indx)
144
- scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
145
- return routing_data, gather_idx, scatter_idx
146
-
147
-
148
-def legacy_routing_from_sparsematrix(
149
- sparse_logits: "SparseMatrix",
150
- n_expts_tot: int,
151
- n_expts_act: int,
152
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
153
- """
154
- Creates routing data from a SparseMatrix representation.
155
- """
156
- dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
157
- combine_indx = sparse_logits.mask_metadata.col_sorted_indx
158
- ragged_batch_metadata = make_ragged_tensor_metadata(
159
- sparse_logits.mask_metadata.col_sum,
160
- dispatch_indx.shape[0],
161
- )
162
- gate_scal = sparse_logits.vals.flatten()[combine_indx]
163
- routing_data = RoutingData(
164
- gate_scal,
165
- ragged_batch_metadata.block_sizes,
166
- n_expts_tot,
167
- n_expts_act,
168
- ragged_batch_metadata,
169
- )
170
- gather_idx = GatherIndx(combine_indx, dispatch_indx)
171
- scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
172
- return routing_data, gather_idx, scatter_idx
173
-
174
-
175
-def legacy_routing(
176
- logits: torch.Tensor,
177
- n_expts_act: int,
178
- sm_first: bool = False,
179
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
180
- """
181
- Replacement for the removed triton_kernels.routing.routing function.
182
- Computes routing data from gating logits.
183
- """
184
- if use_legacy_triton_kernels:
185
- from triton_kernels.routing import routing
186
-
187
- return routing(logits, n_expts_act, sm_first=sm_first)
188
- if sm_first:
189
- logits = torch.softmax(logits, dim=-1)
190
- sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
191
- return legacy_routing_from_sparsematrix(
192
- sparse_logits,
193
- logits.shape[-1],
194
- n_expts_act,
195
- )
196
-
197
-
116
def triton_kernel_moe_forward(
117
hidden_states: torch.Tensor,
118
w1, # Tensor or triton_kernels.Tensor
@@ -241,26 +159,22 @@ def triton_kernel_moe_forward(
159
unpadded_K_w2=unpadded_K_w2,
160
)
161
162
+ from triton_kernels.topk import topk as topk_fn
163
+
164
+ sm_first = not renormalize
165
+ logits = gating_output
166
+ if sm_first:
167
+ logits = torch.softmax(logits, dim=-1)
168
+ topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
169
+ # topk may return a tuple (vals, indx, bitmatrix) or a
170
+ # SparseMatrix depending on the triton_kernels version.
171
+ if isinstance(topk_result, tuple):
172
+ topk_weights, topk_ids_raw, _ = topk_result
173
+ else:
174
+ topk_weights = topk_result.vals
175
+ topk_ids_raw = topk_result.indx
176
+
177
if expert_map is not None:
245
- # With expert parallelism, legacy_routing produces routing data
246
- # using global expert IDs which don't correspond to local weight
247
- # indices. Split the routing into topk selection + expert_map
248
- # remapping + local routing data construction (matching the
249
- # approach used by OAITritonExperts.apply).
250
- from triton_kernels.topk import topk as topk_fn
251
-
252
- sm_first = not renormalize
253
- logits = gating_output
254
- if sm_first:
255
- logits = torch.softmax(logits, dim=-1)
256
- topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
257
- # topk may return a tuple (vals, indx, bitmatrix) or a
258
- # SparseMatrix depending on the triton_kernels version.
259
- if isinstance(topk_result, tuple):
260
- topk_weights, topk_ids_raw, _ = topk_result
261
- else:
262
- topk_weights = topk_result.vals
263
- topk_ids_raw = topk_result.indx
178
# topk_ids_raw contains global expert IDs - remap to local.
179
topk_ids = expert_map[topk_ids_raw.to(torch.long)]
180
local_num_experts = w1.shape[0]
@@ -271,8 +185,9 @@ def triton_kernel_moe_forward(
185
effective_expert_map = None
186
effective_global_num_experts = local_num_experts
187
else:
274
- routing_data, gather_idx, scatter_idx = legacy_routing(
275
- gating_output, topk, sm_first=not renormalize
188
+ topk_ids = topk_ids_raw.to(torch.long)
189
+ routing_data, gather_idx, scatter_idx = make_routing_data(
190
+ topk_ids, topk_weights, gating_output.shape[-1]
191
)
192
effective_expert_map = expert_map
193
effective_global_num_experts = global_num_experts
@@ -539,10 +454,31 @@ def make_routing_data(
454
455
# matmul_ogs expects invalid topk_weights to be -1s
456
topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
542
- routing_data, gather_indx, scatter_indx = legacy_routing_from_bitmatrix(
543
- bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
544
- )
457
458
+ if use_legacy_triton_kernels:
459
+ from triton_kernels.routing import routing_from_bitmatrix
460
+
461
+ return routing_from_bitmatrix(
462
+ bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
463
+ )
464
+
465
+ sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix)
466
+ dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
467
+ combine_indx = sparse_logits.mask_metadata.col_sorted_indx
468
+ ragged_batch_metadata = make_ragged_tensor_metadata(
469
+ sparse_logits.mask_metadata.col_sum,
470
+ dispatch_indx.shape[0],
471
+ )
472
+ gate_scal = sparse_logits.vals.flatten()[combine_indx]
473
+ routing_data = RoutingData(
474
+ gate_scal,
475
+ ragged_batch_metadata.block_sizes,
476
+ num_local_experts,
477
+ num_topk,
478
+ ragged_batch_metadata,
479
+ )
480
+ gather_indx = GatherIndx(combine_indx, dispatch_indx)
481
+ scatter_indx = ScatterIndx(dispatch_indx, combine_indx)
482
return routing_data, gather_indx, scatter_indx
483
484