[Perf][MXFP4] Optimize activation quant for mxfp4. 3x kernel speedup, 5% throughput improvement (#63101)

Signed-off-by: ilmarkov <markovilya197@gmail.com>
Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>

benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py

@@ -0,0 +1,538 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Benchmark MXFP4 MoE activation quant at GLM-5.3-shaped sizes.
+
+Sweeps expert count at a fixed hidden width (default 6144). Two timed regions
+(``--timed-region``):
+
+* ``quant``: ``mxfp4_experts_quant`` / ``silu_and_mul_mxfp4_experts_quant`` only,
+  with synthetic routing offsets (kernel isolation).
+* ``metadata_and_quant``: ``get_cutlass_moe_mm_data`` → ``shuffle_rows`` → quant.
+* ``both``: run ``quant`` then ``metadata_and_quant``.
+
+Routing uses 32 rows (4 tokens, topk 8, two picks per expert 0..15). GEMMs are
+outside the timed region.
+
+Run from a CUDA-built vLLM checkout, for example:
+
+    .venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py
+    .venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py \\
+        --timed-region both --experts 16,32,64,128,256
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO_ROOT))
+
+import vllm._custom_ops  # noqa: E402, F401 -- registers torch.ops._C kernels
+from vllm import _custom_ops as ops  # noqa: E402
+from vllm.triton_utils import triton  # noqa: E402
+
+ROUTED_ROWS = 32
+TOPK = 8
+DEFAULT_HIDDEN_SIZE = 6144
+DEFAULT_MOE_INTERMEDIATE_SIZE = 2048
+DEFAULT_TP_SIZE = 4
+DEFAULT_EXPERT_COUNTS = (16, 32, 64, 128, 256)
+
+
+class TimedRegion(str, Enum):
+    QUANT = "quant"
+    METADATA_AND_QUANT = "metadata_and_quant"
+    BOTH = "both"
+
+
+@dataclass(frozen=True)
+class Case:
+    kernel: str
+    output_width: int
+    fused_silu_mul: bool
+    experts: int
+    mm_n: int
+    mm_k: int
+    is_gated: bool
+
+
+def routing_offsets(
+    num_experts: int,
+    rows: int = ROUTED_ROWS,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """32 rows, two assignments on experts 0..15 (production-like small decode)."""
+    if rows != ROUTED_ROWS:
+        raise ValueError(f"expected {ROUTED_ROWS} routed rows, got {rows}")
+    if num_experts < 16:
+        raise ValueError("expert sweep must use at least 16 experts")
+
+    counts = [0] * num_experts
+    counts[:16] = [2] * 16
+
+    row_offsets = [0]
+    scale_offsets = [0]
+    expert_ids: list[int] = []
+    for expert, count in enumerate(counts):
+        expert_ids.extend([expert] * count)
+        row_offsets.append(row_offsets[-1] + count)
+        scale_offsets.append(scale_offsets[-1] + ((count + 127) // 128) * 128)
+    assert row_offsets[-1] == rows
+    return (
+        torch.tensor(row_offsets, dtype=torch.int32, device="cuda"),
+        torch.tensor(scale_offsets, dtype=torch.int32, device="cuda"),
+        torch.tensor(expert_ids, dtype=torch.int32, device="cuda"),
+        scale_offsets[-1],
+    )
+
+
+def decode_topk_ids(rows: int = ROUTED_ROWS) -> torch.Tensor:
+    """Two picks per expert 0..15 across 4 tokens (topk=8); shuffled (fixed seed)."""
+    if rows % TOPK != 0:
+        raise ValueError(f"routed rows {rows} must be divisible by topk {TOPK}")
+    expert_ids: list[int] = []
+    for expert in range(16):
+        expert_ids.extend([expert, expert])
+    assert len(expert_ids) == rows
+    gen = torch.Generator(device="cuda")
+    gen.manual_seed(0)
+    perm = torch.randperm(rows, generator=gen, device="cuda")
+    ordered = torch.tensor(expert_ids, device="cuda", dtype=torch.int32)
+    return ordered[perm].view(rows // TOPK, TOPK)
+
+
+def benchmark_invoke(invoke, rep: int) -> float:
+    """Median latency in microseconds (CUDA graph replay via Triton)."""
+    quantiles = [0.5, 0.2, 0.8]
+    ms, _, _ = triton.testing.do_bench_cudagraph(invoke, rep=rep, quantiles=quantiles)
+    return ms * 1_000
+
+
+def _quant_tensors(
+    case: Case,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    scale_rows = int(blockscale_offsets[-1].item())
+    scale_cols_int32 = (case.output_width // 32 + 3) // 4
+    output = torch.empty(
+        (ROUTED_ROWS, case.output_width // 2), dtype=torch.uint8, device="cuda"
+    )
+    output_scales = torch.zeros(
+        (scale_rows, scale_cols_int32), dtype=torch.int32, device="cuda"
+    )
+    return output, output_scales
+
+
+def benchmark_mxfp4_quant_only(
+    case: Case,
+    inputs: torch.Tensor,
+    rep: int,
+    sorted_first: bool,
+) -> tuple[float, float]:
+    expert_offsets, blockscale_offsets, expert_ids, _ = routing_offsets(case.experts)
+    output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets)
+    op = (
+        torch.ops._C.silu_and_mul_mxfp4_experts_quant
+        if case.fused_silu_mul
+        else torch.ops._C.mxfp4_experts_quant
+    )
+
+    def invoke_offset() -> None:
+        op(
+            output,
+            output_scales,
+            inputs,
+            expert_offsets,
+            blockscale_offsets,
+            case.experts,
+            None,
+        )
+
+    def invoke_sorted() -> None:
+        op(
+            output,
+            output_scales,
+            inputs,
+            expert_offsets,
+            blockscale_offsets,
+            case.experts,
+            expert_ids,
+        )
+
+    if sorted_first:
+        direct_id_us = benchmark_invoke(invoke_sorted, rep)
+        baseline_us = benchmark_invoke(invoke_offset, rep)
+    else:
+        baseline_us = benchmark_invoke(invoke_offset, rep)
+        baseline_output = output.clone()
+        baseline_scales = output_scales.clone()
+        direct_id_us = benchmark_invoke(invoke_sorted, rep)
+        if not (
+            torch.equal(output, baseline_output)
+            and torch.equal(output_scales, baseline_scales)
+        ):
+            raise AssertionError(
+                f"MXFP4 output mismatch for {case.kernel} E={case.experts}"
+            )
+    return baseline_us, direct_id_us
+
+
+def benchmark_mxfp4_metadata_and_quant(
+    case: Case,
+    inputs: torch.Tensor,
+    rep: int,
+    sorted_first: bool,
+) -> tuple[float, float]:
+    num_experts = case.experts
+    topk_ids = decode_topk_ids()
+    expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda")
+    blockscale_offsets = torch.empty_like(expert_offsets)
+    problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda")
+    problem_sizes2 = torch.empty_like(problem_sizes1)
+    input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda")
+    output_permutation = torch.empty_like(input_permutation)
+    sorted_expert_ids = torch.empty_like(input_permutation)
+
+    torch.ops._C.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        case.mm_n,
+        case.mm_k,
+        blockscale_offsets,
+        case.is_gated,
+        sorted_expert_ids,
+    )
+    output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets)
+    op = (
+        torch.ops._C.silu_and_mul_mxfp4_experts_quant
+        if case.fused_silu_mul
+        else torch.ops._C.mxfp4_experts_quant
+    )
+
+    def invoke_offset() -> None:
+        torch.ops._C.get_cutlass_moe_mm_data(
+            topk_ids,
+            expert_offsets,
+            problem_sizes1,
+            problem_sizes2,
+            input_permutation,
+            output_permutation,
+            num_experts,
+            case.mm_n,
+            case.mm_k,
+            blockscale_offsets,
+            case.is_gated,
+            None,
+        )
+        sorted_inputs = ops.shuffle_rows(inputs, input_permutation)
+        op(
+            output,
+            output_scales,
+            sorted_inputs,
+            expert_offsets,
+            blockscale_offsets,
+            num_experts,
+            None,
+        )
+
+    def invoke_sorted() -> None:
+        torch.ops._C.get_cutlass_moe_mm_data(
+            topk_ids,
+            expert_offsets,
+            problem_sizes1,
+            problem_sizes2,
+            input_permutation,
+            output_permutation,
+            num_experts,
+            case.mm_n,
+            case.mm_k,
+            blockscale_offsets,
+            case.is_gated,
+            sorted_expert_ids,
+        )
+        sorted_inputs = ops.shuffle_rows(inputs, input_permutation)
+        op(
+            output,
+            output_scales,
+            sorted_inputs,
+            expert_offsets,
+            blockscale_offsets,
+            num_experts,
+            sorted_expert_ids,
+        )
+
+    if sorted_first:
+        direct_id_us = benchmark_invoke(invoke_sorted, rep)
+        baseline_us = benchmark_invoke(invoke_offset, rep)
+    else:
+        baseline_us = benchmark_invoke(invoke_offset, rep)
+        baseline_output = output.clone()
+        baseline_scales = output_scales.clone()
+        direct_id_us = benchmark_invoke(invoke_sorted, rep)
+        if not (
+            torch.equal(output, baseline_output)
+            and torch.equal(output_scales, baseline_scales)
+        ):
+            raise AssertionError(
+                f"MXFP4 output mismatch for {case.kernel} E={case.experts}"
+            )
+    return baseline_us, direct_id_us
+
+
+def verify_sorted_id_writer(num_experts: int) -> None:
+    """Check the routing sort writes the expected sorted expert ID at each row."""
+    topk_ids = torch.arange(ROUTED_ROWS, device="cuda", dtype=torch.int32) % num_experts
+    topk_ids = topk_ids.view(4, 8)
+    expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda")
+    blockscale_offsets = torch.empty_like(expert_offsets)
+    problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda")
+    problem_sizes2 = torch.empty_like(problem_sizes1)
+    input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda")
+    output_permutation = torch.empty_like(input_permutation)
+    sorted_expert_ids = torch.empty_like(input_permutation)
+    torch.ops._C.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        64,
+        64,
+        blockscale_offsets,
+        True,
+        sorted_expert_ids,
+    )
+    expected = torch.sort(topk_ids.flatten()).values
+    torch.accelerator.synchronize()
+    if not torch.equal(sorted_expert_ids, expected):
+        raise AssertionError("routing sort produced incorrect sorted expert IDs")
+
+
+def parse_expert_counts(raw: str) -> tuple[int, ...]:
+    counts = tuple(int(x.strip()) for x in raw.split(",") if x.strip())
+    if not counts or any(c < 16 for c in counts):
+        raise ValueError(
+            "expert counts must be a comma-separated list of integers >= 16"
+        )
+    return counts
+
+
+def parse_timed_region(raw: str) -> TimedRegion:
+    try:
+        return TimedRegion(raw)
+    except ValueError:
+        choices = ", ".join(r.value for r in TimedRegion)
+        raise ValueError(f"timed-region must be one of: {choices}") from None
+
+
+def intermediate_dims(moe_intermediate_size: int, tp_size: int) -> tuple[int, int]:
+    if moe_intermediate_size % tp_size != 0:
+        raise ValueError(
+            f"moe intermediate size {moe_intermediate_size} must be divisible "
+            f"by tp size {tp_size}"
+        )
+    n_per_rank = moe_intermediate_size // tp_size
+    return 2 * n_per_rank, n_per_rank
+
+
+def build_cases(
+    experts: tuple[int, ...],
+    hidden_size: int,
+    moe_intermediate_size: int,
+    tp_size: int,
+) -> list[tuple[Case, torch.Tensor]]:
+    intermediate_in, intermediate_out = intermediate_dims(
+        moe_intermediate_size, tp_size
+    )
+    cases: list[tuple[Case, torch.Tensor]] = []
+    hidden_input = torch.randn(
+        (ROUTED_ROWS, hidden_size), device="cuda", dtype=torch.bfloat16
+    )
+    intermediate_input = torch.randn(
+        (ROUTED_ROWS, intermediate_in), device="cuda", dtype=torch.bfloat16
+    )
+    w1_n = 2 * intermediate_out
+    for expert_count in experts:
+        cases.append(
+            (
+                Case(
+                    "mxfp4_experts_quant",
+                    hidden_size,
+                    False,
+                    expert_count,
+                    mm_n=w1_n,
+                    mm_k=hidden_size,
+                    is_gated=True,
+                ),
+                hidden_input,
+            )
+        )
+        cases.append(
+            (
+                Case(
+                    "silu_and_mul_mxfp4_experts_quant",
+                    intermediate_out,
+                    True,
+                    expert_count,
+                    mm_n=hidden_size,
+                    mm_k=intermediate_out,
+                    is_gated=True,
+                ),
+                intermediate_input,
+            )
+        )
+    return cases
+
+
+def run_region(
+    region: TimedRegion,
+    cases: list[tuple[Case, torch.Tensor]],
+    rep: int,
+    sorted_first: bool,
+) -> None:
+    print(f"timed_region={region.value}")
+    print("kernel,experts,variant,median_us,speedup_vs_offset_search")
+    bench_fn = (
+        benchmark_mxfp4_metadata_and_quant
+        if region == TimedRegion.METADATA_AND_QUANT
+        else benchmark_mxfp4_quant_only
+    )
+    for case, inputs in cases:
+        baseline_us, direct_id_us = bench_fn(case, inputs, rep, sorted_first)
+        print(f"{case.kernel},{case.experts},offset_search,{baseline_us:.4f},1.0000")
+        print(
+            f"{case.kernel},{case.experts},sorted_expert_id,"
+            f"{direct_id_us:.4f},{baseline_us / direct_id_us:.4f}"
+        )
+    print()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--rep",
+        type=int,
+        default=100,
+        help="Repeat count (default: 100)",
+    )
+    parser.add_argument(
+        "--hidden-size",
+        type=int,
+        default=DEFAULT_HIDDEN_SIZE,
+        help=f"hidden activation width K (default: {DEFAULT_HIDDEN_SIZE})",
+    )
+    parser.add_argument(
+        "--tp-size",
+        type=int,
+        default=DEFAULT_TP_SIZE,
+        help=(
+            "tensor-parallel size; shards moe intermediate per rank "
+            f"(default: {DEFAULT_TP_SIZE})"
+        ),
+    )
+    parser.add_argument(
+        "--moe-intermediate-size",
+        type=int,
+        default=DEFAULT_MOE_INTERMEDIATE_SIZE,
+        help=(
+            "total MoE intermediate size before TP; per-rank N is this divided "
+            f"by --tp-size (default: {DEFAULT_MOE_INTERMEDIATE_SIZE})"
+        ),
+    )
+    parser.add_argument(
+        "--experts",
+        type=parse_expert_counts,
+        default=DEFAULT_EXPERT_COUNTS,
+        help=(
+            "comma-separated expert counts (default: "
+            f"{','.join(map(str, DEFAULT_EXPERT_COUNTS))})"
+        ),
+    )
+    parser.add_argument(
+        "--timed-region",
+        type=parse_timed_region,
+        default=TimedRegion.QUANT,
+        help=(
+            "quant: activation quant only; metadata_and_quant: mm_data + shuffle + "
+            "quant; both: run each region"
+        ),
+    )
+    parser.add_argument(
+        "--sorted-id-first",
+        action="store_true",
+        help="time sorted-expert-ID path before offset search (default: offset first)",
+    )
+    parser.add_argument("--skip-metadata-check", action="store_true")
+    args = parser.parse_args()
+    if args.rep < 1:
+        parser.error("--rep must be positive")
+    if args.hidden_size < 128 or args.hidden_size % 128 != 0:
+        parser.error("--hidden-size must be a multiple of 128 and >= 128")
+    if args.tp_size < 1:
+        parser.error("--tp-size must be positive")
+    if args.moe_intermediate_size < args.tp_size:
+        parser.error("--moe-intermediate-size must be >= --tp-size")
+    if args.moe_intermediate_size % args.tp_size != 0:
+        parser.error("--moe-intermediate-size must be divisible by --tp-size")
+    try:
+        intermediate_in, intermediate_out = intermediate_dims(
+            args.moe_intermediate_size, args.tp_size
+        )
+    except ValueError as exc:
+        parser.error(str(exc))
+    if intermediate_out < 32 or intermediate_out % 32 != 0:
+        parser.error(
+            "per-rank MoE intermediate (moe-intermediate-size / tp-size) "
+            "must be a multiple of 32 and >= 32"
+        )
+    if not torch.cuda.is_available():
+        parser.error("requires a CUDA GPU")
+    capability = torch.cuda.get_device_capability()
+    if capability[0] not in (10, 11):
+        parser.error(f"requires an SM10x/11x GPU, got SM{capability[0]}{capability[1]}")
+
+    torch.manual_seed(0)
+    if not args.skip_metadata_check:
+        verify_sorted_id_writer(num_experts=256)
+
+    cases = build_cases(
+        args.experts,
+        args.hidden_size,
+        args.moe_intermediate_size,
+        args.tp_size,
+    )
+    print(f"GPU: {torch.cuda.get_device_name()} (SM{capability[0]}{capability[1]})")
+    print(
+        f"hidden={ROUTED_ROWS}x{args.hidden_size}, "
+        f"intermediate={ROUTED_ROWS}x{intermediate_in}->{intermediate_out} "
+        f"(moe_intermediate={args.moe_intermediate_size}, tp={args.tp_size}), "
+        f"topk={TOPK}"
+    )
+
+    regions: list[TimedRegion]
+    if args.timed_region == TimedRegion.BOTH:
+        regions = [TimedRegion.QUANT, TimedRegion.METADATA_AND_QUANT]
+    else:
+        regions = [args.timed_region]
+
+    for region in regions:
+        run_region(region, cases, args.rep, args.sorted_id_first)
+
+
+if __name__ == "__main__":
+    main()

csrc/ops.h

@@ -85,7 +85,8 @@ void get_cutlass_moe_mm_data(
     torch::stable::Tensor& output_permutation, const int64_t num_experts,
     const int64_t n, const int64_t k,
     const std::optional<torch::stable::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const bool is_gated,
+    const std::optional<torch::stable::Tensor>& sorted_expert_ids);
 
 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
     const torch::stable::Tensor& expert_first_token_offset,

csrc/quantization/fp4/mxfp4_experts_quant.cu

@@ -22,6 +22,7 @@
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>
 #include <cuda_fp8.h>
+#include <optional>
 
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
@@ -66,7 +67,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
                           fp4_packed_t* out, uint32_t* SFout,
                           uint32_t* input_offset_by_experts,
                           uint32_t* output_scale_offset_by_experts,
-                          int n_experts, bool low_latency) {
+                          int n_experts, int32_t const* sorted_expert_ids,
+                          bool low_latency) {
   using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
   static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
                 "Vec size is not matched.");
@@ -89,7 +91,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
     int rowIdx_in_expert = 0;
     int expert_idx = 0;
 
-    if constexpr (SMALL_NUM_EXPERTS) {
+    if (sorted_expert_ids != nullptr) {
+      expert_idx = __ldca(&sorted_expert_ids[rowIdx]);
+      rowIdx_in_expert = rowIdx - __ldca(&input_offset_by_experts[expert_idx]);
+    } else if constexpr (SMALL_NUM_EXPERTS) {
       for (int i = 0; i < n_experts; i++) {
         uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
         uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
@@ -168,7 +173,7 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
                           fp4_packed_t* out, uint32_t* SFout,
                           uint32_t* input_offset_by_experts,
                           uint32_t* output_scale_offset_by_experts,
-                          int n_experts) {
+                          int n_experts, int32_t const* sorted_expert_ids) {
   using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
   static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
                 "Vec size is not matched.");
@@ -208,21 +213,26 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
     int rowIdx_in_expert = 0;
     int expert_idx = 0;
 
-    // Binary search through experts using shared memory
-    int left = 0, right = n_experts - 1;
-    while (left <= right) {
-      int mid = (left + right) / 2;
-      uint32_t mid_offset = shared_input_offsets[mid];
-      uint32_t next_offset = shared_input_offsets[mid + 1];
+    if (sorted_expert_ids != nullptr) {
+      expert_idx = __ldca(&sorted_expert_ids[rowIdx]);
+      rowIdx_in_expert = rowIdx - shared_input_offsets[expert_idx];
+    } else {
+      // Binary search through experts using shared memory
+      int left = 0, right = n_experts - 1;
+      while (left <= right) {
+        int mid = (left + right) / 2;
+        uint32_t mid_offset = shared_input_offsets[mid];
+        uint32_t next_offset = shared_input_offsets[mid + 1];
 
-      if (rowIdx >= mid_offset && rowIdx < next_offset) {
-        rowIdx_in_expert = rowIdx - mid_offset;
-        expert_idx = mid;
-        break;
-      } else if (rowIdx < mid_offset) {
-        right = mid - 1;
-      } else {
-        left = mid + 1;
+        if (rowIdx >= mid_offset && rowIdx < next_offset) {
+          rowIdx_in_expert = rowIdx - mid_offset;
+          expert_idx = mid;
+          break;
+        } else if (rowIdx < mid_offset) {
+          right = mid - 1;
+        } else {
+          left = mid + 1;
+        }
       }
     }
 
@@ -259,7 +269,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
 template <typename T, bool FUSE_SILU_MUL = false>
 void mxfp4_quant_impl(void* output, void* output_scale, void* input,
                       void* input_offset_by_experts,
-                      void* output_scale_offset_by_experts, int m_topk, int k,
+                      void* output_scale_offset_by_experts,
+                      int32_t const* sorted_expert_ids, int m_topk, int k,
                       int n_experts, cudaStream_t stream) {
   int multiProcessorCount =
       get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
@@ -288,7 +299,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
               reinterpret_cast<uint32_t*>(output_scale),
               reinterpret_cast<uint32_t*>(input_offset_by_experts),
               reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
+              n_experts, sorted_expert_ids);
     } else {
       mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
           <<<grid, block, shared_mem_size, stream>>>(
@@ -297,7 +308,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
               reinterpret_cast<uint32_t*>(output_scale),
               reinterpret_cast<uint32_t*>(input_offset_by_experts),
               reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts);
+              n_experts, sorted_expert_ids);
     }
   } else {
     if (n_experts >= 16) {
@@ -308,7 +319,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
               reinterpret_cast<uint32_t*>(output_scale),
               reinterpret_cast<uint32_t*>(input_offset_by_experts),
               reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-              n_experts, /* bool low_latency */ true);
+              n_experts, sorted_expert_ids,
+              /* bool low_latency */ true);
     } else {
       mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
           m_topk, k, reinterpret_cast<T*>(input),
@@ -316,7 +328,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input,
           reinterpret_cast<uint32_t*>(output_scale),
           reinterpret_cast<uint32_t*>(input_offset_by_experts),
           reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
-          n_experts, /* bool low_latency */ true);
+          n_experts, sorted_expert_ids,
+          /* bool low_latency */ true);
     }
   }
 }
