[Bugfix] Fix no-op num_blocks division in get_moe_wna16_block_config

[a-yeyang]

Purpose

Fixes #52590.

In get_moe_wna16_block_config (vllm/model_executor/layers/fused_moe/fused_moe.py), the 256-branch reassigns block_size_k to 256 before using it as the divisor:

if size_k % 256 == 0 and num_blocks >= 256 and block_size_k < 256:
    block_size_k = 256
    num_blocks = num_blocks // (256 // block_size_k)  # 256 // 256 == 1, always a no-op

Since block_size_k is already 256 by the time it's used as the divisor, 256 // block_size_k is always 1, so num_blocks never actually shrinks. The very next branch in the same function does this correctly — it divides using the pre-doubling block size before doubling it:

if (num_m_blocks <= 16 and ... and num_blocks >= 512):
    block_size_k = block_size_k * 2
    num_blocks = num_blocks // 2

This PR reorders the two statements in the 256-branch to match that pattern, so the division actually uses the old block_size_k (128, the common initial default):

if size_k % 256 == 0 and num_blocks >= 256 and block_size_k < 256:
    num_blocks = num_blocks // (256 // block_size_k)
    block_size_k = 256

An artificially-inflated num_blocks can stay above the num_blocks > 1024 threshold a few lines later when it shouldn't, which flips BLOCK_SIZE_N to 1024 instead of 256 for shapes that should cross that boundary — i.e. this silently picks the wrong CUDA MoE WNA16 kernel tiling for affected shapes.

Why this is not duplicating an existing PR

Checked gh pr list --repo vllm-project/vllm --state open against get_moe_wna16_block_config, moe_wna16 block_size_k, and the issue number in-body. Four open PRs touch this function or file but none touches this specific no-op-division defect:

• [Bugfix] Remove duplicate size_k divisibility check in get_moe_wna16_block_config #40547 / fix(fused_moe): remove duplicated condition in wna16 block config #48751 — both remove a duplicated if condition elsewhere in the same function (a redundant repeated check, not this divisor bug).
• [Bugfix][Kernel] Clamp moe_wna16 BLOCK_SIZE_K so BLOCK_SIZE_K // group_size stays in {1,2,4,8} #44563 — clamps BLOCK_SIZE_K so BLOCK_SIZE_K // group_size stays in {1,2,4,8}; a different, unrelated constraint on the final block size.
• [Bugfix] Fix K-tile handling in the Triton MoE and block-FP8 GEMMs #52577 — fixes an out-of-bounds weight load in the Triton fused_moe_kernel_gptq_awq kernel and a scale-indexing bug in w8a8_triton_block_scaled_mm; neither touches get_moe_wna16_block_config.

None of these change the 256-branch's divisor calculation, so this is not duplicate work.

Test commands run and results

$ .venv/bin/python -m pytest tests/kernels/moe/test_moe_wna16_block_config.py -v
tests/kernels/moe/test_moe_wna16_block_config.py::test_get_moe_wna16_block_config_256_branch_actually_halves_num_blocks PASSED
tests/kernels/moe/test_moe_wna16_block_config.py::test_get_moe_wna16_block_config_256_branch_affects_block_size_n PASSED
2 passed in 0.80s

Added tests/kernels/moe/test_moe_wna16_block_config.py, a new pure-Python (no GPU required) regression test covering get_moe_wna16_block_config directly. The second test's shape was chosen so the bug and fix diverge not just internally but in the returned BLOCK_SIZE_N (1024 under the bug vs. 256 under the fix) — verified by temporarily reverting the fix locally and confirming both tests fail against the unfixed code with exactly that mismatch (assert 1024 == 256).

Also ran, with no regressions:

$ .venv/bin/python -m pytest tests/quantization/test_moe_wna16.py -v   # 1 skipped (CUDA-gated, expected on this CPU-only macOS box)
$ pre-commit run --files vllm/model_executor/layers/fused_moe/fused_moe.py tests/kernels/moe/test_moe_wna16_block_config.py   # ruff check/format, typos, mypy: all Passed
$ pre-commit run mypy-3.10 --files vllm/model_executor/layers/fused_moe/fused_moe.py --hook-stage manual   # Passed

Built vLLM from source for CPU on macOS/arm64 per .github/workflows/macos-smoke-test.yml (uv pip install -r requirements/build/cpu.txt && uv pip install -r requirements/cpu.txt && uv pip install -e . --no-build-isolation) to run the above against the real package; import vllm succeeds afterward.

AI assistance disclosure

AI assistance (Claude) was used to locate, diagnose, and fix this bug, and to write the regression test. I reviewed the diff line-by-line, independently hand-derived and cross-checked the arithmetic for both test shapes against the function's actual branch order (including the _ensure_block_size_k_divisible post-processing step), and verified the test fails against the unfixed code before submitting.

[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. Reviewers with write access and configured trusted contributors can comment /ci run whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use /ci run, /ci retry, or /ci cancel. New commits do not start CI automatically.

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.

🚀