Skip to content

ssm_conv: channels-major input mode to drop the delta-net transpose #52

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
roberteg16 wants to merge 2 commits into
base: gfx11
Choose a base branch
from
+202 −78
Draft

ssm_conv: channels-major input mode to drop the delta-net transpose #52

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e96478
delta-net: force contiguous conv-state concat input
roberteg16 Jul 15, 2026
d0c4efc
ssm_conv: add channels-major input mode; drop delta-net transpose
roberteg16 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion ggml/src/ggml-cann/ggml-cann.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2684,7 +2684,8 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
return true;
}
case GGML_OP_SSM_CONV:
return true;
// the channels-major input layout is only implemented on CPU/CUDA
return 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:
Expand Down
32 changes: 21 additions & 11 deletions ggml/src/ggml-cpu/ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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
Expand All @@ -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;
}
Expand Down
6 changes: 4 additions & 2 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4923,8 +4923,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;
Expand Down
109 changes: 72 additions & 37 deletions ggml/src/ggml-cuda/ssm-conv.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,11 @@
#include "ssm-conv.cuh"
#include "unary.cuh"

template <bool apply_silu, size_t split_d_inner, size_t d_conv>
// 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 <bool apply_silu, bool channels_major, size_t split_d_inner, size_t d_conv>
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,
Expand All @@ -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 };
Expand All @@ -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
Expand All @@ -57,7 +65,7 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p
}
}

template <bool apply_silu, size_t split_d_inner, size_t d_conv, int64_t split_n_t>
template <bool apply_silu, bool channels_major, size_t split_d_inner, size_t d_conv, int64_t split_n_t>
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,
Expand All @@ -68,36 +76,51 @@ 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;

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();
Expand All @@ -123,7 +146,7 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0,
}
}

template <bool apply_silu>
template <bool apply_silu, bool channels_major>
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,
Expand All @@ -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<apply_silu, threads, kNC>, launch_params, src0, src1, bias, src0_nb0, src0_nb1,
ggml_cuda_kernel_launch(ssm_conv_f32<apply_silu, channels_major, threads, kNC>, 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<apply_silu, threads, kNC, split_n_t><<<blocks, threads, smem_size, stream>>>(
ssm_conv_long_token_f32<apply_silu, channels_major, threads, kNC, split_n_t><<<blocks, threads, smem_size, stream>>>(
src0, src1, bias, src0_nb0, src0_nb1, src0_nb2, src1_nb1, dst, dst_nb0, dst_nb1, dst_nb2, n_t);
}
};
Expand Down Expand Up @@ -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));

Expand All @@ -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<true>(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<true, true>(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<true, false>(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<false>(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<false, true>(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<false, false>(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);
}
}
}
5 changes: 5 additions & 0 deletions ggml/src/ggml-hexagon/ggml-hexagon.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
2 changes: 2 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-device.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
3 changes: 2 additions & 1 deletion ggml/src/ggml-opencl/ggml-opencl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}.
Expand Down
4 changes: 3 additions & 1 deletion ggml/src/ggml-sycl/ggml-sycl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Loading
Loading