Skip to content

fix(moe_vec): chunk launches when tokens*top_k exceeds gridDim.z limit - #125

Open
BruceLoveDecimal wants to merge 1 commit into
vllm-project:mainfrom
BruceLoveDecimal:fix/moe_vec_grid_z
Open

fix(moe_vec): chunk launches when tokens*top_k exceeds gridDim.z limit#125
BruceLoveDecimal wants to merge 1 commit into
vllm-project:mainfrom
BruceLoveDecimal:fix/moe_vec_grid_z

Conversation

@BruceLoveDecimal

Copy link
Copy Markdown

Problem

The MMVQ MoE kernels in csrc/gguf/moe_vec.cuh place tokens * top_k in the
z dimension of the launch grid:

const dim3 block_nums(block_num_y, 1, tokens * top_k);

gridDim.z is hardware-capped at 65535. vLLM's default
max_num_batched_tokens is 8192, so any MMVQ-path MoE model with top_k >= 8
hits 8192 * 8 = 65536 > 65535 inside profile_run and the engine dies
during startup with:

RuntimeError: CUDA error: invalid argument

This was found while integrating the Hy3 (Hunyuan) architecture (#88), whose
IQ1_M experts (top_k 8) hit this limit in profile_run on every startup.

This is 100% reproducible on startup for such models with default settings —
it is not a load-dependent or intermittent failure. Decode-time batches are
small, so the limit is only reachable through the profile/dummy run or large
prefill batches.

The vec path is selected when both expert weight types are in
MMVQ_QUANT_TYPES but not on the MMQ gemm path
(quantization/fused_moe.py), which in practice means i-matrix (IQ*) quants —
the formats commonly used for large MoE GGUFs.

Fix

Split the launch into per-token-range chunks so each launch keeps
gridDim.z <= 65535, mirroring the existing chunking pattern in
quantize_row_q8_1_cuda (gguf_kernel.cu):

const int max_tokens_per_launch = 65535 / top_k;
for (int tok_off = 0; tok_off < tokens; tok_off += max_tokens_per_launch) {
  ...
}

Chunk boundaries align to whole tokens (multiples of top_k in grid z), and
the vy / dst / topk_ids pointers are offset accordingly, so kernel
indexing is unchanged. Single-launch cases take the same path as before with
one iteration.

The 19 per-type launcher bodies were byte-identical except for template
arguments, so they are consolidated into one MOE_VEC_Q8_1_LAUNCHER macro
(429 → 158 lines) while touching the launch logic.

Test Plan / Test Result

New regression test test_moe_large_batch exercises the chunked launch path
directly: num_tokens=9000, top_k=8 gives 9000 * 8 = 72000 > 65535, which
crashed with CUDA error: invalid argument before the fix. It runs for both
IQ1_M (the motivating i-matrix case) and Q4_0, and compares against the
dequantized fused_experts reference.

Full run on an RTX PRO 6000 Blackwell (sm_120), torch 2.13.0+cu130, CUDA 13
toolchain, vLLM 0.27.0:

pytest tests/test_kernels.py::test_moe_large_batch -v
2 passed in 287.04s

pytest tests/test_kernels.py
1452 passed, 68 failed, 576 skipped in 3338.46s

pytest tests/test_plugin.py tests/test_gguf_utils.py
44 passed in 8.62s

The 68 failures are all test_moe[*-dtype2-...] cases failing inside the
triton reference kernel with OutOfResources: shared memory, Required: 122880, Hardware limit: 101376 — a hardware/triton limitation of this GPU unrelated
to the change. The identical failure set (same test IDs, same error) exists in
the baseline run on the parent commit recorded before this branch
(69 failed there = the same 68 kernel cases plus one unrelated
test_plugin.py ordering flake that passes in isolation and passes here).

ruff check on the full repo also passes.

The MMVQ MoE kernels put tokens*top_k in grid dimension z, which is
capped at 65535 by CUDA. Profile/dummy runs with large batched token
counts (e.g. 8192 tokens x top_k 8 = 65536) exceeded this and crashed
with CUDA invalid argument during engine init on Hy3 IQ1_M.

Split the launch into per-token-range chunks like quantize_row_q8_1_cuda
already does, and add test_moe_large_batch covering tokens*top_k > 65535.
@BruceLoveDecimal
BruceLoveDecimal deleted the fix/moe_vec_grid_z branch September 1, 2026 08:25
@BruceLoveDecimal
BruceLoveDecimal restored the fix/moe_vec_grid_z branch September 1, 2026 09:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant