[Bugfix][Kernel] Clamp moe_wna16 BLOCK_SIZE_K so BLOCK_SIZE_K // group_size stays in {1,2,4,8}

[Sunt-ing]

Purpose

• Bug: GPTQ/AWQ int4 MoE models with group_size=32 crash on the CUDA moe_wna16_gemm path during serving with RuntimeError: BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8] ([Bug]: invoke_fused_moe_wna16_*_kernel calls get_moe_wna16_block_config with bad parameters #36008, reported on Volta/V100 where GPTQ MoE falls back to moe_wna16 instead of Marlin).
• Root cause: get_moe_wna16_block_config can hand _ensure_block_size_k_divisible a BLOCK_SIZE_K whose BLOCK_SIZE_K // group_size is outside the {1, 2, 4, 8} set the kernel is instantiated for. It fires on small decode batches (num_valid_tokens <= 16) in two ways: the fast path returns the heuristic's value directly, which reaches 512 for size_k=2048 (ratio 16, the reporter's V100 case); and the divisor search steps down by group_size and accepts the first divisor regardless of ratio, landing on BLOCK_SIZE_K=96 (ratio 3) when size_k is not a power-of-two multiple of the group (e.g. 1344).
• Fix (4 lines): clamp the candidate to group_size * 8 (caps the fast path at ratio 8) and only accept search results whose ratio is in {1, 2, 4, 8}. Every config the heuristic returns is then kernel-legal.
• Safety: a sweep over group_size x size_k x tokens x experts is legal for every config (0 illegal vs 2304 before) and returns the same BLOCK_SIZE_K on already-legal configs, so there is no tiling or perf change where it was already fine.

Test Plan

Both crash shapes were driven through the real moe_wna16 CUDA kernel (the path the reporter's V100 selects automatically), forced with --quantization moe_wna16 / int4_w4a16 on an sm120 box that would otherwise auto-select Marlin:

• the reporter's group_size=32, size_k=2048 (ratio 16 via the fast path), on the reporter's exact model;
• a constructed group_size=32, size_k=1344 int4 MoE forward (ratio 3 via the divisor search).

Test Result

• Reporter shape (size_k=2048, model btbtyler09/Qwen3.5-35B-A3B-GPTQ-4bit): before, decode crashes with the exact reported error (BLOCK_SIZE_K=512, ratio 16); after, BLOCK_SIZE_K=256 (ratio 8) and generation succeeds. The fixed moe_wna16 output matches the independent gptq_marlin path token for token on 6 of 8 greedy prompts; the other 2 agree until the tail and diverge only there, as expected from floating-point accumulation-order differences between two kernels.
• size_k=1344: before, the heuristic returns BLOCK_SIZE_K=96 (ratio 3) and moe_wna16_gemm aborts with the same error; after, BLOCK_SIZE_K=64 (ratio 2) and the forward runs and matches the dequantized reference. Reverting only the fix on the same script reproduces the crash, so the fix is causal.

Environment and commands

• GPU: RTX PRO 6000 / RTX 4090 (Blackwell sm120 / Ada sm89); the {1, 2, 4, 8} constraint is hardware-independent (a property of the four moe_wna16_gemm_kernel template instantiations in csrc/moe/moe_wna16.cu, not an arch-gated check).
• Reporter shape: LLM(model="btbtyler09/Qwen3.5-35B-A3B-GPTQ-4bit", quantization="moe_wna16", enforce_eager=True, max_num_seqs=1, max_model_len=2048, trust_remote_code=True), greedy. VLLM_USE_FLASHINFER_SAMPLER=0 to skip flashinfer's JIT sampler arch check on this box (unrelated to moe_wna16).
• size_k=1344 shape: constructed uint4b8 int4 MoE weights (group_size=32, e=8, single-token decode) driven through fused_moe so the real heuristic picks BLOCK_SIZE_K and the real moe_wna16_gemm runs, compared against the dequantized torch_moe reference.

Fixed moe_wna16 vs gptq_marlin (reporter model, greedy, max_tokens=32, 8 prompts)

prompt	match
"The capital of France is"	identical (32 tokens)
"The chemical symbol for gold is"	identical
"Water is made of hydrogen and"	identical
"The opposite of hot is"	identical
"The first president of the United States was"	identical
"The speed of light is approximately"	identical
"2 + 2 ="	identical first ~28 tokens, then tail diverges
"The largest planet in the solar system is"	identical first ~30 tokens, then tail diverges

Note on num_experts (question for maintainers)

• The issue and the now-closed [Bugfix] Fix wrong num_experts in invoke_fused_moe_wna16 kernels #36026 also proposed changing num_experts from B.size(1) to B.size(0) in the two invoke_fused_moe_wna16_* call sites. B.size(0) is the expert count and matches what should_moe_wna16_use_cuda already receives, so it is conceptually correct. I left it out of this crash fix for two measured reasons.
• It does not fix the crash. In small-batch decode num_m_blocks is clamped to num_valid_tokens and num_experts drops out, so the illegal BLOCK_SIZE_K still appears regardless of its value; the ratio guard is what fixes it.
• It is a small measured perf regression where it changes anything. A config sweep shows it changes the chosen config for only one model class in my grid (Mixtral-8x7B, E=8), flipping BLOCK_SIZE_K 256 to 512 at the one cuda-eligible batch; a microbench measured that 512 choice as ~2.7% slower.
• Given that, and the deeper oddities in this same heuristic noted in [Bugfix] Remove duplicate size_k divisibility check in get_moe_wna16_block_config #40547 that need a tuning pass and a maintainer call on the dimension convention, the num_experts correction seems better handled there than bundled here. Would you prefer I also include the num_experts change in this PR, or leave it for the heuristic cleanup?

End-to-end reproduction scripts (real LLM path and num_experts microbench)

The crash reproduction below drives the reporter model through a real LLM(...) engine and forces the same moe_wna16 CUDA path that the reporter's V100 selected automatically.

cd /path/to/vllm
export PYTHONPATH="$PWD"
export CUDA_VISIBLE_DEVICES=0
export VLLM_USE_FLASHINFER_SAMPLER=0
export MODEL=btbtyler09/Qwen3.5-35B-A3B-GPTQ-4bit

cat >/tmp/moe_wna16_36008_e2e.py <<'PY'
import os
import traceback

os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")

from vllm import LLM, SamplingParams

MODEL = os.environ.get("MODEL", "btbtyler09/Qwen3.5-35B-A3B-GPTQ-4bit")

try:
    llm = LLM(
        model=MODEL,
        quantization="moe_wna16",
        trust_remote_code=True,
        tensor_parallel_size=1,
        gpu_memory_utilization=0.85,
        max_model_len=2048,
        enforce_eager=True,
        max_num_seqs=1,
    )
    out = llm.generate(
        ["The capital of France is"],
        SamplingParams(max_tokens=16, temperature=0.0),
    )
    print("PASS")
    print(repr(out[0].outputs[0].text))
except Exception as exc:
    print("CRASH")
    print(type(exc).__name__)
    print(str(exc)[:1000])
    if "BLOCK_SIZE_K // group_size must be one of" in str(exc):
        print("MATCHED_REPORTER_ERROR")
    traceback.print_exc()
    raise
PY

python /tmp/moe_wna16_36008_e2e.py

Observed output:

# stock main, reporter model, quantization="moe_wna16"
CRASH
RuntimeError
BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]
MATCHED_REPORTER_ERROR

# this PR, same model and flags
PASS
' Paris.\nThe capital of France is Paris'

The num_experts note in the PR body is supplementary to the crash fix. The script below is the microbench used for the measured Mixtral BLOCK_SIZE_K=256 vs 512 timing mentioned there.

cat >/tmp/moe_wna16_num_experts_smoke.py <<'PY'
import os
import time

import torch

os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

import vllm.model_executor.layers.fused_moe.fused_moe as fm
from vllm.model_executor.layers.fused_moe import fused_experts
from vllm.model_executor.layers.fused_moe.config import int4_w4a16_moe_quant_config

DEV = "cuda"
DT = torch.bfloat16

# Mixtral-8x7B gate_up/down shapes.
E, HID, INTER, TOPK, GS = 8, 4096, 14336, 2, 64
M = 16

w13 = torch.randint(0, 255, (E, 2 * INTER, HID // 2), dtype=torch.uint8, device=DEV)
w2 = torch.randint(0, 255, (E, HID, INTER // 2), dtype=torch.uint8, device=DEV)
w13_s = torch.randn(E, 2 * INTER, HID // GS, dtype=DT, device=DEV)
w2_s = torch.randn(E, HID, INTER // GS, dtype=DT, device=DEV)
qc = int4_w4a16_moe_quant_config(
    w1_scale=w13_s,
    w2_scale=w2_s,
    w1_zp=None,
    w2_zp=None,
    block_shape=[0, GS],
)

x = torch.randn(M, HID, dtype=DT, device=DEV)
topk_ids = torch.randint(0, E, (M, TOPK), dtype=torch.int32, device=DEV)
topk_w = torch.rand(M, TOPK, dtype=torch.float32, device=DEV)

orig = fm.get_moe_wna16_block_config
force_k = {"v": None}


def forced(config, use_moe_wna16_cuda, **kw):
    if not use_moe_wna16_cuda:
        return orig(config=config, use_moe_wna16_cuda=use_moe_wna16_cuda, **kw)
    if "BLOCK_SIZE_N" in config and "BLOCK_SIZE_K" in config:
        return {}
    return {"BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": force_k["v"]}


fm.get_moe_wna16_block_config = forced


def bench(block_size_k, iters=200):
    force_k["v"] = block_size_k
    for _ in range(20):
        fused_experts(
            x, w13, w2, topk_w, topk_ids, global_num_experts=E, quant_config=qc
        )
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fused_experts(
            x, w13, w2, topk_w, topk_ids, global_num_experts=E, quant_config=qc
        )
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters * 1e6


print(f"Mixtral-8x7B gate_up+down, E={E}, gs={GS}, M={M}")
for block_size_k in (256, 512, 256, 512):
    print(f"BLOCK_SIZE_K={block_size_k} -> {bench(block_size_k):.1f} us/iter")
PY

python /tmp/moe_wna16_num_experts_smoke.py

Observed output:

BLOCK_SIZE_K=256 -> 1584 us/iter
BLOCK_SIZE_K=512 -> 1627 us/iter

AI assistance was used to prepare this PR.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

Hi @Sunt-ing, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[ptempier]

Thanks for fixing this — we hit the same BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8] crash independently (AWQ-4bit checkpoint, group_size=32, MoE shared-expert layer with size_k=1344 under TP=4).

I think this fix doesn't fully cover the bug yet. The clamp (block_size_k = min(block_size_k, group_size * 8)) only bounds the starting candidate, but the search loop below it still steps downward by plain group_size increments and accepts the first value that divides size_k, without checking that the resulting ratio is actually one of the four values the kernel supports (1, 2, 4, 8). Ratios of 3, 5, 6, 7 can still come out of that loop.

Concretely, with our shapes:

size_k = 1344
group_size = 32
block_size_k = min(512, group_size * 8)  # = 256, as in this PR

start = (min(block_size_k, size_k) // group_size) * group_size  # 256
for candidate in range(start, group_size - 1, -group_size):
    if size_k % candidate == 0:
        print(candidate, candidate // group_size)
        break
# -> 224, ratio 7   (still invalid -> still crashes with the exact same
#    "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]" error)

So this PR as written wouldn't have fixed our crash.

The simplest correct fix is to only ever consider the four valid ratios (8, 4, 2, 1) as candidates, picking the largest one that still divides size_k, instead of clamping-then-searching-by-group_size:

def _ensure_block_size_k_divisible(
    size_k: int, block_size_k: int, group_size: int
) -> int:
    """Ensure block_size_k divides size_k and block_size_k // group_size
    is one of {1, 2, 4, 8} -- the only ratios moe_wna16_gemm_kernel is
    instantiated for (csrc/moe/moe_wna16.cu).
    """
    max_ratio = max(1, min(block_size_k // group_size, 8))
    for ratio in (8, 4, 2, 1):
        if ratio > max_ratio:
            continue
        candidate = ratio * group_size
        if size_k % candidate == 0:
            return candidate

    # group_size itself always satisfies size_k % group_size == 0 by
    # construction of group_size from the model's quantization config.
    return group_size

For size_k=1344, group_size=32 this returns 64 (ratio 2), which the kernel accepts.

I verified the {1,2,4,8} constraint is hardware-independent (it's a property of the four template instantiations of moe_wna16_gemm_kernel in moe_wna16.cu, not a __CUDA_ARCH__-gated check), so the fix applies the same way on Turing/Ampere/Hopper/Blackwell.

Happy to push this as a commit to this branch or open a separate PR, whichever you'd prefer — didn't want to duplicate effort given this PR already exists.

(Disclosure: this analysis and the suggested fix were produced with AI assistance — Claude — based on a real crash reproduced on our own deployment; I've reviewed the reasoning and verified the counter-example above myself.)

[ptempier]

Follow-up to my earlier comment: this looks like part of a broader pattern, not a one-off. Two other open PRs hit the same root cause — Nemotron's MoE intermediate_size not dividing evenly across TP ranks — in different Marlin/WNA16 quant paths, both also reproduced on Nemotron-3 checkpoints:

• [Bugfix] Pad Marlin FP8 MoE weight dims to tile alignment under TP > 1 #36807 — Marlin FP8 MoE, w13 N / w2 K tile-alignment (tile_n=64, tile_k=16)
• fix marlin fp4 kernel N-dimension alignment #37296 — Marlin FP4 (NVFP4) N-dimension alignment

Cross-linking so reviewers/maintainers looking at one of these see the others — might be worth a shared tile-alignment utility across all three quant paths instead of three independent fixes, given how consistently this one model's intermediate sizes trip it.

[Sunt-ing]

Good catch. The clamp fixed the original size_k=2048 case, but the search could still return an unsupported ratio.

One small correction: with the actual heuristic, size_k=1344 reaches this helper with block_size_k=128, so it returns 96 (ratio 3), not 224. The crash is the same.

I updated the helper to cap the fast path at ratio 8 and only accept ratios {1, 2, 4, 8} in the search. The 2048 case now picks 256, and the 1344 case picks 64. I ran both through the actual moe_wna16 CUDA path; both crash before the change and pass after it. I also added you as a co-author. Thanks for spotting this.

I don't think #36807 and #37296 share an implementation with this one. They pad Marlin weight dimensions to tile boundaries. Here the dimensions are valid; the heuristic only chose a launch config that the CUDA kernel was never instantiated for.