feat(gguf): apply the weight scales of nvfp4 gguf tensors

python/freetoken/engine/engine.py

@@ -1770,6 +1770,9 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             "--dtype float16 with MXFP8 resident weights is unsupported (the "
             "W8A16 fold is only validated exact in bfloat16); use bfloat16."
         )
+    if _dtype == torch.float16 and getattr(getattr(model_config, "quant", None), "weight_scales", False):
+        # GGUF NVFP4 results are stored in the activation dtype before weight_scale_2 (~1e-4) multiplies them, which overflows fp16.
+        raise ValueError("--dtype float16 with NVFP4 GGUF weight scales is unsupported (the unscaled results overflow float16); use bfloat16.")
     if config.attention_backend == "auto":
         override(
             "attention_backend",

python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp

@@ -48,6 +48,7 @@ std::unique_ptr<MoeFormat> make_mxfp4_format(const MoeParams& p, const FormatArg
 std::unique_ptr<MoeFormat> make_ds_fp4_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_fp8_block_format(const MoeParams& p, const FormatArgs& a);
 std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a);
+std::unique_ptr<MoeFormat> make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a);
 void set_ggml_cpu(uintptr_t addr);
 
 namespace {
@@ -97,6 +98,7 @@ inline MoeTask from_word(uint64_t w) {
 const std::pair<const char*, std::unique_ptr<MoeFormat> (*)(const MoeParams&, const FormatArgs&)> kFormats[] = {
     {"bf16", make_bf16_format},     {"nvfp4", make_nvfp4_format},         {"mxfp4_triton", make_mxfp4_format},
     {"ds_fp4", make_ds_fp4_format}, {"fp8_block", make_fp8_block_format}, {"gguf", make_gguf_format},
+    {"gguf_scaled", make_gguf_scaled_format},
 };
 
 std::unique_ptr<MoeFormat> make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) {

python/freetoken/kernel/csrc/cpu_moe/kernels/gguf.cpp

@@ -31,10 +31,11 @@ void set_ggml_cpu(uintptr_t addr) {
 namespace {
 
 // Packed ggml rows of the layers whose experts have this gate_up type and this down type, dotted by ggml-cpu against the activation converted to each type's vec_dot input (q8_0, q8_1 or q8_K, or the float type itself).
+// With scales, each expert's (gate, up, down, 1) NVFP4 weight_scale_2 multiplies its gate and up rows before the glu and its down rows before the routing weight, as llama.cpp's build_moe_ffn does.
 class GgufFormat final : public MoeFormat {
  public:
-  GgufFormat(const MoeParams& p, const FormatArgs& a)
-      : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {
+  GgufFormat(const MoeParams& p, const FormatArgs& a, bool scaled)
+      : MoeFormat(p), ggml_(ggml_api()), gate_up_(a.bank("gate_up")), down_(a.bank("down")), scales_(scaled ? a.bank("global") : nullptr) {
     static std::once_flag tables;
     std::call_once(tables, ggml_.init);
     gu_ = kernel(static_cast<int>(a.option("gate_up_type")), p.H, "gate_up");
@@ -63,10 +64,11 @@ class GgufFormat final : public MoeFormat {
   void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, bf16_t* g) const override {
     const uint8_t* w = static_cast<const uint8_t*>(tbl_at(gate_up_, layer)) + static_cast<size_t>(e) * 2 * p.I * gu_.row_bytes;
     const uint8_t* xq = xq_.data() + static_cast<size_t>(tok) * gu_.act_row_bytes;
+    const float* sc = expert_scales(layer, e);
     gated_rows(p, i0, i1, w_in, g, [&](int row) {
       float s;
       gu_.vec_dot(p.H, &s, 0, w + static_cast<size_t>(row) * gu_.row_bytes, 0, xq, 0, 1);
-      return s;
+      return sc == nullptr ? s : s * sc[row < p.I ? 0 : 1];
     });
   }
 
@@ -76,9 +78,19 @@ class GgufFormat final : public MoeFormat {
       for (int j = 0; j < n; ++j)
         dn_.vec_dot(p.I, &out[j * kRowBlock + h - h0], 0, w + (static_cast<size_t>(routes[j].e) * p.H + h) * dn_.row_bytes, 0,
                     gq_.data() + static_cast<size_t>(routes[j].r) * dn_.act_row_bytes, 0, 1);
+    if (scales_ == nullptr) return;
+    for (int j = 0; j < n; ++j) {
+      const float s = expert_scales(layer, routes[j].e)[2];
+      for (int h = h0; h < h1; ++h) out[j * kRowBlock + h - h0] *= s;
+    }
   }
 
  private:
+  // the expert's (gate, up, down, 1) row of the layer's [E, 4] fp32 scale bank, or nullptr without scales
+  const float* expert_scales(int layer, int e) const {
+    return scales_ == nullptr ? nullptr : static_cast<const float*>(tbl_at(scales_, layer)) + static_cast<size_t>(e) * 4;
+  }
+
   ft_ggml_row_kernel kernel(int type, int k, const char* role) const {
     ft_ggml_row_kernel out{};
     const int rc = ggml_.row_kernel(type, k, &out);
@@ -92,6 +104,7 @@ class GgufFormat final : public MoeFormat {
   const ft_ggml_cpu_api& ggml_;
   const uint64_t* gate_up_;
   const uint64_t* down_;
+  const uint64_t* scales_;
   ft_ggml_row_kernel gu_{}, dn_{};
   std::vector<float> xf_;
   // the activation in the vec_dot input type: x per token, g per route
@@ -101,7 +114,11 @@ class GgufFormat final : public MoeFormat {
 }  // namespace
 
 std::unique_ptr<MoeFormat> make_gguf_format(const MoeParams& p, const FormatArgs& a) {
-  return std::make_unique<GgufFormat>(p, a);
+  return std::make_unique<GgufFormat>(p, a, false);
+}
+
+std::unique_ptr<MoeFormat> make_gguf_scaled_format(const MoeParams& p, const FormatArgs& a) {
+  return std::make_unique<GgufFormat>(p, a, true);
 }
 
 }  // namespace freetoken::cpu_moe

python/freetoken/kernel/csrc/ggml/VERSION

@@ -10,5 +10,6 @@ Local changes, kept in freetoken.patch (already applied):
 - mmvq.cu/.cuh: MMVQ takes bf16 or fp16 activations, can write bf16 or fp16 results, and can reuse activations already quantized to q8_1 (src1_q8_1_pre).
 - mmq.cu/.cuh, mmid.cu/.cuh: MoE MMQ runs a flat grid over only the (expert, column tile) pairs that hold tokens, listed by mm_ids_tile_map; ncols_max > 0 keeps the regular grid.
 - common.cuh: the native fp4 MMA (BLACKWELL_MMA_AVAILABLE, blackwell_mma_available) is off unless GGML_CUDA_FP4_MMA is defined; it needs an sm_120a/sm_120f build.
+- 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.
 
 To update: copy the same files from a newer llama.cpp, apply freetoken.patch, and update this file.

python/freetoken/kernel/csrc/ggml/freetoken.patch

@@ -24,6 +24,16 @@
  }
 
  // Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+@@ -1582,6 +1588,9 @@
+     const void * gate_bias = nullptr;
+     const void * x_scale = nullptr;
+     const void * gate_scale = nullptr;
++    // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank
++    uint32_t x_scale_stride = 1;
++    uint32_t gate_scale_stride = 1;
+     ggml_glu_op glu_op;
+     float glu_limit = 0.0f;
+     const void * shared_up = nullptr;
 --- a/ggml/src/ggml-cuda/mmid.cu
 +++ b/ggml/src/ggml-cuda/mmid.cu
 @@ -177,3 +177,45 @@
@@ -572,6 +582,19 @@
          stride_col_dst = fusion.shared_stride_col_dst;
      }
 
+@@ -702,10 +706,10 @@
+             }
+             if constexpr (type == GGML_TYPE_NVFP4) {
+                 if (use_scale) {
+-                    x_scales = x_scale[ids ? channel_x : 0];
++                    x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride];
+                 }
+                 if (use_gate_scale) {
+-                    gate_scales = gate_scale[ids ? channel_x : 0];
++                    gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride];
+                 }
+             }
+         }
 @@ -837,7 +841,7 @@
                          }
                      }
