Skip to content

fix(gguf): chunk the moe_vec z grid past the 65535 cap - #211

Open
avlp12 wants to merge 1 commit into
FlashML-org:mainfrom
avlp12:fix/moe-vec-z-grid-cap
Open

fix(gguf): chunk the moe_vec z grid past the 65535 cap#211
avlp12 wants to merge 1 commit into
FlashML-org:mainfrom
avlp12:fix/moe-vec-z-grid-cap

Conversation

@avlp12

@avlp12 avlp12 commented Aug 26, 2026

Copy link
Copy Markdown

The bug

moe_vec_q puts the routed-row count (tokens * top_k) on gridDim.z, which the
hardware caps at 65535 — only gridDim.x reaches 2^31-1. Any grouped MoE GEMV with
more than 65535 routed rows therefore fails to launch outright with
cudaErrorInvalidValue.

At top_k = 6 the largest chunk that fits is 10922 tokens (10923 * 6 = 65538). A
DeepSeek-V4 deployment at --max-prefill-length 16384 (z = 98304) dies on its very
first prefill chunk, while 8192 (z = 49152) works. The limit is inherited from the
vendored kernel and was replicated across all 19 per-type launchers.

The fix

The row count has to stay on z. blockIdx.x is the fastest-varying axis and
consecutive x-blocks are consecutive weight rows of the same (token, expert)
pair, which is what keeps one expert's rows co-resident in L2 for this
bandwidth-bound GEMV. Moving the count to x would scatter the weight streaming.

So slice z instead, and pass each launch the base of its chunk in a new
z_offset argument — the same shape as quantize_row_q8_1_cuda's existing
MAX_BLOCK_SIZE loop. An explicit offset rather than a bumped base pointer,
because the kernel derives token = z / topk from the absolute z, and a
pointer scheme would also force every chunk stride to be a multiple of topk.

Also promoted to 64-bit, so this does not merely turn a hard crash into silent
corruption:

  • dst[z * nrows + row], which wrapped around T ~ 174762
  • token * token_stride (token is now int64, so the product widens)
  • tokens / top_k in the launcher signatures, so the product cannot overflow at
    the call boundary
  • &x[off * kx] in quantize_row_q8_1_cuda, which wraps for off > 524288 at
    kx = 4096 and genuinely sees large ky (= tokens * top_k on the down GEMV)

The 19 per-type launchers differed only in their template arguments and all 19 had
copied the capped launch, so they are now generated from one shared
moe_vec_launch helper via a macro instead of hand-maintained. The generated
launchers are argument-for-argument identical to the ones they replace.

Test

tests/kernels/test_moe_vec_large_rows.py covers 12288 and 16384 tokens at
top_k = 6 (z = 73728 / 98304) plus the top_k = 1 down-projection shape, and
anchors the sub-cap case against a dense ggml_dequantize reference.

The assertion has teeth. "Does not throw" is the cheap half — a wrong z_offset
still launches fine, it just reads the wrong token, the wrong expert id, or writes
the wrong dst row. So the tests compare one over-cap call against the
concatenation of sub-cap slices and assert bit-identical results (every routed
row is an independent dot product with an identical reduction order in both paths).

Mutation-tested: with the kernel reverted to z = blockIdx.z, i.e. ignoring
z_offset, all four value comparisons fail while all four "does not throw" cases
stay green.

FAILED test_over_cap_matches_sub_cap_slices[12288]
FAILED test_over_cap_matches_sub_cap_slices[16384]
FAILED test_rows_straddling_the_old_cap_are_individually_correct
FAILED test_down_projection_over_cap_matches_slices
4 failed, 4 passed

With the fix in place: 8 passed. tests/kernels goes from 204 to 212 passed with
an unchanged failure set (5 pre-existing test_pinned_tensor.py failures in this
environment); tests/moe is unchanged.

Verified on RTX 5090 (sm_120), CUDA 13.3 / nvcc 13.3.73, torch 2.11.0+cu130.

Note: on a host without clang++, building this branch also needs #159 (the
-std=c++20 flag) — the two are independent fixes but #159 is what lets the
extension compile with gcc at all.

moe_vec_q puts the routed-row count (tokens * top_k) on gridDim.z, which the
hardware caps at 65535 -- only gridDim.x reaches 2^31-1. Any grouped MoE GEMV
with more than 65535 routed rows therefore fails to launch outright with
cudaErrorInvalidValue. At top_k = 6 the largest chunk that fitted was 10922
tokens (10923 * 6 = 65538), so a DeepSeek-V4 deployment at
--max-prefill-length 16384 (z = 98304) died on its very first prefill chunk
while 8192 (z = 49152) worked. Inherited from the vendored llama.cpp/vLLM
kernel, and replicated across all 19 per-type launchers.

The row count has to stay on z. blockIdx.x is the fastest-varying axis and
consecutive x-blocks are consecutive weight rows of the SAME (token, expert)
pair, which keeps one expert's rows co-resident in L2 for this bandwidth-bound
GEMV; moving the count to x would scatter the weight streaming. So slice z
instead and pass each launch the base of its chunk in a new z_offset argument
-- the same shape as quantize_row_q8_1_cuda's existing MAX_BLOCK_SIZE loop. An
explicit offset rather than a bumped base pointer, because the kernel derives
token = z / topk from the absolute z, and a pointer scheme would also force
every chunk stride to be a multiple of topk.

Also promoted to 64-bit, so that this does not merely turn a hard crash into
silent corruption:
  - dst[z * nrows + row], which wrapped around T ~ 174762
  - token * token_stride (token is now int64, so the product widens)
  - tokens/top_k in the launcher signatures, so the product cannot overflow at
    the call boundary
  - &x[off * kx] in quantize_row_q8_1_cuda, which wraps for off > 524288 at
    kx = 4096 and genuinely sees large ky (= tokens * top_k on the down GEMV)

The 19 per-type launchers differed only in their template arguments and all 19
had copied the capped launch, so they are now generated from one shared
moe_vec_launch helper via a macro instead of hand-maintained.

tests/kernels/test_moe_vec_large_rows.py covers 12288 and 16384 tokens at
top_k = 6 (z = 73728 / 98304) plus the top_k = 1 down-projection shape, and
asserts bit-identical results against the concatenation of sub-cap slices.
The assertion has teeth: with the kernel reverted to `z = blockIdx.z` (i.e.
ignoring z_offset) the four value comparisons all fail while the "does not
throw" cases stay green -- which is why the test asserts equality against a
sliced reference rather than merely checking that the launch succeeds.
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