Skip to content

fix(moe): skip inactive gguf routes in hybrid decode - #823

Merged
jason-fxz merged 4 commits into
FlashML-org:mainfrom
ischencheng:fix/gguf-hybrid-inactive-routes
Oct 29, 2026
Merged

jason-fxz merged 4 commits into
FlashML-org:mainfrom
ischencheng:fix/gguf-hybrid-inactive-routes

Conversation

@ischencheng

@ischencheng ischencheng commented Oct 25, 2026 •

Copy link
Copy Markdown
Contributor

Fixes the first known issue in #671: --moe-strategy hybrid on GGUF experts crashes once a decode step has more tokens than MMVQ takes (8 on Ada and Blackwell, down to 5 on Ampere for Q3_K). With --max-running-requests 16 it already crashes at startup, while capturing the bs=16 decode graph.

Hybrid hands the expert kernel the CPU's routes as slot -1. The GGUF kernel did not declare supports_inactive_slots, so those routes were parked on slot 0 with weight 0, and a token with two or more CPU routes routed to slot 0 more than once. MMVQ reads the ids per route, so it only computed slot 0 again. MoE MMQ does not handle a repeated expert: mm_ids_helper gives the token one compact row for slot 0 but counts every copy in the offsets of the experts above it, so part of ids_dst / ids_src1 is never written and the quantize and MMQ kernels index through it.

