From 3ca0ccd0365cfe803c959221b7ea442e7167c8f2 Mon Sep 17 00:00:00 2001 From: Erich Stuntebeck Date: Tue, 25 Aug 2026 18:53:14 +0000 Subject: [PATCH] fix(kernel): launch GGUF MoE GEMV token axis on grid.x, not grid.z moe_vec_*_q8_1_cuda launched the flat (token, top-k) index as grid.z: const dim3 block_nums(block_num_y, 1, tokens * top_k); maxGridDimZ is 65535 on every CUDA compute capability, so any batch with tokens * top_k > 65535 fails the launch with cudaErrorInvalidValue. With top_k=8 that caps a prefill batch at 8191 tokens, and --max-extend-tokens defaults to 8192, so a single long prompt -- or a few concurrent ones the scheduler packs into one batch -- reproducibly killed the engine. Move the flat index to grid.x (2^31-1) and rows to grid.y, whose extent is ceil(nrows / GGML_CUDA_MMV_Y) and stays far below the cap. The neighbouring quantize_row_q8_1_cuda already tiles its y axis at 65535 for the same reason. None of the 19 launch sites checked the launch return code, so the failure only set the error flag and was reported by whatever unrelated CUDA call ran next -- flashinfer's gelu_tanh_and_mul, the following ggml_moe_a8_vec, even a torch::zeros. Add FT_MOE_VEC_LAUNCH_CHECK() after each launch so a rejected configuration reports itself, with the geometry that caused it. Co-Authored-By: Claude Opus 5 (1M context) --- python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 85 ++++++++++++++----- tests/kernels/test_gguf_moe_vec.py | 67 +++++++++++++++ 2 files changed, 129 insertions(+), 23 deletions(-) create mode 100644 tests/kernels/test_gguf_moe_vec.py diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e08..269f5612 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -2,6 +2,23 @@ // https://github.com/vllm-project/vllm/blob/4492e3a55428e161ca8db381edc28263e5da4c8d/csrc/quantization/gguf/moe_vec.cuh // copied and adapted from // https://github.com/ggerganov/llama.cpp/blob/b2899/ggml-cuda/mmvq.cu + +// A kernel launch rejected for its *configuration* (a grid dimension over the device +// limit, an oversized block) does not raise -- it only sets the error flag. Nothing here +// checked it, so the flag latched and was reported by whatever unrelated CUDA call ran +// next: the grid.z overflow fixed below surfaced as a failure inside flashinfer's +// activation kernel, inside the next ggml_moe_a8_vec, and even inside a torch::zeros. +#ifndef FT_MOE_VEC_LAUNCH_CHECK +#define FT_MOE_VEC_LAUNCH_CHECK() \ + do { \ + cudaError_t err_ = cudaGetLastError(); \ + TORCH_CHECK( \ + err_ == cudaSuccess, \ + "moe_vec launch failed: ", cudaGetErrorString(err_), \ + " (grid=", block_nums.x, ",", block_nums.y, ",", block_nums.z, \ + " block=", block_dims.x, ",", block_dims.y, ",", block_dims.z, ")"); \ + } while (0) +#endif template static __global__ void moe_vec_q( const void* __restrict__ vx, @@ -12,10 +29,13 @@ static __global__ void moe_vec_q( const int ncols, const int nrows, const int token_stride) { - const auto row = blockIdx.x * blockDim.y + threadIdx.y; + // The flat (token, top-k) index rides blockIdx.x, not .z: it reaches tokens*top_k, + // which exceeds the 65535 cap on grid.y/.z (grid.x allows 2^31-1). Rows move to .y, + // whose extent is ceil(nrows / GGML_CUDA_MMV_Y) and stays far below the cap. + const auto row = blockIdx.y * blockDim.y + threadIdx.y; - const auto token = blockIdx.z / topk; - const auto expert = (topk_ids)[blockIdx.z]; + const auto token = blockIdx.x / topk; + const auto expert = (topk_ids)[blockIdx.x]; if (row >= nrows) { return; @@ -47,7 +67,7 @@ static __global__ void moe_vec_q( } if (threadIdx.x == 0) { - dst[blockIdx.z * nrows + row] = tmp; + dst[blockIdx.x * nrows + row] = tmp; } } @@ -64,10 +84,11 @@ static void moe_vec_q4_0_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -83,10 +104,11 @@ static void moe_vec_q4_1_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -102,10 +124,11 @@ static void moe_vec_q5_0_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -121,10 +144,11 @@ static void moe_vec_q5_1_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -140,10 +164,11 @@ static void moe_vec_q8_0_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -159,10 +184,11 @@ static void moe_vec_q2_K_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -178,10 +204,11 @@ static void moe_vec_q3_K_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -197,10 +224,11 @@ static void moe_vec_q4_K_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -216,10 +244,11 @@ static void moe_vec_q5_K_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -235,10 +264,11 @@ static void moe_vec_q6_K_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -254,10 +284,11 @@ static void moe_vec_iq2_xxs_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -273,10 +304,11 @@ static void moe_vec_iq2_xs_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -292,10 +324,11 @@ static void moe_vec_iq2_s_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -311,10 +344,11 @@ static void moe_vec_iq3_xxs_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -330,10 +364,11 @@ static void moe_vec_iq1_s_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -349,10 +384,11 @@ static void moe_vec_iq1_m_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -368,10 +404,11 @@ static void moe_vec_iq4_nl_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -387,10 +424,11 @@ static void moe_vec_iq4_xs_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } template @@ -406,8 +444,9 @@ static void moe_vec_iq3_s_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); + const dim3 block_nums(tokens * top_k, block_num_y, 1); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); moe_vec_q <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + FT_MOE_VEC_LAUNCH_CHECK(); } diff --git a/tests/kernels/test_gguf_moe_vec.py b/tests/kernels/test_gguf_moe_vec.py new file mode 100644 index 00000000..8cba6838 --- /dev/null +++ b/tests/kernels/test_gguf_moe_vec.py @@ -0,0 +1,67 @@ +import numpy as np +import pytest +import torch + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + +# maxGridDim{Y,Z} is 65535 on every CUDA compute capability; grid.x reaches 2**31-1. +MAX_GRID_YZ = 65535 + + +def _q4_0_weights(num_experts: int, nrows: int, ncols: int, seed: int = 0) -> torch.Tensor: + """Structurally valid Q4_0 banks: each 32-value block is {half d; uint8 qs[16]}. + + The scale has to be a real finite half -- random bytes decode to inf/nan and make any + comparison vacuous. + """ + rng = np.random.default_rng(seed) + nblocks = ncols // 32 + blocks = np.zeros((num_experts, nrows, nblocks, 18), dtype=np.uint8) + scale = np.array([0.01], dtype=np.float16).view(np.uint8) + blocks[..., 0], blocks[..., 1] = scale[0], scale[1] + blocks[..., 2:] = rng.integers(0, 256, blocks[..., 2:].shape, dtype=np.uint8) + return torch.from_numpy(blocks.reshape(num_experts, nrows, nblocks * 18)).cuda() + + +def test_moe_vec_above_grid_z_limit_matches_split_batches(): + """A batch whose ``tokens * top_k`` exceeds 65535 must still compute the right thing. + + ``moe_vec_*_q8_1_cuda`` used to launch that product as grid.z, which CUDA rejects with + ``cudaErrorInvalidValue`` above 65535. With top_k=8 that made any prefill batch of 8192 + tokens fail -- and ``--max-extend-tokens`` defaults to exactly 8192, so a single long + prompt, or a few concurrent ones packed into one batch, hit it. The launch return code + was unchecked, so the failure surfaced at whatever unrelated CUDA call ran next. + + Splitting the same batch in half keeps each launch under the old limit, so the halves + are a reference the pre-fix kernel could actually produce. + """ + from freetoken.kernel.gguf import ggml_moe_a8_vec + from freetoken.models.gguf.dequant import GGML_Q4_0 + + num_experts, hidden, nrows, top_k = 4, 256, 64, 8 + tokens = (MAX_GRID_YZ // top_k) + 1 # 8192 -> 65536 > 65535, one over the old cap + assert tokens * top_k > MAX_GRID_YZ + + torch.manual_seed(0) + weight = _q4_0_weights(num_experts, nrows, hidden) + x = torch.randn((tokens, hidden), dtype=torch.bfloat16, device="cuda") + topk_ids = torch.randint(0, num_experts, (tokens, top_k), dtype=torch.int32, device="cuda") + + out = ggml_moe_a8_vec(x, weight, topk_ids, top_k, int(GGML_Q4_0), nrows, tokens) + torch.cuda.synchronize() + + half = tokens // 2 + assert half * top_k <= MAX_GRID_YZ + ref = torch.cat([ + ggml_moe_a8_vec( + x[s:e].contiguous(), weight, topk_ids[s:e].contiguous(), + top_k, int(GGML_Q4_0), nrows, e - s, + ) + for s, e in ((0, half), (half, tokens)) + ]) + torch.cuda.synchronize() + + assert out.shape == (tokens * top_k, nrows) + assert torch.isfinite(out).all() + torch.testing.assert_close(out, ref, rtol=0, atol=0)