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

python/freetoken/kernel/csrc/ggml/VERSION

@@ -13,5 +13,6 @@ Local changes, kept in freetoken.patch (already applied):
 - mmvq.cu, common.cuh: the NVFP4 scale fusion (x_scale, gate_scale) reads strided scale vectors, e.g. columns of a per-expert scale bank.
 - getrows.cu, dequantize.cuh: get_rows reads NVFP4 rows, e.g. an NVFP4 token_embd.
 - mmq.cu/.cuh, mmq-load-tiles.cuh, template-instances/mmq-instance-iq1_m.cu: MMQ runs IQ1_M, which llama.cpp has no MMQ kernel for. Its values times 8 are int8, so it loads into IQ2_XS's tile layout and takes IQ2_XS's configs.
+- mmvq.cu, mmid.cu, mmq.cu/.cuh, quantize.cu: a negative expert id is an inactive route (hybrid decode): MMVQ writes zero rows for it without reading an expert, and MMQ gives it no compact row when called with inactive_ids.
 
 To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file.

python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu

@@ -243,9 +243,42 @@ void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Opt
     }
 }
 
+namespace {
+
+// One 32-thread group per route: zero row r of out[routes, n] when ids[r] < 0.
+template <typename T>
+__global__ void zero_inactive_routes_kernel(const int32_t * __restrict__ ids, T * __restrict__ out, const int64_t routes, const int64_t n) {
+    const int64_t r = (int64_t) blockIdx.x * (blockDim.x / 32) + threadIdx.x / 32;
+    if (r >= routes || ids[r] >= 0) {
+        return;
+    }
+    for (int64_t i = threadIdx.x % 32; i < n; i += 32) {
+        out[r * n + i] = ggml_cuda_cast<T>(0.0f);
+    }
+}
+
+// The matrix kernel writes only the rows of active routes; the inactive ones (negative ids, hybrid decode) read as zero like on the vector kernel.
+void zero_inactive_routes(TensorView ids, TensorView out, ggml_type otype, cudaStream_t stream) {
+    constexpr int threads = 256;
+    const int64_t routes = ids.numel(), n = out.size(2);
+    const dim3 blocks((routes + threads / 32 - 1) / (threads / 32));
+    const int32_t * ip = (const int32_t *) ids.data_ptr();
+    if (otype == GGML_TYPE_BF16) {
+        zero_inactive_routes_kernel<<<blocks, threads, 0, stream>>>(ip, (nv_bfloat16 *) out.data_ptr(), routes, n);
+    } else if (otype == GGML_TYPE_F16) {
+        zero_inactive_routes_kernel<<<blocks, threads, 0, stream>>>(ip, (half *) out.data_ptr(), routes, n);
+    } else {
+        zero_inactive_routes_kernel<<<blocks, threads, 0, stream>>>(ip, (float *) out.data_ptr(), routes, n);
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
+}  // namespace
+
 // out[T, k, N]: row (t, j) multiplies x[t, j or 0] by expert ids[t, j] of W[E, N, K]; the matrix kernel runs a flat grid over the tiles that hold routes
 // scale[ids[t, j]] multiplies row (t, j); only the vector kernel takes it
-void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, Optional<TensorView> scale) {
+// inactive: ids may hold negative (inactive) routes, which the matrix kernel then skips and zeroes; the vector kernel always does
+void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, Optional<TensorView> scale, bool inactive) {
     RuntimeCheck(w.ndim() == 3, "mul_mat_id takes w[E, N, row_bytes]");
     RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
     const DeviceGuard guard(x.device());
@@ -285,7 +318,10 @@ void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids,
         }
         ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst, scale.has_value() ? &fusion : nullptr);
     } else {
-        ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0);
+        ggml_cuda_mul_mat_q(ctx, &src0, &src1, &tids, &dst, 0, inactive);
+        if (inactive) {
+            zero_inactive_routes(ids, out, otype, stream);
+        }
     }
 }
 

python/freetoken/kernel/csrc/ggml/ggml-cuda/mmid.cu

