From a42a2c0a99e95ad5456fdf11b0cd2c454e5685e3 Mon Sep 17 00:00:00 2001 From: avlp12 Date: Wed, 26 Aug 2026 15:50:04 +0900 Subject: [PATCH] fix(gguf): chunk the moe_vec z grid past the 65535 cap 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. --- .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 7 +- python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 439 ++++-------------- tests/kernels/test_moe_vec_large_rows.py | 221 +++++++++ 3 files changed, 316 insertions(+), 351 deletions(-) create mode 100644 tests/kernels/test_moe_vec_large_rows.py diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5f..f12ca6a62 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -67,7 +67,12 @@ static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx, co const dim3 num_blocks(block_num_x, num_blocks_y, 1); const dim3 block_size(CUDA_DEQUANTIZE_BLOCK_SIZE, 1, 1); quantize_q8_1<<>>( - &x[off * kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded); + // int64: ``off`` walks up to ``ky``, which on the MoE down-projection is + // ``tokens * top_k``. ``off * kx`` wraps int32 past off = 524288 at + // kx = 4096, which would turn a launch that now succeeds into a silently + // wrong read. (``kx_padded`` is already int64, so the ``vy`` term is not + // affected.) + &x[(int64_t)off * (int64_t)kx], (int32_t*)vy + off * (kx_padded / 32 * 9), kx, kx_padded); } } diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e080..66c7b53ce 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -11,11 +11,20 @@ static __global__ void moe_vec_q( const int topk, const int ncols, const int nrows, - const int token_stride) { + const int token_stride, + const int64_t z_offset) { const auto row = blockIdx.x * blockDim.y + threadIdx.y; - const auto token = blockIdx.z / topk; - const auto expert = (topk_ids)[blockIdx.z]; + // Routed-row index within the FULL [tokens * topk] range. gridDim.z is capped + // at 65535 by the hardware, so the launcher slices that range and hands each + // launch the base of its own chunk; blockIdx.z is the offset inside the chunk. + // int64 throughout: at top_k = 6 a 16K-token prefill chunk is already ~98K + // routed rows, and dst below is indexed by z * nrows, which wraps int32 well + // before that. + const int64_t z = z_offset + (int64_t)blockIdx.z; + + const int64_t token = z / topk; + const auto expert = (topk_ids)[z]; if (row >= nrows) { return; @@ -28,6 +37,7 @@ static __global__ void moe_vec_q( float tmp = 0.0f; const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row; + // ``token`` is int64, so this product no longer wraps at large token counts. const block_q8_1* y = (const block_q8_1*)(((const int*)vy) + token * token_stride); for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { @@ -47,367 +57,96 @@ static __global__ void moe_vec_q( } if (threadIdx.x == 0) { - dst[blockIdx.z * nrows + row] = tmp; + dst[z * (int64_t)nrows + (int64_t)row] = tmp; } } -template -static void moe_vec_q4_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_1_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q8_0_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q2_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q3_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q4_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q5_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_q6_K_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} +// The routed-row count `tokens * top_k` rides on gridDim.z, and the hardware +// caps gridDim.y/z 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 that is every prefill chunk past 10922 +// tokens (10923 * 6 = 65538). +// +// 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 row count to x would scatter the weight +// streaming. +// +// So past the cap we SLICE z instead, handing each launch the base of its own +// chunk in `z_offset` -- the same shape as `quantize_row_q8_1_cuda`'s existing +// loop over MAX_BLOCK_SIZE in gguf_kernel.cu. The offset is an explicit kernel +// argument rather than a bumped base pointer because the kernel derives +// `token = z / topk` from the absolute z; a pointer scheme would also force +// every chunk stride to be a multiple of topk. +#define MOE_VEC_MAX_GRID_Z 65535 -template -static void moe_vec_iq2_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq2_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq3_xxs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq1_m_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} - -template -static void moe_vec_iq4_nl_q8_1_cuda( +template +static void moe_vec_launch( const void* vx, const void* vy, scalar_t* dst, const int* topk_ids, - const int top_k, - const int tokens, + const int64_t top_k, + const int64_t tokens, const int ncols, const int nrows, const int token_stride, cudaStream_t stream) { + TORCH_CHECK(top_k > 0, "moe_vec: top_k must be positive (got ", top_k, ")"); + TORCH_CHECK(tokens >= 0, "moe_vec: negative token count (", tokens, ")"); 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + // int64: tokens and top_k are int64 at the call boundary precisely so this + // product cannot overflow before it is compared against the cap. + const int64_t routed_rows = tokens * top_k; + for (int64_t z_offset = 0; z_offset < routed_rows; z_offset += MOE_VEC_MAX_GRID_Z) { + const int64_t remaining = routed_rows - z_offset; + const unsigned int z_span = + (unsigned int)(remaining < MOE_VEC_MAX_GRID_Z ? remaining : (int64_t)MOE_VEC_MAX_GRID_Z); + const dim3 block_nums(block_num_y, 1, z_span); + moe_vec_q<<>>( + vx, vy, dst, topk_ids, (int)top_k, ncols, nrows, token_stride, z_offset); + } } -template -static void moe_vec_iq4_xs_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} +// One launcher per quant type. They differ only in their template arguments, so +// they are generated rather than hand-written: 19 hand-copied bodies is exactly +// how the gridDim.z cap came to be replicated 19 times in the first place. +#define MOE_VEC_LAUNCHER(name, qk, qi, block_q_t, vdr, vec_dot) \ + template \ + static void name( \ + const void* vx, \ + const void* vy, \ + scalar_t* dst, \ + const int* topk_ids, \ + const int64_t top_k, \ + const int64_t tokens, \ + const int ncols, \ + const int nrows, \ + const int token_stride, \ + cudaStream_t stream) { \ + moe_vec_launch( \ + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, \ + token_stride, stream); \ + } -template -static void moe_vec_iq3_s_q8_1_cuda( - const void* vx, - const void* vy, - scalar_t* dst, - const int* topk_ids, - const int top_k, - const int tokens, - const int ncols, - const int nrows, - 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_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); -} +MOE_VEC_LAUNCHER(moe_vec_q4_0_q8_1_cuda, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q4_1_q8_1_cuda, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q5_0_q8_1_cuda, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q5_1_q8_1_cuda, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q8_0_q8_1_cuda, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q2_K_q8_1_cuda, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q3_K_q8_1_cuda, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q4_K_q8_1_cuda, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q5_K_q8_1_cuda, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1) +MOE_VEC_LAUNCHER(moe_vec_q6_K_q8_1_cuda, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq2_xxs_q8_1_cuda, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq2_xs_q8_1_cuda, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq2_s_q8_1_cuda, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq3_xxs_q8_1_cuda, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq1_s_q8_1_cuda, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq1_m_q8_1_cuda, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq4_nl_q8_1_cuda, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq4_xs_q8_1_cuda, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1) +MOE_VEC_LAUNCHER(moe_vec_iq3_s_q8_1_cuda, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1) diff --git a/tests/kernels/test_moe_vec_large_rows.py b/tests/kernels/test_moe_vec_large_rows.py new file mode 100644 index 000000000..27f7420c4 --- /dev/null +++ b/tests/kernels/test_moe_vec_large_rows.py @@ -0,0 +1,221 @@ +"""MMVQ grouped-expert GEMV past the 65535 gridDim.z cap. + +``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. At top_k = 6 that +capped a prefill chunk at 10922 tokens (10923 * 6 = 65538), so anything larger +died on its very first launch with cudaErrorInvalidValue. The launcher now +slices z and passes each launch the base of its chunk in ``z_offset``. + +"does not throw" is the cheap half of this test. The half that actually pins the +indexing down is the EQUIVALENCE check: the same routed rows computed in one +over-cap call must come out bit-identical to the concatenation of sub-cap +slices. A wrong ``z_offset`` still launches fine -- it just reads the wrong +token, the wrong expert id, or writes the wrong dst row, and only a value +comparison catches that. + +The geometry is deliberately tiny (1024 hidden, 64 output rows, 8 experts) so +that the token counts, not the weights, are what makes the test big. +""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +GGML_IQ2_XS = 17 +IQ2_XS_BLOCK_BYTES = 74 +QK_K = 256 + +H = 1024 # ncols / hidden -- 4 IQ2_XS blocks per row +ROW_BYTES = (H // QK_K) * IQ2_XS_BLOCK_BYTES # 296 +NROWS = 64 # output features per expert +NUM_EXPERTS = 8 +TOP_K = 6 + +# Largest token count whose routed rows still fit one launch at TOP_K = 6. +OLD_CAP_TOKENS = 65535 // TOP_K # 10922 +SUB_CAP_CHUNK = 8192 # 8192 * 6 = 49152 routed rows + + +def _make_bank(seed: int = 0) -> torch.Tensor: + """``[E, NROWS, ROW_BYTES]`` uint8 IQ2_XS expert bank. + + Everything is random bytes -- the IQ2_XS grid/sign/scale nibbles are all + total over their lookup tables -- except the fp16 super-block scale ``d`` at + the head of each block, which is written explicitly so that random bit + patterns cannot hand us NaN/Inf and make an exact comparison meaningless. + """ + nb = H // QK_K + g = torch.Generator().manual_seed(seed) + bank = torch.randint(0, 256, (NUM_EXPERTS, NROWS, ROW_BYTES), generator=g, dtype=torch.uint8) + d = (0.01 + 0.04 * torch.rand(NUM_EXPERTS, NROWS, nb, generator=g)).to(torch.float16) + d_bytes = d.view(torch.uint8).reshape(NUM_EXPERTS, NROWS, nb, 2) + for b in range(nb): + bank[:, :, b * IQ2_XS_BLOCK_BYTES : b * IQ2_XS_BLOCK_BYTES + 2] = d_bytes[:, :, b] + return bank.cuda() + + +def _make_inputs(tokens: int, seed: int = 1): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(tokens, H, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + ids = torch.randint( + 0, NUM_EXPERTS, (tokens, TOP_K), generator=g, device="cuda", dtype=torch.int32 + ) + return x, ids + + +def _run(x, bank, ids, tokens, top_k=TOP_K): + from freetoken.kernel.gguf import ggml_moe_a8_vec + + y = ggml_moe_a8_vec(x, bank, ids, top_k, GGML_IQ2_XS, NROWS, tokens) + torch.cuda.synchronize() # surface any async launch failure here + return y + + +def _run_sliced(x, bank, ids, tokens, chunk=SUB_CAP_CHUNK): + """Same computation, but every launch stays under the old 65535 z cap.""" + assert chunk * TOP_K < 65536 + parts = [] + for lo in range(0, tokens, chunk): + hi = min(lo + chunk, tokens) + parts.append(_run(x[lo:hi].contiguous(), bank, ids[lo:hi].contiguous(), hi - lo)) + return torch.cat(parts, dim=0) + + +def _sane(y: torch.Tensor) -> None: + assert torch.isfinite(y).all(), "kernel produced non-finite output" + assert y.abs().sum().item() > 0, "kernel output is entirely zero -- test is vacuous" + + +# --------------------------------------------------------------------------- # +# 1. below the cap: behaviour must be unchanged, and absolutely correct. +# --------------------------------------------------------------------------- # + + +def test_small_token_count_matches_dequant_reference(): + """128 tokens (z = 768, one launch). Anchor the kernel to a dense reference + so that the equivalence tests above the cap are comparing against something + known to be right, not just self-consistent.""" + from freetoken.kernel.gguf import ggml_dequantize + + tokens = 128 + bank = _make_bank() + x, ids = _make_inputs(tokens) + y = _run(x, bank, ids, tokens) + assert y.shape == (tokens * TOP_K, NROWS) + _sane(y) + + flat = bank.reshape(NUM_EXPERTS * NROWS, ROW_BYTES) + deq = ggml_dequantize(flat, GGML_IQ2_XS, NUM_EXPERTS * NROWS, H, torch.float32) + deq = deq.reshape(NUM_EXPERTS, NROWS, H) + + flat_ids = ids.reshape(-1).long() + xf = x.to(torch.float32) + ref = torch.zeros(tokens * TOP_K, NROWS, device="cuda", dtype=torch.float32) + row_token = torch.arange(tokens * TOP_K, device="cuda") // TOP_K + for e in range(NUM_EXPERTS): + sel = (flat_ids == e).nonzero(as_tuple=True)[0] + if sel.numel(): + ref[sel] = xf[row_token[sel]] @ deq[e].T + + # The kernel quantizes activations to Q8_1 (8 bit / 32-element block), the + # reference does not, so this is a closeness check, not an exact one. + err = (y.to(torch.float32) - ref).norm() / ref.norm() + assert err < 0.05, f"relative L2 error vs dequant reference too large: {err.item():.4f}" + + +def test_small_token_count_slice_equivalent(): + """Sub-cap counts take exactly one launch with z_offset == 0; slicing them + must still be bit-identical, i.e. the fix changed nothing below the cap.""" + tokens = 128 + bank = _make_bank() + x, ids = _make_inputs(tokens) + full = _run(x, bank, ids, tokens) + sliced = _run_sliced(x, bank, ids, tokens, chunk=32) + _sane(full) + assert torch.equal(full, sliced) + + +# --------------------------------------------------------------------------- # +# 2. above the cap: the whole point. +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("tokens", [12288, 16384]) +def test_over_cap_token_counts_launch(tokens): + """z = 73728 / 98304, both past the old 65535 cap -- these used to raise + cudaErrorInvalidValue on the very first launch.""" + assert tokens * TOP_K > 65535 + bank = _make_bank() + x, ids = _make_inputs(tokens) + y = _run(x, bank, ids, tokens) + assert y.shape == (tokens * TOP_K, NROWS) + _sane(y) + + +@pytest.mark.parametrize("tokens", [12288, 16384]) +def test_over_cap_matches_sub_cap_slices(tokens): + """THE indexing test. + + One over-cap call (which the launcher internally splits at 65535 routed + rows, a boundary that does NOT line up with the token slicing below) versus + the concatenation of runs that each stay under the cap. Every routed row is + an independent dot product with an identical reduction order in both paths, + so a correct ``z_offset`` gives bit-identical results. Any drift in the token + index, the expert id or the dst row shows up here immediately. + """ + bank = _make_bank() + x, ids = _make_inputs(tokens) + full = _run(x, bank, ids, tokens) + sliced = _run_sliced(x, bank, ids, tokens) + _sane(full) + assert full.shape == sliced.shape + if not torch.equal(full, sliced): + bad = (full != sliced).any(dim=1).nonzero(as_tuple=True)[0] + raise AssertionError( + f"{bad.numel()} of {full.shape[0]} routed rows differ; " + f"first at row {bad[0].item()} (token {bad[0].item() // TOP_K}), " + f"max abs delta {(full.to(torch.float32) - sliced.to(torch.float32)).abs().max().item()}" + ) + + +def test_rows_straddling_the_old_cap_are_individually_correct(): + """Pin specific tokens either side of the old 65535-row boundary against a + single-token run of the same weights -- an independent check that does not + rely on the slicing helper at all.""" + tokens = 12288 + bank = _make_bank() + x, ids = _make_inputs(tokens) + full = _run(x, bank, ids, tokens) + + # 10922 is the last token that fitted; 10923 is the one that broke it. + for t in (0, OLD_CAP_TOKENS - 1, OLD_CAP_TOKENS, OLD_CAP_TOKENS + 1, tokens - 1): + one = _run(x[t : t + 1].contiguous(), bank, ids[t : t + 1].contiguous(), 1) + got = full[t * TOP_K : (t + 1) * TOP_K] + assert torch.equal(got, one), f"token {t} mismatches its standalone run" + + +def test_down_projection_over_cap_matches_slices(): + """The down projection calls this same kernel with tokens = T * top_k and + top_k = 1, so it reaches the identical z. Cover that call shape past the cap + (z = 70000) -- and note that top_k = 1 makes ``token = z / topk`` an + identity, a different ``z_offset`` arithmetic path from the top_k = 6 tests + above.""" + routed = 70000 + bank = _make_bank(seed=3) + x, _ = _make_inputs(routed, seed=5) + g = torch.Generator(device="cuda").manual_seed(7) + ids = torch.randint(0, NUM_EXPERTS, (routed, 1), generator=g, device="cuda", dtype=torch.int32) + + full = _run(x, bank, ids, routed, top_k=1) + parts = [] + for lo in range(0, routed, 32768): + hi = min(lo + 32768, routed) + parts.append( + _run(x[lo:hi].contiguous(), bank, ids[lo:hi].contiguous(), hi - lo, top_k=1) + ) + sliced = torch.cat(parts, dim=0) + _sane(full) + assert torch.equal(full, sliced)