@@ -615,6 +638,24 @@
          stride_col_dst = fusion.shared_stride_col_dst;
      }
 
+@@ -963,7 +969,7 @@
+
+             if constexpr (type == GGML_TYPE_NVFP4) {
+                 if (x_scale) {
+-                    result *= x_scale[channel_x];
++                    result *= x_scale[channel_x * fusion.x_scale_stride];
+                 }
+             }
+             if (x_bias) {
+@@ -973,7 +979,7 @@
+                 float gate_value = tmp_gate[threadIdx.x];
+                 if constexpr (type == GGML_TYPE_NVFP4) {
+                     if (gate_scale) {
+-                        gate_value *= gate_scale[channel_x];
++                        gate_value *= gate_scale[channel_x * fusion.gate_scale_stride];
+                     }
+                 }
+                 if (gate_bias) {
 @@ -998,7 +1004,7 @@
                  }
              }
@@ -718,7 +759,30 @@
      GGML_ASSERT(!ids || ids->type  == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
 
      GGML_TENSOR_BINARY_OP_LOCALS;
-@@ -1526,12 +1535,22 @@
+@@ -1498,17 +1507,20 @@
+             GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
+             fusion_local.gate_bias = fusion->gate_bias->data;
+         }
++        // FreeToken: a scale may be a strided vector, read with its element stride
+         if (fusion->x_scale) {
+             GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
+-            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
++            GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0);
+             GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
+             fusion_local.x_scale = fusion->x_scale->data;
++            fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float);
+         }
+         if (fusion->gate_scale) {
+             GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
+-            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
++            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0);
+             GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
+             fusion_local.gate_scale = fusion->gate_scale->data;
++            fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float);
+         }
+         fusion_local.glu_op = fusion->glu_op;
+         fusion_local.glu_limit = fusion->glu_limit;
+@@ -1526,12 +1538,22 @@
      }
 
      const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING);
