[Bugfix] Fix K-tile handling in the Triton MoE and block-FP8 GEMMs

[truong-v]

Purpose

Fixes #52576.

Two Triton GEMMs on the quantized-MoE path get the K tile wrong, and in both cases a sibling code path already handles it correctly:

• fused_moe_kernel_gptq_awq applies k_mask to the scale and both zero-point loads but leaves the weight load unmasked, so the final K iteration reads past the end of the packed weight tensor; the sibling fused_moe_kernel masks the identical load, and a masked-to-zero scale multiplies the values read, so this is a memory-safety fix, not a numerical one.
• w8a8_triton_block_scaled_mm reads one scale pair per K tile, indexed (k * BLOCK_SIZE_K) // group_k, which is correct only when BLOCK_SIZE_K divides group_k; nothing enforced that, the comment above the default config stated the requirement inverted, and invoke_fused_moe_kernel already clamps this while this launcher did not.

Neither change affects the paths that are correct today: k_mask is None when block_k_diviable is true, so the generated code for the divisible case is unchanged, and the block-FP8 default config already sets BLOCK_SIZE_K = block_k, so only a tuned config asking for a larger tile behaves differently.

Not a duplicate. gh pr list --state open was checked for fused_moe_kernel_gptq_awq, moe wna16 out of bounds, w8a8_triton_block_scaled_mm, and BLOCK_SIZE_K block_shape. The nearby open PRs are all in different code: #44563 clamps BLOCK_SIZE_K inside get_moe_wna16_block_config for the CUDA moe_wna16_gemm instantiation set, #45539 and #46209 fix bounds and index widths in csrc/.../moe_wna16.cu. None touches the Triton fused_moe_kernel_gptq_awq weight load or fp8_utils.py. Re-checked 2026-08-17: the two newer open PRs in this file, #51515 (ROCm WNA16 block assignment) and #42193 (WNA16 warmup), leave the b_ptrs load alone, and none of the four open PRs touching w8a8_triton_block_scaled_mm (#40925, #47988, #48588, #41834) changes the config lookup or the per-tile scale index.

AI assistance was used to find and diagnose both defects and to draft this change.

Test Plan

• compute-sanitizer --tool memcheck over an int4 w4a16 MoE matmul driven through invoke_fused_moe_kernel, with K divisible and not divisible by BLOCK_SIZE_K. The weight tensor is allocated last so the overrun reaches unallocated memory rather than a neighbouring tensor.
• w8a8_triton_block_scaled_mm against a dequantize-and-matmul reference at block_shape = [128, 128], with the config lookup standing in for a tuned configs/*.json at BLOCK_SIZE_K of 64, 128, 256 and 512.

Both scripts are in the linked issue. They drive this repo's own code (VLLM_USE_PRECOMPILED=1 pip install -e .), so the line numbers below are this file's.

Test Result

Out-of-bounds read, before:

  K= 128 BLOCK_SIZE_K= 64: qweight holds 64 bytes per row, the loop covers 64  -> exact fit
      ERROR SUMMARY: 0 errors
  K= 128 BLOCK_SIZE_K= 32: qweight holds 64 bytes per row, the loop covers 64  -> exact fit
      ERROR SUMMARY: 0 errors
  K=  96 BLOCK_SIZE_K= 64: qweight holds 48 bytes per row, the loop covers 64  -> 16 bytes past the end
      reported at vllm/model_executor/layers/fused_moe/fused_moe.py:236
      ERROR SUMMARY: 8 errors
  K=  48 BLOCK_SIZE_K= 32: qweight holds 24 bytes per row, the loop covers 32  -> 8 bytes past the end
      reported at vllm/model_executor/layers/fused_moe/fused_moe.py:236
      ERROR SUMMARY: 8 errors

Each reported access lands past the end of qweight (e.g. 16 bytes after the nearest allocation ... of size 12288 bytes), and the launch then dies with CUDA error: unspecified launch failure. After: ERROR SUMMARY: 0 errors for all four, and the two non-divisible cases run to completion. The error count is not a stable property of the defect — how many accesses the driver reports before it kills the context varies — so treat the fault site, not the number, as the signal.

Block-FP8 scales, max relative error against the reference:

BLOCK_SIZE_K	group_k	before	after
64	128	2.226e-07	2.226e-07
128	128	2.082e-07	2.082e-07
256	128	2.284e-01	2.082e-07
512	128	2.633e-01	2.082e-07

No new tests. The existing suites cannot reach either path. test_fused_moe_wn16 parametrizes k over 128 and 1024 and requires k to be a multiple of group_size (64 or 128), while the untuned path caps BLOCK_SIZE_K at 64 — so K % BLOCK_SIZE_K is always 0 there. For the block-FP8 GEMM, no shipped N=…,K=… config carries a BLOCK_SIZE_K larger than block_k, so the path is only reachable by adding one. Covering either would mean a kernel-level test that pins the config explicitly; happy to add one if you would like it in this PR.

No model evals. Neither change alters output on any path that is correct today: the mask is a no-op when block_k_diviable is true, and the clamp is a no-op unless a tuned config asks for BLOCK_SIZE_K > block_k. I do not have a vLLM serving setup on this machine to run tests/evals/ against.

Environment

• vllm-project/vllm at 6664d397bf091cb9371cba481d4efb8233436fe6, installed with VLLM_USE_PRECOMPILED=1 pip install -e .
• NVIDIA B200 (sm_100), driver 595.71.05
• torch 2.13.0+cu130, triton 3.7.1, Python 3.12

[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.

🚀

[afierka-intel]

Heads-up: #52652 adds a tensor-descriptor path to this same kernel and rewrites the b_ptrs load region, so the two changes will conflict textually.

One fact worth having on record for the masked-B change: on the descriptor path B comes from b_desc.load(...), which takes no mask/other argument, so k_mask cannot be carried over to it. There the K tail is neutralised instead by b_scale (and b_zp) being loaded with mask=k_mask, other=0.0, which forces the dequantized weight to exactly 0 regardless of which bytes the descriptor read.

[afierka-intel]

Correction to my note above: I left out the qualifier that matters for you.

#52652 force-disables its descriptor path off-XPU whenever A.size(1) % BLOCK_SIZE_K != 0, so on CUDA the TD load never executes a K tail at all. The b_scale/b_zp masking I described is how the tail is neutralised on XPU only — it is not an argument that the descriptor path handles K tails on your hardware, and I should not have written it without saying so.

Your masking concern therefore does not carry over to the TD branch; it applies to the pointer branch, which #52652 keeps in the else. Whoever lands second preserves your mask=k_mask, other=k_other there.