[Performance]: The default MXFP4 quant backend of FlashInfer has slow performance, add options for backend "cute-dsl".

[wheird-lee]

Summary

When benchmarking Qwen3-8B-like GEMM shapes with FlashInfer, I observed that MXFP4 GEMM itself provides consistent speedup over BF16, and is significantly faster than MXFP8. However, once activation quantization time is included, the default quantization backend ("cuda") causes MXFP4 performance to drop sharply. In most shapes, MXFP4 even becomes much slower than MXFP8.

Switching the activation quantization backend from "cuda" to "cute-dsl" largely fixes this issue: although activation quantization still introduces overhead, MXFP4 again becomes clearly faster than MXFP8 and often remains faster than BF16.

This suggests that the "cuda" backend of MXFP4 quantization kernel is likely the main bottleneck. (Related issue: flashinfer-ai/flashinfer#2496)

vLLM currently (v0.25) uses the "cuda" backend as the default backend for MXFP4 quantization when using FlashInfer and does not provide an option to use "cute-dsl".

Benchmark setup

The following benchmark uses FlashInfer's mxfp gemm kernels and quantization kernels, including:

• mm_fp4 and mm_mxfp8
• mxfp4_quantize and mxfp8_quantize

The benchmark covers four Qwen3-8B GEMM shapes:

• K=4096, N=24576 — mlp.gateup_proj
• K=12288, N=4096 — mlp.down_proj
• K=4096, N=6144 — self_attn.qkv_proj
• K=4096, N=4096 — self_attn.o_proj

For each shape, multiple batch sizes were tested. Results were compared across three modes:

1. no-quant-a: GEMM time only, excluding activation quantization overhead
2. cuda-backend: GEMM + activation quantization using CUDA backend
3. cute-dsl-backend: GEMM + activation quantization using cute-dsl backend

Formats compared:

• BF16
• MXFP8
• MXFP4

All benchmarks were done independently on NVIDIA B200 GPU.

Observations

1. GEMM-only results show MXFP4 is consistently beneficial

When activation quantization overhead is excluded, MXFP4 is faster than BF16 in almost all cases.

This indicates that the MXFP4 GEMM path itself is efficient.

2. Activation quantization with "cuda" backend severely hurts MXFP4

When activation quantization is included and the CUDA backend is used, MXFP4 performance drops dramatically.

The only shape where CUDA-backend MXFP4 still shows clear benefit is K=4096, N=24576. For the other three shapes, MXFP4 is not only slower than BF16, but also significantly slower than MXFP8.

This is unexpected because MXFP4 should generally have lower memory bandwidth and compute cost than MXFP8.

3. cute-dsl backend restores MXFP4 performance

After switching activation quantization to the cute-dsl backend, MXFP4 performance improves substantially.

Although activation quantization overhead still reduces the overall speed compared with GEMM-only results, MXFP4 remains clearly faster than MXFP8 under the cute-dsl backend.

Conclusion

The benchmark results suggest that the "cuda" quant backend appears to be the main bottleneck for inference of MXFP4 quantized models when using FlashInfer backend.

In particular:

• MXFP4 GEMM-only performance is strong.
• "cuda"-backend activation quantization makes MXFP4 much slower for most shapes.
• "cute-dsl"-backend activation quantization restores the expected MXFP4 advantage.

Suggestions

vLLM currently (v0.25) uses the "cuda" backend as the default backend for FlashInfer MXFP4 quantization (see this code). It is suggested to add the "cute-dsl" backend as an alternative.

Report of performance regression

No response

Misc discussion on performance

No response

Your current environment (if you think it is necessary)

The output of `python collect_env.py`

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[wheird-lee]

benchmarking code:

"""Benchmark BF16, MXFP8, and MXFP4 performance on fixed GEMM shapes.

Shapes use the MxKxN format, corresponding to
A[M, K] @ B[N, K].T -> C[M, N].

Example:
    python benchmark_llm_gemm_pure.py \
        --formats bf16,mxfp8,mxfp4
"""

import argparse
import csv
import json
import os
from datetime import datetime

import torch
from flashinfer import mxfp4_quantize, mxfp8_quantize
from flashinfer.gemm import mm_fp4, mm_mxfp8
from vllm.triton_utils import triton


QUANT_BACKEND = "cute-dsl"  # "cuda"
BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
Qwen3_8B_GEMM_SHAPES = [
    # (K, N)
    (4096,  24567), # mlp.gateup_proj
    (12288, 4096),  # mlp.down_proj
    (4096,  6144),  # self_attn.qkv_proj
    (4096,  4096),  # self_attn.o_proj
]
GEMM_SHAPES = [
    (batch_size, k, n)
    for batch_size in sorted(BATCH_SIZES)
    for k, n in Qwen3_8B_GEMM_SHAPES
]
FORMAT_LABELS = {
    "bf16": "BF16 (w16a16)",
    "mxfp8": "MXFP8 (w8a8)",
    "mxfp4": "MXFP4 (w4a4)",
}


def build_bf16_runner(a, b, no_a_quant=False):
    """Build the BF16 baseline runner."""
    del no_a_quant
    return lambda: torch.nn.functional.linear(a, b)


def build_mxfp8_runner(a, b, no_a_quant=False):
    """Build an MXFP8 GEMM runner with pre-quantized weights.

    Activations can be quantized either on the fly or in advance.
    """
    b_q, b_scale = mxfp8_quantize(
        b, is_sf_swizzled_layout=True, backend=QUANT_BACKEND
    )

    def quantize_a():
        return mxfp8_quantize(
            a, is_sf_swizzled_layout=True, backend=QUANT_BACKEND
        )

    if no_a_quant:
        a_q, a_scale = quantize_a()
        get_a = lambda: (a_q, a_scale)
    else:
        get_a = quantize_a

    def run():
        current_a, current_scale = get_a()
        return mm_mxfp8(
            current_a,
            b_q.t(),
            current_scale,
            b_scale,
            out_dtype=torch.bfloat16,
            backend="auto",
        )

    return run


def build_mxfp4_runner(a, b, no_a_quant=False):
    """Build an MXFP4 GEMM runner with pre-quantized weights.

    Activations can be quantized either on the fly or in advance.
    """
    b_q, b_scale = mxfp4_quantize(b, backend=QUANT_BACKEND)
    alpha = torch.tensor([1.0], device=a.device)

    def quantize_a():
        return mxfp4_quantize(a, backend=QUANT_BACKEND)

    if no_a_quant:
        a_q, a_scale = quantize_a()
        get_a = lambda: (a_q, a_scale)
    else:
        get_a = quantize_a

    def run():
        current_a, current_scale = get_a()
        return mm_fp4(
            current_a,
            b_q.t(),
            current_scale,
            b_scale.t(),
            alpha=alpha,
            out_dtype=torch.bfloat16,
            backend="auto",
            use_nvfp4=False,
            block_size=32,
        )

    return run


RUNNER_BUILDERS = {
    "bf16": build_bf16_runner,
    "mxfp8": build_mxfp8_runner,
    "mxfp4": build_mxfp4_runner,
}


def parse_formats(value):
    formats = [x.strip().lower() for x in value.split(",")]
    unknown = [x for x in formats if x not in RUNNER_BUILDERS]
    if not formats or unknown:
        raise argparse.ArgumentTypeError(
            f"Formats must be selected from: {', '.join(RUNNER_BUILDERS)}"
        )
    return formats


def get_args():
    parser = argparse.ArgumentParser(
        description="Benchmark GEMM performance for fixed MxKxN shapes.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--formats",
        type=parse_formats,
        default=["bf16", "mxfp8", "mxfp4"],
        help="comma-separated formats",
    )
    activation_group = parser.add_mutually_exclusive_group()
    activation_group.add_argument(
        "--quant_act_on_the_fly",
        action="store_true",
        dest="quant_act_on_the_fly",
        default=True,
        help="include activation quantization in latency",
    )
    activation_group.add_argument(
        "--no_quant_act_on_the_fly",
        action="store_false",
        dest="quant_act_on_the_fly",
        help="pre-quantize activation and exclude quantization from latency",
    )
    parser.add_argument("--rep", type=int, default=200)
    parser.add_argument("--output_dir", default="./benchmark_results")
    args = parser.parse_args()
    if args.rep <= 0:
        parser.error("--rep must be positive")
    return args


def run_single_benchmark(m, k, n, fmt, rep, quant_act_on_the_fly):
    """Benchmark one shape/format pair and return latency and TFLOP/s."""
    a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
    b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
    runner = RUNNER_BUILDERS[fmt](
        a, b, no_a_quant=not quant_act_on_the_fly
    )
    p50, p20, p80 = triton.testing.do_bench_cudagraph(
        runner, rep=rep, quantiles=[0.5, 0.2, 0.8]
    )

    flops = 2 * m * k * n
    to_tflops = lambda ms: flops * 1e-9 / ms
    return {
        "latency_ms": round(p50, 6),
        "latency_ms_p20": round(p20, 6),
        "latency_ms_p80": round(p80, 6),
        "tflops": round(to_tflops(p50), 3),
        # Throughput percentiles invert the corresponding latency percentiles.
        "tflops_p20": round(to_tflops(p80), 3),
        "tflops_p80": round(to_tflops(p20), 3),
    }


def collect_results(formats, rep, quant_act_on_the_fly):
    results = []
    total = len(GEMM_SHAPES) * len(formats)

    for index, ((m, k, n), fmt) in enumerate(
        ((shape, fmt) for shape in GEMM_SHAPES for fmt in formats), start=1
    ):
        print(
            f"\r[{index}/{total}] M={m} K={k} N={n} "
            f"{FORMAT_LABELS[fmt]}",
            end="",
            flush=True,
        )
        row = {
            "shape": f"{m}x{k}x{n}",
            "M": m,
            "K": k,
            "N": n,
            "format": FORMAT_LABELS[fmt],
            "format_key": fmt,
        }
        try:
            row.update(
                run_single_benchmark(
                    m, k, n, fmt, rep, quant_act_on_the_fly
                )
            )
        except torch.cuda.OutOfMemoryError:
            print(f"\n[OOM] {row['shape']} {fmt}")
            torch.cuda.empty_cache()
            row.update(latency_ms=None, tflops=None, error="OOM")
        except Exception as error:
            print(f"\n[ERR] {row['shape']} {fmt}: {error}")
            row.update(latency_ms=None, tflops=None, error=str(error))
        results.append(row)

    print()
    return results


def save_results(results, formats, output_dir, quant_act_on_the_fly):
    os.makedirs(output_dir, exist_ok=True)
    csv_path = os.path.join(output_dir, "gemm_all_results.csv")
    fields = [
        "shape", "M", "K", "N", "format", "format_key",
        "latency_ms", "latency_ms_p20", "latency_ms_p80",
        "tflops", "tflops_p20", "tflops_p80", "error",
    ]
    with open(csv_path, "w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(file, fieldnames=fields)
        writer.writeheader()
        writer.writerows(
            {field: row.get(field, "") for field in fields}
            for row in results
        )

    summary = {
        "benchmark_time": datetime.now().isoformat(),
        "shapes": [
            {"M": m, "K": k, "N": n} for m, k, n in GEMM_SHAPES
        ],
        "formats": formats,
        "quant_act_on_the_fly": quant_act_on_the_fly,
        "num_results": len(results),
        "num_errors": sum(bool(row.get("error")) for row in results),
    }
    json_path = os.path.join(output_dir, "gemm_summary.json")
    with open(json_path, "w", encoding="utf-8") as file:
        json.dump(summary, file, indent=2, ensure_ascii=False)

    print(f"[SAVED] {csv_path}")
    print(f"[SAVED] {json_path}")


def main():
    args = get_args()
    print(f"GEMM (K, N): {Qwen3_8B_GEMM_SHAPES}")
    print(f"Batch Sizes: {BATCH_SIZES}")
    print(f"Formats: {args.formats}")
    print(
        "Activation quantization: "
        + ("on-the-fly" if args.quant_act_on_the_fly else "pre-quantized")
    )

    results = collect_results(
        args.formats,
        args.rep,
        args.quant_act_on_the_fly,
    )
    save_results(
        results,
        args.formats,
        args.output_dir,
        args.quant_act_on_the_fly,
    )


if __name__ == "__main__":
    main()

[BWAAEEEK]

I reproduced this on an NVIDIA B200 (SM100) using vLLM main at 83762b77, FlashInfer 0.6.13, PyTorch 2.11.0+cu130, and CUDA 13.0.

The current FlashInfer MXFP4 linear path uses the default cuda backend for activation quantization, while the following MXFP4 GEMM is already fixed to cute-dsl.

I first ran an independent A/B benchmark that changed only the activation quantization backend:

• 16 quantization shapes covering M=1..4096, K={4096, 12288}, including the M=127/128/129 padding boundary
• packed FP4 values and scales were bitwise identical for all shapes, in both eager execution and CUDA Graph replay
• quantization latency geometric-mean speedup: 4.74x eager, 6.13x CUDA Graph
• corrected Qwen3-8B-like linear shapes:
  • (K=4096, N=4096)
  • (K=4096, N=6144)
  • (K=4096, N=24576)
  • (K=12288, N=4096)
• quantization + the same cute-dsl GEMM had bitwise-identical outputs across all 20 tested shape/M combinations
• combined pipeline geometric-mean speedup: 1.78x eager, 2.72x CUDA Graph

I also validated the actual vLLM FlashInferMxFp4LinearKernel path with nm-testing/Meta-Llama-3-8B-Instruct-MXFP4-GPTQ at revision 43e77d3bd4586daba9f29925afabfebad0ea3871.

Configuration: 64 prompts, 512 input tokens, 64 output tokens, greedy decoding, synchronous scheduling, prefix caching disabled, FlashInfer autotuning disabled, linear_backend=flashinfer_cutlass, five trials. Median results:

Mode	CUDA quant	CuTe-DSL quant	Change
Eager output throughput	1061.6 tok/s	1714.5 tok/s	+61.5%
Eager mean TTFT	311.4 ms	175.0 ms	-43.8%
CUDA Graph output throughput	4282.4 tok/s	6823.0 tok/s	+59.3%
CUDA Graph mean TTFT	292.6 ms	153.9 ms	-47.4%

Correctness checks:

• CUDA and CuTe-DSL quantized values/scales: bitwise identical
• combined quantization + GEMM outputs: bitwise identical
• 8B single-request greedy generation: all 128 generated token IDs identical
• existing test_compressed_tensors_mxfp4: passed with the CuTe-DSL quantizer
• concurrent 8B runs showed the existing batch-numerics variation in both the baseline and changed runs, so I did not use concurrent token equality as correctness evidence

Since FlashInferMxFp4LinearKernel is already selected only when SM100+ and FlashInfer CuTe-DSL are available, and its GEMM already uses backend="cute-dsl", the smallest change appears to be explicitly using cute-dsl for its activation quantization as well, rather than adding a new public backend option.

I can prepare a PR with the implementation, regression coverage, and these benchmark results. Please let me know if a configurable backend is preferred instead.