GGUF MoE GEMV crashes with cudaErrorInvalidValue once tokens*top_k exceeds 65535 (grid.z limit)

[erichstuntebeck]

Summary

moe_vec_*_q8_1_cuda (python/freetoken/kernel/csrc/gguf/moe_vec.cuh) launches the flat
(token, top-k) index as grid.z:

const dim3 block_nums(block_num_y, 1, tokens * top_k);   // x19 launchers

maxGridDimZ is 65535 on every CUDA compute capability, so any call with
tokens * top_k > 65535 fails the launch with cudaErrorInvalidValue.

For a GGUF MoE model with top_k = 8 this caps a prefill batch at 8191 tokens — and
--max-extend-tokens defaults to exactly 8192. Serving gemma-4-26B-A4B-it-Q4_0 therefore
dies on the first prompt long enough to fill one prefill chunk. It is not limited to a single
long prompt: the scheduler's token budget is per batch, so several concurrent medium prompts
packed into one batch hit it identically.

Proposed fix in #185: moves the flat (token, top-k) index to grid.x (limit 2³¹−1), adds a launch return-code check at each of the 19 launch sites, and adds a regression test that fails on main and passes with the change.

Reproduction

Model-free — synthetic Q4_0 banks, so it needs no checkpoint download and runs in seconds.
Threshold is exact and deterministic: 8191 tokens works, 8192 fails.

"""Minimal repro: ggml_moe_a8_vec fails once tokens*top_k exceeds 65535 (maxGridDimZ)."""
import ctypes, numpy as np, torch
from freetoken.kernel.gguf import ggml_moe_a8_vec
from freetoken.models.gguf.dequant import GGML_Q4_0

rt = ctypes.CDLL("libcudart.so")
rt.cudaDeviceGetAttribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
def attr(a):
    v = ctypes.c_int(0); rt.cudaDeviceGetAttribute(ctypes.byref(v), a, 0); return v.value
print(f"GPU={torch.cuda.get_device_name(0)}  maxGridDimX={attr(5)} Y={attr(6)} Z={attr(7)}")

