sycl: ssm_conv: fuse the SiLU epilogue into the ssm_conv kernel (#28929)

* opencl: fix warnings

* opencl: fix warnings for non adreno

.github/workflows/check-vendor.yml

@@ -19,7 +19,7 @@ on:
 
 jobs:
   check-vendor:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
 
     steps:
       - name: Checkout

.github/workflows/pre-tokenizer-hashes.yml

@@ -12,7 +12,7 @@ on:
 
 jobs:
     pre-tokenizer-hashes:
-        runs-on: [self-hosted, fast]
+        runs-on: ubuntu-slim
 
         steps:
         - name: Checkout repository

.github/workflows/python-check-requirements.yml

@@ -20,7 +20,7 @@ concurrency:
 
 jobs:
   python-check-requirements:
-    runs-on: [self-hosted, CPU, fast]
+    runs-on: ${{ 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
     name: check-requirements
     steps:
       - name: Check out source repository

.github/workflows/python-lint.yml

@@ -21,7 +21,7 @@ concurrency:
 
 jobs:
   flake8-lint:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     name: Lint
     steps:
       - name: Check out source repository

.github/workflows/python-type-check.yml

@@ -22,7 +22,7 @@ concurrency:
 
 jobs:
   python-type-check:
-    runs-on: [self-hosted, fast]
+    runs-on: ubuntu-slim
     name: python type-check
     steps:
       - name: Check out source repository

.github/workflows/update-ops-docs.yml

@@ -16,7 +16,7 @@ on:
 
 jobs:
     update-ops-docs:
-        runs-on: [self-hosted, fast, ARM64]
+        runs-on: ubuntu-slim
 
         steps:
         - name: Checkout repository

CODEOWNERS

@@ -97,7 +97,6 @@
 /src/models/                            @CISC
 /tests/                                 @ggerganov
 /tests/test-chat.*                      @pwilkin
-/tests/test-llama-archs.cpp             @JohannesGaessler
 /tools/batched-bench/                   @ggerganov
 /tools/cli/                             @ngxson
 /tools/completion/                      @ggerganov

ggml/include/ggml-sycl.h

@@ -25,7 +25,7 @@ GGML_BACKEND_API bool ggml_backend_is_sycl(ggml_backend_t backend);
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_buffer_type(int device);
 
 // split tensor buffer that splits matrices by rows across multiple devices
-GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type(const float * tensor_split);
+GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type([[maybe_unused]] int main_device, const float * tensor_split);
 
 // Tensor parallelism (--split-mode tensor): comm_init/free/allreduce_tensor
 // trio queried by the meta-backend via ggml_backend_reg_get_proc_address.

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -1440,6 +1440,7 @@ static void load_cl_kernels_argsort(ggml_backend_opencl_context *backend_ctx) {
 
 static bool use_adreno_bin_kernels(ggml_backend_opencl_context * backend_ctx) {
 #ifndef GGML_OPENCL_USE_ADRENO_BIN_KERNELS
+    GGML_UNUSED(backend_ctx);
     return false;
 #else
     if (backend_ctx->gpu_family != GPU_FAMILY::ADRENO) {
@@ -6309,6 +6310,8 @@ static void ggml_opencl_print_backend_info(ggml_backend_opencl_device_context *
 
     auto * backend_ctx = dev_ctx->backend_ctx;
 
+    GGML_LOG_INFO("ggml_opencl: OpenCL device: %s\n",
+        backend_ctx->device_name.c_str());
     GGML_LOG_INFO("ggml_opencl: OpenCL driver: %s\n",
         backend_ctx->driver_version.c_str());
     GGML_LOG_INFO("ggml_opencl: vector subgroup broadcast support: %s\n",
@@ -6325,11 +6328,11 @@ static void ggml_opencl_print_backend_info(ggml_backend_opencl_device_context *
         backend_ctx->global_mem_size/1024/1024);
     GGML_LOG_INFO("ggml_opencl: max mem alloc size: %zu MB\n",
         backend_ctx->max_alloc_size/1024/1024);
-    GGML_LOG_INFO("ggml_opencl: device max image buffer size (pixels): %lu\n",
+    GGML_LOG_INFO("ggml_opencl: device max image buffer size (pixels): %zu\n",
         backend_ctx->image_max_buffer_size);
-    GGML_LOG_INFO("ggml_opencl: device max image2d size: %lu x %lu\n",
+    GGML_LOG_INFO("ggml_opencl: device max image2d size: %zu x %zu\n",
         backend_ctx->image2d_max_width, backend_ctx->image2d_max_height);
-    GGML_LOG_INFO("ggml_opencl: device max workgroup size: %lu\n",
+    GGML_LOG_INFO("ggml_opencl: device max workgroup size: %zu\n",
         backend_ctx->max_workgroup_size);
     GGML_LOG_INFO("ggml_opencl: SVM coarse grain buffer support: %s\n",
         backend_ctx->svm_caps & CL_DEVICE_SVM_COARSE_GRAIN_BUFFER ? "true" : "false");
@@ -7627,7 +7630,7 @@ static void ggml_cl_moe_bias_glu_fused(ggml_backend_t backend, ggml_tensor * gat
     size_t global_work_size[] = { (size_t)glu->ne[1]*nth, (size_t)glu->ne[2], 1 };
     size_t local_work_size[]  = { (size_t)nth, 1, 1 };
 
-    backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, (ggml_tensor *)glu);
+    backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, glu);
 }
 
 // Fusion B: the MoE down-projection bias add feeding the combine.
@@ -7771,7 +7774,7 @@ static void ggml_cl_moe_bias_combine_fused(ggml_backend_t backend, const ggml_te
 
     size_t lws[2] = { 64, 1 };
     size_t gws[2] = { (size_t)(((n_embd4 + 63) / 64) * 64), (size_t)nt };
-    backend_ctx->enqueue_ndrange_kernel(kernel, 2, gws, lws, (ggml_tensor *)dst);
+    backend_ctx->enqueue_ndrange_kernel(kernel, 2, gws, lws, dst);
 }
 
 
@@ -8114,7 +8117,6 @@ static void ggml_cl_mul_mat_q4_k_glu_fused(ggml_backend_t backend, ggml_tensor *
     GGML_UNUSED(gate_tensor);
     GGML_UNUSED(up_tensor);
     GGML_UNUSED(glu_tensor);
-    GGML_ABORT("q4_K GLU fusion requires GGML_OPENCL_USE_ADRENO_KERNELS");
 #endif
 }
 
@@ -11336,7 +11338,6 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
         if (use_adreno_moe_kernels(backend_ctx, tensor)) {
-            cl_int err;
             cl_kernel kernel = backend_ctx->kernel_restore_block_q4_0_trans4_ns;
 
             cl_mem data_device = ggml_cl_create_temp_download_buffer(context, queue, ggml_nbytes(tensor), tensor->name);
@@ -11535,7 +11536,6 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
         if (use_adreno_moe_kernels(backend_ctx, tensor)) {
-            cl_int err;
             // TODO: use ggml_cl_buffer to manage this temporary buffer
             cl_mem data_device = ggml_cl_create_temp_download_buffer(context, queue, ggml_nbytes(tensor), tensor->name);
             GGML_ASSERT(data_device != NULL && "get_tensor: temp download buffer alloc failed");
@@ -11638,7 +11638,6 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer,
 
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
         if (use_adreno_moe_kernels(backend_ctx, tensor)) {
-            cl_int err;
             // TODO: use ggml_cl_buffer to manage this temporary buffer
             cl_mem data_device = ggml_cl_create_temp_download_buffer(context, queue, ggml_nbytes(tensor), tensor->name);
             GGML_ASSERT(data_device != NULL && "get_tensor: temp download buffer alloc failed");

ggml/src/ggml-sycl/fusion.cpp

@@ -208,5 +208,53 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ
         return true;
     }
 
+    if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY &&
+        unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) {
+        const ggml_tensor * ssm_conv = cgraph->nodes[node_idx];
+        const ggml_tensor * silu     = cgraph->nodes[node_idx + 1];
+
+        if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) {
+            return false;
+        }
+        if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) {
+            return false;
+        }
+        // the fused kernel writes the SiLU output with dense strides, so it must be contiguous
+        if (!ggml_is_contiguous(silu)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_ADD &&
+        ops.begin()[2] == GGML_OP_UNARY && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) {
+        const ggml_tensor * ssm_conv = cgraph->nodes[node_idx];
+        const ggml_tensor * add      = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * silu     = cgraph->nodes[node_idx + 2];
+
+        if (ggml_get_unary_op(silu) != unary_ops.begin()[0]) {
+            return false;
+        }
+        if (ssm_conv->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) {
+            return false;
+        }
+        // the fused kernel writes the SiLU output with dense strides, so it must be contiguous
+        if (!ggml_is_contiguous(silu)) {
+            return false;
+        }
+
+        // ADD must consume ssm_conv's output and broadcast a 1-D channel-wise bias
+        const ggml_tensor * bias = (add->src[0] == ssm_conv) ? add->src[1] : add->src[0];
+        if (bias->type != GGML_TYPE_F32 || !ggml_is_contiguous(bias)) {
+            return false;
+        }
+        if (ggml_nelements(bias) != ssm_conv->ne[0] || bias->ne[0] != ssm_conv->ne[0]) {
+            return false;
+        }
+
+        return true;
+    }
+
     return false;
 }

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -1484,7 +1484,7 @@ static ggml_backend_buffer_type_i ggml_backend_sycl_split_buffer_type_interface
     /* .is_host          = */ ggml_backend_sycl_split_buffer_type_is_host,
 };
 
-ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type(const float * tensor_split) {
+ggml_backend_buffer_type_t ggml_backend_sycl_split_buffer_type([[maybe_unused]] int main_device, const float * tensor_split) {
     GGML_SYCL_DEBUG("[SYCL] call ggml_backend_sycl_split_buffer_type\n");
 
     static std::mutex mutex;
@@ -3243,7 +3243,7 @@ inline void ggml_sycl_op_scale(ggml_backend_sycl_context & ctx, ggml_tensor * ds
     SYCL_CHECK(0);
 }
 
-static void ggml_sycl_set_peer_access(const int n_tokens, int main_device) {
+static void ggml_sycl_set_peer_access(const int n_tokens, [[maybe_unused]] int main_device) {
     static bool peer_access_enabled = false;
 
     const bool enable_peer_access = n_tokens <= GGML_SYCL_PEER_MAX_BATCH_SIZE;
@@ -6034,6 +6034,20 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             }
         }
 
+        if (node->op == GGML_OP_SSM_CONV &&
+            ggml_sycl_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_ADD, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) {
+            ggml_sycl_ssm_conv_fused(*sycl_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
+            i += 2;
+            continue;
+        }
+
+        if (node->op == GGML_OP_SSM_CONV &&
+            ggml_sycl_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) {
+            ggml_sycl_ssm_conv_fused(*sycl_ctx, node, nullptr, cgraph->nodes[i + 1]);
+            i++;
+            continue;
+        }
+
         if (node->op == GGML_OP_MUL_MAT && ggml_sycl_mul_mat_glu_mmvq_fused(*sycl_ctx, cgraph, i)) {
             i += 2;
             continue;

ggml/src/ggml-sycl/ssm_conv.cpp

@@ -1,11 +1,71 @@
 #include "ssm_conv.hpp"
 #include "common.hpp"
+#include "element_wise.hpp"
 
 #include <cstdio>
 
 using namespace sycl;
 
-static void kernel_ssm_conv(
+// One output element of the conv. DC is d_conv as a compile-time constant (0 keeps the
+// runtime loop); unfused callers pass literal false/nullptr so the epilogue folds away.
+template <int DC>
+static __dpct_inline__ void ssm_conv_element(
+    size_t idx,
+    const float *src_data,
+    const float *weights,
+    float *dst_data,
+    int d_conv,
+    int d_inner,
+    int n_t,
+    int src_stride_inner,
+    int src_stride_seq,
+    int dst_stride_token,
+    int dst_stride_seq,
+    bool apply_silu,
+    const float *bias
+) {
+    // src is token-contiguous per channel, dst is channel-contiguous per token,
+    // so indexing token-fastest coalesces the d_conv loads.
+    const int token   = static_cast<int>(idx % n_t);
+    const int channel = static_cast<int>((idx / n_t) % d_inner);
+    const int seq     = static_cast<int>(idx / (static_cast<size_t>(n_t) * static_cast<size_t>(d_inner)));
+
+    const float *s = src_data
+        + static_cast<size_t>(seq) * static_cast<size_t>(src_stride_seq)
+        + static_cast<size_t>(channel) * static_cast<size_t>(src_stride_inner)
+        + static_cast<size_t>(token);
+
+    const float *c = weights + static_cast<size_t>(channel) * static_cast<size_t>(d_conv);
+
+    float sumf = 0.0f;
+    if constexpr (DC > 0) {
+#pragma unroll
+        for (int i0 = 0; i0 < DC; ++i0) {
+            sumf += s[i0] * c[i0];
+        }
+    } else {
+        for (int i0 = 0; i0 < d_conv; ++i0) {
+            sumf += s[i0] * c[i0];
+        }
+    }
+
+    // fused bias add: the ADD node broadcasts a 1-D channel bias over tokens
+    if (bias != nullptr) {
+        sumf += bias[channel];
+    }
+
+    const size_t dst_idx =
+        static_cast<size_t>(seq) * static_cast<size_t>(dst_stride_seq) +
+        static_cast<size_t>(token) * static_cast<size_t>(dst_stride_token) +
+        static_cast<size_t>(channel);
+
+    dst_data[dst_idx] = apply_silu ? op_silu(sumf) : sumf;
+}
+
+// FUSED=false keeps apply_silu/bias out of the kernel capture list, so the unfused launch
+// takes the pre-fusion argument list; matters at n_t == 1, where the op is launch-bound.
+template <int DC, bool FUSED>
+static void kernel_ssm_conv_impl(
     queue &q,
     const float *src_data,
     const float *weights,
@@ -18,7 +78,9 @@ static void kernel_ssm_conv(
     int src_stride_inner,
     int src_stride_seq,
     int dst_stride_token,
-    int dst_stride_seq
+    int dst_stride_seq,
+    bool apply_silu,
+    const float *bias
 ) {
     const size_t total_work = static_cast<size_t>(d_inner) * static_cast<size_t>(n_t) * static_cast<size_t>(n_s);
     const size_t work_group_size = 256;
@@ -27,53 +89,199 @@ static void kernel_ssm_conv(
     const range<1> global_range(num_work_groups * work_group_size);
     const range<1> local_range(work_group_size);
 
-    q.submit([&](handler &h) {
-        h.parallel_for(
-            nd_range<1>(global_range, local_range),
-            [=](nd_item<1> item) {
-                const size_t idx = item.get_global_id(0);
-                if (idx >= total_work) {
-                    return;
+    if constexpr (FUSED) {
+        q.submit([&](handler &h) {
+            h.parallel_for(
+                nd_range<1>(global_range, local_range),
+                [=](nd_item<1> item) {
+                    const size_t idx = item.get_global_id(0);
+                    if (idx >= total_work) {
+                        return;
+                    }
+
+                    ssm_conv_element<DC>(idx, src_data, weights, dst_data, d_conv, d_inner, n_t,
+                                         src_stride_inner, src_stride_seq, dst_stride_token,
+                                         dst_stride_seq, apply_silu, bias);
                 }
+            );
+        });
+    } else {
+        GGML_UNUSED(apply_silu);
+        GGML_UNUSED(bias);
 
-                // src has the tokens of one channel contiguous, dst has the channels of one
-                // token contiguous, so either the loads or the store must be strided. Indexing
-                // token-fastest coalesces the d_conv loads, which measured faster except for
-                // short, cache-resident rows.
-                const int token   = static_cast<int>(idx % n_t);
-                const int channel = static_cast<int>((idx / n_t) % d_inner);
-                const int seq     = static_cast<int>(idx / (static_cast<size_t>(n_t) * static_cast<size_t>(d_inner)));
+        q.submit([&](handler &h) {
+            h.parallel_for(
+                nd_range<1>(global_range, local_range),
+                [=](nd_item<1> item) {
+                    const size_t idx = item.get_global_id(0);
+                    if (idx >= total_work) {
+                        return;
+                    }
 
-                const float *s = src_data
-                    + static_cast<size_t>(seq) * static_cast<size_t>(src_stride_seq)
-                    + static_cast<size_t>(channel) * static_cast<size_t>(src_stride_inner)
-                    + static_cast<size_t>(token);
-
-                const float *c = weights + static_cast<size_t>(channel) * static_cast<size_t>(d_conv);
-
-                float sumf = 0.0f;
-                for (int i0 = 0; i0 < d_conv; ++i0) {
-                    sumf += s[i0] * c[i0];
+                    ssm_conv_element<DC>(idx, src_data, weights, dst_data, d_conv, d_inner, n_t,
+                                         src_stride_inner, src_stride_seq, dst_stride_token,
+                                         dst_stride_seq, false, nullptr);
                 }
-
-                const size_t dst_idx =
-                    static_cast<size_t>(seq) * static_cast<size_t>(dst_stride_seq) +
-                    static_cast<size_t>(token) * static_cast<size_t>(dst_stride_token) +
-                    static_cast<size_t>(channel);
-
-                dst_data[dst_idx] = sumf;
-            }
-        );
-    });
+            );
+        });
+    }
 }
 
-inline void ggml_sycl_op_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+// SLM transpose tile: coalesces both the loads and the stores. The +1 pad makes the row
+// stride 33, coprime with 32 banks, so both phases are bank-conflict-free.
+template <int DC, int TT, int TC, int WG>
+static __dpct_inline__ void ssm_conv_tile(
+    nd_item<1> it, local_accessor<float, 1> tile, const float *src_data, const float *weights,
+    float *dst_data, int n_t, int nt_tiles, int nc_tiles, int src_stride_inner,
+    int src_stride_seq, int dst_stride_token, int dst_stride_seq, bool apply_silu,
+    const float *bias
+) {
+    const int    lid = static_cast<int>(it.get_local_id(0));
+    const size_t g   = it.get_group(0);
+    const int    tt  = static_cast<int>(g % nt_tiles);
+    const int    ct  = static_cast<int>((g / nt_tiles) % nc_tiles);
+    const int    seq = static_cast<int>(g / (static_cast<size_t>(nt_tiles) * nc_tiles));
+    const int    t0 = tt * TT, c0 = ct * TC;
+
+    const int ti = lid % TT;
+    const int cj = lid / TT;
+#pragma unroll
+    for (int r = 0; r < TC / (WG / TT); ++r) {
+        const int c   = cj + r * (WG / TT);
+        const int tok = t0 + ti;
+        float sumf = 0.0f;
+        if (tok < n_t) {
+            const float *s = src_data + static_cast<size_t>(seq) * src_stride_seq
+                           + static_cast<size_t>(c0 + c) * src_stride_inner + tok;
+            const float *cw = weights + static_cast<size_t>(c0 + c) * DC;
+#pragma unroll
+            for (int i = 0; i < DC; ++i) sumf += s[i] * cw[i];
+            if (bias != nullptr) sumf += bias[c0 + c];
+            if (apply_silu) sumf = op_silu(sumf);
+        }
+        tile[c * (TT + 1) + ti] = sumf;
+    }
+    it.barrier(access::fence_space::local_space);
+
+    const int cc = lid % TC;
+    const int tj = lid / TC;
+#pragma unroll
+    for (int r = 0; r < TT / (WG / TC); ++r) {
+        const int t   = tj + r * (WG / TC);
+        const int tok = t0 + t;
+        if (tok < n_t) {
+            dst_data[static_cast<size_t>(seq) * dst_stride_seq
+                     + static_cast<size_t>(tok) * dst_stride_token + c0 + cc]
+                = tile[cc * (TT + 1) + t];
+        }
+    }
+}
+
+// Same FUSED split as kernel_ssm_conv_impl. The fused instantiation keeps the runtime
+// apply_silu/bias branches: at n_t >= 32 they are amortized over the whole tile.
+template <int DC, bool FUSED>
+static void kernel_ssm_conv_tiled(
+    queue &q, const float *src_data, const float *weights, float *dst_data,
+    int d_inner, int n_t, int n_s, int src_stride_inner, int src_stride_seq,
+    int dst_stride_token, int dst_stride_seq, bool apply_silu, const float *bias
+) {
+    constexpr int TT = 32, TC = 32, WG = 256;
+    const int nt_tiles = (n_t + TT - 1) / TT;
+    const int nc_tiles = d_inner / TC;
+    const size_t groups = static_cast<size_t>(nt_tiles) * nc_tiles * n_s;
+
+    if constexpr (FUSED) {
+        q.submit([&](handler &h) {
+            local_accessor<float, 1> tile(range<1>(TC * (TT + 1)), h);
+            h.parallel_for(nd_range<1>(range<1>(groups * WG), range<1>(WG)), [=](nd_item<1> it) {
+                ssm_conv_tile<DC, TT, TC, WG>(it, tile, src_data, weights, dst_data, n_t, nt_tiles,
+                                              nc_tiles, src_stride_inner, src_stride_seq,
+                                              dst_stride_token, dst_stride_seq, apply_silu, bias);
+            });
+        });
+    } else {
+        GGML_UNUSED(apply_silu);
+        GGML_UNUSED(bias);
+
+        q.submit([&](handler &h) {
+            local_accessor<float, 1> tile(range<1>(TC * (TT + 1)), h);
+            h.parallel_for(nd_range<1>(range<1>(groups * WG), range<1>(WG)), [=](nd_item<1> it) {
+                ssm_conv_tile<DC, TT, TC, WG>(it, tile, src_data, weights, dst_data, n_t, nt_tiles,
+                                              nc_tiles, src_stride_inner, src_stride_seq,
+                                              dst_stride_token, dst_stride_seq, false, nullptr);
+            });
+        });
+    }
+}
+
+static void kernel_ssm_conv(
+    queue &q,
+    const float *src_data,
+    const float *weights,
+    float *dst_data,
+    int d_conv,
+    int d_inner,
+    int n_t,
+    int n_s,
+    int ncs,
+    int src_stride_inner,
+    int src_stride_seq,
+    int dst_stride_token,
+    int dst_stride_seq,
+    bool apply_silu,
+    const float *bias
+) {
+    // Only the fused instantiations carry apply_silu/bias as kernel arguments; the plain
+    // ssm_conv launch keeps the argument list it had before the fusion landed.
+    const bool fused = apply_silu || bias != nullptr;
+
+    // d_inner must be a multiple of 32 so the channel tiles are exact; the transpose is only
+    // worth it for n_t >= 32. d_conv == 4 is the only window with a DC-specialized kernel.
+    if (d_conv == 4 && n_t >= 32 && (d_inner % 32) == 0) {
+        if (fused) {
+            kernel_ssm_conv_tiled<4, true>(q, src_data, weights, dst_data, d_inner, n_t, n_s,
+                                           src_stride_inner, src_stride_seq, dst_stride_token,
+                                           dst_stride_seq, apply_silu, bias);
+        } else {
+            kernel_ssm_conv_tiled<4, false>(q, src_data, weights, dst_data, d_inner, n_t, n_s,
+                                            src_stride_inner, src_stride_seq, dst_stride_token,
+                                            dst_stride_seq, apply_silu, bias);
+        }
+        return;
+    }
+
+    if (d_conv == 4) {
+        if (fused) {
+            kernel_ssm_conv_impl<4, true>(q, src_data, weights, dst_data, d_conv, d_inner, n_t, n_s,
+                                          ncs, src_stride_inner, src_stride_seq, dst_stride_token,
+                                          dst_stride_seq, apply_silu, bias);
+        } else {
+            kernel_ssm_conv_impl<4, false>(q, src_data, weights, dst_data, d_conv, d_inner, n_t, n_s,
+                                           ncs, src_stride_inner, src_stride_seq, dst_stride_token,
+                                           dst_stride_seq, apply_silu, bias);
+        }
+        return;
+    }
+
+    if (fused) {
+        kernel_ssm_conv_impl<0, true>(q, src_data, weights, dst_data, d_conv, d_inner, n_t, n_s,
+                                      ncs, src_stride_inner, src_stride_seq, dst_stride_token,
+                                      dst_stride_seq, apply_silu, bias);
+    } else {
+        kernel_ssm_conv_impl<0, false>(q, src_data, weights, dst_data, d_conv, d_inner, n_t, n_s,
+                                       ncs, src_stride_inner, src_stride_seq, dst_stride_token,
+                                       dst_stride_seq, apply_silu, bias);
+    }
+}
+
+inline void ggml_sycl_op_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * silu_dst = nullptr, const float * bias = nullptr) {
     ggml_tensor * src0 = dst->src[0];
     ggml_tensor * src1 = dst->src[1];
 
     GGML_ASSERT(src0->type == GGML_TYPE_F32);
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
     GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(bias == nullptr || silu_dst != nullptr);
 
     const int d_conv   = src1->ne[0];
     const int ncs      = src0->ne[0];
@@ -104,7 +312,8 @@ inline void ggml_sycl_op_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor *
 
         const float *src_data = static_cast<const float *>(src0->data);
         const float *weights  = static_cast<const float *>(src1->data);
-        float *dst_data       = static_cast<float *>(dst->data);
+        const bool apply_silu = silu_dst != nullptr;
+        float *dst_data       = static_cast<float *>((silu_dst ? silu_dst : dst)->data);
 
         GGML_ASSERT(src_data && weights && dst_data);
 
@@ -121,7 +330,9 @@ inline void ggml_sycl_op_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor *
             src_stride_inner,
             src_stride_seq,
             dst_stride_token,
-            dst_stride_seq
+            dst_stride_seq,
+            apply_silu,
+            bias
         );
 
     } catch (const std::exception &e) {
@@ -134,3 +345,17 @@ void ggml_sycl_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
     ggml_sycl_op_ssm_conv(ctx, dst);
 }
+
+// Fused ssm_conv + ADD + SiLU: write silu(conv(x) + b) straight into silu_dst, eliding the
+// standalone SiLU launch and its HBM round-trip of the conv output.
+void ggml_sycl_ssm_conv_fused(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * add, ggml_tensor * silu_dst) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
+    GGML_ASSERT(silu_dst && ggml_are_same_shape(dst, silu_dst) && silu_dst->type == GGML_TYPE_F32);
+    // the fused kernel reads only the ADD's bias operand; the ADD result is never written
+    const float * bias = nullptr;
+    if (add != nullptr) {
+        const ggml_tensor * bias_t = (add->src[0] == dst) ? add->src[1] : add->src[0];
+        bias = static_cast<const float *>(bias_t->data);
+    }
+    ggml_sycl_op_ssm_conv(ctx, dst, silu_dst, bias);
+}

ggml/src/ggml-sycl/ssm_conv.hpp

@@ -3,3 +3,4 @@
 #include "common.hpp"
 
 void ggml_sycl_ssm_conv(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+void ggml_sycl_ssm_conv_fused(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * add, ggml_tensor * silu_dst);

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1130,7 +1130,9 @@ struct vk_device_struct {
     vk_pipeline pipeline_count_equal_i32;
     vk_pipeline pipeline_dsv4_hc_comb_f32;
     vk_pipeline pipeline_dsv4_hc_pre_f32;
+    vk_pipeline pipeline_dsv4_hc_pre_gated_f32;
     vk_pipeline pipeline_dsv4_hc_post_f32;
+    vk_pipeline pipeline_dsv4_hc_post_nocomb_f32;
     std::map<vk_solve_tri_pipeline_state, vk_pipeline> pipeline_solve_tri_f32;
     vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16;
     vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16;
@@ -1514,12 +1516,14 @@ struct vk_op_dsv4_hc_pre_push_constants {
     uint32_t n_tokens;
 
     uint32_t nbx0; uint32_t nbx1; uint32_t nbx2;
-    uint32_t nbw0; uint32_t nbw1;
+    uint32_t nbw0; uint32_t nbw1; uint32_t nbw2;
     uint32_t nbd0; uint32_t nbd1;
 
     uint32_t x_offset;
     uint32_t w_offset;
     uint32_t d_offset;
+
+    float scale;
 };
 
 struct vk_op_dsv4_hc_post_push_constants {
@@ -2737,7 +2741,7 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk
     p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
     p.r_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
     p.p_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type);
-    p.c_offset = get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type);
+    p.c_offset = src3 ? get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type) : 0;
     p.d_offset = get_misalign_bytes(ctx, dst)  / ggml_type_size(dst->type);
 }
 
@@ -6173,8 +6177,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
         ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, sizeof(vk_op_dsv4_hc_comb_push_constants), {tokens_per_workgroup, 1, 1}, { device->subgroup_size }, 1, true, true, device->subgroup_size);
     }
 
-    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32,  "dsv4_hc_pre_f32",  dsv4_hc_pre_f32_len,  dsv4_hc_pre_f32_data,  "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants),  {256, 1, 1}, { 256 }, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32,        "dsv4_hc_pre_f32",        dsv4_hc_pre_f32_len,  dsv4_hc_pre_f32_data,  "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants),  {256, 1, 1}, { 256, 0 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_gated_f32,  "dsv4_hc_pre_gated_f32",  dsv4_hc_pre_f32_len,  dsv4_hc_pre_f32_data,  "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants),  {256, 1, 1}, { 256, 1 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32,       "dsv4_hc_post_f32",       dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 1 }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_nocomb_f32,"dsv4_hc_post_nocomb_f32",dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256, 0 }, 1);
 
     for (auto &s : device->pipeline_solve_tri_f32) {
         const vk_solve_tri_pipeline_state &state = s.first;
@@ -10356,7 +10362,10 @@ static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subc
 static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst) {
     VK_LOG_DEBUG("ggml_vk_dsv4_hc_pre(" << x << ", " << weights << ", " << dst << ")");
 
-    vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32;
+    const float scale = ggml_get_op_params_f32(dst, 0);
+    const bool  gated = ggml_get_op_params_i32(dst, 1) != 0;
+
+    vk_pipeline pipeline = gated ? ctx->device->pipeline_dsv4_hc_pre_gated_f32 : ctx->device->pipeline_dsv4_hc_pre_f32;
     GGML_ASSERT(pipeline != nullptr);
 
     const uint32_t n_embd   = (uint32_t)x->ne[0];
@@ -10371,9 +10380,10 @@ static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subct
     vk_op_dsv4_hc_pre_push_constants pc = {
         n_embd, n_tokens,
         ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), ggml_vk_nb_elem(x, 2),
-        ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1),
+        ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1), ggml_vk_nb_elem(weights, 2),
         ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1),
         0, 0, 0,
+        scale,
     };
     init_pushconst_tensor_offsets(ctx, pc, x, weights, nullptr, nullptr, dst);
 
@@ -10383,7 +10393,7 @@ static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subct
 static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst) {
     VK_LOG_DEBUG("ggml_vk_dsv4_hc_post(" << x << ", " << residual << ", " << post << ", " << comb << ", " << dst << ")");
 
-    vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32;
+    vk_pipeline pipeline = comb ? ctx->device->pipeline_dsv4_hc_post_f32 : ctx->device->pipeline_dsv4_hc_post_nocomb_f32;
     GGML_ASSERT(pipeline != nullptr);
 
     const uint32_t n_embd   = (uint32_t)x->ne[0];
@@ -10394,7 +10404,7 @@ static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subc
     const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x,        true);
     const vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, residual, true);
     const vk_subbuffer p_buf = ggml_vk_tensor_subbuffer(ctx, post,     true);
-    const vk_subbuffer c_buf = ggml_vk_tensor_subbuffer(ctx, comb,     true);
+    const vk_subbuffer c_buf = comb ? ggml_vk_tensor_subbuffer(ctx, comb, true) : x_buf;
     const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst,      true);
 
     vk_op_dsv4_hc_post_push_constants pc = {
@@ -10402,7 +10412,7 @@ static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subc
         ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1),
         ggml_vk_nb_elem(residual, 0), ggml_vk_nb_elem(residual, 1), ggml_vk_nb_elem(residual, 2),
         ggml_vk_nb_elem(post, 0), ggml_vk_nb_elem(post, 1),
-        ggml_vk_nb_elem(comb, 0), ggml_vk_nb_elem(comb, 1), ggml_vk_nb_elem(comb, 2),
+        comb ? ggml_vk_nb_elem(comb, 0) : 0, comb ? ggml_vk_nb_elem(comb, 1) : 0, comb ? ggml_vk_nb_elem(comb, 2) : 0,
         ggml_vk_nb_elem(dst,  0), ggml_vk_nb_elem(dst,  1), ggml_vk_nb_elem(dst,  2),
         0, 0, 0, 0, 0,
     };
@@ -19692,10 +19702,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 }
                 // hc is hardcoded to 4 in the shaders. ggml only constrains it
                 // to 4 for COMB, so PRE/POST have to be checked here.
-                if (op->op == GGML_OP_DSV4_HC_PRE && (op->src[0]->ne[1] != 4 || ggml_get_op_params_i32(op, 1) != 0)) {
+                if (op->op == GGML_OP_DSV4_HC_PRE && op->src[0]->ne[1] != 4) {
                     return false;
                 }
-                if (op->op == GGML_OP_DSV4_HC_POST && (op->src[1]->ne[1] != 4 || op->src[3] == nullptr)) {
+                if (op->op == GGML_OP_DSV4_HC_POST && op->src[1]->ne[1] != 4) {
                     return false;
                 }
                 if (op->op == GGML_OP_DSV4_HC_COMB) {
@@ -20695,7 +20705,11 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
             tensor_clone = ggml_dsv4_hc_comb(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
                 ggml_get_op_params_f32(tensor, 0), ggml_get_op_params_i32(tensor, 1));
         } else if (tensor->op == GGML_OP_DSV4_HC_PRE) {
-            tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]);
+            if (ggml_get_op_params_i32(tensor, 1) != 0) {
+                tensor_clone = ggml_dsv4_hc_pre_gated(ggml_ctx, src_clone[0], src_clone[1], ggml_get_op_params_f32(tensor, 0));
+            } else {
+                tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]);
+            }
         } else if (tensor->op == GGML_OP_DSV4_HC_POST) {
             tensor_clone = ggml_dsv4_hc_post(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
         } else if (tensor->op == GGML_OP_MEAN) {

ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp

@@ -27,6 +27,8 @@ layout (push_constant) uniform parameter {
     uint inner_end;
 } p;
 
+shared int s;
+
 void argsort(bool needs_bounds_check, const uint row) {
     // bitonic sort
     int col = int(gl_GlobalInvocationID.x);
@@ -35,6 +37,12 @@ void argsort(bool needs_bounds_check, const uint row) {
     const uint row_offset = row * p.ncols;
     uint idx_offset = row * p.ncols_padded;
 
+    // workaround for NV driver/compiler bug - dummy use of shared memory
+    if (gl_LocalInvocationIndex == 0) {
+        s = 0;
+        col += s;
+    }
+
     bool need_barrier = false;
 
     // initialize indices

ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp

@@ -7,8 +7,13 @@
 //
 //   dst[i0, idst, it] = x[i0, it]*post[idst, it]
 //                     + sum_isrc residual[i0, isrc, it]*comb[idst, isrc, it]
+//
+// HAS_COMB == 0: identity mixing, each stream keeps its own residual:
+//
+//   dst[i0, idst, it] = x[i0, it]*post[idst, it] + residual[i0, idst, it]
 
 layout(constant_id = 0) const uint BLOCK_SIZE = 256;
+layout(constant_id = 1) const uint HAS_COMB = 1;
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
@@ -48,7 +53,7 @@ void main() {
     if (tid < hc) {
         post_s[tid] = data_p[p_offset + tid * nbp0 + it * nbp1];
     }
-    if (tid < hc * hc) {
+    if (HAS_COMB == 1 && tid < hc * hc) {
         const uint idst = tid & 3;
         const uint isrc = tid >> 2;
         comb_s[tid] = data_c[c_offset + idst * nbc0 + isrc * nbc1 + it * nbc2];
@@ -74,9 +79,13 @@ void main() {
     [[unroll]]
     for (uint idst = 0; idst < hc; ++idst) {
         float result = xv * post_s[idst];
-        [[unroll]]
-        for (uint isrc = 0; isrc < hc; ++isrc) {
-            result = fma(r[isrc], comb_s[idst + hc * isrc], result);
+        if (HAS_COMB == 1) {
+            [[unroll]]
+            for (uint isrc = 0; isrc < hc; ++isrc) {
+                result = fma(r[isrc], comb_s[idst + hc * isrc], result);
+            }
+        } else {
+            result += r[idst];
         }
         data_d[d_offset + i0 * nbd0 + idst * nbd1 + it * nbd2] = result;
     }

ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp

@@ -4,9 +4,14 @@
 
 // Collapse the hc residual streams of a token into one, weighted per stream:
 //
-//   dst[i0, it] = sum_ih x[i0, ih, it] * weights[ih, it]
+//   dst[i0, it] = scale * sum_ih x[i0, ih, it] * weights[ih, it]
+//
+// GATED: weights is a per-element gate [n_embd, hc, n_tokens], applied as sigmoid:
+//
+//   dst[i0, it] = scale * sum_ih x[i0, ih, it] * sigmoid(gate[i0, ih, it])
 
 layout(constant_id = 0) const uint BLOCK_SIZE = 256;
+layout(constant_id = 1) const uint GATED = 0;
 
 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
 
@@ -16,12 +21,14 @@ layout(push_constant) uniform parameter
     uint n_tokens;
 
     uint nbx0; uint nbx1; uint nbx2;   // x
-    uint nbw0; uint nbw1;              // weights
+    uint nbw0; uint nbw1; uint nbw2;   // weights / gate
     uint nbd0; uint nbd1;              // dst
 
     uint x_offset;
     uint w_offset;
     uint d_offset;
+
+    float scale;
 };
 
 layout(binding = 0, std430) readonly buffer X { float data_x[]; };
@@ -36,10 +43,12 @@ void main() {
     const uint tid = gl_LocalInvocationID.x;
     const uint it  = gl_WorkGroupID.y;
 
-    if (tid < hc) {
-        w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1];
+    if (GATED == 0) {
+        if (tid < hc) {
+            w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1];
+        }
+        barrier();
     }
-    barrier();
 
     // After the barrier, so every invocation reaches it.
     const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid;
@@ -48,12 +57,19 @@ void main() {
     }
 
     const uint xb = x_offset + i0 * nbx0 + it * nbx2;
+    const uint wb = w_offset + i0 * nbw0 + it * nbw2;
 
     float result = 0.0f;
     [[unroll]]
     for (uint ih = 0; ih < hc; ++ih) {
-        result = fma(data_x[xb + ih * nbx1], w[ih], result);
+        float wv;
+        if (GATED == 1) {
+            wv = 1.0f / (1.0f + exp(-data_w[wb + ih * nbw1]));
+        } else {
+            wv = w[ih];
+        }
+        result = fma(data_x[xb + ih * nbx1], wv, result);
     }
 
-    data_d[d_offset + i0 * nbd0 + it * nbd1] = result;
+    data_d[d_offset + i0 * nbd0 + it * nbd1] = scale * result;
 }

ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp

@@ -31,7 +31,7 @@ layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
 layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
 
 #if BDA
-layout (buffer_reference) buffer D_ptr {D_TYPE d;};
+layout (buffer_reference, buffer_reference_align = D_SIZE) buffer D_ptr {D_TYPE d;};
 #endif
 
 void im2col(const uint ow, const uint z_idx) {

ggml/src/ggml-vulkan/vulkan-shaders/im2col_3d.comp

@@ -50,7 +50,7 @@ layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
 layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};
 
 #if BDA
-layout (buffer_reference) buffer D_ptr {D_TYPE d;};
+layout (buffer_reference, buffer_reference_align = D_SIZE) buffer D_ptr {D_TYPE d;};
 #endif
 
 void main() {

src/llama-context.cpp

@@ -599,8 +599,11 @@ void llama_context::sched_reserve() {
 
     LLAMA_LOG_DEBUG("%s: max_nodes = %zu\n", __func__, max_nodes);
 
-    gf_res_prev.reset(new llm_graph_result(max_nodes));
+    for (auto & res : gf_res_prev) {
+        res.reset();
+    }
     gf_res_reserve.reset(new llm_graph_result(max_nodes));
+    gf_res_prev_active = nullptr;
 
     sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload));
 
@@ -816,10 +819,14 @@ bool llama_context::memory_update(bool optimize) {
                 }
         }
 
-        // reset the previous graph result to make sure that it won't be reused
-        // TODO: change the mctx->apply() to return information if a graph reserve is needed
-        //       reset the graph result only if the memory module did reset the scheduler
-        gf_res_prev->reset();
+        // reset the previous graph results to make sure that they won't be reused
+        // TODO: make mctx->apply() report if a graph reserve is needed, then reset graph results only if the memory module reset the scheduler
+        for (auto & res : gf_res_prev) {
+            if (res) {
+                res->reset();
+            }
+        }
+        gf_res_prev_active = nullptr;
 
         if (!mctx->apply()) {
             LLAMA_LOG_ERROR("%s: failed to apply memory update\n", __func__);
@@ -1340,14 +1347,14 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
         return nullptr;
     }
 
-    auto * res = gf_res_prev.get();
+    auto * res = get_gf_res_prev();
     auto * gf  = res->get_gf();
 
     // the new graph parameters
     // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters
     const auto gparams = graph_params(res, ubatch, mctx, gtype);
 
-    if (!graph_reuse_disable && res->can_reuse(gparams)) {
+    if (!graph_reuse_disable && gf_res_prev_active == res && res->can_reuse(gparams)) {
         //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__);
 
         // with pipeline parallelism, the previous graph_compute_async may still be running
@@ -1359,6 +1366,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
 
         n_reused++;
     } else {
+        gf_res_prev_active = nullptr;
         res->reset();
 
         ggml_backend_sched_reset(sched.get());
@@ -1381,6 +1389,8 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
             ret = GGML_STATUS_ALLOC_FAILED;
             return nullptr;
         }
+
+        gf_res_prev_active = res;
     }
 
     // set the input data for the input tensors
@@ -2361,6 +2371,14 @@ llm_graph_result * llama_context::get_gf_res_reserve() const {
     return static_cast<llm_graph_result *>(gf_res_reserve.get());
 }
 
+llm_graph_result * llama_context::get_gf_res_prev() {
+    auto & res = gf_res_prev[n_outputs > 0];
+    if (!res) {
+        res.reset(new llm_graph_result(gf_res_reserve->get_max_nodes()));
+    }
+    return res.get();
+}
+
 // pack sampler outputs into as few sequences as possible before using sequences without samplers
 static void ubatch_prepare_reserve(
               llama_ubatch                            & ubatch,
@@ -2430,8 +2448,13 @@ ggml_cgraph * llama_context::graph_reserve(
 
     ggml_backend_sched_reset(sched.get());
 
-    // when the scheduler is reset, we cannot reuse the old graph, so we reset the previous graph result to prevent that
-    gf_res_prev->reset();
+    // when the scheduler is reset, we cannot reuse old graphs, so we reset the previous graph results
+    for (auto & res : gf_res_prev) {
+        if (res) {
+            res->reset();
+        }
+    }
+    gf_res_prev_active = nullptr;
 
     // store the n_outputs as it is, and restore it afterwards
     // TODO: not sure if needed, might simplify in the future by removing this
@@ -3521,10 +3544,12 @@ void llama_context::opt_epoch_iter(
                 break;
             }
 
-            auto * res = gf_res_prev.get();
+            auto * res = get_gf_res_prev();
 
             const auto gparams = graph_params(res, ubatch, mctx.get(), ctx_type_to_graph_type(cparams.ctx_type));
 
+            // the optimizer graph is allocated outside sched, so the next decode must rebuild
+            gf_res_prev_active = nullptr;
             res->reset();
 
             auto * gf = model.build_graph(gparams);

src/llama-context.h

@@ -11,6 +11,7 @@
 #include "ggml-cpp.h"
 #include "ggml-opt.h"
 
+#include <array>
 #include <map>
 #include <vector>
 
@@ -254,6 +255,8 @@ public:
     bool set_sampler(llama_seq_id seq_id, llama_sampler * sampler);
 
 private:
+    llm_graph_result * get_gf_res_prev();
+
     llm_graph_params graph_params(
                         llm_graph_result * res,
                       const llama_ubatch & ubatch,
@@ -364,9 +367,12 @@ private:
     std::vector<ggml_backend_buffer_type_t> backend_buft;
     std::vector<size_t>                     backend_buf_exp_size; // expected buffer sizes
 
-    llm_graph_result_ptr gf_res_prev;
+    // Separate arenas give batches with and without outputs distinct CUDA graph cache keys.
+    std::array<llm_graph_result_ptr, 2> gf_res_prev;
     llm_graph_result_ptr gf_res_reserve;
 
+    llm_graph_result * gf_res_prev_active = nullptr;
+
     // host buffer for the model output (logits and embeddings)
     ggml_backend_buffer_ptr buf_output;
 

src/llama-model.cpp

@@ -603,8 +603,20 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     };
 
     auto get_split_segments = [&](int axis, uint32_t il) -> std::vector<std::pair<int64_t, uint32_t>> {
+        // TODO: clarify why this is necessary specifically for these models
+        // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN]
         if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
                 ud->model->arch == LLM_ARCH_QWEN4EXP) {
+
+            // fused full attention layers with Q gate tensors that need n_embd doubled:
+            if (!hparams.is_recr(il) && (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias))) {
+                const int64_t n_embd      = hparams.n_head(il) * hparams.n_embd_head_k(il) * 2;
+                const int64_t n_embd_gqa  = hparams.n_embd_v_gqa(il);
+                GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa);
+                GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);
+                return {{n_embd, 1}, {n_embd_gqa, 2}};
+            }
+
             const int64_t head_k_dim = hparams.ssm_d_state;
             const int64_t head_v_dim = hparams.ssm_d_state;
             const int64_t n_k_heads  = hparams.ssm_n_group;
@@ -654,9 +666,9 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
         }
 
         if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
-            const int64_t n_embd      = hparams.n_embd;
+            const int64_t n_embd      = hparams.n_head(il) * hparams.n_embd_head_k(il);
             const int64_t n_embd_gqa  = hparams.n_embd_v_gqa(il);
-            GGML_ASSERT(hparams.n_embd_k_gqa() == n_embd_gqa);
+            GGML_ASSERT(hparams.n_embd_k_gqa(il) == n_embd_gqa);
             GGML_ASSERT(tensor->ne[axis] == n_embd + 2*n_embd_gqa);
             return {{n_embd, 1}, {n_embd_gqa, 2}};
         }
@@ -742,6 +754,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) {
                 GGML_ASSERT(segments.size() == 1);
                 // some models have Q gate tensors, for those cases the granularity needs to be doubled:
+                // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN]
                 if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
                         ud->model->arch == LLM_ARCH_QWEN4EXP) {
                     return {std::lcm(2*n_embd_q, blck_size_perf)};
@@ -769,6 +782,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             }
             if (std::regex_match(tensor_name, pattern_qkv_weight) || std::regex_match(tensor_name, pattern_qkv_bias)) {
                 GGML_ASSERT(segments.size() == 2);
+                // fused full attention layers need Q gate tensors handled like above:
+                // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN]
+                if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE ||
+                        ud->model->arch == LLM_ARCH_QWEN4EXP) {
+                    return {std::lcm(2*n_embd_q, blck_size_perf), granularity_kv};
+                }
                 return {granularity_q, granularity_kv};
             }
         }