[Bugfix] Support -1 (invalid/non-local) slots in topk_ids for Triton MoE (#46408)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com>
Woosuk Kwon committed
Jun 24, 2026 at 13:59 UTC
d7ab9be775526cca1042009e3754f8a4ef14b56c
1 file changed
+88
-7
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+88
-7
@@ -4,7 +4,6 @@
4
import torch
5
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
-from vllm import _custom_ops as ops
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
from vllm.model_executor.layers.fused_moe.config import (
@@ -815,6 +814,85 @@ def make_routing_data(
814
return routing_data, gather_indx, scatter_indx
815
816
817
+@triton.jit
818
+def _masked_topk_sum_kernel(
819
+ inp_ptr, # (M, topk, K) contiguous
820
+ topk_ids_ptr, # (M, topk) int: -1 marks an invalid / non-local slot
821
+ out_ptr, # (M, K), same dtype as inp
822
+ K,
823
+ topk: tl.constexpr,
824
+ BLOCK_K: tl.constexpr,
825
+):
826
+ pid_m = tl.program_id(0).to(tl.int64)
827
+ k = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K)
828
+ k_mask = k < K
829
+ base = pid_m * topk
830
+ acc = tl.zeros((BLOCK_K,), dtype=tl.float32)
831
+ for j in tl.static_range(topk):
832
+ eid = tl.load(topk_ids_ptr + base + j)
833
+ # NOTE: This is NaN-safe because the invalid slots are skipped.
834
+ if eid >= 0:
835
+ x = tl.load(inp_ptr + (base + j) * K + k, mask=k_mask)
836
+ acc += x.to(tl.float32)
837
+ tl.store(out_ptr + pid_m * K + k, acc.to(out_ptr.dtype.element_ty), mask=k_mask)
838
+
839
+
840
+def masked_moe_sum(
841
+ intermediate: torch.Tensor, # (M, topk, K)
842
+ topk_ids: torch.Tensor, # (M, topk) int, -1 = invalid / non-local slot
843
+ output: torch.Tensor, # (M, K)
844
+) -> None:
845
+ M, topk, K = intermediate.shape
846
+ BLOCK_K = 1024
847
+ grid = (M, triton.cdiv(K, BLOCK_K))
848
+ _masked_topk_sum_kernel[grid](
849
+ intermediate, topk_ids, output, K, topk=topk, BLOCK_K=BLOCK_K
850
+ )
851
+
852
+
853
+@triton.jit
854
+def _remap_topk_to_local_kernel(
855
+ topk_ids_ptr, # [n] global expert IDs (-1 = invalid)
856
+ expert_map_ptr, # [num_experts] global->local (-1 for non-local)
857
+ out_ptr, # [n] int64 local expert IDs (-1 for invalid/non-local)
858
+ n_elements,
859
+ BLOCK: tl.constexpr,
860
+):
861
+ pid = tl.program_id(0)
862
+ offs = pid * BLOCK + tl.arange(0, BLOCK)
863
+ mask = offs < n_elements
864
+ tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1)
865
+ # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid
866
+ # rows don't read OOB, then select -1 for them. Matches
867
+ # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a
868
+ # plain expert_map[-1] would wrap to a valid local id and misroute).
869
+ valid = tid >= 0
870
+ idx = tl.where(valid, tid, 0)
871
+ local = tl.load(expert_map_ptr + idx, mask=mask, other=-1)
872
+ out = tl.where(valid, local.to(tl.int64), -1)
873
+ tl.store(out_ptr + offs, out, mask=mask)
874
+
875
+
876
+def remap_topk_to_local(
877
+ topk_ids: torch.Tensor, expert_map: torch.Tensor
878
+) -> torch.Tensor:
879
+ """Fused global->local expert-id mapping over a topk_ids tensor, preserving -1.
880
+
881
+ Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)``
882
+ with one kernel. Returns a NEW int64 tensor -- the caller keeps the original
883
+ ``topk_ids`` as ``global_topk_ids``, so this must not write in place.
884
+
885
+ (Distinct from ``deep_gemm_utils.apply_expert_map``, which is a scalar
886
+ ``@triton.jit`` device helper called from within other kernels.)
887
+ """
888
+ out = torch.empty_like(topk_ids, dtype=torch.int64)
889
+ n = topk_ids.numel()
890
+ BLOCK = 1024
891
+ grid = (triton.cdiv(n, BLOCK),)
892
+ _remap_topk_to_local_kernel[grid](topk_ids, expert_map, out, n, BLOCK=BLOCK)
893
+ return out
894
+
895
+
896
class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
897
@property
898
def expects_unquantized_inputs(self) -> bool:
@@ -946,7 +1024,9 @@ class OAITritonExperts(BaseOAITritonExperts):
1024
self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG
1025
1026
if expert_map is not None:
949
- topk_ids = expert_map[topk_ids]
1027
+ # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch):
1028
+ # make_routing_data treats -1 as the skip sentinel.
1029
+ topk_ids = remap_topk_to_local(topk_ids, expert_map)
1030
1031
local_num_experts = w1.shape[0]
1032
if global_num_experts == -1:
@@ -1018,9 +1098,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1098
output = (M, K)
1099
return (workspace1, workspace2, output)
1100
1021
- def moe_sum(self, input: torch.Tensor, output: torch.Tensor):
1022
- ops.moe_sum(input, output)
1023
-
1101
def activation(
1102
self,
1103
activation: MoEActivation,
@@ -1091,7 +1168,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1168
1169
global_topk_ids = topk_ids
1170
if expert_map is not None:
1094
- topk_ids = expert_map[topk_ids]
1171
+ # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch):
1172
+ # make_routing_data treats -1 as the skip sentinel.
1173
+ topk_ids = remap_topk_to_local(topk_ids, expert_map)
1174
1175
local_num_experts = w1.shape[0]
1176
if global_num_experts == -1:
@@ -1214,7 +1293,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts):
1293
top_k_num=topk,
1294
)
1295
1217
- self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
1296
+ # matmul_ogs leaves invalid (-1 / non-local EP) slots unwritten.
1297
+ # Reduce over topk skipping those slots.
1298
+ masked_moe_sum(intermediate_cache3.view(-1, topk, K), topk_ids, output)
1299
1300
1301
class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):