@@ -345,12 +358,21 @@ static void validate_mxfp4_experts_quant_inputs(
     torch::stable::Tensor const& input,
     torch::stable::Tensor const& input_offset_by_experts,
     torch::stable::Tensor const& output_scale_offset_by_experts,
+    std::optional<torch::stable::Tensor> const& sorted_expert_ids,
     int64_t n_experts, int64_t m_topk, int64_t k) {
   CHECK_INPUT(output, "output");
   CHECK_INPUT(output_scale, "output_scale");
   CHECK_INPUT(input, "input");
   CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
   CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
+  if (sorted_expert_ids.has_value()) {
+    CHECK_INPUT(sorted_expert_ids.value(), "sorted_expert_ids");
+    STD_TORCH_CHECK(sorted_expert_ids->dim() == 1);
+    STD_TORCH_CHECK(sorted_expert_ids->scalar_type() == INT);
+    STD_TORCH_CHECK(sorted_expert_ids->size(0) == m_topk);
+    STD_TORCH_CHECK(sorted_expert_ids->get_device_index() ==
+                    input.get_device_index());
+  }
 
   STD_TORCH_CHECK(output.dim() == 2);
   STD_TORCH_CHECK(output_scale.dim() == 2);
@@ -393,7 +415,8 @@ void mxfp4_experts_quant(
     torch::stable::Tensor const& input,
     torch::stable::Tensor const& input_offset_by_experts,
     torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
+    int64_t n_experts,
+    std::optional<torch::stable::Tensor> const& sorted_expert_ids) {
 #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
   int32_t sm = get_sm_version_num();
   STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
@@ -405,7 +428,7 @@ void mxfp4_experts_quant(
 
   validate_mxfp4_experts_quant_inputs(
       output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
+      output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k);
 
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -417,8 +440,11 @@ void mxfp4_experts_quant(
         vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
             output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
             input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
+            output_scale_offset_by_experts.data_ptr(),
+            sorted_expert_ids.has_value()
+                ? static_cast<int32_t const*>(sorted_expert_ids->data_ptr())
+                : nullptr,
+            m_topk, k, n_experts, stream);
       });
 #else
   STD_TORCH_CHECK_NOT_IMPLEMENTED(false,
@@ -431,7 +457,8 @@ void silu_and_mul_mxfp4_experts_quant(
     torch::stable::Tensor const& input,
     torch::stable::Tensor const& input_offset_by_experts,
     torch::stable::Tensor const& output_scale_offset_by_experts,
-    int64_t n_experts) {
+    int64_t n_experts,
+    std::optional<torch::stable::Tensor> const& sorted_expert_ids) {
 #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED
   int32_t sm = get_sm_version_num();
   STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm),
@@ -445,7 +472,7 @@ void silu_and_mul_mxfp4_experts_quant(
 
   validate_mxfp4_experts_quant_inputs(
       output, output_scale, input, input_offset_by_experts,
-      output_scale_offset_by_experts, n_experts, m_topk, k);
+      output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k);
 
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
@@ -457,8 +484,11 @@ void silu_and_mul_mxfp4_experts_quant(
         vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
             output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
             input_offset_by_experts.data_ptr(),
-            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
-            stream);
+            output_scale_offset_by_experts.data_ptr(),
+            sorted_expert_ids.has_value()
+                ? static_cast<int32_t const*>(sorted_expert_ids->data_ptr())
+                : nullptr,
+            m_topk, k, n_experts, stream);
       });
 #else
   STD_TORCH_CHECK_NOT_IMPLEMENTED(

csrc/quantization/w8a8/cutlass/moe/moe_data.cu

@@ -88,6 +88,7 @@ __global__ void compute_expert_blockscale_offsets(
 __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
                                   int32_t* input_permutation,
                                   int32_t* output_permutation,
+                                  int32_t* sorted_expert_ids,
                                   int32_t* atomic_buffer, const int topk_length,
                                   const int topk) {
   int const blk_expert_id = blockIdx.x;
@@ -108,6 +109,9 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
       int start = atomicAdd(&atomic_buffer[expert_id], 1);
       input_permutation[start] = i / topk;
       output_permutation[i] = start;
+      if (sorted_expert_ids != nullptr) {
+        sorted_expert_ids[start] = expert_id;
+      }
     }
   }
 }