@@ -744,7 +808,7 @@
      }
 
      const int64_t s01 = src0->nb[1] / ts_src0;
-@@ -1556,11 +1575,23 @@
+@@ -1556,11 +1578,23 @@
 
      const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0;
 

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

@@ -115,6 +115,14 @@ int64_t check_weight(TensorView w, ggml_type type, int64_t k) {
     return row_bytes;
 }
 
+// one fp32 scale per expert, e.g. a column of a [E, n] scale bank; the MMVQ epilogue multiplies NVFP4 results by it
+ggml_tensor describe_scale(TensorView s, ggml_type type, int64_t e) {
+    RuntimeCheck(type == GGML_TYPE_NVFP4, "per-expert scales apply to NVFP4 weights, not ", ggml_type_name(type));
+    RuntimeCheck(is_dtype(s.dtype(), kDLFloat, 32) && s.ndim() == 1 && s.size(0) == e && s.stride(0) > 0, "a per-expert scale must be an fp32 [E] view");
+    const size_t nb0 = sizeof(float) * s.stride(0);
+    return describe(GGML_TYPE_F32, s.data_ptr(), {e, 1, 1, 1}, {nb0, nb0 * e, nb0 * e, nb0 * e});
+}
+
 int64_t q8_1_row_bytes(int64_t k) { return GGML_PAD(k, MATRIX_ROW_PADDING) / QK8_1 * (int64_t) sizeof(block_q8_1); }
 
 // MMQ tiles read up to MATRIX_ROW_PADDING values past the last row; the caller keeps this many bytes after it
