vLLM MoE Triton Tuner

This repository contains a small patch to benchmarks/kernels/benchmark_moe.py that makes the MoE Triton tuner practical for single-GPU setups and older vLLM APIs (tested with vLLM 0.10.1.1).

It narrows the search space, avoids API mismatches, and writes the tuned config JSON directly into vLLM’s built-in configs/ directory so vLLM will auto-load it at model startup.

What changed (and why)

Target version

• The script matches vLLM 0.10.1.1 APIs. We avoid newer helpers (FusedMoEQuantConfig.make, _get_config_dtype_str, etc.) that don’t exist in 0.10.1.1 and instead use the functions and call sites available in this release.

Safer dtype handling

• Older configs don’t expose config.dtype. We use config.torch_dtype (or torch.float16 on ROCm) when benchmarking to avoid AttributeError.

Minimal tuning space (fast)

• Reduced CUDA search ranges (BLOCK_M/N/K, warps, stages) and optional “tiny” space to cut total configs from ~1900 down to 8–64 (minutes vs. many hours). This still finds good kernels for common MoE sizes.

Writes to the right place

• Saves the resulting JSON to the installed vLLM package directory, e.g. /usr/local/lib/python3.11/dist-packages/vllm/model_executor/layers/fused_moe/configs/ so the runtime can find it automatically (you’ll see “Using default MoE config” disappear in logs).

No Ray dashboard

• Avoids slow/failed dashboard startup noise in containers.

Single-GPU friendly

• Defaults to one Ray worker bound to one GPU so the JIT is compiled exactly where the kernel will run.

DeepGEMM off by default

• Focuses the tuner on Triton kernels (DeepGEMM can be toggled if you really need it, but was a source of confusion/OOR in this environment).

Expected results

• A file named like: E=<experts>,N=<intermediate/2>,device_name=<autodetected>.json appears in vLLM’s configs/ directory when tuning completes.
• Subsequent vLLM startups for that model+GPU print no “Using default MoE config” warning; throughput and queueing behavior should improve noticeably for MoE models.
• In our tests this change unlocked stable multi-request concurrency and ~tens of tokens/s on problematic prompts.

File to edit/replace

Edit (or replace) this file in your local vLLM checkout:

benchmarks/kernels/benchmark_moe.py

⚠️ Keep the file path the same; the CLI invocation expects it.

Get the right vLLM source

Clone the vLLM repo and check out the branch/tag that matches your runtime (here we use v0.10.1.1 as an example):

git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.10.1.1

Now apply/overwrite benchmarks/kernels/benchmark_moe.py with the patched version from this repo.

How to run (examples)

These examples assume:

• Python at /usr/bin/python3
• You’re inside a container/host that’s already running vLLM 0.10.1.1
• You want the tuned JSON to land in vLLM’s installed package directory

Environment (recommended)

export PYTHONNOUSERSITE=1
export RAY_DISABLE_DASHBOARD=1
CFG_DIR="/usr/local/lib/python3.11/dist-packages/vllm/model_executor/layers/fused_moe/configs"
mkdir -p "$CFG_DIR"

Example 1 — Qwen3-Coder-REAP-25B-A3B (one GPU, tiny search)

/usr/bin/python3 benchmarks/kernels/benchmark_moe.py \
  --model /var/lib/gpustack/cache/huggingface/cerebras/Qwen3-Coder-REAP-25B-A3B \
  --tp-size 1 \
  --dtype auto \
  --tune \
  --batch-size 1 2 4 8 16 32 64 \
  --save-dir "$CFG_DIR"

What you should see

• A line like Start tuning over 8 configurations... (or another small number, depending on GPU)
• Per-batch “Completed tuning…” lines
• A final Writing best config to .../configs/E=...,N=...,device_name=....json

Example 2 — Granite 4.0 H Tiny on RTX 5090

/usr/bin/python3 benchmarks/kernels/benchmark_moe.py \
  --model /var/lib/gpustack/cache/huggingface/ibm-granite/granite-4.0-h-tiny \
  --tp-size 1 \
  --dtype auto \
  --tune \
  --batch-size 1 2 4 8 16 32 64 \
  --save-dir "$CFG_DIR"

After tuning, start vLLM. The startup logs should not contain:

WARNING ... Using default MoE config. Performance might be sub-optimal!

Tips & Notes

• Device name The JSON filename embeds device_name=.... vLLM picks a pretty specific string (e.g., NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition or NVIDIA_GeForce_RTX_4090). If you move the file between machines with slightly different device_names, duplicate/rename accordingly.

• Batch sizes The script tunes across a small set of batch sizes (1 2 4 8 16 32 64) that map well to common scheduler choices. You can change --batch-size to include the sizes you care about.

• Expert parallel The tuner assumes no EP (--tp-size 1 and --enable-expert-parallel off). If you use EP at runtime, the MoE shard shapes change; tune again with the same topology you’ll serve.

• FP8 / quantized paths This tuner uses the release-0.10.1.1 call pattern (use_fp8_w8a8, use_int8_w8a16, per-tensor or block scales) and avoids newer QuantConfig helpers that aren’t present in this tag.

• Ray dashboard noise If Ray prints “Failed to start dashboard”, that’s fine—we disable it and the workers still run. The “Started a local Ray instance.” line indicates success.

• Where configs are loaded from vLLM searches its own configs/ package directory at runtime. Saving the JSON into:

  .../vllm/model_executor/layers/fused_moe/configs/
  

  ensures the engine finds it automatically.

Troubleshooting

• No file written / still “Using default MoE config” Confirm --save-dir points to the installed vLLM package configs/ path (not your repo checkout). Also confirm the filename’s E, N, and device_name match your model & GPU.

License

This patch follows vLLM’s original license headers (Apache-2.0).