@@ -242,7 +246,8 @@ void get_cutlass_moe_mm_data_caller(
     torch::stable::Tensor& output_permutation, const int64_t num_experts,
     const int64_t n, const int64_t k,
     const std::optional<torch::stable::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const bool is_gated,
+    const std::optional<torch::stable::Tensor>& sorted_expert_ids) {
   auto device = topk_ids.device();
   const torch::stable::accelerator::DeviceGuard device_guard(device.index());
   auto stream = get_current_cuda_stream(device.index());
@@ -278,6 +283,9 @@ void get_cutlass_moe_mm_data_caller(
       static_cast<const int32_t*>(topk_ids.data_ptr()),
       static_cast<int32_t*>(input_permutation.data_ptr()),
       static_cast<int32_t*>(output_permutation.data_ptr()),
+      sorted_expert_ids.has_value()
+          ? static_cast<int32_t*>(sorted_expert_ids->data_ptr())
+          : nullptr,
       static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
       topk_ids.size(1));
 }

csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu

@@ -97,7 +97,8 @@ void get_cutlass_moe_mm_data_caller(
     torch::stable::Tensor& output_permutation, const int64_t num_experts,
     const int64_t n, const int64_t k,
     const std::optional<torch::stable::Tensor>& blockscale_offsets,
-    const bool is_gated);
+    const bool is_gated,
+    const std::optional<torch::stable::Tensor>& sorted_expert_ids);
 
 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
     const torch::stable::Tensor& expert_first_token_offset,
