fix(gguf): chunk the moe_vec z grid past the 65535 cap - #211
Open
avlp12 wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The bug
moe_vec_qputs the routed-row count (tokens * top_k) ongridDim.z, which thehardware caps at 65535 — only
gridDim.xreaches 2^31-1. Any grouped MoE GEMV withmore than 65535 routed rows therefore fails to launch outright with
cudaErrorInvalidValue.At
top_k = 6the largest chunk that fits is 10922 tokens (10923 * 6 = 65538). ADeepSeek-V4 deployment at
--max-prefill-length 16384(z = 98304) dies on its veryfirst 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.xis the fastest-varying axis andconsecutive 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
xwould scatter the weight streaming.So slice
zinstead, and pass each launch the base of its chunk in a newz_offsetargument — the same shape asquantize_row_q8_1_cuda's existingMAX_BLOCK_SIZEloop. An explicit offset rather than a bumped base pointer,because the kernel derives
token = z / topkfrom the absolutez, and apointer 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 ~ 174762token * token_stride(tokenis now int64, so the product widens)tokens/top_kin the launcher signatures, so the product cannot overflow atthe call boundary
&x[off * kx]inquantize_row_q8_1_cuda, which wraps foroff > 524288atkx = 4096and genuinely sees largeky(=tokens * top_kon 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_launchhelper via a macro instead of hand-maintained. The generatedlaunchers are argument-for-argument identical to the ones they replace.
Test
tests/kernels/test_moe_vec_large_rows.pycovers 12288 and 16384 tokens attop_k = 6(z = 73728 / 98304) plus thetop_k = 1down-projection shape, andanchors the sub-cap case against a dense
ggml_dequantizereference.The assertion has teeth. "Does not throw" is the cheap half — a wrong
z_offsetstill launches fine, it just reads the wrong token, the wrong expert id, or writes
the wrong
dstrow. So the tests compare one over-cap call against theconcatenation 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. ignoringz_offset, all four value comparisons fail while all four "does not throw" casesstay green.
With the fix in place:
8 passed.tests/kernelsgoes from 204 to 212 passed withan unchanged failure set (5 pre-existing
test_pinned_tensor.pyfailures in thisenvironment);
tests/moeis 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++20flag) — the two are independent fixes but #159 is what lets theextension compile with gcc at all.