diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index ac133665d978..be8ad751dca6 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2438,11 +2438,27 @@ extern "C" { struct ggml_tensor * d, bool masked); + // memory layout of the ssm_conv input (sx), stored in op_params[0] + enum ggml_ssm_conv_layout { + GGML_SSM_CONV_LAYOUT_TIME_MAJOR = 0, // sx = [d_conv-1+n_t, d_inner, n_s] (time contiguous) + GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR = 1, // sx = [d_inner, d_conv-1+n_t, n_s] (channels contiguous) + }; + GGML_API struct ggml_tensor * ggml_ssm_conv( struct ggml_context * ctx, struct ggml_tensor * sx, struct ggml_tensor * c); + // same as ggml_ssm_conv but sx is channels-major: [d_inner, d_conv-1+n_t, n_s] + // (channels contiguous). output layout is unchanged: [d_inner, n_t, n_s]. + GGML_API struct ggml_tensor * ggml_ssm_conv_channels_major( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c); + + // input (sx) memory layout of an SSM_CONV op (see enum ggml_ssm_conv_layout) + GGML_API enum ggml_ssm_conv_layout ggml_ssm_conv_get_layout(const struct ggml_tensor * op); + GGML_API struct ggml_tensor * ggml_ssm_scan( struct ggml_context * ctx, struct ggml_tensor * s, diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 1f29ec86712d..e084932cd94d 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -754,6 +754,15 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + if (ggml_ssm_conv_get_layout(tensor) == GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR) { + // channels-major: d_inner is sx (src0) axis 0 and c (src1) axis 1, so a d_inner + // split lands on different source axes; the output [d_inner, n_t, n_s] splits on axis 0. + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0 && src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_1) { + return {GGML_BACKEND_SPLIT_AXIS_0, {0}, {1}, 1}; + } + return handle_generic(src_ss, /*scalar_only =*/ false); + } + // time-major: d_inner is axis 1 of both sx and c if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5f51ea3bb3c8..ad734b4a4a33 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2684,7 +2684,12 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten return true; } case GGML_OP_SSM_CONV: - return true; + // ggml_cann_ssm_conv() requires F32 src/dst and the time-major layout + // (the channels-major layout is only implemented on CPU/CUDA) + return op->type == GGML_TYPE_F32 && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR; case GGML_OP_CUMSUM: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_TRI: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index df0028cf15e3..84b2b1b40850 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9498,17 +9498,29 @@ static void ggml_compute_forward_ssm_conv_f32( const int ith = params->ith; const int nth = params->nth; - const int nc = src1->ne[0]; // d_conv - const int ncs = src0->ne[0]; // d_conv - 1 + n_t - const int nr = src0->ne[1]; // d_inner - const int n_t = dst->ne[1]; // tokens per sequence - const int n_s = dst->ne[2]; // number of sequences in the batch + // time-major : src0 = [d_conv-1+n_t, d_inner, n_s] (time contiguous) + // channels-major : src0 = [d_inner, d_conv-1+n_t, n_s] (channels contiguous) + // output is [d_inner, n_t, n_s] (channels contiguous) in both cases + const bool channels_major = ggml_ssm_conv_get_layout(dst) == GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR; + + const int nc = src1->ne[0]; // d_conv + const int ncs = channels_major ? src0->ne[1] : src0->ne[0]; // d_conv - 1 + n_t + const int nr = channels_major ? src0->ne[0] : src0->ne[1]; // d_inner + const int n_t = dst->ne[1]; // tokens per sequence + const int n_s = dst->ne[2]; // number of sequences in the batch GGML_ASSERT( dst->ne[0] == nr); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); // input is contiguous (either layout) GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src0->nb[1] == src0->ne[0]*sizeof(float)); + // byte strides of the input to move one channel / one time step, per layout + const size_t chan_nb = channels_major ? src0->nb[0] : src0->nb[1]; + const size_t time_nb = channels_major ? src0->nb[1] : src0->nb[0]; + const int chan_stride = chan_nb / sizeof(float); + const int time_stride = time_nb / sizeof(float); + GGML_UNUSED(ncs); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -9519,13 +9531,11 @@ static void ggml_compute_forward_ssm_conv_f32( for (int i3 = 0; i3 < n_s; ++i3) { for (int i2 = 0; i2 < n_t; ++i2) { - // {d_conv - 1 + n_t, d_inner, n_seqs} - // sliding window - const float * s = (const float *) ((const char *) src0->data + ir0*(src0->nb[1]) + i2*(src0->nb[0]) + i3*(src0->nb[2])); // {d_conv, d_inner, n_s} + // sliding window over the time axis (starting at token i2) + const float * s = (const float *) ((const char *) src0->data + ir0*chan_nb + i2*time_nb + i3*(src0->nb[2])); const float * c = (const float *) ((const char *) src1->data + ir0*(src1->nb[1])); // {d_conv, d_inner} float * x = (float *) ((char *) dst->data + ir0*(dst->nb[0]) + i2*(dst->nb[1]) + i3*(dst->nb[2])); // {d_inner, n_t, n_s} - // TODO: transpose the output for smaller strides for big batches? // d_inner for (int i1 = 0; i1 < ir; ++i1) { // rowwise dot product @@ -9534,7 +9544,7 @@ static void ggml_compute_forward_ssm_conv_f32( // d_conv for (int i0 = 0; i0 < nc; ++i0) { - sumf += s[i0 + i1*ncs] * c[i0 + i1*nc]; + sumf += s[i0*time_stride + i1*chan_stride] * c[i0 + i1*nc]; } x[i1] = sumf; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index ca424c75ad2b..0ac6239f7480 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5079,8 +5079,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } } case GGML_OP_SSM_CONV: { - // assumes d_inner % threads == 0 - return op->src[0]->ne[1] % 128 == 0; + // assumes d_inner % threads == 0; d_inner is on ne[1] (time-major) or ne[0] (channels-major) + const bool channels_major = ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR; + const int64_t d_inner = channels_major ? op->src[0]->ne[0] : op->src[0]->ne[1]; + return d_inner % 128 == 0; } case GGML_OP_CONT: return true; diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu index 1463169cf78b..d1052772de39 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu @@ -2,7 +2,11 @@ #include "ssm-conv.cuh" #include "unary.cuh" -template +// channels_major selects the input (src0) memory layout: +// false = time-major : src0 = [d_conv-1+n_t, d_inner, n_s], time contiguous (nb0) +// true = channels-major : src0 = [d_inner, d_conv-1+n_t, n_s], channels contiguous (nb0) +// The output (dst) is [d_inner, n_t, n_s] (channels contiguous) in both cases. +template static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_ptr, const float * bias_ptr, const int src0_nb0, const int src0_nb1, const int src0_nb2, const int src1_nb1, @@ -13,18 +17,22 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p const float * GGML_CUDA_RESTRICT src1 = src1_ptr; const float * GGML_CUDA_RESTRICT bias = bias_ptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; - GGML_UNUSED(src0_nb0); const int tid = threadIdx.x; const int bidx = blockIdx.x; const int bidy = blockIdx.y; - const float * x_block = (const float *) ((const char *) src0 + bidx * src0_nb2 + bidy * split_d_inner * src0_nb1); + // byte strides of the input to move one channel / one time step, per layout + const int chan_nb = channels_major ? src0_nb0 : src0_nb1; + const int time_nb = channels_major ? src0_nb1 : src0_nb0; + + const float * x_block = (const float *) ((const char *) src0 + bidx * src0_nb2 + bidy * split_d_inner * chan_nb); const float * w_block = (const float *) ((const char *) src1 + bidy * split_d_inner * src1_nb1); float * y_block = (float *) ((char *) dst + bidx * dst_nb2 + bidy * split_d_inner * dst_nb0); - const int stride_x = src0_nb1 / sizeof(float); - const int stride_w = src1_nb1 / sizeof(float); - const int stride_y = dst_nb1 / sizeof(float); + const int stride_xc = chan_nb / sizeof(float); // stride to the next channel + const int stride_xt = time_nb / sizeof(float); // stride to the next time step + const int stride_w = src1_nb1 / sizeof(float); + const int stride_y = dst_nb1 / sizeof(float); float x[d_conv] = { 0.0f }; float w[d_conv] = { 0.0f }; @@ -42,10 +50,10 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p if (i == 0) { for (size_t j = 0; j < d_conv; j++) { - x[j] = x_block[tid * stride_x + j]; + x[j] = x_block[tid * stride_xc + j * stride_xt]; } } else { - x[(i - 1) % d_conv] = x_block[tid * stride_x + i + d_conv - 1]; + x[(i - 1) % d_conv] = x_block[tid * stride_xc + (i + d_conv - 1) * stride_xt]; } #pragma unroll @@ -57,7 +65,7 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p } } -template +template static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, const float * __restrict__ src1, const float * __restrict__ bias, const int src0_nb0, const int src0_nb1, const int src0_nb2, @@ -68,15 +76,19 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, const int bidy = blockIdx.y; const int bidz = blockIdx.z; - const float * x_block = (const float *) ((const char *) src0 + bidx * src0_nb2 + bidy * split_d_inner * src0_nb1 + - bidz * split_n_t * src0_nb0); + const int chan_nb = channels_major ? src0_nb0 : src0_nb1; + const int time_nb = channels_major ? src0_nb1 : src0_nb0; + + const float * x_block = (const float *) ((const char *) src0 + bidx * src0_nb2 + bidy * split_d_inner * chan_nb + + bidz * split_n_t * time_nb); const float * w_block = (const float *) ((const char *) src1 + bidy * split_d_inner * src1_nb1); float * y_block = (float *) ((char *) dst + bidx * dst_nb2 + bidz * split_n_t * dst_nb1 + bidy * split_d_inner * dst_nb0); - const int stride_x = src0_nb1 / sizeof(float); - const int stride_w = src1_nb1 / sizeof(float); - const int stride_y = dst_nb1 / sizeof(float); + const int stride_xc = chan_nb / sizeof(float); // stride to the next channel + const int stride_xt = time_nb / sizeof(float); // stride to the next time step + const int stride_w = src1_nb1 / sizeof(float); + const int stride_y = dst_nb1 / sizeof(float); const int64_t local_n_t = min(split_n_t, n_t - bidz * split_n_t); const int n_cols = d_conv - 1 + split_n_t; @@ -84,20 +96,31 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, extern __shared__ float smem[]; constexpr int load_cols = d_conv - 1 + split_n_t; - constexpr int total_elems = split_d_inner * load_cols; - int row = tid / load_cols; - int col = tid % load_cols; + if (channels_major) { + // channels are contiguous: each thread streams its own channel over the + // time tile, so consecutive threads read consecutive addresses (coalesced). #pragma unroll - for (int idx = 0; idx < total_elems; idx += split_d_inner) { - if (row < (int)split_d_inner) { - smem[row * n_cols + col] = x_block[row * stride_x + col]; + for (int col = 0; col < load_cols; col++) { + smem[tid * n_cols + col] = x_block[tid * stride_xc + col * stride_xt]; } + } else { + // time is contiguous: distribute the [channel, time] tile across threads + // so that consecutive threads read consecutive (time) addresses (coalesced). + constexpr int total_elems = split_d_inner * load_cols; + int row = tid / load_cols; + int col = tid % load_cols; +#pragma unroll + for (int idx = 0; idx < total_elems; idx += split_d_inner) { + if (row < (int)split_d_inner) { + smem[row * n_cols + col] = x_block[row * stride_xc + col * stride_xt]; + } - col += split_d_inner; - row += col / load_cols; - col = col % load_cols; - if (idx >= total_elems - tid - split_d_inner) { - break; + col += split_d_inner; + row += col / load_cols; + col = col % load_cols; + if (idx >= total_elems - tid - split_d_inner) { + break; + } } } __syncthreads(); @@ -123,7 +146,7 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, } } -template +template static void ssm_conv_f32_cuda(const float * src0, const float * src1, const float * bias, const int src0_nb0, const int src0_nb1, const int src0_nb2, const int src1_nb1, float * dst, const int dst_nb0, const int dst_nb1, const int dst_nb2, const int64_t nc, const int64_t nr, const int64_t n_t, @@ -136,13 +159,13 @@ static void ssm_conv_f32_cuda(const float * src0, const float * src1, const floa if (n_t <= 32) { const dim3 blocks(n_s, (nr + threads - 1) / threads, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream); - ggml_cuda_kernel_launch(ssm_conv_f32, launch_params, src0, src1, bias, src0_nb0, src0_nb1, + ggml_cuda_kernel_launch(ssm_conv_f32, launch_params, src0, src1, bias, src0_nb0, src0_nb1, src0_nb2, src1_nb1, dst, dst_nb0, dst_nb1, dst_nb2, n_t); } else { const int64_t split_n_t = 32; dim3 blocks(n_s, (nr + threads - 1) / threads, (n_t + split_n_t - 1) / split_n_t); const size_t smem_size = threads * (kNC - 1 + split_n_t) * sizeof(float); - ssm_conv_long_token_f32<<>>( + ssm_conv_long_token_f32<<>>( src0, src1, bias, src0_nb0, src0_nb1, src0_nb2, src1_nb1, dst, dst_nb0, dst_nb1, dst_nb2, n_t); } }; @@ -172,13 +195,15 @@ void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, g // When fusing, write to silu_dst (the node downstream references). const struct ggml_tensor * out = fuse_silu ? silu_dst : dst; - const int64_t nc = src1->ne[0]; // d_conv - const int64_t nr = src0->ne[1]; // d_inner - const int64_t n_t = out->ne[1]; // tokens per sequence - const int64_t n_s = out->ne[2]; // number of sequences in the batch + const bool channels_major = ggml_ssm_conv_get_layout(dst) == GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR; + + const int64_t nc = src1->ne[0]; // d_conv + const int64_t nr = channels_major ? src0->ne[0] : src0->ne[1]; // d_inner + const int64_t n_t = out->ne[1]; // tokens per sequence + const int64_t n_s = out->ne[2]; // number of sequences in the batch GGML_ASSERT(out->ne[0] == nr); - GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); // input is contiguous (either layout) GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src0->nb[1] == src0->ne[0] * sizeof(float)); @@ -197,10 +222,20 @@ void ggml_cuda_op_ssm_conv(ggml_backend_cuda_context & ctx, ggml_tensor * dst, g } if (fuse_silu) { - ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], - out->nb[2], nc, nr, n_t, n_s, stream); + if (channels_major) { + ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], + out->nb[2], nc, nr, n_t, n_s, stream); + } else { + ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], + out->nb[2], nc, nr, n_t, n_s, stream); + } } else { - ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], - out->nb[2], nc, nr, n_t, n_s, stream); + if (channels_major) { + ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], + out->nb[2], nc, nr, n_t, n_s, stream); + } else { + ssm_conv_f32_cuda(src0_d, src1_d, bias_d, src0->nb[0], src0->nb[1], src0->nb[2], src1->nb[1], dst_d, out->nb[0], out->nb[1], + out->nb[2], nc, nr, n_t, n_s, stream); + } } } diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 76c71d7ee70e..3285e0933347 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3250,6 +3250,11 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session * const struct ggml_tensor * src1 = op->src[1]; const struct ggml_tensor * dst = op; + // the channels-major input layout is only implemented on CPU/CUDA + if (ggml_ssm_conv_get_layout(op) != GGML_SSM_CONV_LAYOUT_TIME_MAJOR) { + return false; + } + // Only support FP32 for now if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { return false; diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 5d29250f654b..aa881fd90b25 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1265,6 +1265,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te } return has_simdgroup_mm; // TODO: over-restricted for vec-kernels case GGML_OP_SSM_CONV: + // the channels-major input layout is only implemented on CPU/CUDA + return has_simdgroup_reduction && ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR; case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; case GGML_OP_RWKV_WKV6: diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 5c96b9a9f6de..8d023a201456 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -6645,7 +6645,8 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) || (op->src[0]->type == GGML_TYPE_F16 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32); case GGML_OP_SSM_CONV: - return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32); + // the channels-major input layout is only implemented on CPU/CUDA + return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR); case GGML_OP_GATED_DELTA_NET: { // Match the Vulkan backend: only F32 -> F32, S_v in {16, 32, 64, 128}. diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 5226fb184a98..7bf03446bc6a 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5924,9 +5924,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_GATED_DELTA_NET: return true; case GGML_OP_SSM_CONV: + // the channels-major input layout is only implemented on CPU/CUDA return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && - op->src[1]->type == GGML_TYPE_F32; + op->src[1]->type == GGML_TYPE_F32 && + ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR; case GGML_OP_ROLL: return op->type == GGML_TYPE_F32; case GGML_OP_ARANGE: diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a483d22c1a26..a4fe4dbb7f20 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17683,7 +17683,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm return true; } case GGML_OP_SSM_CONV: - return op->src[0]->type == GGML_TYPE_F32; + // the channels-major input layout is only implemented on CPU/CUDA + return op->src[0]->type == GGML_TYPE_F32 && ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR; case GGML_OP_CONV_TRANSPOSE_1D: return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; case GGML_OP_COL2IM_1D: diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 29025e9ba4e3..92b54e361311 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4354,7 +4354,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); break; case GGML_OP_SSM_CONV: - supports_op = op->type == GGML_TYPE_F32; + // the channels-major input layout is only implemented on CPU/CUDA + supports_op = op->type == GGML_TYPE_F32 && ggml_ssm_conv_get_layout(op) == GGML_SSM_CONV_LAYOUT_TIME_MAJOR; break; case GGML_OP_SSM_SCAN: supports_op = op->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index de0321d9ffd9..241695a28b96 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5531,25 +5531,38 @@ struct ggml_tensor * ggml_flash_attn_back( // ggml_ssm_conv -struct ggml_tensor * ggml_ssm_conv( +// The input (sx) memory layout is selected by `layout` and stored in op_params[0]. +// The output is [d_inner, n_t, n_s] in both cases. +static struct ggml_tensor * ggml_ssm_conv_impl( struct ggml_context * ctx, struct ggml_tensor * sx, - struct ggml_tensor * c) { + struct ggml_tensor * c, + enum ggml_ssm_conv_layout layout) { GGML_ASSERT(ggml_is_3d(sx)); GGML_ASSERT(ggml_is_matrix(c)); const int64_t d_conv = c->ne[0]; const int64_t d_inner = c->ne[1]; - const int64_t n_t = sx->ne[0] - d_conv + 1; // tokens per sequence const int64_t n_s = sx->ne[2]; + // the time dimension index within sx depends on the layout + const int64_t time_ne = layout == GGML_SSM_CONV_LAYOUT_TIME_MAJOR ? sx->ne[0] : sx->ne[1]; + const int64_t n_t = time_ne - d_conv + 1; // tokens per sequence + // TODO: maybe support other strides than 1? - GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); - GGML_ASSERT(sx->ne[1] == d_inner); + if (layout == GGML_SSM_CONV_LAYOUT_TIME_MAJOR) { + GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); + GGML_ASSERT(sx->ne[1] == d_inner); + } else { + GGML_ASSERT(sx->ne[0] == d_inner); + GGML_ASSERT(sx->ne[1] == d_conv - 1 + n_t); + } GGML_ASSERT(n_t >= 0); struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + ggml_set_op_params_i32(result, 0, (int32_t) layout); + result->op = GGML_OP_SSM_CONV; result->src[0] = sx; result->src[1] = c; @@ -5557,6 +5570,25 @@ struct ggml_tensor * ggml_ssm_conv( return result; } +struct ggml_tensor * ggml_ssm_conv( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c) { + return ggml_ssm_conv_impl(ctx, sx, c, GGML_SSM_CONV_LAYOUT_TIME_MAJOR); +} + +struct ggml_tensor * ggml_ssm_conv_channels_major( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c) { + return ggml_ssm_conv_impl(ctx, sx, c, GGML_SSM_CONV_LAYOUT_CHANNELS_MAJOR); +} + +enum ggml_ssm_conv_layout ggml_ssm_conv_get_layout(const struct ggml_tensor * op) { + GGML_ASSERT(op->op == GGML_OP_SSM_CONV); + return (enum ggml_ssm_conv_layout) ggml_get_op_params_i32(op, 0); +} + // ggml_ssm_scan struct ggml_tensor * ggml_ssm_scan( diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad6612647736..9fab0572972c 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -463,28 +463,33 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_tensor * conv_states = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs); cb(conv_states, "conv_states", il); - conv_states = ggml_reshape_3d(ctx0, conv_states, conv_kernel_size - 1, conv_channels, n_seqs); + // channels-major layout: [conv_channels, time, n_seqs] (channels contiguous). + // qkv_mixed already arrives channels-major as [conv_channels, n_tokens, n_seqs], + // so no transpose is needed: we prepend the recurrent conv state along the time + // axis (dim 1) and feed ggml_ssm_conv_channels_major directly. This avoids the + // physical transpose (a CONT kernel) that the time-major path required. + conv_states = ggml_reshape_3d(ctx0, conv_states, conv_channels, conv_kernel_size - 1, n_seqs); cb(conv_states, "conv_states_reshaped", il); - qkv_mixed = ggml_transpose(ctx0, qkv_mixed); - cb(qkv_mixed, "qkv_mixed_transposed", il); - - ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 0); + ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 1); cb(conv_input, "conv_input", il); const int64_t row_count = (conv_kernel_size - 1) * conv_channels; const size_t row_size = ggml_row_size(conv_states_all->type, row_count); + // number of time steps in conv_input (state + new tokens) + const int64_t n_time = conv_input->ne[1]; + if (cparams.n_rs_seq == 0) { - const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0]; + const int64_t s_idx = n_time - (conv_kernel_size - 1); const int64_t s_slot = 0; ggml_tensor * conv_state_last = ggml_view_3d(ctx0, conv_input, - conv_kernel_size - 1, conv_channels, n_seqs, + conv_channels, conv_kernel_size - 1, n_seqs, conv_input->nb[1], conv_input->nb[2], - ggml_row_size(conv_input->type, s_idx)); + s_idx * conv_input->nb[1]); cb(conv_state_last, "conv_state_last", il); ggml_tensor * conv_state_update = @@ -502,14 +507,14 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const int64_t K = (int64_t) cparams.n_rs_seq + 1; for (int64_t t = 1; t <= K; ++t) { - const int64_t s_idx = std::max(0, conv_input->ne[0] - conv_states->ne[0] - K + t); + const int64_t s_idx = std::max(0, n_time - (conv_kernel_size - 1) - K + t); const int64_t s_slot = K - t; ggml_tensor * conv_state_last = ggml_view_3d(ctx0, conv_input, - conv_kernel_size - 1, conv_channels, n_seqs, + conv_channels, conv_kernel_size - 1, n_seqs, conv_input->nb[1], conv_input->nb[2], - ggml_row_size(conv_input->type, s_idx)); + s_idx * conv_input->nb[1]); ggml_tensor * conv_state_update = ggml_view_2d(ctx0, diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index d8ffe43ae76c..5a5092372efa 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -391,7 +391,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); - ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + ggml_tensor * conv_output_proper = ggml_ssm_conv_channels_major(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index 7b0876cbb04b..27a79b31cfca 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -415,7 +415,7 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); - ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + ggml_tensor * conv_output_proper = ggml_ssm_conv_channels_major(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index 09b66423d5a8..7abadad7dcb1 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -442,7 +442,7 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn_linear( state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); - ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + ggml_tensor * conv_output_proper = ggml_ssm_conv_channels_major(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index d7b909103084..4edaba125079 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3754,20 +3754,23 @@ struct test_ssm_conv : public test_case { const ggml_type type; const std::array ne_a; const std::array ne_b; + const bool channels_major; // if true, ne_a is [d_inner, d_conv-1+n_t, n_s, 1] std::string vars() override { - return VARS_TO_STR3(type, ne_a, ne_b); + return VARS_TO_STR4(type, ne_a, ne_b, channels_major); } test_ssm_conv(ggml_type type = GGML_TYPE_F32, std::array ne_a = {10, 10, 10, 1}, - std::array ne_b = {3, 3, 1, 1}) - : type(type), ne_a(ne_a), ne_b(ne_b) {} + std::array ne_b = {3, 3, 1, 1}, + bool channels_major = false) + : type(type), ne_a(ne_a), ne_b(ne_b), channels_major(channels_major) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne_b.data()); - ggml_tensor * out = ggml_ssm_conv(ctx, a, b); + ggml_tensor * out = channels_major ? ggml_ssm_conv_channels_major(ctx, a, b) + : ggml_ssm_conv(ctx, a, b); return out; } }; @@ -8480,6 +8483,14 @@ static std::vector> make_test_cases_eval() { // long token (n_t > 32, exercises the long_token kernel path) test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv - 1 + 64, d_inner, 1, 1}, {d_conv, d_inner, 1, 1})); test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv - 1 + 64, d_inner, 4, 1}, {d_conv, d_inner, 1, 1})); + + // channels-major input: ne_a is [d_inner, d_conv-1+n_t, n_s, 1] + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_inner, d_conv, 1, 1}, {d_conv, d_inner, 1, 1}, true)); + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_inner, 2 * d_conv, 1, 1}, {d_conv, d_inner, 1, 1}, true)); + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_inner, d_conv, 4, 1}, {d_conv, d_inner, 1, 1}, true)); + // long token (n_t > 32, exercises the long_token kernel path) + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_inner, d_conv - 1 + 64, 1, 1}, {d_conv, d_inner, 1, 1}, true)); + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_inner, d_conv - 1 + 64, 4, 1}, {d_conv, d_inner, 1, 1}, true)); } }