@@ -234,7 +242,8 @@ void op_mul_mat(TensorView w, int64_t type_id, TensorView x, TensorView out, Opt
 }
 
 // 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
-void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out) {
+// 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) {
     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());
@@ -264,8 +273,15 @@ void op_mul_mat_id(TensorView w, int64_t type_id, TensorView x, TensorView ids,
     ggml_tensor src1 = describe(a.type, a.data, {k, ky, t, 1}, {a.esz, a.esz * k, a.esz * k * ky, a.esz * k * ky * t});
     const size_t oe = ggml_type_size(otype);
     ggml_tensor dst = describe(otype, out.data_ptr(), {n, top_k, t, 1}, {oe, oe * n, oe * n * top_k, oe * n * top_k * t});
+    RuntimeCheck(mmvq || !scale.has_value(), "mul_mat_id takes a per-expert scale only on the vector kernel, T=", t);
     if (mmvq) {
-        ggml_cuda_mul_mat_vec_q(ctx, &src0, &src1, &tids, &dst);
+        ggml_cuda_mm_fusion_args_host fusion{};
+        ggml_tensor tscale;
+        if (scale.has_value()) {
+            tscale = describe_scale(scale.value(), type, e);
+            fusion.x_scale = &tscale;
+        }
+        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);
     }
@@ -291,7 +307,9 @@ ggml_glu_op fused_glu_op(const String & name) {
 }  // namespace
 
 // out[T, k, I] = glu(x @ gate^T, x @ up^T) for small T in one MMVQ pass; w[E, 2I, row_bytes] holds each expert's gate rows then its up rows; limit is swiglu_clamp's
-void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit) {
+// gate_scale / up_scale multiply the expert's gate / up results before the glu
+void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView ids, TensorView out, String glu, double limit,
+                       Optional<TensorView> gate_scale, Optional<TensorView> up_scale) {
     const ggml_glu_op glu_op = fused_glu_op(glu);
     RuntimeCheck(w.ndim() == 3 && w.size(1) % 2 == 0, "mul_mat_id_glu takes w[E, 2I, row_bytes]");
     RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 2 && ids.IsContiguous(), "ids must be a contiguous int32 [T, k] tensor");
@@ -325,6 +343,15 @@ void op_mul_mat_id_glu(TensorView w, int64_t type_id, TensorView x, TensorView i
     fusion.gate = &gate;
     fusion.glu_op = glu_op;
     fusion.glu_limit = (float) limit;
+    ggml_tensor tgate_scale, tup_scale;
+    if (gate_scale.has_value()) {
+        tgate_scale = describe_scale(gate_scale.value(), type, e);
+        fusion.gate_scale = &tgate_scale;
+    }
+    if (up_scale.has_value()) {
+        tup_scale = describe_scale(up_scale.value(), type, e);
+        fusion.x_scale = &tup_scale;
+    }
     ggml_cuda_mul_mat_vec_q(ctx, &up, &src1, &tids, &dst, &fusion);
 }
 

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

@@ -1588,6 +1588,9 @@ struct ggml_cuda_mm_fusion_args_device {
     const void * gate_bias = nullptr;
     const void * x_scale = nullptr;
     const void * gate_scale = nullptr;
+    // FreeToken: element strides of x_scale / gate_scale, so they can be columns of a per-expert scale bank
+    uint32_t x_scale_stride = 1;
+    uint32_t gate_scale_stride = 1;
     ggml_glu_op glu_op;
     float glu_limit = 0.0f;
     const void * shared_up = nullptr;

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

@@ -706,10 +706,10 @@ static __global__ void mul_mat_vec_q(
             }
             if constexpr (type == GGML_TYPE_NVFP4) {
                 if (use_scale) {
-                    x_scales = x_scale[ids ? channel_x : 0];
+                    x_scales = x_scale[(ids ? channel_x : 0) * fusion.x_scale_stride];
                 }
                 if (use_gate_scale) {
-                    gate_scales = gate_scale[ids ? channel_x : 0];
+                    gate_scales = gate_scale[(ids ? channel_x : 0) * fusion.gate_scale_stride];
                 }
             }
         }