This makes the GGUF kernel skip inactive routes, the same contract as the block-FP8 (#665) and ds_fp4 Triton kernels: a negative id gives a zero row and reads no expert.

  • mmvq.cu: both MoE kernels write zero rows for a negative id and return.
  • mmid.cu: mm_ids_helper compares ids unsigned, so a negative id sorts after every expert and gets no compact row.
  • mmq.cu, mmq.cuh, quantize.cu: with inactive_ids, ids_src1 starts at -1; the q8_1 and fp4 quantize kernels skip such rows (gather) and slots (scatter).
  • ggml_ops.cu, ggml.py: mul_mat_id takes inactive. With it, MMQ gets inactive_ids and the rows it leaves unwritten are zeroed, so both kernels return zeros for inactive routes. Without it, the MMQ call is the same as on main.
  • gguf.py: GgmlMoEKernel.supports_inactive_slots = True, and apply passes inactive only in hybrid decode (not is_prefill and cfg.decode_target == "hybrid"), the one path that hands over -1 routes. A float GGUF half still runs on the bf16 fused MoE, which reads every route, so it gets the ids clamped to 0 as before.
  • freetoken.patch is regenerated against llama.cpp 8c2134a8 (the same git diff --no-index gives today's patch byte for byte on main), and VERSION lists the change.

test_q4_moe_routes_match_cpu gets inactive and per_route_x cases and a batch of 4: batch 1 runs the single-token MMVQ kernel, 4 the multi-token one, 17 MMQ, against gguf-py's dequantization with the inactive rows set to zero.

Tested on top of 0bf7dbb with driver 610.57.04, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0. The hybrid runs use L4 (sm_89) with 8 CPU threads.

Checkpoint unsloth/Qwen3.6-35B-A3B-GGUF, file Qwen3.6-35B-A3B-UD-Q4_K_M.gguf (Q4_K gate/up, Q5_K or Q6_K down):

ft serve --model Qwen3.6-35B-A3B-UD-Q4_K_M.gguf --moe-strategy hybrid --moe-cache-rate 0.25 \
  --moe-cpu-threads 8 --max-running-requests 16 --moe-collect-stats --max-seq-len-override 4096

then one streamed request, then 16 greedy /v1/completions sent together (Question: what is a * b?\nAnswer: a * b =, 24 tokens), so the decode steps carry 16 tokens.

offload hybrid
main 16/16 correct crashes while capturing the bs=16 graph: illegal memory access was encountered in mul_mat_id (gate/up MMQ)
this PR 16/16 correct starts, 16/16 correct, no CUDA errors in the log

The hybrid run logs 4.5 to 4.8 CPU routes per layer-step (--moe-collect-stats), so most steps hand inactive routes to the GGUF kernel.

ggml.mul_mat_id alone (run on ce2d3c9), 16 experts, 8 routes per token with half of them inactive, Q4_0 / Q4_K / NVFP4, activations shared by the routes and per route, against gguf-py's dequantization:

  • main, inactive routes parked on slot 0 as hybrid does today: right at 4 tokens (MMVQ), illegal memory access at 16 tokens (MMQ) for all three types.
  • this PR, inactive routes as -1: right at 4 and 16 tokens for all three types and both activation layouts; the inactive rows are exactly zero, and the active rows are as close to the reference as in a run with no inactive routes (q8_1 activation rounding).

pytest tests/kernels/test_ggml.py::test_q4_moe_routes_match_cpu, each case in its own process: on main (only the test applied, L4) the 18 cases without inactive routes pass and the 18 with them fail, 17 on an illegal memory access and 1 on a mismatch; with this PR all 36 pass on L4 and on A10 (sm_86).

pytest tests/kernels tests/moe tests/models/test_flint.py tests/models/test_gemma4_gguf_rope.py -m "not slow" on A10: main 2 failed / 495 passed / 117 skipped, this PR 2 failed / 520 passed / 122 skipped. The 30 new cases pass; 5 AVX-512 CPU MoE tests ran on the main run's host and skip on this one's (AVX2). The two failures are the same on main: the test container has no UVA (cudaHostGetDevicePointer fails in test_host_device_ptr_is_identity_under_uva), and test_swiglu_clamp.py::test_triton_matches_reference fails in the same run.

Prefill and fused / offload decode never see -1 routes, so they skip the extra MMQ work. L40S (sm_89), main, ce2d3c9 and 2ff4040 in one container, alternating:

main ce2d3c9 2ff4040
mul_mat_id gate_up + down, 16 tokens 317.7 us 323.9 us 317.8 us
mul_mat_id gate_up + down, 32 tokens 494.0 us 500.0 us 495.5 us
--moe-strategy fused decode, 16 streams 942.5 tok/s 924.5 tok/s 940.7 tok/s
--moe-strategy fused decode, 32 streams 1409.8 tok/s 1396.1 tok/s 1410.7 tok/s

The kernel rows use Qwen3.6-35B-A3B's MoE shapes (256 experts, hidden 2048, intermediate 512, top-8) with random Q4_K blocks, timed in a CUDA graph, median of 30 replays, 3 rounds. The decode rows use Qwen3.6-35B-A3B-UD-Q4_K_M.gguf with --max-running-requests 32, greedy ignore_eos streams of 192 tokens, tok/s over the window in which every stream is decoding, 4 runs each.

a45f0b1 then compares the ids unsigned in mm_ids_helper instead of remapping negative ids. Same harness, main / 2ff4040 / a45f0b1, with nvcc 13.0:

main 2ff4040 a45f0b1
mm_ids_helper, 512 tokens, RTX PRO 6000 (sm_120) 19.1 us 19.0 us 19.1 us
mm_ids_helper, 512 tokens, L40S 15.8 us 15.7 us 15.8 us
mul_mat_id gate_up + down, 32 tokens, RTX PRO 6000 226.7 us 229.8 us 229.8 us
mul_mat_id gate_up + down, 32 tokens, L40S 495.0 us 496.3 us 496.2 us

The helper times are kernel times from a profiler trace of 20 eager gate_up calls at 512 tokens. On a45f0b1, and again after merging main with #835 (69d55db), the 36 test_q4_moe_routes_match_cpu cases and the suite above pass on A10, with the same two failures.

Commands

Decode, the same on each tree:

ft serve --model Qwen3.6-35B-A3B-UD-Q4_K_M.gguf --moe-strategy fused --max-running-requests 32 --max-seq-len-override 4096

then one burst of 32 to warm up, and per run N = 16 or 32 concurrent streamed /v1/completions with {"prompt": "Write a long essay about the history of <topic> (part i).", "max_tokens": 192, "temperature": 0, "ignore_eos": true, "stream": true}, 16 different topics. tok/s = text events with timestamps in [max(first event), min(last event)] over all streams, divided by that window.

Kernel:

import torch
from freetoken.kernel import ggml

E, H, I, K, Q4_K = 256, 2048, 512, 8, 12
def weights(n, k):
    b = torch.randint(0, 256, (E, n, k // 256, 144), dtype=torch.uint8)
    b[..., :4] = torch.tensor([0.01, 0.002], dtype=torch.float16).view(torch.uint8)
    return b.view(E, n, k // 256 * 144).cuda()
gate_up, down = weights(2 * I, H), weights(H, I)
for t in (16, 32):
    ids = torch.stack([torch.randperm(E)[:K] for _ in range(t)]).cuda().int()
    x, inter = torch.randn(t, H).cuda().bfloat16(), torch.randn(t, K, I).cuda().bfloat16()
    step = lambda: (ggml.mul_mat_id(gate_up, Q4_K, x, ids), ggml.mul_mat_id(down, Q4_K, inter, ids))
    step(); torch.cuda.synchronize()
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for _ in range(10):
            step()
    times = []
    for _ in range(30):
        a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        a.record(); g.replay(); b.record(); b.synchronize()
        times.append(a.elapsed_time(b) * 100)  # us per gate_up + down
    print(t, sorted(times)[15])

@jason-fxz jason-fxz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, the fix is correct, and hybrid now serves batches above 8 on every GGUF type I tried.

One ask: only hybrid decode passes negative ids, but the new ids_src1 memset in ggml_cuda_mul_mat_q and the zero_inactive_routes launch in op_mul_mat_id run on every MMQ call, so prefill and fused / offload decode pay for them too. That is about 4 us per call; Qwen3.6-35B-A3B UD-Q4_K_M with --moe-strategy fused decodes 3.3% slower at batch 16 and 2.5% slower at batch 32.

Please add an inactive flag to ggml.mul_mat_id, set it in GgmlMoEKernel.apply from not is_prefill and cfg.decode_target == "hybrid", and run the memset and zero_inactive_routes only when it is set. The MMVQ early return can stay as it is.

…ecode

Only hybrid decode hands the GGUF kernel -1 routes, so mul_mat_id takes an inactive flag
and MMQ does the ids_src1 memset and the zeroing pass only when it is set. Prefill and
fused / offload decode go back to the main-branch MMQ calls.
@jason-fxz

Copy link
Copy Markdown
Collaborator

@ischencheng

Thanks, the fix is correct, and hybrid now serves batches above 8 on every GGUF type I tried.

One ask: only hybrid decode passes negative ids, but the new ids_src1 memset in ggml_cuda_mul_mat_q and the zero_inactive_routes launch in op_mul_mat_id run on every MMQ call, so prefill and fused / offload decode pay for them too. That is about 4 us per call; Qwen3.6-35B-A3B UD-Q4_K_M with --moe-strategy fused decodes 3.3% slower at batch 16 and 2.5% slower at batch 32.

Please add an inactive flag to ggml.mul_mat_id, set it in GgmlMoEKernel.apply from not is_prefill and cfg.decode_target == "hybrid", and run the memset and zero_inactive_routes only when it is set. The MMVQ early return can stay as it is.

@ischencheng

ischencheng commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor Author

@jason-fxz fixed in 2ff4040: ggml.mul_mat_id takes an inactive flag, GgmlMoEKernel.apply sets it from not is_prefill and cfg.decode_target == "hybrid", and the ids_src1 memset and zero_inactive_routes run only when it's set. On L40S fused decode is back to main, 940.7 vs 942.5 tok/s at 16 streams and 1410.7 vs 1409.8 at 32 (ce2d3c9: 924.5 / 1396.1), and hybrid with 16 concurrent requests still answers 16/16 right on L4. Updated the description with the numbers.

@jason-fxz jason-fxz left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ischencheng

Thanks, 2ff4040 does what I asked: fused decode is back to main, and hybrid still runs clean on Q4_K_M, Q4_0, Q8_K_XL and NVFP4 GGUFs.

One more thing, found while timing it: the expert_used < 0 ? INT_MAX : expert_used line in mm_ids_helper makes nvcc unroll its token loop 2x instead of 4x (sm_120 SASS: 2 loads in flight instead of 4). The helper waits on those loads, so it gets about 1.7x slower on every MMQ call with ids, inactive or not, and prefill gets about 3% slower. (On ce2d3c9 this was hidden: zero_inactive_routes had just read ids, so the next helper found it in L2.)

Comparing as unsigned gives the same order without the remap, because a negative id then sorts after every expert and never equals one:

const int expert_used = ...;  // as on main
const int iex_used = expert_used == expert ? iex : -1;
nex_prev += (unsigned) expert_used < (unsigned) expert;

and the same nex_prev line in the generic loop. test_ggml.py and my real-bank checks with -1 routes still pass with it.

Qwen3.6-35B-A3B-UD-Q4_K_M.gguf, --moe-strategy fused, RTX PRO 6000:

main 2ff4040 2ff4040 + unsigned compare
mm_ids_helper alone, 512 tokens 19.5 us 33.5 us 19.5 us
prefill, 1959-token prompt 93.7 ms 96.7 ms 93.2 ms
prefill, 8477-token prompt 374.3 ms 386.6 ms 373.7 ms
decode, 16 streams 1590.1 tok/s 1584.3 tok/s 1586.2 tok/s
decode, 32 streams 2437.3 tok/s 2439.1 tok/s 2437.7 tok/s

A negative (inactive) id compared unsigned sorts after every expert and equals none, so the
helper needs no remap to INT_MAX; the remap made nvcc unroll its token loop half as far.
@ischencheng

Copy link
Copy Markdown
Contributor Author

done in a45f0b1, both loops compare unsigned as you wrote and the INT_MAX remap is gone. test_ggml.py and the kernels/moe suite pass on A10 with the same two container failures as main. With nvcc 13.0 I couldn't reproduce the slower helper: on an RTX PRO 6000 mm_ids_helper at 512 tokens is 19.1 us on main, 2ff4040 and a45f0b1 alike, numbers are in the description.

Resolves freetoken.patch and VERSION after FlashML-org#835: the patch is regenerated against llama.cpp
8c2134a8 with both changes, and VERSION lists both.
@jason-fxz

Copy link
Copy Markdown
Collaborator

Thanks, a45f0b1 is what I meant, LGTM.

@jason-fxz
jason-fxz merged commit ecbbce4 into FlashML-org:main Oct 29, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants