cuda: add opt-in SM75 SwiGLU input fusion

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -3511,6 +3511,36 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
 
     ggml_tensor * node = cgraph->nodes[i];
 
+    static const bool fuse_swiglu_q8 = getenv("GGML_CUDA_SM75_SWIGLU_Q8_1") != nullptr && std::atoi(getenv("GGML_CUDA_SM75_SWIGLU_Q8_1"));
+    static const bool trace_swiglu_q8 = getenv("GGML_CUDA_SM75_SWIGLU_Q8_1_TRACE") != nullptr;
+    if (fuse_swiglu_q8 && node->op == GGML_OP_GLU && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU &&
+            node->src[1] != nullptr && node->type == GGML_TYPE_F32 && node->ne[1] > 8 &&
+            node->ne[2] == 1 && node->ne[3] == 1 && node->ne[0] % 4 == 0 &&
+            ggml_is_contiguous(node) && ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(node->src[1]) &&
+            ggml_cuda_is_aligned(node->src[0], 16) && ggml_cuda_is_aligned(node->src[1], 16) &&
+            node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_F32 &&
+            ggml_are_same_shape(node, node->src[0]) && ggml_are_same_shape(node, node->src[1]) &&
+            ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_GLU, GGML_OP_MUL_MAT }, { i + 1 })) {
+        ggml_tensor * dst = cgraph->nodes[i + 1];
+        const ggml_tensor * weight = dst->src[0];
+        const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+        const bool cublas_batch = node->ne[1] >= 1024 && node->ne[1] <= 2048 &&
+                (weight->type == GGML_TYPE_Q4_K || weight->type == GGML_TYPE_IQ4_XS);
+        const int output_idx = i + 1;
+        if (cc == GGML_CUDA_CC_TURING && !cublas_batch && dst->src[1] == node && dst->type == GGML_TYPE_F32 &&
+                weight->ne[2] == 1 && weight->ne[3] == 1 && ggml_is_contiguous(weight) && ggml_is_contiguous(dst) &&
+                !ggml_backend_buffer_is_host(weight->buffer) &&
+                !ggml_cuda_should_use_mmvq(weight->type, cc, node->ne[1]) &&
+                ggml_cuda_should_use_mmq(weight->type, cc, node->ne[1], 0) &&
+                ggml_cuda_check_fusion_memory_ranges(cgraph, i, 2, &output_idx, 1)) {
+            if (trace_swiglu_q8) {
+                GGML_LOG_INFO("SM75_SWIGLU_Q8_1 type=%s n=%lld k=%lld\n", ggml_type_name(weight->type), (long long) node->ne[1], (long long) node->ne[0]);
+            }
+            ggml_cuda_mul_mat_q(*cuda_ctx, weight, node, nullptr, dst, true);
+            return 1;
+        }
+    }
+
     if (node->op == GGML_OP_MUL) {
         ggml_cuda_moe_weighted_reduction_match match;
         if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {

ggml/src/ggml-cuda/mmq.cu

@@ -83,7 +83,7 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
 }
 
 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, bool swiglu) {
     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.
@@ -145,7 +145,11 @@ void ggml_cuda_mul_mat_q(
             const int64_t s11 = src1->nb[1] / ts_src1;
             const int64_t s12 = src1->nb[2] / ts_src1;
             const int64_t s13 = src1->nb[3] / ts_src1;
-            if (use_native_fp4) {
+            if (swiglu) {
+                GGML_ASSERT(!use_native_fp4 && ne12 == 1 && ne13 == 1);
+                quantize_swiglu_mmq_q8_1_cuda((const float *) src1->src[0]->data, (const float *) src1->src[1]->data,
+                        src1_q8_1.get(), src0->type, ne10, ne10_padded, ne11, stream);
+            } else if (use_native_fp4) {
                 static constexpr size_t align_float8 = 32;
                 const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
                 static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));

ggml/src/ggml-cuda/mmq.cuh

@@ -1600,6 +1600,6 @@ extern DECL_MMQ_CASE(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, bool swiglu = false);
 
 bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts);

ggml/src/ggml-cuda/quantize.cu

@@ -1,4 +1,5 @@
 #include "quantize.cuh"
+#include "unary.cuh"
 #include <cstdint>
 
 #if defined(BLACKWELL_MMA_AVAILABLE)
@@ -454,11 +455,11 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
 }
 
 // scatter: grid over tokens, quantize once, write to all the token's compact rows
