Repository: vllm-project/vllm
File: vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py (Python)

Code before:
    w1,  # Tensor or triton_kernels.Tensor
    w2,  # Tensor or triton_kernels.Tensor
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    activation: str = "silu",
    quant_config: FusedMoEQuantConfig | None = None,
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
    expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
    if expert_map is not None:
        # With expert parallelism, legacy_routing produces routing data
        # using global expert IDs which don't correspond to local weight
        # indices.  Split the routing into topk selection + expert_map
        # remapping + local routing data construction (matching the
        # approach used by OAITritonExperts.apply).
        from triton_kernels.topk import topk as topk_fn

        sm_first = not renormalize
        logits = gating_output
        if sm_first:
            logits = torch.softmax(logits, dim=-1)
        sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first)
        # sparse_logits.indx contains global expert IDs – remap to local.
        topk_ids = expert_map[sparse_logits.indx]
        topk_weights = sparse_logits.vals
        local_num_experts = w1.size(0)
        routing_data, gather_idx, scatter_idx = make_routing_data(
            topk_ids, topk_weights, local_num_experts
        )
        # expert_map already applied; pass None downstream.
        effective_expert_map = None
        effective_global_num_experts = local_num_experts
    else:
        routing_data, gather_idx, scatter_idx = legacy_routing(
            gating_output, topk, sm_first=not renormalize
        )
        effective_expert_map = expert_map
        effective_global_num_experts = global_num_experts

    output = torch.empty_like(hidden_states)

    return triton_kernel_fused_experts(
        output,
        hidden_states,
        w1,
        w2,
        routing_data,
        gather_idx,
        scatter_idx,

Review comment (suggestion):
Hey @elizabetht  - I tried this PR on an H100 and I ran into 
```
(Worker_TP1_EP1 pid=716570) ERROR 02-13 17:09:54 [multiproc_executor.py:863]   File "/home/varun-sundar-rabindranath/code/vllm/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py", line 196, in triton_kernel_moe_forward
(Worker_TP1_EP1 pid=716570) ERROR 02-13 17:09:54 [multiproc_executor.py:863]     topk_ids = expert_map[sparse_logits.indx]
(Worker_TP1_EP1 pid=716570) ERROR 02-13 17:09:54 [multiproc_executor.py:863]                ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^
(Worker_TP1_EP1 pid=716570) ERROR 02-13 17:09:54 [multiproc_executor.py:863] IndexError: tensors used as indices must be long, int, byte or bool tensors
```
`topk_ids = expert_map[sparse_logits.indx.to(torch.long)]`
fixed it for me. Can you make this change please. Thanks 🙌 

Code after:
    hidden_states: torch.Tensor,
    w1,  # Tensor or triton_kernels.Tensor
    w2,  # Tensor or triton_kernels.Tensor
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
    activation: MoEActivation = MoEActivation.SWIGLUOAI,
    quant_config: FusedMoEQuantConfig | None = None,
    apply_router_weight_on_input: bool = False,
    global_num_experts: int = -1,
    expert_map: torch.Tensor | None = None,
) -> torch.Tensor:
    if expert_map is not None:
        # With expert parallelism, legacy_routing produces routing data
        # using global expert IDs which don't correspond to local weight
        # indices.  Split the routing into topk selection + expert_map
        # remapping + local routing data construction (matching the
        # approach used by OAITritonExperts.apply).
        from triton_kernels.topk import topk as topk_fn

        sm_first = not renormalize
        logits = gating_output
        if sm_first:
            logits = torch.softmax(logits, dim=-1)
        sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first)
        # sparse_logits.indx contains global expert IDs – remap to local.
        topk_ids = expert_map[sparse_logits.indx.to(torch.long)]
        topk_weights = sparse_logits.vals
        local_num_experts = w1.size(0)
        routing_data, gather_idx, scatter_idx = make_routing_data(
            topk_ids, topk_weights, local_num_experts
        )
        # expert_map already applied; pass None downstream.
        effective_expert_map = None
        effective_global_num_experts = local_num_experts
    else:
        routing_data, gather_idx, scatter_idx = legacy_routing(
            gating_output, topk, sm_first=not renormalize
        )
        effective_expert_map = expert_map
        effective_global_num_experts = global_num_experts

    output = torch.empty_like(hidden_states)

    return triton_kernel_fused_experts(
        output,
        hidden_states,
        w1,
        w2,
        routing_data,
        gather_idx,