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

[erichstuntebeck]

Fixes #186.

Problem

moe_vec_*_q8_1_cuda launches the flat (token, top-k) index as grid.z:

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

maxGridDimZ is 65535 on every CUDA compute capability, so a 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 exactly 8192. So a GGUF MoE model reliably dies on the first prompt long
enough to fill one chunk. It is not limited to one long prompt either: the scheduler's token
budget is per batch, so a few concurrent medium prompts packed together trip it identically
(observed here as Prefill batch, #new-seq: 3, #new-token: 7299 — three streams, one batch).

This cost me a long time to find, because none of the 19 launch sites check the launch
return code. A rejected launch only sets the error flag, so it was reported by whatever
unrelated CUDA call ran next. I saw the same single fault blamed on three different places:

• flashinfer.activation.gelu_tanh_and_mul (the very next CUDA call, via act_fn in fused_q4_0.py)
• the following ggml_moe_a8_vec (once flashinfer was gated off)
• a bare torch.zeros in an unrelated allocation

CUDA_LAUNCH_BLOCKING=1 does not help here — the failure is in the launch configuration,
not in execution.

Fix

Move the flat index to grid.x (limit 2³¹−1) and rows to grid.y, whose extent is
ceil(nrows / GGML_CUDA_MMV_Y) and stays far below the cap. quantize_row_q8_1_cuda next
door already tiles its y axis at 65535 for the same reason, so the limit is known in this file's
neighbourhood — moe_vec.cuh just missed it.

Also adds FT_MOE_VEC_LAUNCH_CHECK() after each of the 19 launches, so a rejected
configuration reports itself with the geometry that caused it instead of latching.

In short, three things: the flat (token, top-k) index moves to grid.x; every launch site
gains a return-code check; and a regression test is added that fails on main and passes here.

Test

tests/kernels/test_gguf_moe_vec.py runs a batch one token over the old cap
(8192 * 8 = 65536) and compares against the same batch split in half, each half under the
old limit — so it checks the result is correct, not merely that nothing raised.

# on main (2757bb5)
FAILED tests/kernels/test_gguf_moe_vec.py::test_moe_vec_above_grid_z_limit_matches_split_batches
E   torch.AcceleratorError: CUDA error: invalid argument
1 failed in 92.68s

# on this branch
1 passed in 89.20s

Both runs used a fresh JIT build cache. Worth flagging for reviewers: ninja compares
mtimes, so swapping moe_vec.cuh in place over a warm ~/.cache/torch_extensions silently
reuses the old .so and the test passes when it should fail.

Tested on

• GPU: Tesla T4 (sm_75, 16 GiB), driver 580.178.04, CUDA 13.0.88, torch 2.11.0+cu130
• OS/toolchain: Ubuntu 24.04 container, g++ 13.3, -std=c++17 (built clean)
• 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), served via --moe-backend offload
• Command:

  ft serve --model=<ckpt> --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
  

• End to end: a 13,962-token prompt that killed the engine on every prior build now
  completes, and a 3-needle recall test passes 3/3 at 14,329 tokens. Server logs show
  Prefill batch, #new-token: 8192 succeeding — the exact geometry that used to fail.

Performance

Not the goal, but it measures faster rather than slower. Isolated kernel benchmark, median of
10 after 3 warmups, two independent runs per arm, each arm built from a clean cache
(nrows=1408, hidden=2816, top_k=8):

tokens	main	this branch	delta
2048	187.1 / 193.5 ms	174.0 / 176.5 ms	≈ −8%
4096	380.0 / 394.3 ms	346.1 / 355.1 ms	≈ −9%
8191	768.6 / 808.3 ms	705.8 / 711.6 ms	≈ −10%

The mechanism is not root-caused, and this PR claims none. Swapping the axes changes which
blocks are co-scheduled, but the per-token expert varies, so the obvious weight-reuse story does
not straightforwardly hold. The numbers are recorded here as a measurement, not as a
justification for the change — the fix stands on the correctness bug alone.

Notes

• Filed as GGUF MoE GEMV crashes with cudaErrorInvalidValue once tokens*top_k exceeds 65535 (grid.z limit) #186, which carries the model-free repro script and full environment details.
• Only q4_0 is exercised on my hardware, but all 19 launchers in the file shared the
  identical grid expression and are changed identically.
