[Perf] Reuse topk SparseMatrix routing metadata in GPT-OSS MoE forward (#45457)
Signed-off-by: Shenglei Fu <sfu@confluent.io> Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Shenglei Fu committed
Aug 29, 2026 at 09:14 UTC
129087ddab2b197b067e54e130d2cf5fcb094bd5
2 files changed
+102
-15
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+57
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass, fields
3
+from dataclasses import dataclass, fields, is_dataclass
4
5
import pytest
6
import torch
@@ -371,3 +371,59 @@ def test_unit_shuffle():
371
)
372
373
assert_close(ref=out_ref, tri=out)
374
+
375
+
376
+@pytest.mark.parametrize("n_tokens", [1, 33, 512])
377
+@pytest.mark.parametrize("n_experts,topk", [(32, 4), (128, 4)])
378
+def test_routing_data_from_sparse_topk_parity(n_tokens, n_experts, topk):
379
+ """routing_data_from_sparse_topk must produce routing structures
380
+ identical to make_routing_data for the same topk result."""
381
+ from vllm.model_executor.layers.fused_moe.experts import (
382
+ gpt_oss_triton_kernels_moe as gptoss_moe,
383
+ )
384
+
385
+ make_routing_data = gptoss_moe.make_routing_data
386
+ routing_data_from_sparse_topk = gptoss_moe.routing_data_from_sparse_topk
387
+ use_legacy_triton_kernels = gptoss_moe.use_legacy_triton_kernels
388
+
389
+ if use_legacy_triton_kernels:
390
+ pytest.skip("SparseMatrix path requires triton_kernels v3.6.0+")
391
+
392
+ from triton_kernels.topk import topk as topk_fn
393
+
394
+ torch.manual_seed(0)
395
+ logits = torch.randn(n_tokens, n_experts, dtype=torch.bfloat16, device="cuda")
396
+ sparse_topk = topk_fn(logits, topk, apply_softmax=True)
397
+ assert not isinstance(sparse_topk, tuple)
398
+
399
+ rd_new, gather_new, scatter_new = routing_data_from_sparse_topk(
400
+ sparse_topk, n_experts
401
+ )
402
+ rd_ref, gather_ref, scatter_ref = make_routing_data(
403
+ sparse_topk.indx.to(torch.long), sparse_topk.vals, n_experts
404
+ )
405
+
406
+ def assert_equivalent(new, ref, path="RoutingData"):
407
+ """Recursively compare tensor-valued leaves of routing structures."""
408
+ if isinstance(ref, torch.Tensor):
409
+ torch.testing.assert_close(new, ref, msg=lambda m: f"{path}: {m}")
410
+ elif isinstance(ref, (int, float, bool, str)) or ref is None:
411
+ assert new == ref, f"{path}: {new!r} != {ref!r}"
412
+ elif is_dataclass(ref):
413
+ for f in fields(ref):
414
+ assert_equivalent(
415
+ getattr(new, f.name), getattr(ref, f.name), f"{path}.{f.name}"
416
+ )
417
+ elif callable(ref):
418
+ assert callable(new), f"{path}: callable vs {type(new)}"
419
+ elif hasattr(ref, "__dict__"):
420
+ for k, v in vars(ref).items():
421
+ assert_equivalent(getattr(new, k), v, f"{path}.{k}")
422
+ else:
423
+ assert type(new) is type(ref), f"{path}: {type(new)} vs {type(ref)}"
424
+
425
+ assert_equivalent(rd_new, rd_ref)
426
+ torch.testing.assert_close(gather_new.src_indx, gather_ref.src_indx)
427
+ torch.testing.assert_close(gather_new.dst_indx, gather_ref.dst_indx)
428
+ torch.testing.assert_close(scatter_new.src_indx, scatter_ref.src_indx)
429
+ torch.testing.assert_close(scatter_new.dst_indx, scatter_ref.dst_indx)
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+45
-14
@@ -578,13 +578,14 @@ def triton_kernel_moe_forward(
578
if sm_first:
579
logits = torch.softmax(logits, dim=-1)
580
topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
581
+ is_sparse_topk = not isinstance(topk_result, tuple)
582
# topk may return a tuple (vals, indx, bitmatrix) or a
583
# SparseMatrix depending on the triton_kernels version.
583
- if isinstance(topk_result, tuple):
584
- topk_weights, topk_ids_raw, _ = topk_result
585
- else:
584
+ if is_sparse_topk:
585
topk_weights = topk_result.vals
586
topk_ids_raw = topk_result.indx
587
+ else:
588
+ topk_weights, topk_ids_raw, _ = topk_result
589
590
if expert_map is not None:
591
# topk_ids_raw contains global expert IDs - remap to local.
@@ -597,11 +598,23 @@ def triton_kernel_moe_forward(
598
effective_expert_map = None
599
effective_global_num_experts = local_num_experts
600
else:
600
- topk_ids = topk_ids_raw.to(torch.long)
601
- routing_data, gather_idx, scatter_idx = make_routing_data(
602
- topk_ids, topk_weights, gating_output.shape[-1]
603
- )
604
- effective_expert_map = expert_map
601
+ if is_sparse_topk:
602
+ # v3.6+ topk returns a SparseMatrix whose construction already
603
+ # computed the bitmatrix and its metadata; reuse it instead of
604
+ # re-packing the bitmatrix and recomputing the metadata in
605
+ # make_routing_data. Only the monolithic path reaches this reuse;
606
+ # the modular OAITritonExperts (expert-parallel) path calls
607
+ # make_routing_data directly.
608
+ routing_data, gather_idx, scatter_idx = routing_data_from_sparse_topk(
609
+ topk_result, gating_output.shape[-1]
610
+ )
611
+ else:
612
+ topk_ids = topk_ids_raw.to(torch.long)
613
+ routing_data, gather_idx, scatter_idx = make_routing_data(
614
+ topk_ids, topk_weights, gating_output.shape[-1]
615
+ )
616
+
617
+ effective_expert_map = None
618
effective_global_num_experts = global_num_experts
619
620
output = torch.empty_like(hidden_states)
@@ -782,18 +795,36 @@ def make_routing_data(
795
)
796
797
sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix)
785
- dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
786
- combine_indx = sparse_logits.mask_metadata.col_sorted_indx
798
+ return routing_data_from_sparse_topk(sparse_logits, num_local_experts)
799
+
800
+
801
+def routing_data_from_sparse_topk(
802
+ sparse_topk, # SparseMatrix
803
+ n_expts: int,
804
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
805
+ """Build routing structures from the ``SparseMatrix`` returned by
806
+ ``triton_kernels.topk``, reusing its precomputed bitmatrix metadata.
807
+
808
+ ``make_routing_data`` re-packs a bitmatrix from the topk ids and
809
+ constructs a second ``SparseMatrix``, recomputing metadata the topk
810
+ call already produced; when the topk ids are used as-is (no expert-map
811
+ remapping) that work is redundant.
812
+ """
813
+ dispatch_indx = sparse_topk.mask_metadata.row_sorted_indx
814
+ combine_indx = sparse_topk.mask_metadata.col_sorted_indx
815
ragged_batch_metadata = make_ragged_tensor_metadata(
788
- sparse_logits.mask_metadata.col_sum,
816
+ sparse_topk.mask_metadata.col_sum,
817
dispatch_indx.shape[0],
818
)
791
- gate_scal = sparse_logits.vals.flatten()[combine_indx]
819
+ n_expts_act = sparse_topk.indx.shape[1]
820
+ # Unlike make_routing_data, no bfloat16 cast: gate_scal inherits the
821
+ # gating dtype, which is bfloat16 for all callers of this path.
822
+ gate_scal = sparse_topk.vals.flatten()[combine_indx]
823
routing_data = RoutingData(
824
gate_scal,
825
ragged_batch_metadata.block_sizes,
795
- num_local_experts,
796
- num_topk,
826
+ n_expts,
827
+ n_expts_act,
828
ragged_batch_metadata,
829
)
830
gather_indx = GatherIndx(combine_indx, dispatch_indx)