E, H, NROWS, TOPK = 4, 256, 64, 8
rng = np.random.default_rng(0)
blocks = np.zeros((E, NROWS, H // 32, 18), dtype=np.uint8)      # Q4_0: {half d; uint8 qs[16]}
d = np.array([0.01], dtype=np.float16).view(np.uint8)
blocks[..., 0], blocks[..., 1] = d[0], d[1]
blocks[..., 2:] = rng.integers(0, 256, blocks[..., 2:].shape, dtype=np.uint8)
w = torch.from_numpy(blocks.reshape(E, NROWS, H // 32 * 18)).cuda()

for tokens in (8191, 8192):
    x = torch.randn((tokens, H), dtype=torch.bfloat16, device="cuda")
    ids = torch.randint(0, E, (tokens, TOPK), dtype=torch.int32, device="cuda")
    print(f"\n--- tokens={tokens}  grid.z = tokens*top_k = {tokens*TOPK} "
          f"({'<=' if tokens*TOPK <= 65535 else '>'} 65535) ---")
    y = ggml_moe_a8_vec(x, w, ids, TOPK, int(GGML_Q4_0), NROWS, tokens)
    torch.cuda.synchronize()
    print(f"    ok, out={tuple(y.shape)} finite={bool(torch.isfinite(y.float()).all())}")

Full output, on a clean checkout of main @ 2757bb5

GPU=Tesla T4  maxGridDimX=2147483647 Y=65535 Z=65535

--- tokens=8191  grid.z = tokens*top_k = 65528 (<= 65535) ---
    ok, out=(65528, 64) finite=True

--- tokens=8192  grid.z = tokens*top_k = 65536 (> 65535) ---
Traceback (most recent call last):
  File "/t.py", line 27, in <module>
    print(f"    ok, out={tuple(y.shape)} finite={bool(torch.isfinite(y.float()).all())}")
                                                                     ^^^^^^^^^
torch.AcceleratorError: CUDA error: invalid argument
Search for `cudaErrorInvalidValue' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

The reported location is never the real one

None of the 19 launch sites in moe_vec.cuh checks the launch return code. A rejected launch
only sets the error flag, so it is reported by whatever CUDA call runs next. In the trace
above that is a torch.isfinite on the following line. While debugging a live server I saw the
same single fault attributed to three different places:

• flashinfer.activation.gelu_tanh_and_mul — the next CUDA call, via act_fn in moe/fused_q4_0.py:44
• the following ggml_moe_a8_vec at moe/fused_q4_0.py:46, once flashinfer was gated off
• a bare torch.zeros in an unrelated allocation

CUDA_LAUNCH_BLOCKING=1 does not help, because the failure is in the launch configuration
rather than in execution. This cost me a lot of time and is why #185 also adds a check at each
launch site.

Environment

FreeToken	0.1.2, source build, tested at 2757bb5 (git rev-parse --short HEAD)
GPU	Tesla T4, 16 GiB, sm_75
Driver / CUDA	580.178.04 / CUDA 13.0.88, torch 2.11.0+cu130
CPU / RAM	AMD EPYC 7V12 (4 vCPU) / 27 GiB
OS	Ubuntu 24.04.4 LTS, kernel 6.17.0-1022-azure (Azure VM), no WSL, no proxy
Checkpoint	ggml-org/gemma-4-26B-A4B-it-GGUF → gemma-4-26B-A4B-it-Q4_0.gguf (128 experts/layer, top_k=8, 30 layers)
Other GPUs	none

Serving command under which this first showed up:

ft serve --model=<path>/gemma-4-26B-A4B-it-Q4_0.gguf --host=0.0.0.0 --port=1919 \
         --attention-backend=triton --memory-ratio=0.55 --max-running-requests=3 \
         --max-seq-len-override=16384 --num-tokens=49152 --moe-prefill-hit-d2d

Any prompt long enough to produce a full 8192-token prefill chunk killed the engine
(Backend worker is gone and cannot be restarted).

One caveat, stated plainly: the repro above is on an unmodified main, but the serving
path is not — a stock main cannot load this checkpoint on a T4 for two unrelated reasons (the
gemma4 GGUF loader hardcodes the token-embedding quant as Q6_K where this file stores Q8_0, and
the Triton decode attention tile overflows Turing's 64 KiB shared memory at head_dim=512). I
carry local patches for both. Neither touches moe_vec.cuh or anything on the path in the repro
script, and the script itself runs on a completely clean tree. Happy to open separate issues for
those two if useful.

Not a duplicate of

• Unhandled CUDA OOM in prefill expert workspace kills the engine; server keeps accepting requests that never return (root cause behind #20-style headless state; likely also #72) #110 / Backend worker is gone and cannot be restarted; stopping the API server #20 — those are torch.OutOfMemoryError in the prefill expert workspace. This is
  cudaErrorInvalidValue from a grid dimension over the device limit, at a fixed token count,
  with VRAM to spare.
• Request hang (silent, zero-log) with hybrid_radix cache + moe-backend offload + nvfp4 triton backend — reproduced on 2 independent models, not explained by #110 #123 — that is a silent hang; this is a hard crash with a deterministic threshold.

Related, but not reachable here

moe.cuh puts tokens_post_padded / mmq_x in grid.y and mmvq.cuh puts nvecs in grid.y.
Both can overflow the same 65535 cap in principle, though neither is reachable on my config
(ggml_moe_a8 has no Python callers, and nvecs is bounded by max_seq_len). Note
quantize_row_q8_1_cuda in gguf_kernel.cu already tiles its y axis at 65535 — so the limit is
handled a few lines away, moe_vec.cuh just missed it. I left those alone to keep #185 to one change.