• Sibling kernels were checked for the same mistake: moe.cuh puts tokens/mmq_x in grid.y
  and mmvq.cuh puts nvecs in grid.y, both of which can overflow in principle but neither
  is reachable on my config. Left alone to keep this to one change.

[salekseev]

Independent confirmation on Ada (sm_89) with a different model, plus a data point on the
misattribution you describe.

Before the patch

FreeToken 0.1.2 (PyPI wheel), RTX 4080 SUPER 16 GB, driver 610.57.04 / CUDA UMD 13.3,
serving unsloth/gemma-4-26B-A4B-it-qat-GGUF (Q4_0, top_k = 8, 30 layers x 128 experts).

Any prompt long enough to fill a default 8192-token prefill chunk killed the worker, and it
never came back:

RuntimeError: Check failed: (err == cudaSuccess) is false: Failed to launch kernel: invalid argument
  __tvm_ffi_gelu_tanh_and_mul -> gelu_tanh_and_mul.cu:44
Backend supervisor: backend worker freetoken-TP0-scheduler exited
Backend worker is gone and cannot be restarted; stopping the API server

I spent an afternoon on this before finding #186, and landed in exactly the trap your PR
description calls out: the fault is reported by the next CUDA call, so I had it filed in my
own notes as "a bad launch geometry in the AOT gelu_tanh_and_mul kernel" — the innocent
bystander. Your explanation (no launch-return check at any of the 19 sites, so the flag
latches) is the missing piece. The FT_MOE_VEC_LAUNCH_CHECK() addition is arguably the more
valuable half of this PR for anyone who hits it next.

Client-visible behaviour is also worth noting: the request does not error, it hangs — the
process exits underneath an open connection, so a caller sees a 300 s timeout, and
docker ps still reports the container up while nvidia-smi shows the VRAM gone.

Measured boundary, one prompt per server since the first failure is terminal:

prompt tokens	default 8192 chunk	--max-prefill-length 4096
32215	OK, 4072 tok/s	—
35415	worker died	OK
61175	—	OK, 4437 tok/s

4096 * 8 = 32768 stays under the 65535 cap, which is why halving the chunk was a complete
workaround — consistent with your tokens * top_k analysis.

After the patch

Applied to the installed tree (the .cuh is JIT-compiled, so no wheel rebuild is needed —
though the TORCH_EXTENSIONS_DIR cache has to be busted, since ninja can otherwise keep a
.so built from the pre-patch header). --max-prefill-length back to its 8192 default:

prompt tokens	result	prefill tok/s
35415	OK	773
51543	OK	1958
77303	OK	1570
103063	OK	1786

Engine still {"status":"ok"} afterwards, uptime 14180 s. Decode unaffected (150 tok/s at
400 tokens, 142 at 1200). So on this box the patch removes the cap outright rather than
moving it.

Caveat on attribution: that run also carried #103 (--kv-cache-dtype q8_0) because I was
testing both, and a 131072-token KV pool. Neither touches the launch geometry, and the
before/after boundary above was measured on bf16 KV with a 65536 pool, but flagging it so
nobody reads more into the numbers than they support.

Also worth flagging for anyone hitting this on a GGUF MoE: the practical trigger is not
"someone sent a huge prompt". --max-extend-tokens defaults to 8192 and top_k = 8 is
ordinary, so the default configuration of a GGUF MoE model is one long prompt away from a
dead server — and agent traffic (system prompt plus tool schemas) clears 8k routinely.

[salekseev]

Follow-up to remove the attribution caveat in my comment above: I re-ran with only this
patch active — --kv-cache-dtype left at auto (bf16), --kv-reserve-tokens 65536
(pool: 65560 tokens / 5.43 GiB), --max-prefill-length at its 8192 default. That is
byte-for-byte the configuration where 35415 prompt tokens killed the worker before.

prompt tokens	pre-patch	post-patch
32215	OK	OK
35415	worker died, server shut down	OK, 819 tok/s
51543	untestable (server already gone)	OK, 2282 tok/s

Engine {"status":"ok"} afterwards. So the fix stands on its own here — #103 and the
131072-token pool in my earlier numbers were incidental.