@@ -969,7 +969,7 @@ static __global__ void mul_mat_vec_q_moe(
 
             if constexpr (type == GGML_TYPE_NVFP4) {
                 if (x_scale) {
-                    result *= x_scale[channel_x];
+                    result *= x_scale[channel_x * fusion.x_scale_stride];
                 }
             }
             if (x_bias) {
@@ -979,7 +979,7 @@ static __global__ void mul_mat_vec_q_moe(
                 float gate_value = tmp_gate[threadIdx.x];
                 if constexpr (type == GGML_TYPE_NVFP4) {
                     if (gate_scale) {
-                        gate_value *= gate_scale[channel_x];
+                        gate_value *= gate_scale[channel_x * fusion.gate_scale_stride];
                     }
                 }
                 if (gate_bias) {
@@ -1507,17 +1507,20 @@ void ggml_cuda_mul_mat_vec_q(
             GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]);
             fusion_local.gate_bias = fusion->gate_bias->data;
         }
+        // FreeToken: a scale may be a strided vector, read with its element stride
         if (fusion->x_scale) {
             GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32);
-            GGML_ASSERT(ggml_is_contiguous(fusion->x_scale));
+            GGML_ASSERT(ggml_nelements(fusion->x_scale) == fusion->x_scale->ne[0] && fusion->x_scale->nb[0] % sizeof(float) == 0);
             GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1));
             fusion_local.x_scale = fusion->x_scale->data;
+            fusion_local.x_scale_stride = fusion->x_scale->nb[0] / sizeof(float);
         }
         if (fusion->gate_scale) {
             GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32);
-            GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale));
+            GGML_ASSERT(ggml_nelements(fusion->gate_scale) == fusion->gate_scale->ne[0] && fusion->gate_scale->nb[0] % sizeof(float) == 0);
             GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1));
             fusion_local.gate_scale = fusion->gate_scale->data;
+            fusion_local.gate_scale_stride = fusion->gate_scale->nb[0] / sizeof(float);
         }
         fusion_local.glu_op = fusion->glu_op;
         fusion_local.glu_limit = fusion->glu_limit;

python/freetoken/kernel/ggml.py

@@ -95,24 +95,37 @@ 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) -> 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]``."""
+def mul_mat_id(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, scale: torch.Tensor | None = None) -> 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``.
+    """
     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)
+    _module().mul_mat_id(w, ggml_type, (x.unsqueeze(1) if x.dim() == 2 else x).contiguous(), ids, out, scale)
     return out
 
 
-def mul_mat_id_glu(w: torch.Tensor, ggml_type: int, x: torch.Tensor, ids: torch.Tensor, glu: str, limit: float = 0.0) -> torch.Tensor:
+def mul_mat_id_glu(
+    w: torch.Tensor,
+    ggml_type: int,
+    x: torch.Tensor,
+    ids: torch.Tensor,
+    glu: str,
+    limit: float = 0.0,
+    gate_scale: torch.Tensor | None = None,
+    up_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
     """``glu(gate, up)`` per route in one pass for small ``T``: ``w[E, 2I, row_bytes]`` holds gate rows then up rows; returns ``[T, k, I]`` in ``x``'s dtype.
 
     ``glu`` is one of ``swiglu``, ``geglu`` (tanh GELU), ``swiglu_oai`` (alpha 1.702, limit 7) and ``swiglu_clamp`` (alpha 1, clamped at ``limit``).
+    ``gate_scale`` / ``up_scale`` are fp32 ``[E]`` views that multiply each NVFP4 expert's gate / up results before the glu.
     """
     t, top_k = ids.shape
     out = torch.empty(t, top_k, w.shape[1] // 2, dtype=x.dtype, device=x.device)
-    _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit)
+    _module().mul_mat_id_glu(w, ggml_type, x.contiguous(), ids, out, glu, limit, gate_scale, up_scale)
     return out