@@ -49,7 +49,7 @@ static __global__ void mm_ids_helper(
             int iex_used = -1; // The index at which the expert is used, if any.
             for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) {
                 const int expert_used = ids[it*si1 + iex];
-                nex_prev += expert_used < expert;
+                nex_prev += (unsigned) expert_used < (unsigned) expert; // FreeToken: unsigned, so an inactive route (negative id) sorts after every expert
                 if (expert_used == expert) {
                     iex_used = iex;
                 }
@@ -75,7 +75,7 @@ static __global__ void mm_ids_helper(
             const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ?
                 ids[it*si1 + iex] : INT_MAX;
             const int iex_used = expert_used == expert ? iex : -1;
-            nex_prev += expert_used < expert;
+            nex_prev += (unsigned) expert_used < (unsigned) expert; // FreeToken: unsigned, so an inactive route (negative id) sorts after every expert
 
             // Whether the threads at this token position have used the expert:
             const int it_compact_add_self = warp_reduce_any<neu_padded>(iex_used != -1);

python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cu

@@ -136,7 +136,7 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm
 
 void ggml_cuda_mul_mat_q(
         ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
-        int64_t ncols_max) {
+        int64_t ncols_max, bool inactive_ids) {
     // FreeToken: bf16 and fp16 activations are quantized directly, except for the native fp4 paths; dst is fp32, bf16 or fp16
     GGML_ASSERT(        src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_BF16 || src1->type == GGML_TYPE_F16);
     GGML_ASSERT(        dst->type  == GGML_TYPE_F32 || dst->type  == GGML_TYPE_BF16 || dst->type  == GGML_TYPE_F16);
@@ -299,6 +299,10 @@ void ggml_cuda_mul_mat_q(
     ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
     ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access.
     ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
+    // FreeToken: inactive routes (negative ids, only with inactive_ids) get no compact row, so their ids_src1 entries stay -1
+    if (inactive_ids) {
+        CUDA_CHECK(cudaMemsetAsync(ids_src1.get(), 0xFF, ne_get_rows*sizeof(int32_t), stream));
+    }
 
     // gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
     // scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).

python/freetoken/kernel/csrc/ggml/ggml-cuda/mmq.cuh

@@ -1686,6 +1686,6 @@ extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
 // FreeToken: under ids, ncols_max > 0 bounds the tokens any one expert gets and sizes the tile grid; otherwise the grid is flat, see ggml_cuda_launch_mm_ids_tile_map
 void ggml_cuda_mul_mat_q(
         ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
-        int64_t ncols_max = -1);
+        int64_t ncols_max = -1, bool inactive_ids = false);
 
 bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);

python/freetoken/kernel/csrc/ggml/ggml-cuda/mmvq.cu

@@ -646,6 +646,15 @@ static __global__ void mul_mat_vec_q(
     channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst;
     sample_dst = blockIdx.z;
 
+    // FreeToken: a negative id is an inactive route (hybrid decode): zero rows, no expert read
+    if (ncols_dst == 1 && ids && !shared_expert && int32_t(channel_x) < 0) {
+        if (threadIdx.y == 0 && threadIdx.x < rows_per_cuda_block &&
+            (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) {
+            dst[sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row0 + threadIdx.x] = ggml_cuda_cast<dst_t>(0.0f);
+        }
+        return;
+    }
+
     const uint32_t sample_x    = fastdiv(sample_dst, sample_ratio);
     const uint32_t sample_y    = sample_dst;
 
@@ -926,6 +935,14 @@ static __global__ void mul_mat_vec_q_moe(
     const uint32_t channel_x = shared_expert ? 0 : ids[channel_dst + token_idx * ids_stride];
     const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y);
 
+    // FreeToken: a negative id is an inactive route (hybrid decode): zero rows, no expert read
+    if (!shared_expert && int32_t(channel_x) < 0) {
+        if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
+            dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = ggml_cuda_cast<dst_t>(0.0f);
+        }
+        return;
+    }
+
     const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y;
     const int kbx_offset  = channel_x*stride_channel_x + row0*stride_row_x;
 

python/freetoken/kernel/csrc/ggml/ggml-cuda/quantize.cu

@@ -139,6 +139,9 @@ static __global__ void quantize_mmq_nvfp4(
     } else {
         const int64_t i2  = blockIdx.y % ne2;
         const int64_t i3  = blockIdx.y / ne2;
+        if (ids && ids[blockIdx.x] < 0) {
+            return; // FreeToken: a compact row past the active routes (inactive routes, hybrid decode)
+        }
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
         base_idx = i3 * s03 + i2 * s02 + i01 * s01;
     }
@@ -184,6 +187,9 @@ static __global__ void quantize_mmq_nvfp4(
 #pragma unroll
                 for (int slot = 0; slot < n_expert_used; ++slot) {
                     const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                    if (i < 0) {
+                        continue; // FreeToken: an inactive route has no compact row
+                    }
                     scale[i] = warp_amax[0];
                 }
             } else {
@@ -310,6 +316,9 @@ static __global__ void quantize_mmq_nvfp4(
 #pragma unroll
             for (int slot = 0; slot < n_expert_used; ++slot) {
                 const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                if (i < 0) {
+                    continue; // FreeToken: an inactive route has no compact row
+                }
                 block_fp4_mmq * yb = y + (k_block * ne1 + i);
                 uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
                 yqs[2 * sub + 0] = q0;
@@ -376,6 +385,9 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
     } else {
         const int64_t i2  = blockIdx.z % ne2;
         const int64_t i3  = blockIdx.z / ne2;
+        if (ids && ids[blockIdx.x] < 0) {
+            return; // FreeToken: a compact row past the active routes (inactive routes, hybrid decode)
+        }
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
         base_pos = i3 * s03 + i2 * s02 + i01 * s01;
     }
@@ -429,6 +441,9 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
 #pragma unroll
         for (int slot = 0; slot < n_expert_used; ++slot) {
             const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            if (i < 0) {
+                continue; // FreeToken: an inactive route has no compact row
+            }
             block_fp4_mmq * yb = y + (k_block * ne1 + i);
             char2 * yqs2 = (char2 *) yb->qs;
             if (lane_in_group == 0) {
@@ -498,6 +513,9 @@ static __global__ void quantize_mmq_q8_1(
     } else {
         const int64_t i2  = blockIdx.z % ne2;
         const int64_t i3  = blockIdx.z / ne2;
+        if (ids && ids[blockIdx.x] < 0) {
+            return; // FreeToken: a compact row past the active routes (inactive routes, hybrid decode)
+        }
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
         base_idx = i3*s03 + i2*s02 + i01*s01;
     }
@@ -546,6 +564,9 @@ static __global__ void quantize_mmq_q8_1(
         int64_t ib;
         if constexpr (scatter) {
             const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+            if (i < 0) {
+                continue; // FreeToken: an inactive route has no compact row
+            }
             ib = k_block*ne1 + i;
         } else {
             const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel

python/freetoken/kernel/ggml.py

@@ -95,16 +95,19 @@ def uses_mmvq(ggml_type: int, m: int) -> bool:
     return _module().uses_mmvq(ggml_type, m)
 
 
-def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None) -> torch.Tensor:
+def mul_mat_id(
+    w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None, inactive: bool = False
+) -> torch.Tensor:
     """Per-route expert GEMM over ``w[E, N, row_bytes]``: returns ``[T, k, N]`` in ``x``'s dtype; ``x`` is ``[T, K]`` (shared by the k routes) or ``[T, k, K]``.
 
     ``scale`` is an fp32 ``[E]`` view (strided is fine) that multiplies each NVFP4 expert's rows; it needs ``T <= mmvq_max_tokens``.
+    ``inactive`` says ``ids`` may hold -1 (hybrid decode's CPU routes): those rows come back zero. Above ``mmvq_max_tokens`` it costs a memset and a zeroing pass, so pass it only then.
     """
     t, top_k = ids.shape
     out = torch.empty(t, top_k, w.shape[1], dtype=x.dtype, device=x.device)
     if t > mmvq_max_tokens(ggml_type):
         _check_mmq_tail(w, ggml_type, x.shape[-1])
-    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale)
+    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale, inactive)
     return out
 
 

python/freetoken/layers/quantization/moe/gguf.py

@@ -69,6 +69,7 @@ def float_mul_mat_id(w: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, topk_w
 class GgmlMoEKernel(MoEKernel):
     name = "ggml"
     cpu_format = "gguf"
+    supports_inactive_slots = True
 
     def unusable_reason(self, cfg: MoEConfig) -> str | None:
         if not is_quantized(expert_types(cfg)[0]):
@@ -103,6 +104,7 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
             return self._apply_float(layer, x, topk_weights, topk_ids, view, gate_up_type, is_prefill=is_prefill)
         assert not layer.apply_router_weight_on_input, "GGUF experts apply the router weights on the output"
         ids = topk_ids.to(torch.int32)
+        inactive = not is_prefill and cfg.decode_target == "hybrid"  # only hybrid decode hands over -1 routes
         gate_up = view.tensors["gate_up"]
         # a float half runs on the bf16 fused MoE's grouped GEMM, with its config for this layer's shape
         config = None
@@ -114,8 +116,9 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 
         def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor:
             if is_quantized(ggml_type):
-                return ggml.mul_mat_id(w, ggml_type, a, ids)
-            return float_mul_mat_id(w, a, ids, topk_weights, config)
+                return ggml.mul_mat_id(w, ggml_type, a, ids, inactive=inactive)
+            # the bf16 kernels read every route's slot: an inactive one reads slot 0 and gets weight 0
+            return float_mul_mat_id(w, a, ids.clamp_min(0), topk_weights, config)
 
         # llama.cpp's build_moe_ffn: each NVFP4 GEMM result times the expert's weight_scale_2, gate and up before the glu, down before the routing weight
         scales = view.tensors.get("global")
@@ -135,7 +138,7 @@ def mul_mat_id(w: torch.Tensor, ggml_type: int, a: torch.Tensor) -> torch.Tensor
             gated_act_and_mul(layer.activation, h.view(-1, two_i), inter.view(-1, two_i // 2), alpha=float(layer.alpha), limit=limit_or_inf(layer))
             del h  # free the gate_up output before the down GEMM allocates its own
         if down_scales is not None and x.shape[0] <= ggml.mmvq_max_tokens(down_type):
-            y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids, scale=down_scales[:, 2])
+            y = ggml.mul_mat_id(view.tensors["down"], down_type, inter, ids, scale=down_scales[:, 2], inactive=inactive)
         else:
             y = mul_mat_id(view.tensors["down"], down_type, inter)
             if down_scales is not None:
@@ -151,6 +154,7 @@ def _apply_float(layer, x, topk_weights, topk_ids, view: ExpertView, ggml_type:
         if gate_up.dtype != x.dtype:
             raise NotImplementedError(f"GGUF {type_name(ggml_type)} experts need {gate_up.dtype} activations, got {x.dtype}")
         impl = fused_experts_impl if is_prefill or is_resident(layer) else fused_experts_decode_impl
+        topk_ids = topk_ids.clamp_min(0)  # as in apply: inactive routes read slot 0 with weight 0
         return impl(x, gate_up, down, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, float(layer.alpha), limit_or_inf(layer))
 
 

tests/kernels/test_ggml.py

@@ -91,14 +91,20 @@ def test_q4_matrix_matches_cpu(dtype, batch):
 
 
 @pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("batch", [1, 17])
-def test_q4_moe_routes_match_cpu(dtype, batch):
+@pytest.mark.parametrize("batch", [1, 4, 17])
+@pytest.mark.parametrize("per_route_x", [False, True])
+@pytest.mark.parametrize("inactive", [False, True])
+def test_q4_moe_routes_match_cpu(dtype, batch, per_route_x, inactive):
     w, ref = _q4_weights((4, 64, 256))
-    x = _activation((batch, 256), dtype)
+    x = _activation((batch, 2, 256) if per_route_x else (batch, 256), dtype)
     ids = torch.tensor([[i % 4, (i + 1) % 4] for i in range(batch)], dtype=torch.int32, device="cuda")
-    actual = ggml.mul_mat_id(w, int(QType.Q4_0), x, ids)
-    expected = torch.einsum("tk, tjok -> tjo", x.cpu().float(), ref[ids.cpu().long()]).to(dtype)
-    torch.testing.assert_close(actual.cpu(), expected, atol=2e-4, rtol=1e-3)
+    if inactive:  # hybrid decode hands the GPU the CPU's routes as id -1: zero rows, no expert read
+        ids[::2, 0] = -1
+    actual = ggml.mul_mat_id(w, int(QType.Q4_0), x, ids, inactive=inactive)
+    xs = x.cpu().float() if per_route_x else x.cpu().float().unsqueeze(1).expand(-1, 2, -1)
+    expected = torch.einsum("tjk, tjok -> tjo", xs, ref[ids.cpu().long().clamp_min(0)])
+    expected[ids.cpu() < 0] = 0
+    torch.testing.assert_close(actual.cpu(), expected.to(dtype), atol=2e-4, rtol=1e-3)
 
 
 @pytest.mark.parametrize("dtype", DTYPES)