@@ -337,17 +338,18 @@ void get_cutlass_moe_mm_data(
     torch::stable::Tensor& output_permutation, const int64_t num_experts,
     const int64_t n, const int64_t k,
     const std::optional<torch::stable::Tensor>& blockscale_offsets,
-    const bool is_gated) {
+    const bool is_gated,
+    const std::optional<torch::stable::Tensor>& sorted_expert_ids) {
   // This function currently gets compiled only if we have a valid cutlass moe
   // mm to run it for.
   int32_t version_num = get_sm_version_num();
 #if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
     (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
     (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
-  get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
-                                 problem_sizes2, input_permutation,
-                                 output_permutation, num_experts, n, k,
-                                 blockscale_offsets, is_gated);
+  get_cutlass_moe_mm_data_caller(
+      topk_ids, expert_offsets, problem_sizes1, problem_sizes2,
+      input_permutation, output_permutation, num_experts, n, k,
+      blockscale_offsets, is_gated, sorted_expert_ids);
   return;
 #endif
   STD_TORCH_CHECK_NOT_IMPLEMENTED(

csrc/torch_bindings.cpp

@@ -151,7 +151,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "                        Tensor! input_permutation, "
       "                        Tensor! output_permutation, int num_experts, "
       "                        int n, int k, Tensor? blockscale_offsets, "
-      "                        bool is_gated) -> ()");
+      "                        bool is_gated, "
+      "                        Tensor!? sorted_expert_ids=None) -> ()");
 
   // compute per-expert problem sizes from expert_first_token_offset
   // produced by vLLM's moe_permute kernel
@@ -235,14 +236,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
       "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
+      "Tensor output_scale_offset_by_experts, int n_experts, "
+      "Tensor? sorted_expert_ids=None) -> ()");
 
   // Fused SiLU+Mul+MXFP4 experts quantization.
   ops.def(
       "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
       "output_scale,"
       "Tensor input, Tensor input_offset_by_experts,"
-      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
+      "Tensor output_scale_offset_by_experts, int n_experts, "
+      "Tensor? sorted_expert_ids=None) -> ()");
 
   // Fused SiLU+Mul+NVFP4 quantization.
   ops.def(

tests/kernels/moe/test_mxfp4_moe.py

@@ -19,6 +19,73 @@ set_random_seed(42)
 MXFP4_BLOCK_SIZE = 32
 
 
+def _mxfp4_experts_quant_for_test(
+    input_tensor: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    n_experts: int,
+    sorted_expert_ids: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Call the raw quant op with a compact, zero-initialized scale buffer."""
+    m_numtopk, k = input_tensor.shape
+    scales_k = k // MXFP4_BLOCK_SIZE
+    padded_k = (scales_k + 3) // 4
+    num_sf_rows = int(blockscale_offsets[-1].item())
+    output = torch.empty(
+        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
+    )
+    output_scales = torch.zeros(
+        num_sf_rows,
+        padded_k,
+        dtype=torch.int32,
+        device=input_tensor.device,
+    )
+    torch.ops._C.mxfp4_experts_quant(
+        output,
+        output_scales,
+        input_tensor,
+        expert_offsets,
+        blockscale_offsets,
+        n_experts,
+        sorted_expert_ids,
+    )
+    return output, output_scales.view(torch.uint8)
+
+
+def _silu_and_mul_mxfp4_experts_quant_for_test(
+    input_tensor: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    n_experts: int,
+    sorted_expert_ids: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Call the raw fused quant op with a compact, zero-initialized scale buffer."""
+    m_numtopk, k_times_2 = input_tensor.shape
+    k = k_times_2 // 2
+    scales_k = k // MXFP4_BLOCK_SIZE
+    padded_k = (scales_k + 3) // 4
+    num_sf_rows = int(blockscale_offsets[-1].item())
+    output = torch.empty(
+        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
+    )
+    output_scales = torch.zeros(
+        num_sf_rows,
+        padded_k,
+        dtype=torch.int32,
+        device=input_tensor.device,
+    )
+    torch.ops._C.silu_and_mul_mxfp4_experts_quant(
+        output,
+        output_scales,
+        input_tensor,
+        expert_offsets,
+        blockscale_offsets,
+        n_experts,
+        sorted_expert_ids,
+    )
+    return output, output_scales.view(torch.uint8)
+
+
 def align(val: int, alignment: int = 128) -> int:
     return int((val + alignment - 1) // alignment * alignment)
 
@@ -458,5 +525,155 @@ def test_mxfp4_experts_quant_no_saturation():
     )
 
 
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="get_cutlass_moe_mm_data sorted expert IDs require CUDA SM100 MoE",
+)
+def test_get_cutlass_moe_mm_data_populates_sorted_expert_ids():
+    num_experts = 32
+    m, topk = 8, 4
+    topk_ids = torch.randint(
+        0, num_experts, (m, topk), device="cuda", dtype=torch.int32
+    )
+    expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32)
+    blockscale_offsets = torch.empty_like(expert_offsets)
+    problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32)
+    problem_sizes2 = torch.empty_like(problem_sizes1)
+    input_permutation = torch.empty(m * topk, device="cuda", dtype=torch.int32)
+    output_permutation = torch.empty_like(input_permutation)
+    sorted_expert_ids = torch.empty_like(input_permutation)
+
+    ops.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        128,
+        512,
+        blockscale_offsets,
+        is_gated=True,
+        sorted_expert_ids=sorted_expert_ids,
+    )
+
+    expected = torch.sort(topk_ids.flatten(), stable=True).values
+    torch.testing.assert_close(sorted_expert_ids, expected)
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="mxfp4_experts_quant requires CUDA SM100",
+)
+@pytest.mark.parametrize("num_experts", [16, 64, 256])
+def test_mxfp4_experts_quant_sorted_expert_ids_matches_offset_search(
+    num_experts: int,
+):
+    """Sorted expert IDs must match the offset-search quantization path."""
+    set_random_seed(0)
+    m, topk, k, n = 32, 1, 512, 128
+    topk_ids = torch.randint(
+        0, num_experts, (m, topk), device="cuda", dtype=torch.int32
+    )
+    expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32)
+    blockscale_offsets = torch.empty_like(expert_offsets)
+    problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32)
+    problem_sizes2 = torch.empty_like(problem_sizes1)
+    input_permutation = torch.empty(m * topk, device="cuda", dtype=torch.int32)
+    output_permutation = torch.empty_like(input_permutation)
+    sorted_expert_ids = torch.empty_like(input_permutation)
+
+    ops.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        n,
+        k,
+        blockscale_offsets,
+        is_gated=True,
+        sorted_expert_ids=sorted_expert_ids,
+    )
+
+    activations = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) / 5
+    sorted_activations = ops.shuffle_rows(activations, input_permutation)
+
+    baseline_out, baseline_sf = _mxfp4_experts_quant_for_test(
+        sorted_activations,
+        expert_offsets,
+        blockscale_offsets,
+        num_experts,
+    )
+    direct_out, direct_sf = _mxfp4_experts_quant_for_test(
+        sorted_activations,
+        expert_offsets,
+        blockscale_offsets,
+        num_experts,
+        sorted_expert_ids=sorted_expert_ids,
+    )
+
+    assert torch.equal(baseline_out, direct_out)
+    assert torch.equal(baseline_sf, direct_sf)
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="silu_and_mul_mxfp4_experts_quant requires CUDA SM100",
+)
+def test_silu_and_mul_mxfp4_sorted_expert_ids_matches_offset_search():
+    set_random_seed(1)
+    num_experts = 64
+    m, topk, k, n = 32, 1, 256, 128
+    topk_ids = torch.randint(
+        0, num_experts, (m, topk), device="cuda", dtype=torch.int32
+    )
+    expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32)
+    blockscale_offsets = torch.empty_like(expert_offsets)
+    problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32)
+    problem_sizes2 = torch.empty_like(problem_sizes1)
+    input_permutation = torch.empty(m * topk, device="cuda", dtype=torch.int32)
+    output_permutation = torch.empty_like(input_permutation)
+    sorted_expert_ids = torch.empty_like(input_permutation)
+
+    ops.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        input_permutation,
+        output_permutation,
+        num_experts,
+        n,
+        k,
+        blockscale_offsets,
+        is_gated=True,
+        sorted_expert_ids=sorted_expert_ids,
+    )
+
+    gate_up = torch.randn(m, 2 * k, device="cuda", dtype=torch.bfloat16) / 5
+    sorted_gate_up = ops.shuffle_rows(gate_up, input_permutation)
+
+    baseline_out, baseline_sf = _silu_and_mul_mxfp4_experts_quant_for_test(
+        sorted_gate_up,
+        expert_offsets,
+        blockscale_offsets,
+        num_experts,
+    )
+    direct_out, direct_sf = _silu_and_mul_mxfp4_experts_quant_for_test(
+        sorted_gate_up,
+        expert_offsets,
+        blockscale_offsets,
+        num_experts,
+        sorted_expert_ids=sorted_expert_ids,
+    )
+
+    assert torch.equal(baseline_out, direct_out)
+    assert torch.equal(baseline_sf, direct_sf)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

