Skip to content

fix(kernel): launch GGUF MoE GEMV token axis on grid.x, not grid.z #185

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erichstuntebeck wants to merge 1 commit into
base: main
Choose a base branch
from
+129 −23
Open

fix(kernel): launch GGUF MoE GEMV token axis on grid.x, not grid.z #185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 62 additions & 23 deletions python/freetoken/kernel/csrc/gguf/moe_vec.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename scalar_t, int qk, int qi, typename block_q_t, int vdr, vec_dot_q_cuda_t vec_dot_q_cuda>
static __global__ void moe_vec_q(
const void* __restrict__ vx,
Expand All @@ -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;
Expand Down Expand Up @@ -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;
}
}

Expand All @@ -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<scalar_t, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK4_0, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI2_XXS, block_iq2_xxs, 1, vec_dot_iq2_xxs_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI2_XS, block_iq2_xs, 1, vec_dot_iq2_xs_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI2_S, block_iq2_s, 1, vec_dot_iq2_s_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI3_XXS, block_iq3_xxs, 1, vec_dot_iq3_xxs_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI1_S, block_iq1_s, 1, vec_dot_iq1_s_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI1_M, block_iq1_m, 1, vec_dot_iq1_m_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK4_NL, QI4_NL, block_iq4_nl, VDR_Q4_0_Q8_1_MMVQ, vec_dot_iq4_nl_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI4_XS, block_iq4_xs, 1, vec_dot_iq4_xs_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}

template <typename scalar_t>
Expand All @@ -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<scalar_t, QK_K, QI3_XS, block_iq3_s, 1, vec_dot_iq3_s_q8_1>
<<<block_nums, block_dims, 0, stream>>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride);
FT_MOE_VEC_LAUNCH_CHECK();
}
67 changes: 67 additions & 0 deletions tests/kernels/test_gguf_moe_vec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)