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 6, 08:15 UTC
[Perf] Eliminate redundant SparseMatrix creation in gpt_oss_triton_kernels (#37683)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Mar 20, 2026 at 10:28 UTC
d0532bf38da5c8f4758e34e53a3708be0955d2db
2 files changed
+73
-4
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+44
@@ -21,12 +21,16 @@ from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_m
21
from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
22
from triton_kernels.tensor_details import layout
23
from triton_kernels.testing import assert_close
24
+from triton_kernels.topk import topk as topk_fn
25
26
from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
27
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
28
+ legacy_routing,
29
+ make_routing_data,
30
triton_kernel_moe_forward,
31
)
32
from vllm.utils.math_utils import round_up
33
+from vllm.utils.torch_utils import set_random_seed
34
35
from .utils import shuffle_weight
36
@@ -355,3 +359,43 @@ def test_unit_shuffle():
359
)
360
361
assert_close(ref=out_ref, tri=out)
362
+
363
+
364
+@pytest.mark.parametrize("num_tokens", [2, 8, 64])
365
+@pytest.mark.parametrize("num_experts", [32, 128])
366
+@pytest.mark.parametrize("topk", [1, 4])
367
+@pytest.mark.parametrize("renormalize", [True, False])
368
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
369
+def test_legacy_routing(
370
+ num_tokens: int, num_experts: int, topk: int, renormalize: bool, dtype: torch.dtype
371
+):
372
+ set_random_seed(0)
373
+ gating_output = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
374
+
375
+ sm_first = not renormalize
376
+ logits = gating_output
377
+ if sm_first:
378
+ logits = torch.softmax(logits, dim=-1)
379
+ sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first)
380
+ topk_ids = sparse_logits.indx.to(torch.long)
381
+ topk_weights = sparse_logits.vals
382
+ routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data(
383
+ topk_ids, topk_weights, num_experts
384
+ )
385
+
386
+ routing_data, gather_indx, scatter_indx = legacy_routing(
387
+ gating_output, topk, sm_first=sm_first
388
+ )
389
+
390
+ assert_close(
391
+ ref=gather_indx_ref.src_indx, tri=gather_indx.src_indx, maxtol=0, rmstol=0
392
+ )
393
+ assert_close(
394
+ ref=gather_indx_ref.dst_indx, tri=gather_indx.dst_indx, maxtol=0, rmstol=0
395
+ )
396
+ assert_close(
397
+ ref=scatter_indx_ref.src_indx, tri=scatter_indx.src_indx, maxtol=0, rmstol=0
398
+ )
399
+ assert_close(
400
+ ref=scatter_indx_ref.dst_indx, tri=scatter_indx.dst_indx, maxtol=0, rmstol=0
401
+ )
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+29
-4
@@ -142,6 +142,33 @@ def legacy_routing_from_bitmatrix(
142
return routing_data, gather_idx, scatter_idx
143
144
145
+def legacy_routing_from_sparsematrix(
146
+ sparse_logits: "SparseMatrix",
147
+ n_expts_tot: int,
148
+ n_expts_act: int,
149
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
150
+ """
151
+ Creates routing data from a SparseMatrix representation.
152
+ """
153
+ dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
154
+ combine_indx = sparse_logits.mask_metadata.col_sorted_indx
155
+ ragged_batch_metadata = make_ragged_tensor_metadata(
156
+ sparse_logits.mask_metadata.col_sum,
157
+ dispatch_indx.shape[0],
158
+ )
159
+ gate_scal = sparse_logits.vals.flatten()[combine_indx]
160
+ routing_data = RoutingData(
161
+ gate_scal,
162
+ ragged_batch_metadata.block_sizes,
163
+ n_expts_tot,
164
+ n_expts_act,
165
+ ragged_batch_metadata,
166
+ )
167
+ gather_idx = GatherIndx(combine_indx, dispatch_indx)
168
+ scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
169
+ return routing_data, gather_idx, scatter_idx
170
+
171
+
172
def legacy_routing(
173
logits: torch.Tensor,
174
n_expts_act: int,
@@ -158,10 +185,8 @@ def legacy_routing(
185
if sm_first:
186
logits = torch.softmax(logits, dim=-1)
187
sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
161
- return legacy_routing_from_bitmatrix(
162
- sparse_logits.mask,
163
- sparse_logits.vals,
164
- sparse_logits.indx,
188
+ return legacy_routing_from_sparsematrix(
189
+ sparse_logits,
190
logits.shape[-1],
191
n_expts_act,
192
)