vllm/_custom_ops.py

@@ -883,6 +883,7 @@ def get_cutlass_moe_mm_data(
     k: int,
     blockscale_offsets: torch.Tensor | None = None,
     is_gated: bool = True,
+    sorted_expert_ids: torch.Tensor | None = None,
 ):
     """Prepare data necessary to perform CUTLASS grouped matrix multiplications
     used in CUTLASS-based fused MoE.
@@ -907,6 +908,8 @@ def get_cutlass_moe_mm_data(
                           blockscale_offsets[E]
     - is_gated: Whether the activation is gated (gate + up). When True, the
                 first GEMM N dimension is 2*n; when False, it is n.
+    - sorted_expert_ids: Optional expert ID for each row after sorting, used to
+                         skip expert-offset searches in MXFP4 conversion.
     """
     return torch.ops._C.get_cutlass_moe_mm_data(
         topk_ids,
@@ -920,6 +923,7 @@ def get_cutlass_moe_mm_data(
         k,
         blockscale_offsets,
         is_gated,
+        sorted_expert_ids,
     )
 
 
@@ -1736,6 +1740,7 @@ def mxfp4_experts_quant(
     blockscale_offsets: torch.Tensor,
     n_experts: int,
     topk: int,
+    sorted_expert_ids: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Quantize input tensor to MXFP4 for packed MoE inputs.
     Uses 32-element blocks with E8M0 (power-of-two) scale factors.
@@ -1747,6 +1752,9 @@ def mxfp4_experts_quant(
         blockscale_offsets: [n_experts+1] SF row boundaries per expert
         n_experts: number of experts
         topk: number of top-k experts
+        sorted_expert_ids: Optional expert ID per sorted row; skips offset
+            lookup in the CUDA kernel when provided.
+
     Returns:
         output: [m_topk, k//2] packed E2M1 values (uint8)
         output_scales: E8M0 blockscales in swizzled layout (uint8 view)
@@ -1783,6 +1791,7 @@ def mxfp4_experts_quant(
         expert_offsets,
         blockscale_offsets,
         n_experts,
+        sorted_expert_ids,
     )
     # E8M0 SFs are stored as uint8
     output_scales = output_scales.view(torch.uint8)
@@ -1795,9 +1804,24 @@ def silu_and_mul_mxfp4_experts_quant(
     blockscale_offsets: torch.Tensor,
     n_experts: int,
     topk: int,
+    sorted_expert_ids: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
     MXFP4 has no global scale - only block-level E8M0 scale factors.
+
+    Args:
+        input_tensor: [m_topk, 2*k] BF16/FP16 gate||up activations
+        expert_offsets: [n_experts+1] token boundaries per expert
+        blockscale_offsets: [n_experts+1] SF row boundaries per expert
+        n_experts: number of experts
+        topk: number of top-k experts
+        sorted_expert_ids: Optional expert ID per sorted row; skips offset
+            lookup in the CUDA kernel when provided.
+
+    Returns:
+        output: [m_topk, k//2] packed E2M1 values (uint8)
+        output_scales: E8M0 blockscales in swizzled layout (uint8 view)
+
     """
     assert not current_platform.is_rocm()
     assert input_tensor.ndim == 2
@@ -1827,6 +1851,7 @@ def silu_and_mul_mxfp4_experts_quant(
         expert_offsets,
         blockscale_offsets,
         n_experts,
+        sorted_expert_ids,
     )
     output_scales = output_scales.view(torch.uint8)
     return output, output_scales

vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py

@@ -890,6 +890,11 @@ def run_cutlass_moe_mxfp4(
 
     a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
     c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
+    # Below 17 experts the quant kernel uses a cheap linear offset scan; skip
+    # sorted_expert_ids to avoid extra routing writes and ID lookups.
+    sorted_expert_ids: torch.Tensor | None = None
+    if e > 16:
+        sorted_expert_ids = torch.empty_like(a_map)
 
     if apply_router_weight_on_input:
         assert num_topk == 1, (
@@ -909,6 +914,7 @@ def run_cutlass_moe_mxfp4(
         k,
         blockscale_offsets,
         is_gated=is_gated,
+        sorted_expert_ids=sorted_expert_ids,
     )
 
     a = ops.shuffle_rows(a, a_map)
@@ -918,6 +924,7 @@ def run_cutlass_moe_mxfp4(
         blockscale_offsets,
         e,
         num_topk,
+        sorted_expert_ids,
     )
     c1 = _resize_cache(workspace13, (m * topk, w1_n))
     c2 = _resize_cache(workspace2, (m * topk, n))
@@ -938,7 +945,12 @@ def run_cutlass_moe_mxfp4(
         activation_config is None or activation_config.clamp_limit is None
     ):
         int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
-            c1, expert_offsets, blockscale_offsets, e, num_topk
+            c1,
+            expert_offsets,
+            blockscale_offsets,
+            e,
+            num_topk,
+            sorted_expert_ids,
         )
     else:
         apply_moe_activation(
@@ -948,7 +960,12 @@ def run_cutlass_moe_mxfp4(
             activation_config=activation_config,
         )
         int_fp4, int_blockscale = ops.mxfp4_experts_quant(
-            c2, expert_offsets, blockscale_offsets, e, num_topk
+            c2,
+            expert_offsets,
+            blockscale_offsets,
+            e,
+            num_topk,
+            sorted_expert_ids,
         )
 
     ops.cutlass_mxfp4_moe_mm(