From f3bdb0e02515b972372a50844a2d7ca65eed209b Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 26 Oct 2026 21:51:55 +0100 Subject: [PATCH] [Perf][MXFP4] Optimize activation quant for mxfp4. 3x kernel speedup, 5% throughput improvement (#63101) Signed-off-by: ilmarkov Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> --- ...enchmark_mxfp4_experts_activation_quant.py | 538 ++++++++++++++++++ csrc/ops.h | 3 +- csrc/quantization/fp4/mxfp4_experts_quant.cu | 90 ++- .../quantization/w8a8/cutlass/moe/moe_data.cu | 10 +- .../w8a8/cutlass/scaled_mm_entry.cu | 14 +- csrc/torch_bindings.cpp | 9 +- tests/kernels/moe/test_mxfp4_moe.py | 217 +++++++ vllm/_custom_ops.py | 25 + .../layers/fused_moe/experts/cutlass_moe.py | 21 +- 9 files changed, 884 insertions(+), 43 deletions(-) create mode 100644 benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py diff --git a/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py b/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py new file mode 100644 index 00000000000..12f7571e93c --- /dev/null +++ b/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() diff --git a/csrc/ops.h b/csrc/ops.h index b137d7db041..1c5ddc0beb4 100644 --- a/csrc/ops.h +++ b/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& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& sorted_expert_ids); void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( const torch::stable::Tensor& expert_first_token_offset, diff --git a/csrc/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu index 4c6875f0623..1f9fd7c8a16 100644 --- a/csrc/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -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; 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; 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 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(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts); + n_experts, sorted_expert_ids); } else { mxfp4_cvt_fp16_to_fp4 <<>>( @@ -297,7 +308,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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<<>>( m_topk, k, reinterpret_cast(input), @@ -316,7 +328,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(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 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 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( 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(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 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( 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(sorted_expert_ids->data_ptr()) + : nullptr, + m_topk, k, n_experts, stream); }); #else STD_TORCH_CHECK_NOT_IMPLEMENTED( diff --git a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu index b6043e80330..268650ae181 100644 --- a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/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& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& 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(topk_ids.data_ptr()), static_cast(input_permutation.data_ptr()), static_cast(output_permutation.data_ptr()), + sorted_expert_ids.has_value() + ? static_cast(sorted_expert_ids->data_ptr()) + : nullptr, static_cast(atomic_buffer.data_ptr()), topk_ids.numel(), topk_ids.size(1)); } diff --git a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 7e4f942f544..0fd72ebd082 100644 --- a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/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& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& 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& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& 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( diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 4e9530cf3b6..1c84b53ece2 100644 --- a/csrc/torch_bindings.cpp +++ b/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( diff --git a/tests/kernels/moe/test_mxfp4_moe.py b/tests/kernels/moe/test_mxfp4_moe.py index 7c7ca108857..8bc328ba9f7 100644 --- a/tests/kernels/moe/test_mxfp4_moe.py +++ b/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"]) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index effb39f7ac1..09a45982698 100644 --- a/vllm/_custom_ops.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py index 65d7499ffae..95f113015fa 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py +++ b/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(