From b2f8eea93e1c7b7d8602befc23ec177709b43dae Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sat, 3 May 2025 19:25:29 -0500 Subject: [PATCH 1/6] cuda: refactored ssm_scan to use CUB --- ggml/src/ggml-cuda/ssm-scan.cu | 268 +++++++++++++++++++++++---------- 1 file changed, 192 insertions(+), 76 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 37ee208c09d46..741ae57203b1a 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -1,104 +1,220 @@ +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#define USE_CUB +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 + +#ifdef USE_CUB +#include +using namespace cub; +#endif // USE_CUB + #include "ssm-scan.cuh" template __global__ void __launch_bounds__(splitD, 2) - ssm_scan_f32(const float * __restrict__ src0, const float * __restrict__ src1, const float * __restrict__ src2, - const float * __restrict__ src3, const float * __restrict__ src4, const float * __restrict__ src5, - const int src0_nb1, const int src0_nb2, const int src1_nb0, const int src1_nb1, const int src1_nb2, - const int src1_nb3, const int src2_nb0, const int src2_nb1, const int src2_nb2, const int src3_nb1, + ssm_scan_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2, + const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5, + const int src0_nb1, const int src0_nb2, const int src1_nb1, const int src1_nb2, + const int src1_nb3, const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb1, const int src4_nb2, const int src5_nb1, const int src5_nb2, - float * __restrict__ dst, const int64_t L) { - GGML_UNUSED(src1_nb0); - GGML_UNUSED(src2_nb0); - const int bidx = blockIdx.x; // split along B - const int bidy = blockIdx.y; // split along D - const int tid = threadIdx.x; - const int wid = tid / 32; - const int wtid = tid % 32; - - extern __shared__ float smem[]; - const int stride_sA = N + 1; - const int stride_ss0 = N + 1; - float * smem_A = smem; - float * smem_s0 = smem_A + splitD * stride_sA; - - const float * s0_block = (const float *) ((const char *) src0 + bidx * src0_nb2 + bidy * splitD * src0_nb1); - const float * x_block = (const float *) ((const char *) src1 + (bidx * src1_nb2) + bidy * splitD * sizeof(float)); - const float * dt_block = (const float *) ((const char *) src2 + (bidx * src2_nb2) + bidy * splitD * sizeof(float)); - const float * A_block = (const float *) ((const char *) src3 + bidy * splitD * src3_nb1); - const float * B_block = (const float *) ((const char *) src4 + (bidx * src4_nb2)); - const float * C_block = (const float *) ((const char *) src5 + (bidx * src5_nb2)); - float * y_block = (float *) ((char *) dst + (bidx * src1_nb2) + bidy * splitD * sizeof(float)); - float * s_block = (float *) ((char *) dst + src1_nb3 + bidx * src0_nb2 + bidy * splitD * src0_nb1); + float *__restrict__ dst, const int64_t L) +{ - const int stride_s0 = src0_nb1 / sizeof(float); - const int stride_x = src1_nb1 / sizeof(float); + const float *s0_block = (const float *)((const char *)src0 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); + const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); + const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); + const float *A_block = (const float *)((const char *)src3 + blockIdx.y * splitD * src3_nb1); + const float *B_block = (const float *)((const char *)src4 + (blockIdx.x * src4_nb2)); + const float *C_block = (const float *)((const char *)src5 + (blockIdx.x * src5_nb2)); + float *y_block = (float *)((char *)dst + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); + float *s_block = (float *)((char *)dst + src1_nb3 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); + + const int stride_x = src1_nb1 / sizeof(float); const int stride_dt = src2_nb1 / sizeof(float); - const int stride_A = src3_nb1 / sizeof(float); - const int stride_B = src4_nb1 / sizeof(float); - const int stride_C = src5_nb1 / sizeof(float); - const int stride_s = stride_s0; - const int stride_y = stride_x; - - // can N not be 16? for example 32? - if (N == 16) { + const int stride_B = src4_nb1 / sizeof(float); + const int stride_C = src5_nb1 / sizeof(float); + const int stride_y = stride_x; + + float regA[N]; + float regs0[N]; + + __shared__ float smemB[N]; + __shared__ float smemC[N]; + +#ifdef USE_CUB + using BlockLoadA = cub::BlockLoad; + using BlockLoadS0 = cub::BlockLoad; + using BlockStoreS = cub::BlockStore; + + __shared__ typename BlockLoadA::TempStorage block_load_tempA; + __shared__ typename BlockLoadS0::TempStorage block_load_tempS0; + __shared__ typename BlockStoreS::TempStorage block_store_tempS; + + BlockLoadA(block_load_tempA).Load(A_block, regA); + BlockLoadS0(block_load_tempS0).Load(s0_block, regs0); +#else + const int stride_s0 = src0_nb1 / sizeof(float); + const int stride_A = src3_nb1 / sizeof(float); #pragma unroll - for (size_t i = 0; i < splitD / 4; i += 2) { - float value = A_block[(wid * warpSize + i) * stride_A + wtid]; - // todo: bank conflict - // I am always confused with how to use the swizzling method to solve - // bank conflit. Hoping somebody can tell me. - smem_A[(wid * warpSize + i) * stride_sA + wtid + ((wtid / 16) > 0 ? 1 : 0)] = value; + for (int j = 0; j < N; ++j) + { + regA[j] = A_block[threadIdx.x * stride_A + j]; + regs0[j] = s0_block[threadIdx.x * stride_s0 + j]; + } +#endif + + for (int i = 0; i < L; i++) + { + if (threadIdx.x < N) + { + smemB[threadIdx.x] = B_block[i * stride_B + threadIdx.x]; + smemC[threadIdx.x] = C_block[i * stride_C + threadIdx.x]; + } + __syncthreads(); + + float dt_soft_plus = dt_block[i * stride_dt + threadIdx.x]; + if (dt_soft_plus <= 20.0f) + { + dt_soft_plus = log1pf(expf(dt_soft_plus)); } + float x_dt = x_block[i * stride_x + threadIdx.x] * dt_soft_plus; + + float sumf = 0.0f; #pragma unroll - for (size_t i = 0; i < splitD / 4; i += 2) { - float value = s0_block[(wid * warpSize + i) * stride_s0 + wtid]; - smem_s0[(wid * warpSize + i) * stride_ss0 + wtid + ((wtid / 16) > 0 ? 1 : 0)] = value; + for (int j = 0; j < N; j++) + { + float state = regs0[j] * expf(dt_soft_plus * regA[j]) + smemB[j] * x_dt; + sumf += state * smemC[j]; + regs0[j] = state; } + y_block[i * stride_y + threadIdx.x] = sumf; + } + +#ifdef USE_CUB + BlockStoreS(block_store_tempS).Store(s_block, regs0); +#else + const int stride_s = stride_s0; +#pragma unroll + for (int j = 0; j < N; ++j) + { + s_block[threadIdx.x * stride_s + j] = regs0[j]; } +#endif +} + +template +__global__ void __launch_bounds__(splitD, 2) + ssm_scan_single_step_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2, + const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5, + const int src0_nb1, const int src0_nb2, const int src1_nb2, + const int src1_nb3, const int src2_nb2, const int src3_nb1, + const int src4_nb2, const int src5_nb2, + float *__restrict__ dst) +{ + const float *s0_block = (const float *)((const char *)src0 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); + const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); + const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); + const float *A_block = (const float *)((const char *)src3 + blockIdx.y * splitD * src3_nb1); + const float *B_block = (const float *)((const char *)src4 + (blockIdx.x * src4_nb2)); + const float *C_block = (const float *)((const char *)src5 + (blockIdx.x * src5_nb2)); + float *y_block = (float *)((char *)dst + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); + float *s_block = (float *)((char *)dst + src1_nb3 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); + + float regA[N]; + float regs0[N]; + + __shared__ float smemB[N]; + __shared__ float smemC[N]; + +#ifdef USE_CUB + using BlockLoadA = cub::BlockLoad; + using BlockLoadS0 = cub::BlockLoad; + using BlockStoreS = cub::BlockStore; + + __shared__ typename BlockLoadA::TempStorage block_load_tempA; + __shared__ typename BlockLoadS0::TempStorage block_load_tempS0; + __shared__ typename BlockStoreS::TempStorage block_store_tempS; + + BlockLoadA(block_load_tempA).Load(A_block, regA); + BlockLoadS0(block_load_tempS0).Load(s0_block, regs0); +#else + const int stride_s0 = src0_nb1 / sizeof(float); + const int stride_A = src3_nb1 / sizeof(float); +#pragma unroll + for (int j = 0; j < N; ++j) + { + regA[j] = A_block[threadIdx.x * stride_A + j]; + regs0[j] = s0_block[threadIdx.x * stride_s0 + j]; + } +#endif + if (threadIdx.x < N) + { + smemB[threadIdx.x] = B_block[threadIdx.x]; + smemC[threadIdx.x] = C_block[threadIdx.x]; + } __syncthreads(); - for (int64_t i = 0; i < L; i++) { - float dt_soft_plus = dt_block[i * stride_dt + tid]; - if (dt_soft_plus <= 20.0f) { - dt_soft_plus = log1pf(exp(dt_soft_plus)); + { + float dt_soft_plus = dt_block[threadIdx.x]; + if (dt_soft_plus <= 20.0f) + { + dt_soft_plus = log1pf(expf(dt_soft_plus)); } - float x_dt = x_block[i * stride_x + tid] * dt_soft_plus; + float x_dt = x_block[threadIdx.x] * dt_soft_plus; float sumf = 0.0f; #pragma unroll - for (size_t j = 0; j < N; j++) { - float state = (smem_s0[tid * stride_ss0 + j] * expf(dt_soft_plus * smem_A[tid * stride_sA + j])) + - (B_block[i * stride_B + j] * x_dt); - sumf += state * C_block[i * stride_C + j]; - if (i == L - 1) { - s_block[tid * stride_s + j] = state; - } else { - smem_s0[tid * stride_ss0 + j] = state; - } + for (int j = 0; j < N; j++) + { + float state = regs0[j] * expf(dt_soft_plus * regA[j]) + smemB[j] * x_dt; + sumf += state * smemC[j]; + regs0[j] = state; } - __syncthreads(); - y_block[i * stride_y + tid] = sumf; + y_block[threadIdx.x] = sumf; + } + +#ifdef USE_CUB + BlockStoreS(block_store_tempS).Store(s_block, regs0); +#else + const int stride_s = s0; +#pragma unroll + for (int j = 0; j < N; ++j) + { + s_block[threadIdx.x * stride_s + j] = regs0[j]; } +#endif } -static void ssm_scan_f32_cuda(const float * src0, const float * src1, const float * src2, const float * src3, - const float * src4, const float * src5, const int src0_nb1, const int src0_nb2, - const int src1_nb0, const int src1_nb1, const int src1_nb2, const int src1_nb3, - const int src2_nb0, const int src2_nb1, const int src2_nb2, const int src3_nb1, +static void ssm_scan_f32_cuda(const float *src0, const float *src1, const float *src2, const float *src3, + const float *src4, const float *src5, const int src0_nb1, const int src0_nb2, + const int src1_nb1, const int src1_nb2, const int src1_nb3, + const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb1, const int src4_nb2, const int src5_nb1, const int src5_nb2, - float * dst, const int64_t N, const int64_t D, const int64_t L, const int64_t B, - cudaStream_t stream) { + float *dst, const int64_t N, const int64_t D, const int64_t L, const int64_t B, + cudaStream_t stream) +{ const int threads = 128; // todo: consider D cannot be divided,does this situation exist? GGML_ASSERT(D % threads == 0); const dim3 blocks(B, (D + threads - 1) / threads, 1); - const int smem_size = (threads * (N + 1) * 2) * sizeof(float); - if (N == 16) { - ssm_scan_f32<128, 16><<>>( - src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb0, src1_nb1, src1_nb2, src1_nb3, src2_nb0, - src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); - } else { + if (N == 16) + { + if (L > 1) + { + ssm_scan_f32<<>>( + src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + } + else + { + ssm_scan_single_step_f32<<>>( + src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb2, + src1_nb3, src2_nb2, src3_nb1, + src4_nb2, src5_nb2, + dst); + } + } + else + { GGML_ABORT("doesn't support N!=16."); } } @@ -147,7 +263,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_F32); - ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src0->nb[1], src0->nb[2], src1->nb[0], - src1->nb[1], src1->nb[2], src1->nb[3], src2->nb[0], src2->nb[1], src2->nb[2], src3->nb[1], + ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src0->nb[1], src0->nb[2], + src1->nb[1], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[1], src4->nb[2], src5->nb[1], src5->nb[2], dst_d, nc, nr, n_t, n_s, stream); } From c7d4d45fd192da0ff64c7a2579b86c9ed93d60d9 Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sun, 4 May 2025 10:22:02 -0500 Subject: [PATCH 2/6] fixed compilation error when when not using CUB --- ggml/src/ggml-cuda/ssm-scan.cu | 38 +++++++++++++++++----------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 741ae57203b1a..402cbee31d39c 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -55,10 +55,10 @@ __global__ void __launch_bounds__(splitD, 2) const int stride_s0 = src0_nb1 / sizeof(float); const int stride_A = src3_nb1 / sizeof(float); #pragma unroll - for (int j = 0; j < N; ++j) + for (size_t n = 0; n < N; ++n) { - regA[j] = A_block[threadIdx.x * stride_A + j]; - regs0[j] = s0_block[threadIdx.x * stride_s0 + j]; + regA[n] = A_block[threadIdx.x * stride_A + n]; + regs0[n] = s0_block[threadIdx.x * stride_s0 + n]; } #endif @@ -80,11 +80,11 @@ __global__ void __launch_bounds__(splitD, 2) float sumf = 0.0f; #pragma unroll - for (int j = 0; j < N; j++) + for (size_t n = 0; n < N; n++) { - float state = regs0[j] * expf(dt_soft_plus * regA[j]) + smemB[j] * x_dt; - sumf += state * smemC[j]; - regs0[j] = state; + float state = regs0[n] * expf(dt_soft_plus * regA[n]) + smemB[n] * x_dt; + sumf += state * smemC[n]; + regs0[n] = state; } y_block[i * stride_y + threadIdx.x] = sumf; } @@ -94,9 +94,9 @@ __global__ void __launch_bounds__(splitD, 2) #else const int stride_s = stride_s0; #pragma unroll - for (int j = 0; j < N; ++j) + for (size_t n = 0; n < N; ++n) { - s_block[threadIdx.x * stride_s + j] = regs0[j]; + s_block[threadIdx.x * stride_s + n] = regs0[n]; } #endif } @@ -140,10 +140,10 @@ __global__ void __launch_bounds__(splitD, 2) const int stride_s0 = src0_nb1 / sizeof(float); const int stride_A = src3_nb1 / sizeof(float); #pragma unroll - for (int j = 0; j < N; ++j) + for (size_t n = 0; n < N; ++n) { - regA[j] = A_block[threadIdx.x * stride_A + j]; - regs0[j] = s0_block[threadIdx.x * stride_s0 + j]; + regA[n] = A_block[threadIdx.x * stride_A + n]; + regs0[n] = s0_block[threadIdx.x * stride_s0 + n]; } #endif @@ -163,11 +163,11 @@ __global__ void __launch_bounds__(splitD, 2) float x_dt = x_block[threadIdx.x] * dt_soft_plus; float sumf = 0.0f; #pragma unroll - for (int j = 0; j < N; j++) + for (size_t n = 0; n < N; n++) { - float state = regs0[j] * expf(dt_soft_plus * regA[j]) + smemB[j] * x_dt; - sumf += state * smemC[j]; - regs0[j] = state; + float state = regs0[n] * expf(dt_soft_plus * regA[n]) + smemB[n] * x_dt; + sumf += state * smemC[n]; + regs0[n] = state; } y_block[threadIdx.x] = sumf; } @@ -175,11 +175,11 @@ __global__ void __launch_bounds__(splitD, 2) #ifdef USE_CUB BlockStoreS(block_store_tempS).Store(s_block, regs0); #else - const int stride_s = s0; + const int stride_s = stride_s0; #pragma unroll - for (int j = 0; j < N; ++j) + for (size_t n = 0; n < N; ++n) { - s_block[threadIdx.x * stride_s + j] = regs0[j]; + s_block[threadIdx.x * stride_s + n] = regs0[n]; } #endif } From 949e4fa2929cd580537c6967bb338296224185d2 Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sat, 10 May 2025 20:36:03 -0500 Subject: [PATCH 3/6] assign L to constant and use size_t instead of int --- ggml/src/ggml-cuda/ssm-scan.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 402cbee31d39c..8c8b33126de2e 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -16,9 +16,9 @@ __global__ void __launch_bounds__(splitD, 2) const int src0_nb1, const int src0_nb2, const int src1_nb1, const int src1_nb2, const int src1_nb3, const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb1, const int src4_nb2, const int src5_nb1, const int src5_nb2, - float *__restrict__ dst, const int64_t L) + float *__restrict__ dst, const int64_t L_param) { - + const size_t L = L_param; const float *s0_block = (const float *)((const char *)src0 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); @@ -62,7 +62,7 @@ __global__ void __launch_bounds__(splitD, 2) } #endif - for (int i = 0; i < L; i++) + for (size_t i = 0; i < L; i++) { if (threadIdx.x < N) { From 75520d673dac8649edf31872c474fb8836129919 Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sun, 11 May 2025 09:18:51 -0500 Subject: [PATCH 4/6] deduplicated functions --- ggml/src/ggml-cuda/ssm-scan.cu | 174 ++++++++++++--------------------- 1 file changed, 60 insertions(+), 114 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 8c8b33126de2e..ad18b34042362 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -9,7 +9,7 @@ using namespace cub; #include "ssm-scan.cuh" -template +template __global__ void __launch_bounds__(splitD, 2) ssm_scan_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2, const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5, @@ -18,7 +18,7 @@ __global__ void __launch_bounds__(splitD, 2) const int src4_nb1, const int src4_nb2, const int src5_nb1, const int src5_nb2, float *__restrict__ dst, const int64_t L_param) { - const size_t L = L_param; + const size_t L = L_template == 0 ? L_param : L_template; const float *s0_block = (const float *)((const char *)src0 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); @@ -62,6 +62,7 @@ __global__ void __launch_bounds__(splitD, 2) } #endif +#pragma unroll for (size_t i = 0; i < L; i++) { if (threadIdx.x < N) @@ -101,89 +102,6 @@ __global__ void __launch_bounds__(splitD, 2) #endif } -template -__global__ void __launch_bounds__(splitD, 2) - ssm_scan_single_step_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2, - const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5, - const int src0_nb1, const int src0_nb2, const int src1_nb2, - const int src1_nb3, const int src2_nb2, const int src3_nb1, - const int src4_nb2, const int src5_nb2, - float *__restrict__ dst) -{ - const float *s0_block = (const float *)((const char *)src0 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); - const float *x_block = (const float *)((const char *)src1 + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); - const float *dt_block = (const float *)((const char *)src2 + (blockIdx.x * src2_nb2) + blockIdx.y * splitD * sizeof(float)); - const float *A_block = (const float *)((const char *)src3 + blockIdx.y * splitD * src3_nb1); - const float *B_block = (const float *)((const char *)src4 + (blockIdx.x * src4_nb2)); - const float *C_block = (const float *)((const char *)src5 + (blockIdx.x * src5_nb2)); - float *y_block = (float *)((char *)dst + (blockIdx.x * src1_nb2) + blockIdx.y * splitD * sizeof(float)); - float *s_block = (float *)((char *)dst + src1_nb3 + blockIdx.x * src0_nb2 + blockIdx.y * splitD * src0_nb1); - - float regA[N]; - float regs0[N]; - - __shared__ float smemB[N]; - __shared__ float smemC[N]; - -#ifdef USE_CUB - using BlockLoadA = cub::BlockLoad; - using BlockLoadS0 = cub::BlockLoad; - using BlockStoreS = cub::BlockStore; - - __shared__ typename BlockLoadA::TempStorage block_load_tempA; - __shared__ typename BlockLoadS0::TempStorage block_load_tempS0; - __shared__ typename BlockStoreS::TempStorage block_store_tempS; - - BlockLoadA(block_load_tempA).Load(A_block, regA); - BlockLoadS0(block_load_tempS0).Load(s0_block, regs0); -#else - const int stride_s0 = src0_nb1 / sizeof(float); - const int stride_A = src3_nb1 / sizeof(float); -#pragma unroll - for (size_t n = 0; n < N; ++n) - { - regA[n] = A_block[threadIdx.x * stride_A + n]; - regs0[n] = s0_block[threadIdx.x * stride_s0 + n]; - } -#endif - - if (threadIdx.x < N) - { - smemB[threadIdx.x] = B_block[threadIdx.x]; - smemC[threadIdx.x] = C_block[threadIdx.x]; - } - __syncthreads(); - - { - float dt_soft_plus = dt_block[threadIdx.x]; - if (dt_soft_plus <= 20.0f) - { - dt_soft_plus = log1pf(expf(dt_soft_plus)); - } - float x_dt = x_block[threadIdx.x] * dt_soft_plus; - float sumf = 0.0f; -#pragma unroll - for (size_t n = 0; n < N; n++) - { - float state = regs0[n] * expf(dt_soft_plus * regA[n]) + smemB[n] * x_dt; - sumf += state * smemC[n]; - regs0[n] = state; - } - y_block[threadIdx.x] = sumf; - } - -#ifdef USE_CUB - BlockStoreS(block_store_tempS).Store(s_block, regs0); -#else - const int stride_s = stride_s0; -#pragma unroll - for (size_t n = 0; n < N; ++n) - { - s_block[threadIdx.x * stride_s + n] = regs0[n]; - } -#endif -} - static void ssm_scan_f32_cuda(const float *src0, const float *src1, const float *src2, const float *src3, const float *src4, const float *src5, const int src0_nb1, const int src0_nb2, const int src1_nb1, const int src1_nb2, const int src1_nb3, @@ -198,19 +116,46 @@ static void ssm_scan_f32_cuda(const float *src0, const float *src1, const float const dim3 blocks(B, (D + threads - 1) / threads, 1); if (N == 16) { - if (L > 1) + switch (L) { - ssm_scan_f32<<>>( + case 1: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 2: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 3: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 4: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 5: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 6: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 7: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + case 8: + ssm_scan_f32<<>>(src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, + src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); + break; + + default: + ssm_scan_f32<<>>( src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb1, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1, src4_nb1, src4_nb2, src5_nb1, src5_nb2, dst, L); - } - else - { - ssm_scan_single_step_f32<<>>( - src0, src1, src2, src3, src4, src5, src0_nb1, src0_nb2, src1_nb2, - src1_nb3, src2_nb2, src3_nb1, - src4_nb2, src5_nb2, - dst); + break; } } else @@ -219,23 +164,24 @@ static void ssm_scan_f32_cuda(const float *src0, const float *src1, const float } } -void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const struct ggml_tensor * src0 = dst->src[0]; // s - const struct ggml_tensor * src1 = dst->src[1]; // x - const struct ggml_tensor * src2 = dst->src[2]; // dt - const struct ggml_tensor * src3 = dst->src[3]; // A - const struct ggml_tensor * src4 = dst->src[4]; // B - const struct ggml_tensor * src5 = dst->src[5]; // C +void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context &ctx, ggml_tensor *dst) +{ + const struct ggml_tensor *src0 = dst->src[0]; // s + const struct ggml_tensor *src1 = dst->src[1]; // x + const struct ggml_tensor *src2 = dst->src[2]; // dt + const struct ggml_tensor *src3 = dst->src[3]; // A + const struct ggml_tensor *src4 = dst->src[4]; // B + const struct ggml_tensor *src5 = dst->src[5]; // C // const int64_t d_state = src0->ne[0]; // const int64_t d_inner = src0->ne[1]; // const int64_t l = src1->ne[1]; // const int64_t b = src0->ne[2]; - const int64_t nc = src0->ne[0]; // d_state - const int64_t nr = src0->ne[1]; // d_inner - const int64_t n_t = src1->ne[1]; // number of tokens per sequence - const int64_t n_s = src0->ne[2]; // number of sequences in the batch + const int64_t nc = src0->ne[0]; // d_state + const int64_t nr = src0->ne[1]; // d_inner + const int64_t n_t = src1->ne[1]; // number of tokens per sequence + const int64_t n_s = src0->ne[2]; // number of sequences in the batch GGML_ASSERT(ggml_nelements(src1) + ggml_nelements(src0) == ggml_nelements(dst)); GGML_ASSERT(src0->nb[0] == sizeof(float)); @@ -251,14 +197,14 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { // required to get correct offset for state destination (i.e. src1->nb[3]) GGML_ASSERT(src1->nb[3] == src1->ne[0] * src1->ne[1] * src1->ne[2] * sizeof(float)); - const float * src0_d = (const float *) src0->data; - const float * src1_d = (const float *) src1->data; - const float * src2_d = (const float *) src2->data; - const float * src3_d = (const float *) src3->data; - const float * src4_d = (const float *) src4->data; - const float * src5_d = (const float *) src5->data; - float * dst_d = (float *) dst->data; - cudaStream_t stream = ctx.stream(); + const float *src0_d = (const float *)src0->data; + const float *src1_d = (const float *)src1->data; + const float *src2_d = (const float *)src2->data; + const float *src3_d = (const float *)src3->data; + const float *src4_d = (const float *)src4->data; + const float *src5_d = (const float *)src5->data; + float *dst_d = (float *)dst->data; + cudaStream_t stream = ctx.stream(); GGML_ASSERT(src0->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_F32); From 7e559f3e56be62a787b8e709d32785c508fad1fa Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sun, 11 May 2025 09:27:44 -0500 Subject: [PATCH 5/6] change min blocks per mp to 1 --- ggml/src/ggml-cuda/ssm-scan.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index ad18b34042362..65cec767a8117 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -10,7 +10,7 @@ using namespace cub; #include "ssm-scan.cuh" template -__global__ void __launch_bounds__(splitD, 2) +__global__ void __launch_bounds__(splitD, 1) ssm_scan_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2, const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5, const int src0_nb1, const int src0_nb2, const int src1_nb1, const int src1_nb2, From 7d259d9eb3fb0880b0e75a1e855af9f2a8dec404 Mon Sep 17 00:00:00 2001 From: David Zhao Date: Sun, 11 May 2025 10:43:59 -0500 Subject: [PATCH 6/6] Use cub load and store warp transpose --- ggml/src/ggml-cuda/ssm-scan.cu | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 65cec767a8117..0fa697dfa17d0 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -41,16 +41,17 @@ __global__ void __launch_bounds__(splitD, 1) __shared__ float smemC[N]; #ifdef USE_CUB - using BlockLoadA = cub::BlockLoad; - using BlockLoadS0 = cub::BlockLoad; - using BlockStoreS = cub::BlockStore; + using BlockLoad = cub::BlockLoad; + using BlockStore = cub::BlockStore; - __shared__ typename BlockLoadA::TempStorage block_load_tempA; - __shared__ typename BlockLoadS0::TempStorage block_load_tempS0; - __shared__ typename BlockStoreS::TempStorage block_store_tempS; + union CubTempStorage { + typename BlockLoad::TempStorage load_temp; + typename BlockStore::TempStorage store_temp; + }; + __shared__ CubTempStorage cub_temp_storage; - BlockLoadA(block_load_tempA).Load(A_block, regA); - BlockLoadS0(block_load_tempS0).Load(s0_block, regs0); + BlockLoad(cub_temp_storage.load_temp).Load(A_block, regA); + BlockLoad(cub_temp_storage.load_temp).Load(s0_block, regs0); #else const int stride_s0 = src0_nb1 / sizeof(float); const int stride_A = src3_nb1 / sizeof(float); @@ -91,7 +92,7 @@ __global__ void __launch_bounds__(splitD, 1) } #ifdef USE_CUB - BlockStoreS(block_store_tempS).Store(s_block, regs0); + BlockStore(cub_temp_storage.store_temp).Store(s_block, regs0); #else const int stride_s = stride_s0; #pragma unroll