-template <mmq_q8_1_ds_layout ds_layout, bool scatter>
+template <mmq_q8_1_ds_layout ds_layout, bool scatter, bool swiglu = false>
 static __global__ void quantize_mmq_q8_1(
         const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
-        const int64_t ne0, const int ne1, const int ne2, const int n_expert_used) {
+        const int64_t ne0, const int ne1, const int ne2, const int n_expert_used, const float * gate = nullptr) {
 
     constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32;
     constexpr int vals_per_sum   = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32;
@@ -489,7 +490,16 @@ static __global__ void quantize_mmq_q8_1(
     const int64_t iqs     = i0 % QK8_1_MMQ; // quant index in block
 
     // Load 4 floats per thread and calculate max. abs. value between them:
-    const float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    float4 xi = i0 < ne00 ? x4[(base_idx + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
+    if constexpr (swiglu) {
+        if (i0 < ne00) {
+            const float4 gi = ((const float4 *) gate)[(base_idx + i00)/4];
+            xi.x = ggml_cuda_op_silu_single(xi.x) * gi.x;
+            xi.y = ggml_cuda_op_silu_single(xi.y) * gi.y;
+            xi.z = ggml_cuda_op_silu_single(xi.z) * gi.z;
+            xi.w = ggml_cuda_op_silu_single(xi.w) * gi.w;
+        }
+    }
     float amax = fabsf(xi.x);
     amax = fmaxf(amax, fabsf(xi.y));
     amax = fmaxf(amax, fabsf(xi.z));
@@ -602,6 +612,32 @@ void quantize_mmq_q8_1_cuda(
     }
 }
 
+void quantize_swiglu_mmq_q8_1_cuda(
+        const float * x, const float * gate, void * vy, ggml_type type_src0,
+        int64_t ne00, int64_t ne0, int64_t ne1, cudaStream_t stream) {
+    GGML_ASSERT(ne00 % 4 == 0);
+    GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
+    const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
+    const dim3 num_blocks(ne1, block_num_y, 1);
+    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
+    switch (mmq_get_q8_1_ds_layout(type_src0)) {
+        case MMQ_Q8_1_DS_LAYOUT_D4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D4, false, true>
+                <<<num_blocks, block_size, 0, stream>>>(x, nullptr, vy, ne00, ne00, 0, 0, ne0, ne1, 1, 0, gate);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_DS4:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_DS4, false, true>
+                <<<num_blocks, block_size, 0, stream>>>(x, nullptr, vy, ne00, ne00, 0, 0, ne0, ne1, 1, 0, gate);
+            break;
+        case MMQ_Q8_1_DS_LAYOUT_D2S6:
+            quantize_mmq_q8_1<MMQ_Q8_1_DS_LAYOUT_D2S6, false, true>
+                <<<num_blocks, block_size, 0, stream>>>(x, nullptr, vy, ne00, ne00, 0, 0, ne0, ne1, 1, 0, gate);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
 // scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row)
 void quantize_scatter_mmq_q8_1_cuda(
         const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,

ggml/src/ggml-cuda/quantize.cuh

@@ -26,6 +26,10 @@ void quantize_mmq_q8_1_cuda(
         ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
         int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
 
+void quantize_swiglu_mmq_q8_1_cuda(
+        const float * x, const float * gate, void * vy, ggml_type type_src0,
+        int64_t ne00, int64_t ne0, int64_t ne1, cudaStream_t stream);
+
 void quantize_mmq_fp4_cuda(const float *   x,
                              const int32_t * ids,
                              void *          vy,

sm75/run-shared-chat-kernels.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+import json
+import os
+from pathlib import Path
+import random
+import re
+import statistics
+import subprocess
+import time
+
+root = Path(__file__).resolve().parent.parent
+out = root / ('kernel-results-' + time.strftime('%Y%m%d-%H%M%S'))
+out.mkdir()
+binary = root / 'build/bin/test-backend-ops'
+env = dict(os.environ, GGML_TEST_SWIGLU_MMQ='1')
+active = subprocess.run(['systemctl', 'is-active', '--quiet', 'cmp-idle-governor.service']).returncode == 0
+records = []
+
+def run(mode, enabled, label):
+    current_env = dict(env, GGML_CUDA_SM75_SWIGLU_Q8_1=str(enabled))
+    if mode == 'test' and enabled:
+        current_env['GGML_CUDA_SM75_SWIGLU_Q8_1_TRACE'] = '1'
+    command = [str(binary), mode, '-b', 'CUDA0', '-o', 'SWIGLU_MMQ']
+    before = subprocess.check_output(['nvidia-smi', '--query-gpu=pstate,clocks.sm,clocks.mem,temperature.gpu,power.limit,memory.used', '--format=csv,noheader'], text=True)
+    result = subprocess.run(command, env=current_env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    (out / (label + '.log')).write_text(result.stdout)
+    (out / (label + '-state.txt')).write_text(before)
+    if result.returncode != 0 or 'FAIL' in result.stdout:
+        raise RuntimeError(f'{label} failed: see {out}')
+    if mode == 'test':
+        if '50/50 tests passed' not in result.stdout:
+            raise RuntimeError(f'{label}: expected 50/50 test summary')
+        if enabled and 'SM75_SWIGLU_Q8_1 type=' not in result.stdout:
+            raise RuntimeError(f'{label}: fusion did not run')
+        print(label, '50/50 passed', flush=True)
+        return
+    rows = re.findall(r'SWIGLU_MMQ\(([^\n]*?)\):.*?([0-9.]+) us/run', result.stdout, re.DOTALL)
+    if len(rows) != 18:
+        raise RuntimeError(f'{label}: expected 18 timings, got {len(rows)}')
+    for shape, us in rows:
+        records.append(dict(label=label, enabled=enabled, shape=shape, us=float(us)))
+    (out / 'records.json').write_text(json.dumps(records, indent=2))
+    print(label, '18 timings', flush=True)
+
+try:
+    for enabled in [0, 1]:
+        run('test', enabled, f'correctness-{enabled}')
+    if active:
+        subprocess.run(['sudo', '-n', 'systemctl', 'stop', 'cmp-idle-governor.service'], check=True)
+    for enabled in [0, 1]:
+        run('perf', enabled, f'warmup-{enabled}')
+    rng = random.Random(20260909)
+    for repeat in range(5):
+        order = [0, 1]
+        rng.shuffle(order)
+        for enabled in order:
+            run('perf', enabled, f'run-{repeat}-{enabled}')
+    summary = []
+    for shape in sorted({r['shape'] for r in records}):
+        row = dict(shape=shape)
+        for enabled in [0, 1]:
+            values = [r['us'] for r in records if r['shape'] == shape and r['enabled'] == enabled and r['label'].startswith('run-')]
+            row[str(enabled)] = dict(median=statistics.median(values), minimum=min(values), maximum=max(values))
+        row['throughput_change_pct'] = 100 * (row['0']['median'] / row['1']['median'] - 1)
+        summary.append(row)
+    (out / 'summary.json').write_text(json.dumps(summary, indent=2))
+    print('COMPLETE', out, flush=True)
+finally:
+    if active:
+        subprocess.run(['sudo', '-n', 'systemctl', 'start', 'cmp-idle-governor.service'], check=True)

tests/test-backend-ops.cpp

@@ -1649,6 +1649,10 @@ struct test_case {
             n_runs = (int)std::min<int64_t>(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_size / op_size(out)) + 1;
         }
 
+        if (current_op_name == "SWIGLU_MMQ") {
+            n_runs = 1;
+        }
+
         // duplicate the op
         for (int i = 1; i < n_runs; i++) {
             ggml_graph_add_node(gf, out);
@@ -5042,6 +5046,88 @@ struct test_mul_mat : public test_case {
     }
 };
 
+struct test_swiglu_mmq : public test_mul_mat {
+    test_swiglu_mmq(ggml_type type, int64_t rows, int64_t tokens, int64_t width)
+        : test_mul_mat(type, GGML_TYPE_F32, rows, tokens, width, {1, 1}, {1, 1}) {}
+
+    bool run_whole_graph() override { return true; }
+
+    std::string op_desc(ggml_tensor *) override { return "SWIGLU_MMQ"; }
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * weight = ggml_new_tensor_2d(ctx, type_a, k, m);
+        ggml_tensor * gate = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, n);
+        ggml_tensor * up = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, n);
+        return ggml_mul_mat(ctx, weight, ggml_swiglu_split(ctx, gate, up));
+    }
+};
+
+static std::vector<std::unique_ptr<test_case>> make_swiglu_mmq_cases(bool perf) {
+    std::vector<std::unique_ptr<test_case>> cases;
+    for (ggml_type type : {GGML_TYPE_Q4_K, GGML_TYPE_IQ4_XS, GGML_TYPE_Q6_K, GGML_TYPE_Q2_K, GGML_TYPE_Q5_K}) {
+        if (perf && (type == GGML_TYPE_Q2_K || type == GGML_TYPE_Q5_K)) {
+            continue;
+        }
+        if (perf) {
+            for (int64_t tokens : {256, 1024, 2048}) {
+                cases.emplace_back(new test_swiglu_mmq(type, 4096, tokens, 12288));
+                cases.emplace_back(new test_swiglu_mmq(type, 5120, tokens, 17408));
+            }
+        } else {
+            for (int64_t tokens : {1, 9, 32, 256, 1024}) {
+                cases.emplace_back(new test_swiglu_mmq(type, 128, tokens, 256));
+                cases.emplace_back(new test_swiglu_mmq(type, 129, tokens, 768));
+            }
+        }
+    }
+    return cases;
+}
+
+#define P 1.0f
+#define N -1.0f
+
+// constant Hadamard matrix via Paley I construction
+static constexpr float H12[12][12] = {
+    { P, P, P, P, P, P, P, P, P, P, P, P },
+    { P, N, P, N, P, P, P, N, N, N, P, N },
+    { P, N, N, P, N, P, P, P, N, N, N, P },
+    { P, P, N, N, P, N, P, P, P, N, N, N },
+    { P, N, P, N, N, P, N, P, P, P, N, N },
+    { P, N, N, P, N, N, P, N, P, P, P, N },
+    { P, N, N, N, P, N, N, P, N, P, P, P },
+    { P, P, N, N, N, P, N, N, P, N, P, P },
+    { P, P, P, N, N, N, P, N, N, P, N, P },
+    { P, P, P, P, N, N, N, P, N, N, P, N },
+    { P, N, P, P, P, N, N, N, P, N, N, P },
+    { P, P, N, P, P, P, N, N, N, P, N, N }
+};
+
+static constexpr float H20[20][20] = {
+    { P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P },
+    { P, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N },
+    { P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P },
+    { P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P },
+    { P, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N },
+    { P, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N },
+    { P, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N },
+    { P, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N },
+    { P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P },
+    { P, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N },
+    { P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P },
+    { P, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N },
+    { P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P },
+    { P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P },
+    { P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P },
+    { P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P },
+    { P, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N },
+    { P, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N },
+    { P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P },
+    { P, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N }
+};
+
+#undef P
+#undef N
+
 // GGML_HINT_SRC0_IS_HADAMARD
 struct test_mul_mat_hadamard : public test_mul_mat {
     test_mul_mat_hadamard(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
@@ -8969,6 +9055,9 @@ static const ggml_type other_types[] = {
 
 // Test cases for evaluation: should try to cover edge cases while using small input sizes to keep the runtime low
 static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
+    if (getenv("GGML_TEST_SWIGLU_MMQ")) {
+        return make_swiglu_mmq_cases(false);
+    }
     std::vector<std::unique_ptr<test_case>> test_cases;
     std::default_random_engine rng(0);
 
@@ -11083,6 +11172,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
 // Test cases for performance evaluation: should be representative of real-world use cases
 static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
+    if (getenv("GGML_TEST_SWIGLU_MMQ")) {
+        return make_swiglu_mmq_cases(true);
+    }
     std::vector<std::unique_ptr<test_case>> test_cases;
 
     // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands