Commit c49edf096 for llama.cpp
commit c49edf0963a76d5e78b2a7278c9bd5b71d26c84f
Author: anujj <ajalota@nvidia.com>
Date: Sun Oct 25 14:54:38 2026 +0530
CUDA: fuse FFN gate/up matmuls and GLU in MMQ with interleaved gate/up rows (#29948)
* cuda: fuse FFN gate/up and GLU in MMQ with interleaved gate/up rows
For prompt processing, run the gate and up matmuls of a dense FFN as one
mul_mat_q launch and apply the GLU in the write-back.
- The loaders take the gate weights as a second pointer and fill the weight
tile with up and gate rows interleaved in chunks of 8, so each thread holds
an up value and the matching gate value (rows r and r + 8 of an MMA tile).
- The MMA loop is unchanged: same tile, warp layout, vec_dot and one sum[]
per thread. The fused write-back computes the GLU from registers and writes
half the rows.
- The fused path uses the regular tiling (no stream-k) with the column tiles
innermost so that neighbouring blocks share weight tiles.
- Enabled for Q4_K, Q5_K, Q6_K, Q8_0 and IQ4_XS on NVIDIA Turing and newer.
* cuda: assert the fused MMQ gate argument and trim redundant comments
* cuda: move the fused MMQ hint check to should_fuse and add TODOs
* cuda: pass prec_src1 to the fused MMQ write-back
* cuda: fuse only with efficient tiling and no fallback, merge the GLU write-back
* cuda: require has_fusion explicitly, drop the fusion-only src1 precision check
* cuda: launch the fused MMQ kernel through the stream-k kernel with one block per tile
* cuda: inline the MMQ tiling-efficiency check
* tests: one fused gate/up MMQ case per type and GLU op
* cuda: always launch the fused MMQ kernel with one block per tile
* tests: shorten the fused gate/up MMQ test comment
* cuda: select the fused MMQ loader and write-back in ggml_cuda_mmq_get_util_funcs
* cuda: run the fused MMQ kernel on the plain tiling path
* cuda: put the column tiles innermost in the MMQ plain tiling path
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index b6efc8736..aff687a9a 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1973,6 +1973,61 @@ static bool ggml_cuda_match_shared_expert(const ggml_cgraph * graph, int routed_
return true;
}
+// Buffer- and batch-size independent, so graph_optimize produces the same graph topology for every ubatch size (see #29986).
+static bool ggml_cuda_mul_mat_q_fusion_matches(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu, const int device) {
+ const ggml_tensor * x_up = up->src[0];
+ const ggml_tensor * y = up->src[1];
+
+ if (up->op != GGML_OP_MUL_MAT || !ggml_cuda_mmq_fusion_supported(x_up->type) || y->type != GGML_TYPE_F32) {
+ return false;
+ }
+
+ // The fused kernel is only compiled for weights whose rows fill whole tiles (no fallback).
+ if (ggml_cuda_mmq_needs_fallback(x_up->ne[1])) {
+ return false;
+ }
+
+ if (!ggml_cuda_should_fuse_mul_mat(up, gate, glu)) {
+ return false;
+ }
+
+ // The fused write-back applies SWIGLU_OAI with the default alpha and limit of ggml_cuda_op_swiglu_oai_single.
+ if (ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_OAI && (ggml_get_op_params_f32(glu, 2) != 1.702f || ggml_get_op_params_f32(glu, 3) != 7.0f)) {
+ return false;
+ }
+
+ // NVIDIA only: the fused write-back relies on the MMA accumulator layout.
+ return turing_mma_available(ggml_cuda_info().devices[device].cc);
+}
+
+static bool ggml_cuda_should_fuse_mul_mat_q(const ggml_tensor * up, const ggml_tensor * gate, const ggml_tensor * glu) {
+ // A hinted matmul can be fused, but its unfused path may be faster (e.g. FWHT for GGML_HINT_SRC0_IS_HADAMARD).
+ if (ggml_get_op_params_i32(up, 1) != GGML_HINT_NONE || ggml_get_op_params_i32(gate, 1) != GGML_HINT_NONE) {
+ return false;
+ }
+
+ // The padding of a compute-buffer weight that is a view cannot be cleared, see ggml_cuda_mul_mat.
+ const auto bad_padding_clear = [](const ggml_tensor * x) {
+ return ggml_backend_buffer_get_usage(x->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE &&
+ ggml_nbytes(x) != ggml_backend_buffer_get_alloc_size(x->buffer, x) && x->view_src;
+ };
+ const int device = ggml_cuda_get_device();
+ if (!ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, device) || bad_padding_clear(up->src[0]) || bad_padding_clear(gate->src[0])) {
+ return false;
+ }
+
+ // The batch size checks are done here, not in ggml_cuda_mul_mat_q_fusion_matches.
+ const ggml_tensor * x_up = up->src[0];
+ const ggml_tensor * y = up->src[1];
+ const int cc = ggml_cuda_info().devices[device].cc;
+ if (ggml_cuda_should_use_mmvq(x_up->type, cc, y->ne[1]) || !ggml_cuda_should_use_mmq(x_up->type, cc, y->ne[1], 0)) {
+ return false;
+ }
+
+ // The fused kernel needs one block per tile (no stream-k split), so only fuse if the tiling is efficient.
+ return ggml_cuda_mmq_fusion_is_efficient(x_up->type, x_up->ne[1], y->ne[1], y->ne[2]*y->ne[3], device);
+}
+
static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_TENSOR_BINARY_OP_LOCALS
@@ -4378,6 +4433,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
fused_node_count = 3;
break;
}
+
+ if (ggml_cuda_should_fuse_mul_mat_q(up, gate, glu)) {
+ ggml_cuda_mm_fusion_args_host fusion_data{};
+ fusion_data.gate = gate->src[0];
+ fusion_data.glu_op = ggml_get_glu_op(glu);
+ fusion_data.glu_limit = ggml_get_op_params_f32(glu, 3);
+
+ ggml_cuda_mul_mat_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data);
+ fused_mul_mat_vec = true;
+ fused_node_count = 3;
+ break;
+ }
}
}
@@ -4974,6 +5041,19 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
}
}
for (int i = 0; i < cgraph->n_nodes; ++i) {
+ // Keep the inputs of the fused gate/up/GLU MMQ kernel alive until the GLU so that they cannot alias its output (under op offload the weights are buffer copies).
+ if (i + 2 < cgraph->n_nodes && cgraph->nodes[i + 2]->op == GGML_OP_GLU) {
+ ggml_tensor * glu = cgraph->nodes[i + 2];
+ ggml_tensor * gate = glu->src[0];
+ ggml_tensor * up = glu->src[1];
+ const bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) || (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]);
+ if (ok && ggml_cuda_mul_mat_q_fusion_matches(up, gate, glu, cuda_ctx->device)) {
+ add_alloc_deps(i, i + 2);
+ i += 2;
+ continue;
+ }
+ }
+
ggml_cuda_moe_weighted_reduction_match match;
if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);
diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh
index 2c7f4cd0a..4951db72e 100644
--- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh
+++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh
@@ -5,7 +5,9 @@
#include "mmq.cuh"
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q1_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -116,7 +118,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -205,7 +209,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -268,7 +274,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_1(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -331,7 +339,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -411,7 +421,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_1(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -488,8 +500,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0(
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -519,7 +531,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx;
+ const block_q8_0 * bxi = ggml_cuda_mmq_x_row<block_q8_0, has_fusion>(x, x_gate, kbx0, i, stride) + kbx;
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
x_qs[i*sram_stride + 0 + txi] = get_int_b2(bxi[0].qs, kqsx);
@@ -542,7 +554,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd;
+ const block_q8_0 * bxi = ggml_cuda_mmq_x_row<block_q8_0, has_fusion>(x, x_gate, kbx0, i, stride) + kbxd;
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
x_df[i*sram_stride + kbxd] = bxi->d;
@@ -555,7 +567,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
// ---------------------------------------------------------------------------------------------
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -616,7 +630,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q3_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -728,8 +744,8 @@ static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, co
((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q4_K(
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -757,7 +773,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+ const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int qs0 = get_int_b4(bxi->qs, txi);
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -786,7 +802,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+ const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int * scales = (const int *) bxi->scales;
const int ksc = threadIdx.x % 2;
@@ -814,7 +830,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride;
+ const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride);
x_dm[i] = bxi->dm;
}
@@ -827,7 +843,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
+ const block_q4_K * bxi = ggml_cuda_mmq_x_row<block_q4_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8);
const int * scales = (const int *) bxi->scales;
@@ -839,8 +855,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q5_K(
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -868,7 +884,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+ const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int ky = QR5_K*txi;
const int ql = get_int_b4(bxi->qs, txi);
@@ -909,7 +925,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+ const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int * scales = (const int *) bxi->scales;
const int ksc = threadIdx.x % 2;
@@ -937,7 +953,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+ const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
x_dm[i] = bxi->dm;
}
@@ -951,7 +967,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride;
+ const block_q5_K * bxi = ggml_cuda_mmq_x_row<block_q5_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int * scales = (const int *) bxi->scales;
@@ -963,8 +979,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q6_K(
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -993,7 +1009,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+ const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride);
const int ql = get_int_b2(bxi->ql, txi);
const int ql0 = (ql >> 0) & 0x0F0F0F0F;
@@ -1023,7 +1039,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride;
+ const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride);
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
x_df[i*sram_stride] = bxi->d;
@@ -1041,7 +1057,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
+ const block_q6_K * bxi = ggml_cuda_mmq_x_row<block_q6_K, has_fusion>(x, x_gate, kbx0, i, stride) + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4;
#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
x_sc[i*sram_stride + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8));
@@ -1054,7 +1070,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
// ---------------------------------------------------------------------------------------------
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq1_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1116,7 +1134,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xxs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1180,7 +1200,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_xs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1245,7 +1267,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq2_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1313,7 +1337,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_xxs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1377,7 +1403,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq3_s(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1445,8 +1473,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
}
-template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+template <ggml_type type, int J, bool fallback, bool has_fusion> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_xs(
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1473,7 +1501,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+ const block_iq4_xs * bxi = ggml_cuda_mmq_x_row<block_iq4_xs, has_fusion>(x, x_gate, kbx0, i, stride);
const int aux_q4 = get_int_b4(bxi->qs, kqsx);
const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl);
@@ -1497,7 +1525,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
i = min(i, i_max);
}
- const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride;
+ const block_iq4_xs * bxi = ggml_cuda_mmq_x_row<block_iq4_xs, has_fusion>(x, x_gate, kbx0, i, stride);
const float d = __half2float(bxi->d);
@@ -1513,7 +1541,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_iq4_nl(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1582,7 +1612,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
// ---------------------------------------------------------------------------------------------
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
@@ -1649,7 +1681,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_mxfp4_fp4(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
@@ -1691,7 +1725,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
}
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kb0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -1747,7 +1783,9 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
}
template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_nvfp4_nvfp4(
- const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ const char * __restrict__ x, const char * __restrict__ x_gate, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) {
+ GGML_UNUSED(x_gate);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q4) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q4);
diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
index 46efdfe0f..27558ec77 100644
--- a/ggml/src/ggml-cuda/mmq.cu
+++ b/ggml/src/ggml-cuda/mmq.cu
@@ -131,12 +131,51 @@ static ggml_prec ggml_cuda_mmq_get_prec_src1(const ggml_tensor * src0, const ggm
return GGML_PREC_Q4;
}
+// Tile width in ne11 (dense) / ne12 (MoE) direction with the fewest tiles that fits into shared memory, 0 if there is none.
+static int ggml_cuda_mmq_get_J_best(
+ const ggml_type type, const int64_t ncols, const bool fallback, const int cc, const size_t smpbo, const ggml_prec prec_src1,
+ int * nthreads_best = nullptr) {
+ int J_best = 0;
+ int ntiles_J_best = INT_MAX;
+
+ for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
+ const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc, prec_src1);
+ if (config.type == GGML_TYPE_COUNT) {
+ continue;
+ }
+
+ if (mmq_get_nbytes_shared(config, cc) > smpbo) {
+ continue;
+ }
+
+ const int ntiles_x = (ncols + config.J - 1) / config.J;
+
+ if (ntiles_x < ntiles_J_best) {
+ J_best = J;
+ if (nthreads_best) {
+ *nthreads_best = config.nthreads;
+ }
+ ntiles_J_best = ntiles_x;
+ }
+ }
+ return J_best;
+}
+
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) {
+ ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+ const ggml_cuda_mm_fusion_args_host * fusion) {
GGML_ASSERT( src1->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F32);
GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID.
+ // Fused gate/up/GLU: src0 is the up weight, dst the GLU output.
+ const ggml_tensor * gate = fusion ? fusion->gate : nullptr;
+ if (fusion) {
+ // TODO MoE support
+ GGML_ASSERT(gate && !ids && !fusion->x_bias && !fusion->gate_bias && !fusion->x_scale && !fusion->gate_scale);
+ GGML_ASSERT(gate->type == src0->type && ggml_are_same_shape(gate, src0) && ggml_are_same_stride(gate, src0));
+ }
+
GGML_TENSOR_BINARY_OP_LOCALS;
cudaStream_t stream = ctx.stream();
@@ -158,14 +197,16 @@ void ggml_cuda_mul_mat_q(
const float * src1_d = (const float *) src1->data;
float * dst_d = (float *) dst->data;
- // If src0 is a temporary compute buffer, clear any potential padding.
- if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
- const size_t size_data = ggml_nbytes(src0);
- const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0);
- if (size_alloc > size_data) {
- GGML_ASSERT(ggml_is_contiguously_allocated(src0));
- GGML_ASSERT(!src0->view_src);
- CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream));
+ // If src0 (or the fused gate weight) is a temporary compute buffer, clear any potential padding.
+ for (const ggml_tensor * t : { src0, gate }) {
+ if (t && ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+ const size_t size_data = ggml_nbytes(t);
+ const size_t size_alloc = ggml_backend_buffer_get_alloc_size(t->buffer, t);
+ if (size_alloc > size_data) {
+ GGML_ASSERT(ggml_is_contiguously_allocated(t));
+ GGML_ASSERT(!t->view_src);
+ CUDA_CHECK(cudaMemsetAsync((char *) t->data + size_data, 0, size_alloc - size_data, stream));
+ }
}
}
@@ -180,7 +221,9 @@ void ggml_cuda_mul_mat_q(
const bool fallback = ggml_cuda_mmq_needs_fallback(ne01);
- const ggml_prec prec_src1 = ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
+ // With fusion dst is the GLU output, not a MUL_MAT; the fused types always quantize src1 to Q8_1.
+ // TODO support for Q4 precision if both tensors allow it
+ const ggml_prec prec_src1 = gate ? GGML_PREC_Q8 : ggml_cuda_mmq_get_prec_src1(src0, dst, cc);
const bool use_native_fp4 = prec_src1 == GGML_PREC_Q4;
const size_t y_block_size = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
@@ -201,26 +244,7 @@ void ggml_cuda_mul_mat_q(
}
}
- int ntiles_J_best = INT_MAX;
-
- for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) {
- const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(src0->type, J, fallback, cc, prec_src1);
- if (config.type == GGML_TYPE_COUNT) {
- continue;
- }
-
- if (mmq_get_nbytes_shared(config, cc) > smpbo) {
- continue;
- }
-
- const int ntiles_x = (ncols_opt + config.J - 1) / config.J;
-
- if (ntiles_x < ntiles_J_best) {
- J_best = J;
- nthreads_best = config.nthreads;
- ntiles_J_best = ntiles_x;
- }
- }
+ J_best = ggml_cuda_mmq_get_J_best(src0->type, ncols_opt, fallback, cc, smpbo, prec_src1, &nthreads_best);
}
GGML_ASSERT(J_best > 0);
@@ -262,13 +286,20 @@ void ggml_cuda_mul_mat_q(
ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
const int64_t s13 = ne12*s12;
+ ggml_cuda_mm_fusion_args_device fusion_device = {};
+ if (gate) {
+ fusion_device.gate = gate->data;
+ fusion_device.glu_op = fusion->glu_op;
+ fusion_device.glu_limit = fusion->glu_limit;
+ }
+
const mmq_args args = {
src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- ne1, J_best};
+ ne1, J_best, fusion_device};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
return;
}
@@ -485,3 +516,22 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
return (!GGML_CUDA_CC_IS_CDNA(cc)) || ne11 < MMQ_DP4A_MAX_BATCH_SIZE;
}
+
+bool ggml_cuda_mmq_fusion_is_efficient(const ggml_type type, const int64_t nrows_x, const int64_t ncols, const int64_t nchannels, const int device) {
+ const int cc = ggml_cuda_info().devices[device].cc;
+ const int nsm = ggml_cuda_info().devices[device].nsm;
+ const size_t smpbo = ggml_cuda_info().devices[device].smpbo;
+
+ // Same tile width as ggml_cuda_mul_mat_q; the fused kernel has no fallback, quantizes src1 to Q8_1 and writes I/2 output rows per tile.
+ const bool fallback = false;
+ const int J_best = ggml_cuda_mmq_get_J_best(type, ncols, fallback, cc, smpbo, GGML_PREC_Q8);
+ if (J_best == 0) {
+ return false;
+ }
+ const int64_t nrows_tile = ggml_cuda_mmq_get_config(type, J_best, fallback, cc, GGML_PREC_Q8).I / 2;
+ const int64_t ntiles = (ncols + J_best - 1)/J_best * ((nrows_x + nrows_tile - 1)/nrows_tile) * nchannels;
+ // Same condition as launch_mul_mat_q uses to run the stream-k kernel with one block per tile.
+ const int64_t tiles_nwaves = (ntiles + nsm - 1) / nsm;
+ const int64_t tiles_efficiency_percent = 100 * ntiles / (nsm*tiles_nwaves);
+ return tiles_efficiency_percent >= 90;
+}
diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
index 1375ba3e8..565ed1481 100644
--- a/ggml/src/ggml-cuda/mmq.cuh
+++ b/ggml/src/ggml-cuda/mmq.cuh
@@ -1,6 +1,7 @@
#pragma once
#include "common.cuh"
+#include "unary.cuh"
#include <climits>
#include <cstdint>
@@ -10,10 +11,12 @@
#define MMQ_ITER_K_FP4 512
#define MMQ_NWARPS 8
-typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
+typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, const char * __restrict__ x_gate, int * x_tile,
+ const int kbx0, const int i_max, const int stride);
typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
- float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);
+ float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+ const ggml_glu_op glu_op, const float glu_limit);
enum mmq_q8_1_ds_layout {
MMQ_Q8_1_DS_LAYOUT_D4,
@@ -426,12 +429,37 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config
// ------------------------------------------------------------
+// Weight types with a fused gate/up/GLU path (see ggml_cuda_mmq_get_util_funcs).
+static constexpr __host__ __device__ bool ggml_cuda_mmq_fusion_supported(const ggml_type type) {
+ return type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_XS;
+}
+
+// Fused gate/up/GLU: the weight tile alternates chunks of this many up and gate rows, so rows r and r + 8 of an MMA tile (same thread) pair up.
+static constexpr int GGML_CUDA_MMQ_FUSION_ROWS = 8;
+
+// Source of weight tile row i; with has_fusion the up or gate row that the interleaved tile holds there.
+template <typename block_t, bool has_fusion>
+static __device__ __forceinline__ const block_t * ggml_cuda_mmq_x_row(
+ const char * __restrict__ x, const char * __restrict__ x_gate, const int kbx0, const int i, const int stride) {
+ if constexpr (has_fusion) {
+ constexpr int R = GGML_CUDA_MMQ_FUSION_ROWS;
+ const char * x_src = (i / R) % 2 == 0 ? x : x_gate;
+ return (const block_t *) x_src + kbx0 + ((i / (2*R))*R + i % R)*stride;
+ } else {
+ GGML_UNUSED(x_gate);
+ return (const block_t *) x + kbx0 + i*stride;
+ }
+}
+
#include "mmq-load-tiles.cuh"
#include "mmq-vec-dot.cuh"
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
- const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+ const ggml_glu_op glu_op, const float glu_limit) {
+ GGML_UNUSED_VARS(glu_op, glu_limit);
+
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -469,16 +497,19 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
}
}
-template<ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+// has_fusion: accumulators l and l + ne/2 hold the up and gate values of one output element, the GLU of the two is written.
+template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
- const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+ const ggml_glu_op glu_op, const float glu_limit) {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
#else
typedef tile<16, 8, int> tile_C;
#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+ static_assert(!has_fusion || tile_C::I == 2*GGML_CUDA_MMQ_FUSION_ROWS, "each C tile must hold one chunk of up rows and one of gate rows");
constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback, prec_src1);
constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp.
@@ -492,20 +523,23 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
#pragma unroll
for (int n = 0; n < ntx; ++n) {
#pragma unroll
- for (int l = 0; l < tile_C::ne; ++l) {
+ for (int l = 0; l < (has_fusion ? tile_C::ne/2 : tile_C::ne); ++l) {
const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l);
if (j > j_max) {
continue;
}
- const int i = i0 + n*tile_C::I + tile_C::get_i(l);
+ const int i = (i0 + n*tile_C::I)/(has_fusion ? 2 : 1) + tile_C::get_i(l);
if (fallback && i > i_max) {
continue;
}
- if constexpr (type == GGML_TYPE_NVFP4) {
+ if constexpr (has_fusion) {
+ const int sum_index = (j0/tile_C::J + n)*tile_C::ne + l;
+ dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, sum[sum_index + tile_C::ne/2], sum[sum_index], glu_limit);
+ } else if constexpr (type == GGML_TYPE_NVFP4) {
if (y_scale_used) {
dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
} else {
@@ -518,6 +552,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
}
}
}
+
+ GGML_UNUSED_VARS(y_scale_used, glu_op, glu_limit);
}
// -------------------------------------------------------------------------------------------------------------------------------------
@@ -534,8 +570,9 @@ struct ggml_cuda_mmq_util_funcs {
vdr(vdr), load_tiles(load_tiles), vec_dot(vec_dot), write_back(write_back) {}
};
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_funcs() {
+ static_assert(!has_fusion || ggml_cuda_mmq_fusion_supported(type), "no fused loader for this type");
if (!ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).use_mma_data_layout()) {
switch (type) {
case GGML_TYPE_Q1_0:
@@ -577,7 +614,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
case GGML_TYPE_Q8_0:
return ggml_cuda_mmq_util_funcs(
VDR_Q8_0_Q8_1_MMQ,
- ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, false>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
// ---------------------------------------------------------------------------------------------
@@ -596,19 +633,19 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
case GGML_TYPE_Q4_K:
return ggml_cuda_mmq_util_funcs(
VDR_Q4_K_Q8_1_MMQ,
- ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, false>,
ggml_cuda_mmq_vec_dot_q4_K_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
case GGML_TYPE_Q5_K:
return ggml_cuda_mmq_util_funcs(
VDR_Q5_K_Q8_1_MMQ,
- ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, false>,
ggml_cuda_mmq_vec_dot_q5_K_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
case GGML_TYPE_Q6_K:
return ggml_cuda_mmq_util_funcs(
VDR_Q6_K_Q8_1_MMQ,
- ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, false>,
ggml_cuda_mmq_vec_dot_q6_K_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
// ---------------------------------------------------------------------------------------------
@@ -651,7 +688,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
case GGML_TYPE_IQ4_XS:
return ggml_cuda_mmq_util_funcs(
VDR_IQ4_XS_Q8_1_MMQ,
- ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, false>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_dp4a<type, J, fallback>,
ggml_cuda_mmq_write_back_dp4a<type, J, fallback, prec_src1>);
case GGML_TYPE_IQ4_NL:
@@ -688,7 +725,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
-1,
ggml_cuda_mmq_load_tiles_mxfp4_fp4<type, J, fallback>,
ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
}
break;
case GGML_TYPE_NVFP4:
@@ -697,7 +734,7 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
-1,
ggml_cuda_mmq_load_tiles_nvfp4_nvfp4<type, J, fallback>,
ggml_cuda_mmq_vec_dot_fp4_fp4_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
}
break;
default:
@@ -713,136 +750,136 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
-1,
ggml_cuda_mmq_load_tiles_q1_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q2_0:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q2_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q4_0:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q4_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_DS4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q4_1:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q4_1<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q5_0:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q5_0<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q5_1:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q5_1<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q8_0:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q8_0<type, J, fallback, has_fusion>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
// ---------------------------------------------------------------------------------------------
case GGML_TYPE_Q2_K:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q2_K<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q2_K_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q3_K:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_q3_K<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_Q4_K:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q4_K<type, J, fallback, has_fusion>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
case GGML_TYPE_Q5_K:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q5_K<type, J, fallback, has_fusion>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
case GGML_TYPE_Q6_K:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_q6_K<type, J, fallback, has_fusion>,
ggml_cuda_mmq_vec_dot_q6_K_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
// ---------------------------------------------------------------------------------------------
case GGML_TYPE_IQ1_S:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq1_s<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_1_q8_1_mma<type, J, fallback>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ2_XXS:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq2_xxs<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ2_XS:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq2_xs<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ2_S:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq2_s<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ3_XXS:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq3_xxs<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ3_S:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq3_s<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_IQ4_XS:
return ggml_cuda_mmq_util_funcs(
-1,
- ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback>,
+ ggml_cuda_mmq_load_tiles_iq4_xs<type, J, fallback, has_fusion>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, has_fusion>);
case GGML_TYPE_IQ4_NL:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_iq4_nl<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
// ---------------------------------------------------------------------------------------------
case GGML_TYPE_MXFP4:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_mxfp4<type, J, fallback>,
ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma<type, J, fallback, MMQ_Q8_1_DS_LAYOUT_D4>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
case GGML_TYPE_NVFP4:
return ggml_cuda_mmq_util_funcs(
-1,
ggml_cuda_mmq_load_tiles_nvfp4<type, J, fallback, prec_src1>,
ggml_cuda_mmq_vec_dot_q8_0_16_q8_1_mma<type, J, fallback, prec_src1>,
- ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1>);
+ ggml_cuda_mmq_write_back_mma<type, J, fallback, prec_src1, false>);
default:
return ggml_cuda_mmq_util_funcs(1, nullptr, nullptr, nullptr);
}
@@ -850,41 +887,44 @@ static constexpr __device__ ggml_cuda_mmq_util_funcs ggml_cuda_mmq_get_util_func
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
static constexpr __device__ int ggml_cuda_mmq_get_vdr() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vdr;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vdr;
}
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static constexpr __device__ ggml_cuda_mmq_load_tiles_t ggml_cuda_mmq_get_load_tiles() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().load_tiles;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().load_tiles;
}
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
static constexpr __device__ ggml_cuda_mmq_vec_dot_t ggml_cuda_mmq_get_vec_dot() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().vec_dot;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, false>().vec_dot;
}
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_back() {
- return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1>().write_back;
+ return ggml_cuda_mmq_get_util_funcs<type, J, fallback, prec_src1, has_fusion>().write_back;
}
// ---------------------------------------------------------------------------------------------
-template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1>
+// has_fusion: x and x_gate are the up and gate weights.
+template <ggml_type type, int J, bool fallback, bool fixup, ggml_prec prec_src1, bool has_fusion>
static __device__ __forceinline__ void mul_mat_q_process_tile(
- const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
+ const char * __restrict__ x, const char * __restrict__ x_gate, const int offset_x, const int * __restrict__ y,
const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
- const float * __restrict__ y_scale,
+ const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
const int stride_row_x, const int ncols_y, const int stride_col_dst,
const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+ static_assert(!(fixup && has_fusion), "the fused gate/up path does not split K");
+ static_assert(!(fallback && has_fusion), "the fused gate/up path needs nrows_x % 128 == 0");
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int qk = ggml_cuda_type_traits<type>::qk;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
- constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1>();
+ constexpr ggml_cuda_mmq_load_tiles_t load_tiles = ggml_cuda_mmq_get_load_tiles<type, J, fallback, prec_src1, has_fusion>();
constexpr ggml_cuda_mmq_vec_dot_t vec_dot = ggml_cuda_mmq_get_vec_dot<type, J, fallback, prec_src1>();
- constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1>();
+ constexpr ggml_cuda_mmq_write_back_t write_back = ggml_cuda_mmq_get_write_back<type, J, fallback, prec_src1, has_fusion>();
extern __shared__ int data_mul_mat_q[];
int * tile_y = data_mul_mat_q + J;
@@ -907,7 +947,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
- load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+ load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
{
const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
#pragma unroll
@@ -942,28 +982,34 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
}
if (fixup) {
- write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
+ write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J, glu_op, glu_limit);
} else {
- write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
+ write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j, glu_op, glu_limit);
}
}
// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
static __global__ void mul_mat_q(
- const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
+ const char * __restrict__ x, const char * __restrict__ x_gate, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
- const float * __restrict__ y_scale,
+ const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
const uint3 ntx) {
+#ifdef TURING_MMA_AVAILABLE
+ constexpr bool fusion_available = true;
+#else
+ constexpr bool fusion_available = false;
+#endif // TURING_MMA_AVAILABLE
+
// Skip unused template specializations for faster compilation:
- if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT) {
+ if (ggml_cuda_mmq_get_config(type, J, fallback, prec_src1).type == GGML_TYPE_COUNT || (has_fusion && !fusion_available)) {
NO_DEVICE_CODE;
return;
}
@@ -991,12 +1037,16 @@ static __global__ void mul_mat_q(
}
__syncthreads();
- if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+ // The fused kernel needs the whole K range of a tile in one block, so it always uses tiling.
+ if constexpr (has_fusion || !ggml_cuda_mmq_get_stream_k(type, J, fallback, prec_src1)) {
+ // The fused kernel writes the GLU of I/2 up rows and I/2 gate rows per tile.
+ constexpr int I_dst = has_fusion ? I/2 : I;
+
const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
const int wt = tmp2.x;
const int zt = tmp2.y;
- const int jt = blockIdx.y;
- const int it = blockIdx.x;
+ const int jt = blockIdx.x;
+ const int it = blockIdx.y;
// Defaults for regular matrix multiplication:
int col_low = 0;
@@ -1041,22 +1091,22 @@ static __global__ void mul_mat_q(
}
offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
- offset_dst += it*I;
+ offset_dst += it*I_dst;
const float * y_scale_tile = nullptr;
if constexpr (type == GGML_TYPE_NVFP4) {
offset_y_scale += col_low + jt*J;
y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
}
- const int tile_x_max_i = nrows_x - it*I - 1;
+ const int tile_x_max_i = nrows_x - it*I_dst - 1;
const int tile_y_max_j = col_diff - jt*J - 1;
- const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
+ const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I_dst*stride_row_x;
constexpr bool fixup = false;
- mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
- stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, has_fusion>
+ (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
return;
}
@@ -1148,9 +1198,9 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
- mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
- stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
+ (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
kbc += blocks_per_ne00.z;
@@ -1232,9 +1282,9 @@ static __global__ void mul_mat_q(
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
- mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1>
- (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
- stride_row_x, ncols_y, stride_col_dst,
+ mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
+ (x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+ glu_op, glu_limit, stride_row_x, ncols_y, stride_col_dst,
tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
}
@@ -1384,6 +1434,7 @@ struct mmq_args {
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
int64_t ncols_max;
int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
+ ggml_cuda_mm_fusion_args_device fusion = {}; // gate != nullptr selects the fused gate/up/GLU path
};
static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const int cc) {
@@ -1393,7 +1444,7 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i
return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int));
}
-template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
const int id = ggml_cuda_get_device();
const int cc = ggml_cuda_info().devices[id].cc;
@@ -1407,13 +1458,17 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_dims(warp_size, nwarps, 1);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, false, prec_src1>), nbytes_shared);
- CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, true, prec_src1>), nbytes_shared);
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ false, prec_src1, has_fusion>), nbytes_shared);
+ if constexpr (!has_fusion) {
+ CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q<type, J, /*fallback =*/ true, prec_src1, false>), nbytes_shared);
+ }
- const int nty = (args.nrows_x + config.I - 1) / config.I;
+ const int I_dst = has_fusion ? config.I/2 : config.I; // Rows of dst per tile, see mul_mat_q.
+ const int nty = (args.nrows_x + I_dst - 1) / I_dst;
const int ntx = (args.ncols_max + config.J - 1) / config.J;
const int ntzw = args.nchannels_y * args.nsamples_y;
- const dim3 block_nums_xy_tiling(nty, ntx, ntzw);
+ // Column tiles innermost so that the concurrently running blocks share the same tile of x and read it from L2.
+ const dim3 block_nums_xy_tiling(ntx, nty, ntzw);
GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
GGML_ASSERT(args.nsamples_y % args.nsamples_x == 0);
@@ -1427,9 +1482,10 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio);
const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio);
- if (!config.stream_k) {
- mul_mat_q<type, J, fallback, prec_src1><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
+ if (has_fusion || !config.stream_k) {
+ GGML_ASSERT(nty <= 65535); // gridDim.y limit
+ mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
+ (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1459,8 +1515,8 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_nums_fixup(block_nums_stream_k.x, config.I/warp_size, 1);
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
- mul_mat_q<type, J, fallback, prec_src1><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
- (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
+ mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
+ (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1477,56 +1533,56 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
ntx_fd);
}
-template <ggml_type type, bool fallback, ggml_prec prec_src1>
+template <ggml_type type, bool fallback, ggml_prec prec_src1, bool has_fusion>
void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
switch (args.J_best) {
case 8:
- launch_mul_mat_q<type, 8, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 8, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 16:
- launch_mul_mat_q<type, 16, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 16, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 24:
- launch_mul_mat_q<type, 24, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 24, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 32:
- launch_mul_mat_q<type, 32, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 32, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 40:
- launch_mul_mat_q<type, 40, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 40, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 48:
- launch_mul_mat_q<type, 48, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 48, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 56:
- launch_mul_mat_q<type, 56, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 56, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 64:
- launch_mul_mat_q<type, 64, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 64, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 72:
- launch_mul_mat_q<type, 72, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 72, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 80:
- launch_mul_mat_q<type, 80, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 80, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 88:
- launch_mul_mat_q<type, 88, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 88, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 96:
- launch_mul_mat_q<type, 96, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 96, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 104:
- launch_mul_mat_q<type, 104, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 104, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 112:
- launch_mul_mat_q<type, 112, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 112, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 120:
- launch_mul_mat_q<type, 120, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 120, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
case 128:
- launch_mul_mat_q<type, 128, fallback, prec_src1>(ctx, args, stream);
+ launch_mul_mat_q<type, 128, fallback, prec_src1, has_fusion>(ctx, args, stream);
break;
default:
fprintf(stderr, "J_best=%d\n", args.J_best);
@@ -1534,14 +1590,26 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
}
}
+template <ggml_type type, bool fallback, ggml_prec prec_src1>
+static void mul_mat_q_switch_fusion(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
+ if constexpr (!fallback && ggml_cuda_mmq_fusion_supported(type)) {
+ if (args.fusion.gate) {
+ mul_mat_q_switch_J<type, fallback, prec_src1, true>(ctx, args, stream);
+ return;
+ }
+ }
+ GGML_ASSERT(!args.fusion.gate);
+ mul_mat_q_switch_J<type, fallback, prec_src1, false>(ctx, args, stream);
+}
+
template <ggml_type type, ggml_prec prec_src1>
void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) {
if (ggml_cuda_mmq_needs_fallback(args.nrows_x)) {
constexpr bool fallback = true;
- mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+ mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
} else {
constexpr bool fallback = false;
- mul_mat_q_switch_J<type, fallback, prec_src1>(ctx, args, stream);
+ mul_mat_q_switch_fusion<type, fallback, prec_src1>(ctx, args, stream);
}
}
@@ -1583,6 +1651,9 @@ extern DECL_MMQ_CASE_W4A4(GGML_TYPE_NVFP4);
// -------------------------------------------------------------------------------------------------------------------------
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);
+ ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst,
+ const ggml_cuda_mm_fusion_args_host * fusion = nullptr);
bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);
+
+bool ggml_cuda_mmq_fusion_is_efficient(ggml_type type, int64_t nrows_x, int64_t ncols, int64_t nchannels, int device);
diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh
index 04f3af644..c6ffdb9d6 100644
--- a/ggml/src/ggml-cuda/unary.cuh
+++ b/ggml/src/ggml-cuda/unary.cuh
@@ -121,3 +121,18 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl
return ggml_cuda_op_silu_single(gate) * up;
}
+
+__device__ __forceinline__ float ggml_cuda_op_glu_single(ggml_glu_op op, float gate, float up, float limit) {
+ switch (op) {
+ case GGML_GLU_OP_SWIGLU:
+ return up * ggml_cuda_op_silu_single(gate);
+ case GGML_GLU_OP_GEGLU:
+ return up * ggml_cuda_op_gelu_single(gate);
+ case GGML_GLU_OP_SWIGLU_OAI:
+ return ggml_cuda_op_swiglu_oai_single(gate, up);
+ case GGML_GLU_OP_SWIGLU_CLAMP:
+ return ggml_cuda_op_swiglu_clamp_single(gate, up, limit);
+ default:
+ return up * gate;
+ }
+}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index cb68651dc..fe3a67c64 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12365,6 +12365,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
false, 16, 8, false, false, true, false, { 1, 1 }));
}
+ // Fused gate/up + GLU on the MMQ path, only used if its tiling is efficient: ~1280 tiles per case, one case per type and GLU op.
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q4_K, GGML_GLU_OP_SWIGLU, 17, 81920, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q5_K, GGML_GLU_OP_GEGLU, 150, 40960, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q6_K, GGML_GLU_OP_SWIGLU_OAI, 150, 40960, 256, false, 1, 1, false, false, true, false, { 1, 1 }));
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_Q8_0, GGML_GLU_OP_SWIGLU_CLAMP, 150, 6912, 256, false, 1, 1, false, false, true, false, { 2, 3 }));
+ test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_IQ4_XS, GGML_GLU_OP_SWIGLU, 150, 40960, 512, false, 1, 1, false, false, true, false, { 1, 1 }));
+
for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) {
for (bool with_norm : {false, true}) {
for (bool bias_probs : {false, true}) {