From 8180cdb1fd3f5702ce3a3ae237a49dcea44c9aa3 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Mon, 4 May 2026 15:29:57 +0530 Subject: [PATCH 01/10] ggml-cuda: add chunked SSD matmul for Mamba-2 prefill acceleration --- ggml/src/ggml-cuda/ssm-scan.cu | 683 ++++++++++++++++++++++++++++++++- 1 file changed, 671 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 3022249c77d5..b67edb58c631 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -1,13 +1,23 @@ -#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 +#include "ssm-scan.cuh" +#include "mma.cuh" +#include "cp-async.cuh" -#ifdef USE_CUB +#ifdef GGML_CUDA_USE_CUB #include using namespace cub; -#endif // USE_CUB +#endif // GGML_CUDA_USE_CUB -#include "ssm-scan.cuh" +// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan +// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences) +#define SSM_SSD_MIN_TOKENS 64 + +// Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. +#define SSM_SSD_CHUNK_SIZE 256 + +// Fast exp() using hardware exp2: exp(x) = exp2(x * log2(e)). +__device__ __forceinline__ float fast_expf(float x) { + return exp2f(x * 1.4426950408889634f); // 1/ln(2) = log2(e) +} // We would like to keep pragma unroll for cases where L_template is not 0, // so we suppress the clang transformation warning. @@ -56,7 +66,7 @@ __global__ void __launch_bounds__(splitD, 1) __shared__ float smemB[N]; __shared__ float smemC[N]; -#ifdef USE_CUB +#ifdef GGML_CUDA_USE_CUB using BlockLoad = cub::BlockLoad; using BlockStore = cub::BlockStore; @@ -109,7 +119,7 @@ __global__ void __launch_bounds__(splitD, 1) __syncthreads(); } -#ifdef USE_CUB +#ifdef GGML_CUDA_USE_CUB BlockStore(cub_temp_storage.store_temp).Store(s_block, regs0); #else const int stride_s = stride_s0; @@ -316,6 +326,633 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa } } +// ============================================================================ +// SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) +// +// Instead of a sequential scan, SSD reformulates the output as: +// Y = (L ? (C @ B^T)) @ (X * dt) + decay * C @ s_init +// where L is a causal decay mask derived from A and dt. +// +// This converts the O(T*N) sequential scan into parallel matmuls. +// ============================================================================ +// Softplus(dt) and inclusive prefix sum per head using CUB BlockScan. +// Grid: (n_head, n_seqs) +template +__global__ void ssm_ssd_prepare_dt_kernel( + const float * __restrict__ dt_raw, + float * __restrict__ dt_sp_out, + float * __restrict__ cs_out, + const int n_head, const int n_tok, + const int dt_stride_tok, // elements between tokens in dt + const int dt_stride_seq) { // elements between sequences in dt + + const int h = blockIdx.x; + const int s = blockIdx.y; + + const float * dt_seq = dt_raw + s * dt_stride_seq; + + float * dt_sp_seq = dt_sp_out + s * n_tok * n_head; + float * cs_seq = cs_out + s * n_tok * n_head; + + const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; + + // Phase 1: parallel softplus, each thread accumulates a local sum + float local_vals[MAX_ITEMS]; + float local_sum = 0.0f; + + for (int i = 0; i < items_per_thread; i++) { + const int t = threadIdx.x * items_per_thread + i; // blocked distribution + if (t < n_tok) { + float val = dt_seq[h + t * dt_stride_tok]; + float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; + local_vals[i] = sp; + local_sum += sp; + dt_sp_seq[t * n_head + h] = sp; + } else { + local_vals[i] = 0.0f; + } + } + + // Phase 2: parallel prefix sum of per-thread totals +#ifdef GGML_CUDA_USE_CUB + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage scan_temp; + float thread_inclusive; + BlockScan(scan_temp).InclusiveSum(local_sum, thread_inclusive); + float exclusive_prefix = thread_inclusive - local_sum; +#else + // Fallback: sequential prefix sum in shared memory + __shared__ float sdata[BLOCK_SIZE]; + sdata[threadIdx.x] = local_sum; + __syncthreads(); + if (threadIdx.x == 0) { + for (int i = 1; i < BLOCK_SIZE && i * items_per_thread < n_tok; i++) { + sdata[i] += sdata[i - 1]; + } + } + __syncthreads(); + float exclusive_prefix = (threadIdx.x > 0) ? sdata[threadIdx.x - 1] : 0.0f; +#endif + + // Phase 3: write cumsum = exclusive_prefix + local running sum + float running = exclusive_prefix; + for (int i = 0; i < items_per_thread; i++) { + const int t = threadIdx.x * items_per_thread + i; + if (t < n_tok) { + running += local_vals[i]; + cs_seq[t * n_head + h] = running; + } + } +} + +// Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. +// T_matmul controls precision for X_dt, B_weighted (float or half). +// C_scaled is always float (pairs with float s_cur in step 3c). +// Computation is always FP32; only the final store converts to T_matmul. +// Grid: (ceil(max(C*head_dim, d_state*C) / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_pre_matmul_kernel( + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ dt_sp, // {n_tok, n_head} softplus(dt) + const float * __restrict__ A, // {1, n_head} + const float * __restrict__ x, // {head_dim, n_head, n_tok, n_seqs} + const float * __restrict__ B, // {d_state, n_group, n_tok, n_seqs} + const float * __restrict__ C_src, // {d_state, n_group, n_tok, n_seqs} + T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest + T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end + float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) + const int chunk_len, const int head_dim, const int n_head, const int n_group, + const int d_state, const int A_stride, + const int x_stride_tok, const int x_stride_seq, + const int B_stride_tok, const int B_stride_seq, + const int C_stride_tok, const int C_stride_seq, + const int chunk_offset, + const int n_tok_total) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int g = h / (n_head / n_group); + + const float A_h = A[h * A_stride]; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + + // Prepare X_dt = x * dt, stored d-fastest for coalesced reads and writes. + const int n_xdt = chunk_len * head_dim; + if (idx < n_xdt) { + const int d = idx % head_dim; + const int t = idx / head_dim; + + const float x_val = x[s * x_stride_seq + (chunk_offset + t) * x_stride_tok + d + h * head_dim]; + const float dt_val = dt_sp[cs_seq_off + (chunk_offset + t) * n_head + h]; + + X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); + } + + // Prepare B_weighted = B * decay_from_end for state update matmul. + // dt factor is already in X_dt (x * dt), so B_weighted must NOT include dt + // to avoid double-counting: s_cur = B_weighted @ X_dt^T = B*decay * x*dt = B*x*decay*dt ? + // Column-major {d_state, chunk_len} per head, group broadcast to head. + const int n_bw = d_state * chunk_len; + if (idx < n_bw) { + const int n = idx % d_state; + const int t = idx / d_state; + + const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; + const float decay_from_end = fast_expf(A_h * (cs_last - cs_t)); + + const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; + + B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * decay_from_end); + } + + // Prepare C_scaled = C * decay_to_pos for state contribution matmul. + // Column-major {d_state, chunk_len} per head, group broadcast to head. + const int n_cs = d_state * chunk_len; + if (idx < n_cs) { + const int n = idx % d_state; + const int t = idx / d_state; + + const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; + const float decay_to_pos = fast_expf(A_h * cs_t); + + const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; + + C_scaled[n + t * d_state + h * n_cs + s * n_cs * n_head] = C_val * decay_to_pos; + } +} + +// Scale running state in-place: s_cur *= decay_total(chunk). +// Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. +// Eliminates the s_old buffer and D2D memcpy vs the old approach of: +// memcpy(s_old, s_cur) ? cuBLAS(beta=0) ? s_cur += decay * s_old +// Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) +template +__global__ void ssm_ssd_scale_state_kernel( + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums + const float * __restrict__ A, // {1, n_head} + const int d_state, const int head_dim, const int n_head, + const int chunk_offset, const int chunk_len, + const int n_tok_total, const int A_stride) { + + const int h = blockIdx.y; + const int s = blockIdx.z; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + const int state_per_head = d_state * head_dim; + if (idx >= state_per_head) return; + + const float A_h = A[h * A_stride]; + const int cs_seq_off = s * n_tok_total * n_head; + const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; + const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; + const float decay_total = fast_expf(A_h * cs_last); + + const int off = s * state_per_head * n_head + h * state_per_head + idx; + s_cur[off] *= decay_total; +} + +// ============================================================================ +// Fused Y kernel (MMA tensor-core): Y += X_dt @ M_online^T +// +// Computes M on-the-fly from CB + cs + A in shared memory, never writing M to +// global memory. This eliminates ~32 MB DRAM traffic per chunk vs materializing M. +// Uses m16n8k16 FP16?FP32 MMA instructions via mma.cuh. +// +// The operation computes: +// Y[d, t_out, h] += S_{t_in<=t_out} X_dt[d, t_in, h] +// * exp(A[h]*(cs[t_out]-cs[t_in])) * CB[t_out, t_in, g(h)] +// +// Key design: +// - M is computed on-the-fly per warp in shared memory (never touches DRAM) +// - X_dt tile loaded cooperatively into smem with d?t_in transpose +// - Both A and B tiles loaded via hardware ldmatrix from shared memory +// - C (output) accumulated in registers across all k-blocks +// +// Block: dim3(WARP_SIZE, N_WARPS, 1) — threadIdx.x is lane, threadIdx.y is warp +// Grid: dim3(ceil(head_dim/D_TILE), n_head, n_seqs) +// ============================================================================ + +template +__global__ void ssm_ssd_fused_y_mma_kernel( + const half * __restrict__ X_dt, // {head_dim, chunk_len, n_head, n_seqs} d-fastest, FP16 + const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} FP32 + const float * __restrict__ cs, // {n_tok, n_head, n_seqs} cumulative dt sums + const float * __restrict__ A, // {1, n_head} + float * __restrict__ dst, // {d_inner, n_tok, n_seqs} accumulate onto step 3c result + const int head_dim, const int chunk_len, const int n_head, const int n_group, + const int d_inner, const int n_tok, + const int chunk_offset, const int A_stride, + const int64_t stride_X_head) { // = chunk_len * head_dim + + static_assert(D_TILE == 16, "MMA m16n8k16 requires D_TILE=16"); + static_assert(K_TILE == 16, "K_TILE must be 16 for cp.async pipeline"); + constexpr int T_OUT_TILE = 8; // MMA N dimension (t_out per tile) + + // Number of 16-byte cp.async copies per X tile: D_TILE=16 halfs = 32 bytes = 2 copies per t_in row + constexpr int COPIES_PER_ROW = D_TILE / 8; // 16 halfs / 8 halfs per copy = 2 + constexpr int COPIES_TOTAL = K_TILE * COPIES_PER_ROW; // 16 * 2 = 32 + + const int d_block = blockIdx.x; + const int h = blockIdx.y; + const int s = blockIdx.z; + const int g = h / (n_head / n_group); + const int warp_id = threadIdx.y; + + const float A_h = A[h * A_stride]; + const int d_start = d_block * D_TILE; + + // --- Shared memory layout (double-buffered X for cp.async pipeline) --- + // smem_X stored D-FASTEST (matching global layout) for direct cp.async copies. + // A tile loaded via load_ldmatrix_trans (hardware transpose during load). + // smem_cs padded to 16-byte alignment for cp.async destination requirements. + // [smem_cs: chunk_len_padded floats] + // [smem_X: 2 * K_TILE * D_TILE halfs] ? ping-pong buffers, d-fastest + // [smem_M: N_WARPS * T_OUT_TILE * K_TILE halfs] ? t_in-fastest (for B tile) + extern __shared__ char smem_raw[]; + const int chunk_len_padded = (chunk_len + 3) & ~3; // round up to 16-byte boundary + float * smem_cs = (float *)smem_raw; + half * smem_X_pp = (half *)(smem_cs + chunk_len_padded); // double buffer base + half * smem_M = smem_X_pp + 2 * K_TILE * D_TILE; + + // Load cs values for this chunk (cooperative, block-wide) + { + const int n_threads = WARP_SIZE * N_WARPS; + const int cs_seq_off = s * n_tok * n_head; + for (int i = threadIdx.y * WARP_SIZE + threadIdx.x; i < chunk_len; i += n_threads) { + smem_cs[i] = cs[cs_seq_off + (chunk_offset + i) * n_head + h]; + } + } + __syncthreads(); + + // Global memory pointers + const half * X_base = X_dt + (int64_t)s * stride_X_head * n_head + + (int64_t)h * stride_X_head + d_start; + const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group + + (int64_t)g * chunk_len * chunk_len; + float * dst_base = dst + (int64_t)s * d_inner * n_tok + (int64_t)h * head_dim + d_start; + + // C accumulators in registers + using tile_C = ggml_cuda_mma::tile<16, 8, float>; + tile_C C_acc[N_TOUT_PER_WARP]; + + const int tout_base = warp_id * N_TOUT_PER_WARP; + const int n_k_blocks = (chunk_len + K_TILE - 1) / K_TILE; + + // Helper: issue cp.async copies for X_dt tile into smem buffer (d-fastest layout). + // Each copy is 16 bytes = 8 halfs. D_TILE=16 ? 2 copies per t_in row, 32 total. + // k_len_tile: valid rows in this tile (may be < K_TILE for the last k-block). + // Out-of-bounds rows are zeroed to prevent NaN propagation through MMA (0*NaN=NaN). + auto issue_cp_async_X = [&](half * smem_X_buf, const int t_in_start, const int k_len_tile) { +#ifdef CP_ASYNC_AVAILABLE + const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; + constexpr int n_threads = WARP_SIZE * N_WARPS; + // Zero-fill partial tiles to prevent 0*NaN in MMA (pool memory may contain NaN) + if (k_len_tile < K_TILE) { + for (int i = tid; i < D_TILE * K_TILE; i += n_threads) { + smem_X_buf[i] = __float2half(0.0f); + } + __syncthreads(); + } + if (tid < COPIES_TOTAL) { + const int t = tid / COPIES_PER_ROW; // which t_in (0..K_TILE-1) + const int d_half = tid % COPIES_PER_ROW; // which 8-half group (0 or 1) + if (t < k_len_tile) { + const half * src = X_base + (t_in_start + t) * head_dim + d_half * 8; + half * dst_ptr = smem_X_buf + t * D_TILE + d_half * 8; + const unsigned int dst_addr = ggml_cuda_cvta_generic_to_shared(dst_ptr); + cp_async_cg_16<0>(dst_addr, src); + } + } + asm volatile("cp.async.commit_group;"); +#else + // Fallback: synchronous load (for GPUs below Ampere without cp.async) + const int n_threads = WARP_SIZE * N_WARPS; + const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; + for (int i = tid; i < D_TILE * K_TILE; i += n_threads) { + const int t_local = i / D_TILE; + const int d_local = i % D_TILE; + half val = (t_local < k_len_tile && d_start + d_local < head_dim) + ? X_base[(t_in_start + t_local) * head_dim + d_local] + : __float2half(0.0f); + smem_X_buf[t_local * D_TILE + d_local] = val; + } +#endif + }; + + // === cp.async pipelined loop: overlap next X load with current compute === + int pp_cur = 0; + + // Prologue: issue async load of first k-block into buffer 0 + issue_cp_async_X(smem_X_pp, 0, min(K_TILE, chunk_len)); +#ifdef CP_ASYNC_AVAILABLE + cp_async_wait_all(); +#endif + __syncthreads(); + + for (int kb = 0; kb < n_k_blocks; kb++) { + const int t_in_start = kb * K_TILE; + const int k_len = min(K_TILE, chunk_len - t_in_start); + + half * smem_X_cur = smem_X_pp + pp_cur * K_TILE * D_TILE; + + // --- Issue async load of NEXT k-block (overlaps with compute below!) --- + if (kb + 1 < n_k_blocks) { + half * smem_X_next = smem_X_pp + (1 - pp_cur) * K_TILE * D_TILE; + const int next_k_len = min(K_TILE, chunk_len - (kb + 1) * K_TILE); + issue_cp_async_X(smem_X_next, (kb + 1) * K_TILE, next_k_len); + } + + // --- Per-warp: iterate over assigned t_out blocks --- + for (int tb = 0; tb < N_TOUT_PER_WARP; tb++) { + const int tout_block = tout_base + tb; + const int t_out_start = tout_block * T_OUT_TILE; + + if (t_out_start >= chunk_len) break; + if (t_in_start > t_out_start + T_OUT_TILE - 1) continue; + + // --- Compute M tile on-the-fly in warp's smem region --- + { + half * my_M = smem_M + warp_id * T_OUT_TILE * K_TILE; + for (int i = threadIdx.x; i < T_OUT_TILE * K_TILE; i += WARP_SIZE) { + const int tout_local = i / K_TILE; + const int tin_local = i % K_TILE; + const int t_out_abs = t_out_start + tout_local; + const int t_in_abs = t_in_start + tin_local; + + half val; + if (t_out_abs < chunk_len && tin_local < k_len && t_in_abs <= t_out_abs) { + const float cs_out = smem_cs[t_out_abs]; + const float cs_in = smem_cs[t_in_abs]; + const float decay = fast_expf(A_h * (cs_out - cs_in)); + const float cb_val = CB_g[t_out_abs + t_in_abs * chunk_len]; + val = __float2half(decay * cb_val); + } else { + val = __float2half(0.0f); + } + my_M[tout_local * K_TILE + tin_local] = val; + } + } + __syncwarp(); + + // --- MMA tiles --- + // A tile: loaded via ldmatrix_trans from d-fastest smem_X (hardware transpose) + // B tile: loaded via ldmatrix from t_in-fastest smem_M (no transpose needed) + using tile_A = ggml_cuda_mma::tile<16, 8, half2>; + using tile_B = ggml_cuda_mma::tile<8, 8, half2>; + + tile_A A_tile; + tile_B B_tile; + + // A from smem_X (d-fastest): stride = D_TILE/2 half2 between t_in rows + // ldmatrix_trans transposes: smem rows=t_in ? tile rows=d ? + const half2 * X_h2 = (const half2 *)smem_X_cur; + ggml_cuda_mma::load_ldmatrix_trans(A_tile, X_h2, D_TILE / 2); + + // B from smem_M (t_in-fastest): stride = K_TILE/2 half2 between t_out rows + const half2 * my_M_h2 = (const half2 *)(smem_M + warp_id * T_OUT_TILE * K_TILE); + ggml_cuda_mma::load_ldmatrix(B_tile, my_M_h2, K_TILE / 2); + + ggml_cuda_mma::mma(C_acc[tb], A_tile, B_tile); + } + + // Wait for next buffer load + ensure all warps done with current buffer +#ifdef CP_ASYNC_AVAILABLE + cp_async_wait_all(); +#endif + pp_cur = 1 - pp_cur; + __syncthreads(); + } + + // === Store C accumulators to dst === + for (int tb = 0; tb < N_TOUT_PER_WARP; tb++) { + const int tout_block = tout_base + tb; + const int t_out_start = tout_block * T_OUT_TILE; + if (t_out_start >= chunk_len) break; + +#pragma unroll + for (int l = 0; l < tile_C::ne; l++) { + const int d_local = tile_C::get_i(l); + const int tout_local = tile_C::get_j(l); + const int t_out_abs = t_out_start + tout_local; + + if (d_start + d_local < head_dim && t_out_abs < chunk_len) { + dst_base[d_local + (chunk_offset + t_out_abs) * d_inner] += C_acc[tb].x[l]; + } + } + } +} + +// Copy initial state from src0[ids[s]] into s_cur for each sequence. +// Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) +template +__global__ void ssm_ssd_init_state_kernel( + const float * __restrict__ src0, // {d_state, head_dim, n_head, n_rs} + const int32_t * __restrict__ ids, // {n_seqs} + float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} + const int state_size, // d_state * head_dim * n_head + const int s0_stride_seq) { // elements between state rows + const int s = blockIdx.y; + const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; + if (idx >= state_size) return; + + const float * s_src = src0 + ids[s] * s0_stride_seq; + s_cur[s * state_size + idx] = s_src[idx]; +} + +// SSD (State Space Duality) dispatch for Mamba-2 prefill. +// Chunked matmuls: CB, fused_Y (MMA), S@C, B@X_dt. +// All strides are in elements (floats), not bytes. +static void ssm_scan_ssd_f32_cuda( + ggml_backend_cuda_context & ctx, + const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, + const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, + const int s0_stride_seq, // state (src0) stride between seqs + const int x_stride_tok, const int x_stride_seq, // x (src1) strides + const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides + const int A_stride, // A (src3) stride between heads + const int B_stride_tok, const int B_stride_seq, // B (src4) strides + const int C_stride_tok, const int C_stride_seq, // C (src5) strides + const int64_t s_off, const int64_t d_state, const int64_t head_dim, + const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq) { + + cudaStream_t stream = ctx.stream(); + const int64_t d_inner = head_dim * n_head; + + const int64_t chunk_size = SSM_SSD_CHUNK_SIZE; + const int64_t n_chunks = (n_tok + chunk_size - 1) / chunk_size; + + const int64_t state_per_head = d_state * head_dim; + + using matmul_t = half; + static constexpr cudaDataType matmul_dtype = CUDA_R_16F; + + ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); + ggml_cuda_pool_alloc CB_buf(ctx.pool(), chunk_size * chunk_size * n_group * n_seq); + ggml_cuda_pool_alloc X_dt_buf(ctx.pool(), chunk_size * head_dim * n_head * n_seq); + ggml_cuda_pool_alloc B_w_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + ggml_cuda_pool_alloc C_s_buf(ctx.pool(), d_state * chunk_size * n_head * n_seq); + float * dt_sp = dt_sp_buf.get(); + float * cs = cs_buf.get(); + float * CB = CB_buf.get(); + matmul_t * X_dt = X_dt_buf.get(); + matmul_t * B_weighted = B_w_buf.get(); + float * C_scaled = C_s_buf.get(); + float * s_cur = (float *)((char *)dst_d + s_off); // write state directly to dst + + // Step 1: softplus(dt) and parallel prefix sum over full sequence + { + constexpr int DT_BLOCK = 256; + constexpr int DT_MAX_ITEMS = 32; // handles up to n_tok=8192 with BLOCK=256 + GGML_ASSERT(n_tok <= DT_BLOCK * DT_MAX_ITEMS); + dim3 grid(n_head, n_seq); + ssm_ssd_prepare_dt_kernel<<>>( + src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 2: initialize running state from src0[ids[s]] + { + constexpr int BLOCK = 256; + const int64_t state_size = d_state * head_dim * n_head; + dim3 grid((state_size + BLOCK - 1) / BLOCK, n_seq); + ssm_ssd_init_state_kernel<<>>( + src0_d, src6_d, s_cur, state_size, s0_stride_seq); + CUDA_CHECK(cudaGetLastError()); + } + + // Step 3: chunked SSD loop + // Per chunk: pre_matmul + 3 cuBLAS (CB, S@C, state update) + fused MMA Y + scale_state + cublasHandle_t handle = ctx.cublas_handle(); + CUBLAS_CHECK(cublasSetStream(handle, stream)); + const float alpha_one = 1.0f; + const float beta_zero = 0.0f; + const float beta_one = 1.0f; + const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B + const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B + + for (int64_t k = 0; k < n_chunks; k++) { + const int64_t chunk_offset = k * chunk_size; + const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); + + // 3a: CB = C^T @ B per group + for (int64_t s = 0; s < n_seq; s++) { + const float * C_s = src5_d + s * C_stride_seq + chunk_offset * C_stride_tok; + const float * B_s = src4_d + s * B_stride_seq + chunk_offset * B_stride_tok; + float * CB_s = CB + s * chunk_len * chunk_len * n_group; + + if (n_group == 1) { + CUBLAS_CHECK(cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, + &beta_zero, CB_s, (int)chunk_len)); + } else { + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, + chunk_len, chunk_len, d_state, + &alpha_one, + C_s, lda_C_src, d_state, + B_s, ldb_B_src, d_state, + &beta_zero, + CB_s, (int)chunk_len, (long long)(chunk_len * chunk_len), + n_group)); + } + } + + // 3b: prepare X_dt, B_weighted, C_scaled + { + constexpr int BLOCK = 256; + const int64_t n_xdt = chunk_len * head_dim; + const int64_t n_bw = d_state * chunk_len; + int64_t max_work = n_xdt; + if (n_bw > max_work) max_work = n_bw; + dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_pre_matmul_kernel<<>>( + cs, dt_sp, src3_d, src1_d, src4_d, src5_d, + X_dt, B_weighted, C_scaled, + chunk_len, head_dim, n_head, n_group, d_state, A_stride, + x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, + chunk_offset, n_tok); + CUDA_CHECK(cudaGetLastError()); + } + + // 3c: dst = S_cur^T @ C_scaled (state contribution) + { + const int64_t stride_S = state_per_head; + const int64_t stride_Cs = d_state * chunk_len; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + + CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, + head_dim, chunk_len, d_state, + &alpha_one, + s_cur + s * stride_S * n_head, d_state, stride_S, + C_scaled + s * stride_Cs * n_head, d_state, stride_Cs, + &beta_zero, + dst_chunk, d_inner, head_dim, + n_head)); + } + } + + // 3d: dst += fused MMA Y (intra-chunk contribution, adds to 3c) + { + constexpr int D_TILE = 16; + constexpr int K_TILE = 16; + constexpr int N_WARPS = 8; + constexpr int N_TOUT_PER_WARP = SSM_SSD_CHUNK_SIZE / 8 / N_WARPS; // 256/8/8 = 4 + const int n_d_blocks = (head_dim + D_TILE - 1) / D_TILE; + dim3 grid(n_d_blocks, n_head, n_seq); + dim3 block(WARP_SIZE, N_WARPS, 1); + const size_t smem_cs_padded = ((chunk_len + 3) & ~3) * sizeof(float); // 16-byte aligned + const size_t smem_bytes = smem_cs_padded // smem_cs + + 2 * K_TILE * D_TILE * sizeof(half) // smem_X (double-buffered, d-fastest) + + N_WARPS * 8 * K_TILE * sizeof(half); // smem_M + ssm_ssd_fused_y_mma_kernel + <<>>( + X_dt, CB, cs, src3_d, dst_d, + head_dim, chunk_len, n_head, n_group, + d_inner, n_tok, + chunk_offset, A_stride, + (int64_t)chunk_len * head_dim); + CUDA_CHECK(cudaGetLastError()); + } + + // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) + { + // Scale s_cur in-place by per-head decay_total BEFORE cuBLAS overwrites it + constexpr int BLOCK = 256; + dim3 grid((state_per_head + BLOCK - 1) / BLOCK, n_head, n_seq); + ssm_ssd_scale_state_kernel<<>>( + s_cur, cs, src3_d, + d_state, head_dim, n_head, + chunk_offset, chunk_len, n_tok, A_stride); + CUDA_CHECK(cudaGetLastError()); + + // cuBLAS with beta=1: s_cur = B_weighted @ X_dt^T + 1.0 * s_cur (already scaled) + const int64_t stride_Bw = d_state * chunk_len; + const int64_t stride_X = chunk_len * head_dim; + const int64_t stride_S = state_per_head; + + for (int64_t s = 0; s < n_seq; s++) { + // X_dt is d-fastest {hd, C}, read as OP_T to get {C, hd} + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + d_state, head_dim, chunk_len, + &alpha_one, + B_weighted + s * stride_Bw * n_head, matmul_dtype, d_state, stride_Bw, + X_dt + s * stride_X * n_head, matmul_dtype, head_dim, stride_X, + &beta_one, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } + } + } +} + 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 @@ -357,8 +994,30 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src6->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32); - ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, - src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], - src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], - s_off, nc, nr, nh, ng, n_t, n_s, stream); + // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. + // Requires NVIDIA Turing+ for MMA (m16n8k16 + ldmatrix), otherwise fallback to scan. + const bool is_mamba2 = (src3->nb[1] == sizeof(float)); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && GGML_CUDA_CC_IS_NVIDIA(cc) + && cc >= GGML_CUDA_CC_TURING + && nr % 8 == 0; // head_dim must be 8-aligned for cp.async 16-byte copies + + if (use_ssd) { + // Convert byte strides to element strides (all tensors are f32-contiguous on dim 0) + ssm_scan_ssd_f32_cuda(ctx, + src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + src0->nb[3] / sizeof(float), + src1->nb[2] / sizeof(float), src1->nb[3] / sizeof(float), + src2->nb[1] / sizeof(float), src2->nb[2] / sizeof(float), + src3->nb[1] / sizeof(float), + src4->nb[2] / sizeof(float), src4->nb[3] / sizeof(float), + src5->nb[2] / sizeof(float), src5->nb[3] / sizeof(float), + s_off, nc, nr, nh, ng, n_t, n_s); + } else { + ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], + src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], + s_off, nc, nr, nh, ng, n_t, n_s, stream); + } } From e9633ae8fc398c5100ab317b4c5c9d74abab4da9 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 5 May 2026 12:05:44 +0530 Subject: [PATCH 02/10] cuda: added SSD CICD fixes for CUDA / HIP / MUSA / MSVC. --- ggml/src/ggml-cuda/ssm-scan.cu | 73 +++++++++++++++++++--------------- 1 file changed, 41 insertions(+), 32 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index b67edb58c631..9295f9e9fd66 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -1,14 +1,19 @@ +#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" #include "mma.cuh" #include "cp-async.cuh" -#ifdef GGML_CUDA_USE_CUB -#include -using namespace cub; -#endif // GGML_CUDA_USE_CUB -// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan -// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences) +// Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. +// For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). #define SSM_SSD_MIN_TOKENS 64 // Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. @@ -66,7 +71,7 @@ __global__ void __launch_bounds__(splitD, 1) __shared__ float smemB[N]; __shared__ float smemC[N]; -#ifdef GGML_CUDA_USE_CUB +#ifdef USE_CUB using BlockLoad = cub::BlockLoad; using BlockStore = cub::BlockStore; @@ -119,7 +124,7 @@ __global__ void __launch_bounds__(splitD, 1) __syncthreads(); } -#ifdef GGML_CUDA_USE_CUB +#ifdef USE_CUB BlockStore(cub_temp_storage.store_temp).Store(s_block, regs0); #else const int stride_s = stride_s0; @@ -330,7 +335,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa // SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) // // Instead of a sequential scan, SSD reformulates the output as: -// Y = (L ? (C @ B^T)) @ (X * dt) + decay * C @ s_init +// Y = (L (.) (C @ B^T)) @ (X * dt) + decay * C @ s_init // where L is a causal decay mask derived from A and dt. // // This converts the O(T*N) sequential scan into parallel matmuls. @@ -374,7 +379,7 @@ __global__ void ssm_ssd_prepare_dt_kernel( } // Phase 2: parallel prefix sum of per-thread totals -#ifdef GGML_CUDA_USE_CUB +#ifdef USE_CUB using BlockScan = cub::BlockScan; __shared__ typename BlockScan::TempStorage scan_temp; float thread_inclusive; @@ -454,7 +459,7 @@ __global__ void ssm_ssd_pre_matmul_kernel( // Prepare B_weighted = B * decay_from_end for state update matmul. // dt factor is already in X_dt (x * dt), so B_weighted must NOT include dt - // to avoid double-counting: s_cur = B_weighted @ X_dt^T = B*decay * x*dt = B*x*decay*dt ? + // to avoid double-counting: s_cur = B_weighted @ X_dt^T = B*decay * x*dt = B*x*decay*dt // Column-major {d_state, chunk_len} per head, group broadcast to head. const int n_bw = d_state * chunk_len; if (idx < n_bw) { @@ -488,7 +493,7 @@ __global__ void ssm_ssd_pre_matmul_kernel( // Scale running state in-place: s_cur *= decay_total(chunk). // Called BEFORE cuBLAS state update (beta=1) to fuse inter-chunk decay. // Eliminates the s_old buffer and D2D memcpy vs the old approach of: -// memcpy(s_old, s_cur) ? cuBLAS(beta=0) ? s_cur += decay * s_old +// memcpy(s_old, s_cur) -> cuBLAS(beta=0) -> s_cur += decay * s_old // Grid: (ceil(d_state * head_dim / BLOCK), n_head, n_seqs) template __global__ void ssm_ssd_scale_state_kernel( @@ -520,15 +525,15 @@ __global__ void ssm_ssd_scale_state_kernel( // // Computes M on-the-fly from CB + cs + A in shared memory, never writing M to // global memory. This eliminates ~32 MB DRAM traffic per chunk vs materializing M. -// Uses m16n8k16 FP16?FP32 MMA instructions via mma.cuh. +// Uses m16n8k16 FP16->FP32 MMA instructions via mma.cuh. // // The operation computes: -// Y[d, t_out, h] += S_{t_in<=t_out} X_dt[d, t_in, h] +// Y[d, t_out, h] += Sigma_{t_in<=t_out} X_dt[d, t_in, h] // * exp(A[h]*(cs[t_out]-cs[t_in])) * CB[t_out, t_in, g(h)] // // Key design: // - M is computed on-the-fly per warp in shared memory (never touches DRAM) -// - X_dt tile loaded cooperatively into smem with d?t_in transpose +// - X_dt tile loaded cooperatively into smem with d->t_in transpose // - Both A and B tiles loaded via hardware ldmatrix from shared memory // - C (output) accumulated in registers across all k-blocks // @@ -553,8 +558,8 @@ __global__ void ssm_ssd_fused_y_mma_kernel( constexpr int T_OUT_TILE = 8; // MMA N dimension (t_out per tile) // Number of 16-byte cp.async copies per X tile: D_TILE=16 halfs = 32 bytes = 2 copies per t_in row - constexpr int COPIES_PER_ROW = D_TILE / 8; // 16 halfs / 8 halfs per copy = 2 - constexpr int COPIES_TOTAL = K_TILE * COPIES_PER_ROW; // 16 * 2 = 32 + [[maybe_unused]] constexpr int COPIES_PER_ROW = D_TILE / 8; // 16 halfs / 8 halfs per copy = 2 + [[maybe_unused]] constexpr int COPIES_TOTAL = K_TILE * COPIES_PER_ROW; // 16 * 2 = 32 const int d_block = blockIdx.x; const int h = blockIdx.y; @@ -570,8 +575,8 @@ __global__ void ssm_ssd_fused_y_mma_kernel( // A tile loaded via load_ldmatrix_trans (hardware transpose during load). // smem_cs padded to 16-byte alignment for cp.async destination requirements. // [smem_cs: chunk_len_padded floats] - // [smem_X: 2 * K_TILE * D_TILE halfs] ? ping-pong buffers, d-fastest - // [smem_M: N_WARPS * T_OUT_TILE * K_TILE halfs] ? t_in-fastest (for B tile) + // [smem_X: 2 * K_TILE * D_TILE halfs] <- ping-pong buffers, d-fastest + // [smem_M: N_WARPS * T_OUT_TILE * K_TILE halfs] <- t_in-fastest (for B tile) extern __shared__ char smem_raw[]; const int chunk_len_padded = (chunk_len + 3) & ~3; // round up to 16-byte boundary float * smem_cs = (float *)smem_raw; @@ -603,7 +608,7 @@ __global__ void ssm_ssd_fused_y_mma_kernel( const int n_k_blocks = (chunk_len + K_TILE - 1) / K_TILE; // Helper: issue cp.async copies for X_dt tile into smem buffer (d-fastest layout). - // Each copy is 16 bytes = 8 halfs. D_TILE=16 ? 2 copies per t_in row, 32 total. + // Each copy is 16 bytes = 8 halfs. D_TILE=16 -> 2 copies per t_in row, 32 total. // k_len_tile: valid rows in this tile (may be < K_TILE for the last k-block). // Out-of-bounds rows are zeroed to prevent NaN propagation through MMA (0*NaN=NaN). auto issue_cp_async_X = [&](half * smem_X_buf, const int t_in_start, const int k_len_tile) { @@ -696,7 +701,9 @@ __global__ void ssm_ssd_fused_y_mma_kernel( my_M[tout_local * K_TILE + tin_local] = val; } } +#ifndef GGML_USE_HIP __syncwarp(); +#endif // GGML_USE_HIP // --- MMA tiles --- // A tile: loaded via ldmatrix_trans from d-fastest smem_X (hardware transpose) @@ -708,7 +715,7 @@ __global__ void ssm_ssd_fused_y_mma_kernel( tile_B B_tile; // A from smem_X (d-fastest): stride = D_TILE/2 half2 between t_in rows - // ldmatrix_trans transposes: smem rows=t_in ? tile rows=d ? + // ldmatrix_trans transposes: smem rows=t_in -> tile rows=d const half2 * X_h2 = (const half2 *)smem_X_cur; ggml_cuda_mma::load_ldmatrix_trans(A_tile, X_h2, D_TILE / 2); @@ -788,7 +795,7 @@ static void ssm_scan_ssd_f32_cuda( const int64_t state_per_head = d_state * head_dim; using matmul_t = half; - static constexpr cudaDataType matmul_dtype = CUDA_R_16F; + static constexpr cudaDataType_t matmul_dtype = CUDA_R_16F; ggml_cuda_pool_alloc dt_sp_buf(ctx.pool(), n_tok * n_head * n_seq); ggml_cuda_pool_alloc cs_buf(ctx.pool(), n_tok * n_head * n_seq); @@ -851,14 +858,15 @@ static void ssm_scan_ssd_f32_cuda( &alpha_one, C_s, lda_C_src, B_s, ldb_B_src, &beta_zero, CB_s, (int)chunk_len)); } else { - CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, chunk_len, chunk_len, d_state, &alpha_one, - C_s, lda_C_src, d_state, - B_s, ldb_B_src, d_state, + C_s, CUDA_R_32F, lda_C_src, d_state, + B_s, CUDA_R_32F, ldb_B_src, d_state, &beta_zero, - CB_s, (int)chunk_len, (long long)(chunk_len * chunk_len), - n_group)); + CB_s, CUDA_R_32F, (int)chunk_len, (long long)(chunk_len * chunk_len), + n_group, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); } } @@ -887,14 +895,15 @@ static void ssm_scan_ssd_f32_cuda( for (int64_t s = 0; s < n_seq; s++) { float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; - CUBLAS_CHECK(cublasSgemmStridedBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, head_dim, chunk_len, d_state, &alpha_one, - s_cur + s * stride_S * n_head, d_state, stride_S, - C_scaled + s * stride_Cs * n_head, d_state, stride_Cs, + s_cur + s * stride_S * n_head, CUDA_R_32F, d_state, stride_S, + C_scaled + s * stride_Cs * n_head, CUDA_R_32F, d_state, stride_Cs, &beta_zero, - dst_chunk, d_inner, head_dim, - n_head)); + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); } } From 3f93219202bee31f3236aa4e4e1e57b2bebf0987 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Fri, 15 May 2026 17:27:48 +0530 Subject: [PATCH 03/10] ggml-cuda: review comments fixed. --- ggml/src/ggml-cuda/ssm-scan.cu | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 9295f9e9fd66..c4f75cb68263 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -16,14 +16,14 @@ using namespace cub; // For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). #define SSM_SSD_MIN_TOKENS 64 +// Maximum number of tokens the SSD prepare_dt kernel can handle in one launch. +// Bounded by DT_BLOCK (256) * DT_MAX_ITEMS (32) = 8192 items per block. Sequences +// longer than this fall back to the scan path. +#define SSM_SSD_MAX_TOKENS 8192 + // Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. #define SSM_SSD_CHUNK_SIZE 256 -// Fast exp() using hardware exp2: exp(x) = exp2(x * log2(e)). -__device__ __forceinline__ float fast_expf(float x) { - return exp2f(x * 1.4426950408889634f); // 1/ln(2) = log2(e) -} - // We would like to keep pragma unroll for cases where L_template is not 0, // so we suppress the clang transformation warning. #ifdef __clang__ @@ -467,7 +467,7 @@ __global__ void ssm_ssd_pre_matmul_kernel( const int t = idx / d_state; const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; - const float decay_from_end = fast_expf(A_h * (cs_last - cs_t)); + const float decay_from_end = __expf(A_h * (cs_last - cs_t)); const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; @@ -482,7 +482,7 @@ __global__ void ssm_ssd_pre_matmul_kernel( const int t = idx / d_state; const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; - const float decay_to_pos = fast_expf(A_h * cs_t); + const float decay_to_pos = __expf(A_h * cs_t); const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; @@ -514,7 +514,7 @@ __global__ void ssm_ssd_scale_state_kernel( const int cs_seq_off = s * n_tok_total * n_head; const float cs_base = (chunk_offset > 0) ? cs[cs_seq_off + (chunk_offset - 1) * n_head + h] : 0.0f; const float cs_last = cs[cs_seq_off + (chunk_offset + chunk_len - 1) * n_head + h] - cs_base; - const float decay_total = fast_expf(A_h * cs_last); + const float decay_total = __expf(A_h * cs_last); const int off = s * state_per_head * n_head + h * state_per_head + idx; s_cur[off] *= decay_total; @@ -692,7 +692,7 @@ __global__ void ssm_ssd_fused_y_mma_kernel( if (t_out_abs < chunk_len && tin_local < k_len && t_in_abs <= t_out_abs) { const float cs_out = smem_cs[t_out_abs]; const float cs_in = smem_cs[t_in_abs]; - const float decay = fast_expf(A_h * (cs_out - cs_in)); + const float decay = __expf(A_h * (cs_out - cs_in)); const float cb_val = CB_g[t_out_abs + t_in_abs * chunk_len]; val = __float2half(decay * cb_val); } else { @@ -701,9 +701,6 @@ __global__ void ssm_ssd_fused_y_mma_kernel( my_M[tout_local * K_TILE + tin_local] = val; } } -#ifndef GGML_USE_HIP - __syncwarp(); -#endif // GGML_USE_HIP // --- MMA tiles --- // A tile: loaded via ldmatrix_trans from d-fastest smem_X (hardware transpose) @@ -815,7 +812,6 @@ static void ssm_scan_ssd_f32_cuda( { constexpr int DT_BLOCK = 256; constexpr int DT_MAX_ITEMS = 32; // handles up to n_tok=8192 with BLOCK=256 - GGML_ASSERT(n_tok <= DT_BLOCK * DT_MAX_ITEMS); dim3 grid(n_head, n_seq); ssm_ssd_prepare_dt_kernel<<>>( src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); @@ -1008,6 +1004,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const bool is_mamba2 = (src3->nb[1] == sizeof(float)); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && n_t <= SSM_SSD_MAX_TOKENS // prepare_dt block bound (DT_BLOCK * DT_MAX_ITEMS) && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING && nr % 8 == 0; // head_dim must be 8-aligned for cp.async 16-byte copies From 21f374576f43b8c5966706fcae0fa15db3b42d2d Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 9 Jun 2026 18:59:29 +0530 Subject: [PATCH 04/10] ggml-cuda: Fuse M matrix materialization into pre_matmul kernel and enabled test. --- ggml/src/ggml-cuda/ssm-scan.cu | 320 ++++++--------------------------- tests/test-backend-ops.cpp | 18 +- 2 files changed, 74 insertions(+), 264 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index c4f75cb68263..ecefa840dbcb 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -8,13 +8,11 @@ using namespace cub; #endif // USE_CUB #include "ssm-scan.cuh" -#include "mma.cuh" -#include "cp-async.cuh" // Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path. // For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). -#define SSM_SSD_MIN_TOKENS 64 +#define SSM_SSD_MIN_TOKENS 128 // Maximum number of tokens the SSD prepare_dt kernel can handle in one launch. // Bounded by DT_BLOCK (256) * DT_MAX_ITEMS (32) = 8192 items per block. Sequences @@ -414,7 +412,8 @@ __global__ void ssm_ssd_prepare_dt_kernel( // T_matmul controls precision for X_dt, B_weighted (float or half). // C_scaled is always float (pairs with float s_cur in step 3c). // Computation is always FP32; only the final store converts to T_matmul. -// Grid: (ceil(max(C*head_dim, d_state*C) / BLOCK), n_head, n_seqs) +// Also materializes the causal M matrix = exp(A*(cs_out - cs_in)) * CB (fused with prep to save a launch). +// Grid: (ceil(max(C*head_dim, d_state*C, chunk_len^2) / BLOCK), n_head, n_seqs) template __global__ void ssm_ssd_pre_matmul_kernel( const float * __restrict__ cs, // {n_tok, n_head} cumulative dt sums @@ -426,6 +425,8 @@ __global__ void ssm_ssd_pre_matmul_kernel( T_matmul * __restrict__ X_dt, // {head_dim, C, n_head} x * dt, d-fastest T_matmul * __restrict__ B_weighted, // {d_state, C, n_head} B * decay_from_end float * __restrict__ C_scaled, // {d_state, C, n_head} C * decay_to_pos (always float) + const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} + half * __restrict__ M_out, // {chunk_len, chunk_len, n_head, n_seqs} const int chunk_len, const int head_dim, const int n_head, const int n_group, const int d_state, const int A_stride, const int x_stride_tok, const int x_stride_seq, @@ -458,9 +459,6 @@ __global__ void ssm_ssd_pre_matmul_kernel( } // Prepare B_weighted = B * decay_from_end for state update matmul. - // dt factor is already in X_dt (x * dt), so B_weighted must NOT include dt - // to avoid double-counting: s_cur = B_weighted @ X_dt^T = B*decay * x*dt = B*x*decay*dt - // Column-major {d_state, chunk_len} per head, group broadcast to head. const int n_bw = d_state * chunk_len; if (idx < n_bw) { const int n = idx % d_state; @@ -475,7 +473,6 @@ __global__ void ssm_ssd_pre_matmul_kernel( } // Prepare C_scaled = C * decay_to_pos for state contribution matmul. - // Column-major {d_state, chunk_len} per head, group broadcast to head. const int n_cs = d_state * chunk_len; if (idx < n_cs) { const int n = idx % d_state; @@ -488,6 +485,28 @@ __global__ void ssm_ssd_pre_matmul_kernel( C_scaled[n + t * d_state + h * n_cs + s * n_cs * n_head] = C_val * decay_to_pos; } + + // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. + const int n_M = chunk_len * chunk_len; + if (idx < n_M) { + const int t_out = idx % chunk_len; + const int t_in = idx / chunk_len; + + half val; + if (t_in <= t_out) { + const float cs_out = cs[cs_seq_off + (chunk_offset + t_out) * n_head + h] - cs_base; + const float cs_in = cs[cs_seq_off + (chunk_offset + t_in) * n_head + h] - cs_base; + const float decay = __expf(A_h * (cs_out - cs_in)); + const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group + + (int64_t)g * chunk_len * chunk_len; + const float cb_val = CB_g[t_out + t_in * chunk_len]; + val = __float2half(decay * cb_val); + } else { + val = __float2half(0.0f); + } + + M_out[(int64_t)s * n_M * n_head + (int64_t)h * n_M + t_in * chunk_len + t_out] = val; + } } // Scale running state in-place: s_cur *= decay_total(chunk). @@ -520,236 +539,6 @@ __global__ void ssm_ssd_scale_state_kernel( s_cur[off] *= decay_total; } -// ============================================================================ -// Fused Y kernel (MMA tensor-core): Y += X_dt @ M_online^T -// -// Computes M on-the-fly from CB + cs + A in shared memory, never writing M to -// global memory. This eliminates ~32 MB DRAM traffic per chunk vs materializing M. -// Uses m16n8k16 FP16->FP32 MMA instructions via mma.cuh. -// -// The operation computes: -// Y[d, t_out, h] += Sigma_{t_in<=t_out} X_dt[d, t_in, h] -// * exp(A[h]*(cs[t_out]-cs[t_in])) * CB[t_out, t_in, g(h)] -// -// Key design: -// - M is computed on-the-fly per warp in shared memory (never touches DRAM) -// - X_dt tile loaded cooperatively into smem with d->t_in transpose -// - Both A and B tiles loaded via hardware ldmatrix from shared memory -// - C (output) accumulated in registers across all k-blocks -// -// Block: dim3(WARP_SIZE, N_WARPS, 1) — threadIdx.x is lane, threadIdx.y is warp -// Grid: dim3(ceil(head_dim/D_TILE), n_head, n_seqs) -// ============================================================================ - -template -__global__ void ssm_ssd_fused_y_mma_kernel( - const half * __restrict__ X_dt, // {head_dim, chunk_len, n_head, n_seqs} d-fastest, FP16 - const float * __restrict__ CB, // {chunk_len, chunk_len, n_group, n_seqs} FP32 - const float * __restrict__ cs, // {n_tok, n_head, n_seqs} cumulative dt sums - const float * __restrict__ A, // {1, n_head} - float * __restrict__ dst, // {d_inner, n_tok, n_seqs} accumulate onto step 3c result - const int head_dim, const int chunk_len, const int n_head, const int n_group, - const int d_inner, const int n_tok, - const int chunk_offset, const int A_stride, - const int64_t stride_X_head) { // = chunk_len * head_dim - - static_assert(D_TILE == 16, "MMA m16n8k16 requires D_TILE=16"); - static_assert(K_TILE == 16, "K_TILE must be 16 for cp.async pipeline"); - constexpr int T_OUT_TILE = 8; // MMA N dimension (t_out per tile) - - // Number of 16-byte cp.async copies per X tile: D_TILE=16 halfs = 32 bytes = 2 copies per t_in row - [[maybe_unused]] constexpr int COPIES_PER_ROW = D_TILE / 8; // 16 halfs / 8 halfs per copy = 2 - [[maybe_unused]] constexpr int COPIES_TOTAL = K_TILE * COPIES_PER_ROW; // 16 * 2 = 32 - - const int d_block = blockIdx.x; - const int h = blockIdx.y; - const int s = blockIdx.z; - const int g = h / (n_head / n_group); - const int warp_id = threadIdx.y; - - const float A_h = A[h * A_stride]; - const int d_start = d_block * D_TILE; - - // --- Shared memory layout (double-buffered X for cp.async pipeline) --- - // smem_X stored D-FASTEST (matching global layout) for direct cp.async copies. - // A tile loaded via load_ldmatrix_trans (hardware transpose during load). - // smem_cs padded to 16-byte alignment for cp.async destination requirements. - // [smem_cs: chunk_len_padded floats] - // [smem_X: 2 * K_TILE * D_TILE halfs] <- ping-pong buffers, d-fastest - // [smem_M: N_WARPS * T_OUT_TILE * K_TILE halfs] <- t_in-fastest (for B tile) - extern __shared__ char smem_raw[]; - const int chunk_len_padded = (chunk_len + 3) & ~3; // round up to 16-byte boundary - float * smem_cs = (float *)smem_raw; - half * smem_X_pp = (half *)(smem_cs + chunk_len_padded); // double buffer base - half * smem_M = smem_X_pp + 2 * K_TILE * D_TILE; - - // Load cs values for this chunk (cooperative, block-wide) - { - const int n_threads = WARP_SIZE * N_WARPS; - const int cs_seq_off = s * n_tok * n_head; - for (int i = threadIdx.y * WARP_SIZE + threadIdx.x; i < chunk_len; i += n_threads) { - smem_cs[i] = cs[cs_seq_off + (chunk_offset + i) * n_head + h]; - } - } - __syncthreads(); - - // Global memory pointers - const half * X_base = X_dt + (int64_t)s * stride_X_head * n_head - + (int64_t)h * stride_X_head + d_start; - const float * CB_g = CB + (int64_t)s * chunk_len * chunk_len * n_group - + (int64_t)g * chunk_len * chunk_len; - float * dst_base = dst + (int64_t)s * d_inner * n_tok + (int64_t)h * head_dim + d_start; - - // C accumulators in registers - using tile_C = ggml_cuda_mma::tile<16, 8, float>; - tile_C C_acc[N_TOUT_PER_WARP]; - - const int tout_base = warp_id * N_TOUT_PER_WARP; - const int n_k_blocks = (chunk_len + K_TILE - 1) / K_TILE; - - // Helper: issue cp.async copies for X_dt tile into smem buffer (d-fastest layout). - // Each copy is 16 bytes = 8 halfs. D_TILE=16 -> 2 copies per t_in row, 32 total. - // k_len_tile: valid rows in this tile (may be < K_TILE for the last k-block). - // Out-of-bounds rows are zeroed to prevent NaN propagation through MMA (0*NaN=NaN). - auto issue_cp_async_X = [&](half * smem_X_buf, const int t_in_start, const int k_len_tile) { -#ifdef CP_ASYNC_AVAILABLE - const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; - constexpr int n_threads = WARP_SIZE * N_WARPS; - // Zero-fill partial tiles to prevent 0*NaN in MMA (pool memory may contain NaN) - if (k_len_tile < K_TILE) { - for (int i = tid; i < D_TILE * K_TILE; i += n_threads) { - smem_X_buf[i] = __float2half(0.0f); - } - __syncthreads(); - } - if (tid < COPIES_TOTAL) { - const int t = tid / COPIES_PER_ROW; // which t_in (0..K_TILE-1) - const int d_half = tid % COPIES_PER_ROW; // which 8-half group (0 or 1) - if (t < k_len_tile) { - const half * src = X_base + (t_in_start + t) * head_dim + d_half * 8; - half * dst_ptr = smem_X_buf + t * D_TILE + d_half * 8; - const unsigned int dst_addr = ggml_cuda_cvta_generic_to_shared(dst_ptr); - cp_async_cg_16<0>(dst_addr, src); - } - } - asm volatile("cp.async.commit_group;"); -#else - // Fallback: synchronous load (for GPUs below Ampere without cp.async) - const int n_threads = WARP_SIZE * N_WARPS; - const int tid = threadIdx.y * WARP_SIZE + threadIdx.x; - for (int i = tid; i < D_TILE * K_TILE; i += n_threads) { - const int t_local = i / D_TILE; - const int d_local = i % D_TILE; - half val = (t_local < k_len_tile && d_start + d_local < head_dim) - ? X_base[(t_in_start + t_local) * head_dim + d_local] - : __float2half(0.0f); - smem_X_buf[t_local * D_TILE + d_local] = val; - } -#endif - }; - - // === cp.async pipelined loop: overlap next X load with current compute === - int pp_cur = 0; - - // Prologue: issue async load of first k-block into buffer 0 - issue_cp_async_X(smem_X_pp, 0, min(K_TILE, chunk_len)); -#ifdef CP_ASYNC_AVAILABLE - cp_async_wait_all(); -#endif - __syncthreads(); - - for (int kb = 0; kb < n_k_blocks; kb++) { - const int t_in_start = kb * K_TILE; - const int k_len = min(K_TILE, chunk_len - t_in_start); - - half * smem_X_cur = smem_X_pp + pp_cur * K_TILE * D_TILE; - - // --- Issue async load of NEXT k-block (overlaps with compute below!) --- - if (kb + 1 < n_k_blocks) { - half * smem_X_next = smem_X_pp + (1 - pp_cur) * K_TILE * D_TILE; - const int next_k_len = min(K_TILE, chunk_len - (kb + 1) * K_TILE); - issue_cp_async_X(smem_X_next, (kb + 1) * K_TILE, next_k_len); - } - - // --- Per-warp: iterate over assigned t_out blocks --- - for (int tb = 0; tb < N_TOUT_PER_WARP; tb++) { - const int tout_block = tout_base + tb; - const int t_out_start = tout_block * T_OUT_TILE; - - if (t_out_start >= chunk_len) break; - if (t_in_start > t_out_start + T_OUT_TILE - 1) continue; - - // --- Compute M tile on-the-fly in warp's smem region --- - { - half * my_M = smem_M + warp_id * T_OUT_TILE * K_TILE; - for (int i = threadIdx.x; i < T_OUT_TILE * K_TILE; i += WARP_SIZE) { - const int tout_local = i / K_TILE; - const int tin_local = i % K_TILE; - const int t_out_abs = t_out_start + tout_local; - const int t_in_abs = t_in_start + tin_local; - - half val; - if (t_out_abs < chunk_len && tin_local < k_len && t_in_abs <= t_out_abs) { - const float cs_out = smem_cs[t_out_abs]; - const float cs_in = smem_cs[t_in_abs]; - const float decay = __expf(A_h * (cs_out - cs_in)); - const float cb_val = CB_g[t_out_abs + t_in_abs * chunk_len]; - val = __float2half(decay * cb_val); - } else { - val = __float2half(0.0f); - } - my_M[tout_local * K_TILE + tin_local] = val; - } - } - - // --- MMA tiles --- - // A tile: loaded via ldmatrix_trans from d-fastest smem_X (hardware transpose) - // B tile: loaded via ldmatrix from t_in-fastest smem_M (no transpose needed) - using tile_A = ggml_cuda_mma::tile<16, 8, half2>; - using tile_B = ggml_cuda_mma::tile<8, 8, half2>; - - tile_A A_tile; - tile_B B_tile; - - // A from smem_X (d-fastest): stride = D_TILE/2 half2 between t_in rows - // ldmatrix_trans transposes: smem rows=t_in -> tile rows=d - const half2 * X_h2 = (const half2 *)smem_X_cur; - ggml_cuda_mma::load_ldmatrix_trans(A_tile, X_h2, D_TILE / 2); - - // B from smem_M (t_in-fastest): stride = K_TILE/2 half2 between t_out rows - const half2 * my_M_h2 = (const half2 *)(smem_M + warp_id * T_OUT_TILE * K_TILE); - ggml_cuda_mma::load_ldmatrix(B_tile, my_M_h2, K_TILE / 2); - - ggml_cuda_mma::mma(C_acc[tb], A_tile, B_tile); - } - - // Wait for next buffer load + ensure all warps done with current buffer -#ifdef CP_ASYNC_AVAILABLE - cp_async_wait_all(); -#endif - pp_cur = 1 - pp_cur; - __syncthreads(); - } - - // === Store C accumulators to dst === - for (int tb = 0; tb < N_TOUT_PER_WARP; tb++) { - const int tout_block = tout_base + tb; - const int t_out_start = tout_block * T_OUT_TILE; - if (t_out_start >= chunk_len) break; - -#pragma unroll - for (int l = 0; l < tile_C::ne; l++) { - const int d_local = tile_C::get_i(l); - const int tout_local = tile_C::get_j(l); - const int t_out_abs = t_out_start + tout_local; - - if (d_start + d_local < head_dim && t_out_abs < chunk_len) { - dst_base[d_local + (chunk_offset + t_out_abs) * d_inner] += C_acc[tb].x[l]; - } - } - } -} - // Copy initial state from src0[ids[s]] into s_cur for each sequence. // Grid: (ceil(d_state * head_dim * n_head / BLOCK), n_seqs) template @@ -768,7 +557,7 @@ __global__ void ssm_ssd_init_state_kernel( } // SSD (State Space Duality) dispatch for Mamba-2 prefill. -// Chunked matmuls: CB, fused_Y (MMA), S@C, B@X_dt. +// Chunked matmuls: CB, materialize M + cuBLAS Y, S@C, B@X_dt. // All strides are in elements (floats), not bytes. static void ssm_scan_ssd_f32_cuda( ggml_backend_cuda_context & ctx, @@ -829,7 +618,7 @@ static void ssm_scan_ssd_f32_cuda( } // Step 3: chunked SSD loop - // Per chunk: pre_matmul + 3 cuBLAS (CB, S@C, state update) + fused MMA Y + scale_state + // Per chunk: pre_matmul (incl. M) + 4 cuBLAS (CB, Y, S@C, state update) + scale_state cublasHandle_t handle = ctx.cublas_handle(); CUBLAS_CHECK(cublasSetStream(handle, stream)); const float alpha_one = 1.0f; @@ -838,6 +627,11 @@ static void ssm_scan_ssd_f32_cuda( const int lda_C_src = C_stride_tok; // leading dim for C in CB = C^T @ B const int ldb_B_src = B_stride_tok; // leading dim for B in CB = C^T @ B + // Scratch buffer for causal M matrix, reused across chunks (max size at chunk_size) + const int64_t n_M_max = chunk_size * chunk_size; + ggml_cuda_pool_alloc M_buf(ctx.pool(), n_M_max * n_head * n_seq); + half * M_mat = M_buf.get(); + for (int64_t k = 0; k < n_chunks; k++) { const int64_t chunk_offset = k * chunk_size; const int64_t chunk_len = (chunk_offset + chunk_size <= n_tok) ? chunk_size : (n_tok - chunk_offset); @@ -866,17 +660,20 @@ static void ssm_scan_ssd_f32_cuda( } } - // 3b: prepare X_dt, B_weighted, C_scaled + // 3b: prepare X_dt, B_weighted, C_scaled + materialize causal M matrix + const int64_t n_M = chunk_len * chunk_len; { constexpr int BLOCK = 256; const int64_t n_xdt = chunk_len * head_dim; const int64_t n_bw = d_state * chunk_len; int64_t max_work = n_xdt; - if (n_bw > max_work) max_work = n_bw; + if (n_bw > max_work) max_work = n_bw; + if (n_M > max_work) max_work = n_M; dim3 grid((max_work + BLOCK - 1) / BLOCK, n_head, n_seq); ssm_ssd_pre_matmul_kernel<<>>( cs, dt_sp, src3_d, src1_d, src4_d, src5_d, X_dt, B_weighted, C_scaled, + CB, M_mat, chunk_len, head_dim, n_head, n_group, d_state, A_stride, x_stride_tok, x_stride_seq, B_stride_tok, B_stride_seq, C_stride_tok, C_stride_seq, chunk_offset, n_tok); @@ -903,27 +700,24 @@ static void ssm_scan_ssd_f32_cuda( } } - // 3d: dst += fused MMA Y (intra-chunk contribution, adds to 3c) + // 3d: dst += X_dt @ M^T (intra-chunk contribution, adds to 3c result) + // M is stored as M[t_out, t_in] (lower-triangular), transpose needed for Y = X @ M^T. { - constexpr int D_TILE = 16; - constexpr int K_TILE = 16; - constexpr int N_WARPS = 8; - constexpr int N_TOUT_PER_WARP = SSM_SSD_CHUNK_SIZE / 8 / N_WARPS; // 256/8/8 = 4 - const int n_d_blocks = (head_dim + D_TILE - 1) / D_TILE; - dim3 grid(n_d_blocks, n_head, n_seq); - dim3 block(WARP_SIZE, N_WARPS, 1); - const size_t smem_cs_padded = ((chunk_len + 3) & ~3) * sizeof(float); // 16-byte aligned - const size_t smem_bytes = smem_cs_padded // smem_cs - + 2 * K_TILE * D_TILE * sizeof(half) // smem_X (double-buffered, d-fastest) - + N_WARPS * 8 * K_TILE * sizeof(half); // smem_M - ssm_ssd_fused_y_mma_kernel - <<>>( - X_dt, CB, cs, src3_d, dst_d, - head_dim, chunk_len, n_head, n_group, - d_inner, n_tok, - chunk_offset, A_stride, - (int64_t)chunk_len * head_dim); - CUDA_CHECK(cudaGetLastError()); + const int64_t stride_M = n_M; + const int64_t stride_X_h = (int64_t)chunk_len * head_dim; + + for (int64_t s = 0; s < n_seq; s++) { + float * dst_chunk = dst_d + s * d_inner * n_tok + chunk_offset * d_inner; + CUBLAS_CHECK(cublasGemmStridedBatchedEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, + head_dim, chunk_len, chunk_len, + &alpha_one, + X_dt + s * stride_X_h * n_head, matmul_dtype, head_dim, stride_X_h, + M_mat + s * stride_M * n_head, matmul_dtype, chunk_len, stride_M, + &beta_one, + dst_chunk, CUDA_R_32F, d_inner, head_dim, + n_head, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); + } } // 3e: s_cur = B_weighted @ X_dt^T + decay_total * s_cur_old (state update) @@ -1000,7 +794,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(dst->type == GGML_TYPE_F32); // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. - // Requires NVIDIA Turing+ for MMA (m16n8k16 + ldmatrix), otherwise fallback to scan. + // Requires NVIDIA Turing+ otherwise fallback to scan. const bool is_mamba2 = (src3->nb[1] == sizeof(float)); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e3553c5e9ca1..f98957a66df3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4038,9 +4038,11 @@ struct test_ssm_scan : public test_case { void initialize_tensors(ggml_context * ctx) override { std::random_device rd; std::default_random_engine rng(rd()); + // Tensor order: s(0), x(1), dt(2), A(3), B(4), C(5), ids(6) + int tensor_idx = 0; for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } + if (ggml_is_view_op(t->op)) { tensor_idx++; continue; } // ids for (int64_t r = 0; r < ggml_nrows(t); r++) { std::vector data(t->ne[0]); @@ -4050,9 +4052,19 @@ struct test_ssm_scan : public test_case { std::shuffle(data.begin(), data.end(), rng); ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } + } else if (tensor_idx == 3) { + // A: always negative (decay); moderate range + init_tensor_uniform(t, -1.0f, -0.5f); + } else if (tensor_idx == 2) { + // dt (pre-softplus): softplus([-4, -1]) ~= [0.018, 0.31] + init_tensor_uniform(t, -4.0f, -1.0f); + } else if (tensor_idx == 0) { + // s (initial state): zero as in real inference + init_tensor_uniform(t, 0.0f, 0.0f); } else { init_tensor_uniform(t); } + tensor_idx++; } } }; @@ -8768,6 +8780,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 128, 1)); // Mamba-2 SSD path + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 128, 1)); // Nemotron-9B SSD path test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); @@ -9971,6 +9985,8 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_ssm_conv_bias_silu(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1}, true)); // generate test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B prefill + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 1, 1)); // Nemotron-9B generate // acc test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); From b0af8f4225f625c4c0a6cedc02952d405ec4f75a Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Thu, 11 Jun 2026 17:23:53 +0530 Subject: [PATCH 05/10] ggml-cuda: test updates and fixes --- tests/test-backend-ops.cpp | 44 ++++++++++++++++++++++---------------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f98957a66df3..0652a6d3a2ae 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4000,7 +4000,7 @@ struct test_ssm_scan : public test_case { test_ssm_scan(ggml_type type = GGML_TYPE_F32, int64_t d_state = 32, - int64_t head_dim = 1, // non-zero for Mamba-2 + int64_t head_dim = 1, // 1 = Mamba-1; > 1 = Mamba-2 (scalar A per head) int64_t n_head = 32, int64_t n_group = 1, int64_t n_seq_tokens = 32, @@ -4008,6 +4008,11 @@ struct test_ssm_scan : public test_case { bool xbc_overlap = false) : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {} + double max_nmse_err() override { + // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32. + return (head_dim > 1) ? 2e-7 : 1e-7; + } + ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs); ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs); @@ -4034,16 +4039,14 @@ struct test_ssm_scan : public test_case { return out; } - // similar to test_mul_mat_id + void initialize_tensors(ggml_context * ctx) override { std::random_device rd; std::default_random_engine rng(rd()); - // Tensor order: s(0), x(1), dt(2), A(3), B(4), C(5), ids(6) - int tensor_idx = 0; for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { tensor_idx++; continue; } - // ids + if (ggml_is_view_op(t->op)) { continue; } + // ids: permutation of [0..n_seqs) for (int64_t r = 0; r < ggml_nrows(t); r++) { std::vector data(t->ne[0]); for (int i = 0; i < t->ne[0]; i++) { @@ -4052,19 +4055,23 @@ struct test_ssm_scan : public test_case { std::shuffle(data.begin(), data.end(), rng); ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } - } else if (tensor_idx == 3) { - // A: always negative (decay); moderate range - init_tensor_uniform(t, -1.0f, -0.5f); - } else if (tensor_idx == 2) { - // dt (pre-softplus): softplus([-4, -1]) ~= [0.018, 0.31] - init_tensor_uniform(t, -4.0f, -1.0f); - } else if (tensor_idx == 0) { - // s (initial state): zero as in real inference + } else if (ggml_is_view_op(t->op)) { + continue; + } else if (t->ne[0] == d_state && t->ne[1] == head_dim && t->ne[2] == n_head) { + // s {d_state, head_dim, n_head, n_seqs}: zero initial state. + // ne[2] == n_head disambiguates from B/C {d_state, n_group, n_seq_tokens, n_seqs} + // which otherwise collide when n_group == head_dim (e.g. Mamba-1). init_tensor_uniform(t, 0.0f, 0.0f); + } else if (t->ne[0] == n_head && t->ne[1] == n_seq_tokens) { + // dt {n_head, n_seq_tokens, n_seqs}: pre-softplus values + init_tensor_uniform(t, -4.0f, -1.0f); + } else if (t->ne[1] == n_head && t->ne[2] == 1) { + // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors) + init_tensor_uniform(t, -1.0f, -0.5f); } else { - init_tensor_uniform(t); + // x, B, C: [-0.5, 0.5] keeps intermediate calculations from overflowing + init_tensor_uniform(t, -0.5f, 0.5f); } - tensor_idx++; } } }; @@ -8780,8 +8787,9 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap - test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 128, 1)); // Mamba-2 SSD path - test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 128, 1)); // Nemotron-9B SSD path + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks) + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs) test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1)); test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1)); From f6ed77737dc0659b2620a027517c4df58bd3b133 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 16 Jun 2026 18:11:50 +0530 Subject: [PATCH 06/10] ggml-cuda: test updates to remove hardcoding of tensor initialise data limits. --- tests/test-backend-ops.cpp | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0652a6d3a2ae..a8bf3368c9f2 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4057,20 +4057,11 @@ struct test_ssm_scan : public test_case { } } else if (ggml_is_view_op(t->op)) { continue; - } else if (t->ne[0] == d_state && t->ne[1] == head_dim && t->ne[2] == n_head) { - // s {d_state, head_dim, n_head, n_seqs}: zero initial state. - // ne[2] == n_head disambiguates from B/C {d_state, n_group, n_seq_tokens, n_seqs} - // which otherwise collide when n_group == head_dim (e.g. Mamba-1). - init_tensor_uniform(t, 0.0f, 0.0f); - } else if (t->ne[0] == n_head && t->ne[1] == n_seq_tokens) { - // dt {n_head, n_seq_tokens, n_seqs}: pre-softplus values - init_tensor_uniform(t, -4.0f, -1.0f); } else if (t->ne[1] == n_head && t->ne[2] == 1) { // A {1 or d_state, n_head}: negative decay (2-D tensor, ne[2]==1 distinguishes from 3-D/4-D tensors) init_tensor_uniform(t, -1.0f, -0.5f); } else { - // x, B, C: [-0.5, 0.5] keeps intermediate calculations from overflowing - init_tensor_uniform(t, -0.5f, 0.5f); + init_tensor_uniform(t); } } } From f35681f920cf4249290ec536801310b92a732107 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 16 Jun 2026 19:13:15 +0530 Subject: [PATCH 07/10] ggml-cuda: ssd minor review comment fixed. --- ggml/src/ggml-cuda/ssm-scan.cu | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index ecefa840dbcb..fa63a243a248 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -14,10 +14,12 @@ using namespace cub; // For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences). #define SSM_SSD_MIN_TOKENS 128 -// Maximum number of tokens the SSD prepare_dt kernel can handle in one launch. -// Bounded by DT_BLOCK (256) * DT_MAX_ITEMS (32) = 8192 items per block. Sequences -// longer than this fall back to the scan path. -#define SSM_SSD_MAX_TOKENS 8192 +// prepare_dt kernel dimensions: one block per (head, seq), each block handles DT_MAX_ITEMS items. +#define SSM_SSD_DT_BLOCK 256 +#define SSM_SSD_DT_MAX_ITEMS 32 + +// Maximum tokens the SSD path supports — derived from the prepare_dt kernel block capacity. +#define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) // Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. #define SSM_SSD_CHUNK_SIZE 256 @@ -599,10 +601,8 @@ static void ssm_scan_ssd_f32_cuda( // Step 1: softplus(dt) and parallel prefix sum over full sequence { - constexpr int DT_BLOCK = 256; - constexpr int DT_MAX_ITEMS = 32; // handles up to n_tok=8192 with BLOCK=256 dim3 grid(n_head, n_seq); - ssm_ssd_prepare_dt_kernel<<>>( + ssm_ssd_prepare_dt_kernel<<>>( src2_d, dt_sp, cs, n_head, n_tok, dt_stride_tok, dt_stride_seq); CUDA_CHECK(cudaGetLastError()); } @@ -798,7 +798,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const bool is_mamba2 = (src3->nb[1] == sizeof(float)); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS - && n_t <= SSM_SSD_MAX_TOKENS // prepare_dt block bound (DT_BLOCK * DT_MAX_ITEMS) + && n_t <= SSM_SSD_MAX_TOKENS && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING && nr % 8 == 0; // head_dim must be 8-aligned for cp.async 16-byte copies From a942d1b07066a5715c910712e96c45f317593a90 Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 23 Jun 2026 18:51:31 +0530 Subject: [PATCH 08/10] ggml-cuda: ssd minor CICD fixed. --- ggml/src/ggml-cuda/ssm-scan.cu | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index fa63a243a248..8e450c4c33d9 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -331,6 +331,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa } } +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // ============================================================================ // SSD (State Space Duality) kernels for Mamba-2 prefill (n_tok > SSM_SSD_MIN_TOKENS) // @@ -751,6 +752,7 @@ static void ssm_scan_ssd_f32_cuda( } } } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; // s @@ -793,6 +795,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src6->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32); +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. // Requires NVIDIA Turing+ otherwise fallback to scan. const bool is_mamba2 = (src3->nb[1] == sizeof(float)); @@ -814,10 +817,11 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { src4->nb[2] / sizeof(float), src4->nb[3] / sizeof(float), src5->nb[2] / sizeof(float), src5->nb[3] / sizeof(float), s_off, nc, nr, nh, ng, n_t, n_s); - } else { - ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, - src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], - src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], - s_off, nc, nr, nh, ng, n_t, n_s, stream); + return; } +#endif + ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, + src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], + src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], + s_off, nc, nr, nh, ng, n_t, n_s, stream); } From 3c6938c8c722a8efda279d0b7a371988aa6fadbd Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 21 Jul 2026 22:25:35 +0530 Subject: [PATCH 09/10] CUDA SSD: Fixes correctness by promoting s0_stride_seq to int64_t, improves memory coalescing in ssm_ssd_prepare_dt_kernel, and boosts efficiency by merging B_weighted and C_scaled; also addresses prior review comments. --- ggml/src/ggml-cuda/ssm-scan.cu | 125 +++++++++++++++++++-------------- 1 file changed, 71 insertions(+), 54 deletions(-) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index 8e450c4c33d9..de4fa3ba4b7d 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -18,7 +18,7 @@ using namespace cub; #define SSM_SSD_DT_BLOCK 256 #define SSM_SSD_DT_MAX_ITEMS 32 -// Maximum tokens the SSD path supports — derived from the prepare_dt kernel block capacity. +// Maximum tokens the SSD path supports, derived from the prepare_dt kernel block capacity. #define SSM_SSD_MAX_TOKENS (SSM_SSD_DT_BLOCK * SSM_SSD_DT_MAX_ITEMS) // Chunk size for chunked SSD. Caps matmul cost at O(chunk^2) per chunk. @@ -362,53 +362,64 @@ __global__ void ssm_ssd_prepare_dt_kernel( const int items_per_thread = (n_tok + BLOCK_SIZE - 1) / BLOCK_SIZE; - // Phase 1: parallel softplus, each thread accumulates a local sum + // Phase 1: softplus with interleaved distribution (t = i*BLOCK_SIZE + threadIdx.x). + // Each warp reads BLOCK_SIZE consecutive tokens, giving coalesced dt_raw loads + // (stride n_head between threads vs. items_per_thread*n_head in blocked layout). float local_vals[MAX_ITEMS]; - float local_sum = 0.0f; - for (int i = 0; i < items_per_thread; i++) { - const int t = threadIdx.x * items_per_thread + i; // blocked distribution + const int t = i * BLOCK_SIZE + threadIdx.x; if (t < n_tok) { float val = dt_seq[h + t * dt_stride_tok]; float sp = (val <= 20.0f) ? log1pf(expf(val)) : val; local_vals[i] = sp; - local_sum += sp; dt_sp_seq[t * n_head + h] = sp; } else { local_vals[i] = 0.0f; } } - // Phase 2: parallel prefix sum of per-thread totals + // Phase 2+3: per-step inclusive scan to build cs[] in token order. + // With interleaved distribution the per-thread total scan would not give token-order + // prefix sums, so we scan one BLOCK_SIZE slab at a time and carry a running total. #ifdef USE_CUB using BlockScan = cub::BlockScan; __shared__ typename BlockScan::TempStorage scan_temp; - float thread_inclusive; - BlockScan(scan_temp).InclusiveSum(local_sum, thread_inclusive); - float exclusive_prefix = thread_inclusive - local_sum; -#else - // Fallback: sequential prefix sum in shared memory - __shared__ float sdata[BLOCK_SIZE]; - sdata[threadIdx.x] = local_sum; - __syncthreads(); - if (threadIdx.x == 0) { - for (int i = 1; i < BLOCK_SIZE && i * items_per_thread < n_tok; i++) { - sdata[i] += sdata[i - 1]; + __shared__ float step_total; + + float running = 0.0f; + for (int i = 0; i < items_per_thread; i++) { + float inclusive; + BlockScan(scan_temp).InclusiveSum(local_vals[i], inclusive); + const int t = i * BLOCK_SIZE + threadIdx.x; + if (t < n_tok) { + cs_seq[t * n_head + h] = running + inclusive; + } + if (threadIdx.x == BLOCK_SIZE - 1) { + step_total = inclusive; } + __syncthreads(); + running += step_total; } - __syncthreads(); - float exclusive_prefix = (threadIdx.x > 0) ? sdata[threadIdx.x - 1] : 0.0f; -#endif - - // Phase 3: write cumsum = exclusive_prefix + local running sum - float running = exclusive_prefix; +#else + // Fallback: sequential prefix scan in shared memory, one slab at a time. + __shared__ float sdata[BLOCK_SIZE]; + float running = 0.0f; for (int i = 0; i < items_per_thread; i++) { - const int t = threadIdx.x * items_per_thread + i; + const int t = i * BLOCK_SIZE + threadIdx.x; + sdata[threadIdx.x] = local_vals[i]; + __syncthreads(); + if (threadIdx.x == 0) { + for (int j = 1; j < BLOCK_SIZE; j++) { + sdata[j] += sdata[j - 1]; + } + } + __syncthreads(); if (t < n_tok) { - running += local_vals[i]; - cs_seq[t * n_head + h] = running; + cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; } + running += sdata[BLOCK_SIZE - 1]; } +#endif } // Prepare SSD matmul inputs for one chunk: X_dt, B_weighted, C_scaled. @@ -461,32 +472,20 @@ __global__ void ssm_ssd_pre_matmul_kernel( X_dt[d + t * head_dim + h * n_xdt + s * n_xdt * n_head] = (T_matmul)(x_val * dt_val); } - // Prepare B_weighted = B * decay_from_end for state update matmul. + // Prepare B_weighted and C_scaled together: both share the same index space (d_state * chunk_len) + // and the same cs_t load, so merging halves the cs[] global memory traffic. const int n_bw = d_state * chunk_len; if (idx < n_bw) { const int n = idx % d_state; const int t = idx / d_state; const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; - const float decay_from_end = __expf(A_h * (cs_last - cs_t)); const float B_val = B[s * B_stride_seq + (chunk_offset + t) * B_stride_tok + g * d_state + n]; - - B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * decay_from_end); - } - - // Prepare C_scaled = C * decay_to_pos for state contribution matmul. - const int n_cs = d_state * chunk_len; - if (idx < n_cs) { - const int n = idx % d_state; - const int t = idx / d_state; - - const float cs_t = cs[cs_seq_off + (chunk_offset + t) * n_head + h] - cs_base; - const float decay_to_pos = __expf(A_h * cs_t); + B_weighted[n + t * d_state + h * n_bw + s * n_bw * n_head] = (T_matmul)(B_val * __expf(A_h * (cs_last - cs_t))); const float C_val = C_src[s * C_stride_seq + (chunk_offset + t) * C_stride_tok + g * d_state + n]; - - C_scaled[n + t * d_state + h * n_cs + s * n_cs * n_head] = C_val * decay_to_pos; + C_scaled[n + t * d_state + h * n_bw + s * n_bw * n_head] = C_val * __expf(A_h * cs_t); } // Materialize M = exp(A*(cs_out - cs_in)) * CB with causal mask. @@ -550,12 +549,12 @@ __global__ void ssm_ssd_init_state_kernel( const int32_t * __restrict__ ids, // {n_seqs} float * __restrict__ s_cur, // {d_state, head_dim, n_head, n_seqs} const int state_size, // d_state * head_dim * n_head - const int s0_stride_seq) { // elements between state rows + const int64_t s0_stride_seq) { // elements between state rows const int s = blockIdx.y; const int idx = blockIdx.x * BLOCK_SIZE + threadIdx.x; if (idx >= state_size) return; - const float * s_src = src0 + ids[s] * s0_stride_seq; + const float * s_src = src0 + (int64_t)ids[s] * s0_stride_seq; s_cur[s * state_size + idx] = s_src[idx]; } @@ -566,7 +565,7 @@ static void ssm_scan_ssd_f32_cuda( ggml_backend_cuda_context & ctx, const float * src0_d, const float * src1_d, const float * src2_d, const float * src3_d, const float * src4_d, const float * src5_d, const int32_t * src6_d, float * dst_d, - const int s0_stride_seq, // state (src0) stride between seqs + const int64_t s0_stride_seq, // state (src0) stride between seqs const int x_stride_tok, const int x_stride_seq, // x (src1) strides const int dt_stride_tok, const int dt_stride_seq, // dt (src2) strides const int A_stride, // A (src3) stride between heads @@ -795,6 +794,19 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src6->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32); + // Byte strides are narrowed to int for both scan and SSD paths. + GGML_ASSERT(src0->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src0->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src1->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src2->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src3->nb[1] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src4->nb[3] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[2] <= (size_t)INT_MAX); + GGML_ASSERT(src5->nb[3] <= (size_t)INT_MAX); + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // Mamba-2 with scalar A per head: use SSD matmul path for long sequences. // Requires NVIDIA Turing+ otherwise fallback to scan. @@ -804,18 +816,23 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { && n_t <= SSM_SSD_MAX_TOKENS && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING - && nr % 8 == 0; // head_dim must be 8-aligned for cp.async 16-byte copies + && nr % 8 == 0; // cuBLAS requires 8-element (16-byte) alignment if (use_ssd) { - // Convert byte strides to element strides (all tensors are f32-contiguous on dim 0) + // ssm_ssd_init_state_kernel uses flat linear indexing within each sequence, + // so src0 must be fully contiguous across all inner dimensions. + // The scan path handles non-contiguous nb[2] via src0_nb2 but does not handle nb[1]. + GGML_ASSERT(src0->nb[1] == nc * sizeof(float)); + GGML_ASSERT(src0->nb[2] == nc * nr * sizeof(float)); + ssm_scan_ssd_f32_cuda(ctx, src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, - src0->nb[3] / sizeof(float), - src1->nb[2] / sizeof(float), src1->nb[3] / sizeof(float), - src2->nb[1] / sizeof(float), src2->nb[2] / sizeof(float), - src3->nb[1] / sizeof(float), - src4->nb[2] / sizeof(float), src4->nb[3] / sizeof(float), - src5->nb[2] / sizeof(float), src5->nb[3] / sizeof(float), + (int64_t)(src0->nb[3] / sizeof(float)), + (int)(src1->nb[2] / sizeof(float)), (int)(src1->nb[3] / sizeof(float)), + (int)(src2->nb[1] / sizeof(float)), (int)(src2->nb[2] / sizeof(float)), + (int)(src3->nb[1] / sizeof(float)), + (int)(src4->nb[2] / sizeof(float)), (int)(src4->nb[3] / sizeof(float)), + (int)(src5->nb[2] / sizeof(float)), (int)(src5->nb[3] / sizeof(float)), s_off, nc, nr, nh, ng, n_t, n_s); return; } From f96b0e53082a5c6d0d1765ef5be4f16f50c47f7c Mon Sep 17 00:00:00 2001 From: Bhavik Sharda Date: Tue, 28 Jul 2026 12:19:11 +0530 Subject: [PATCH 10/10] cuda: fix sdata read-write race in prepare_dt fallback scan loop --- ggml/src/ggml-cuda/ssm-scan.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index de4fa3ba4b7d..f3418c2af83d 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -418,6 +418,7 @@ __global__ void ssm_ssd_prepare_dt_kernel( cs_seq[t * n_head + h] = running + sdata[threadIdx.x]; } running += sdata[BLOCK_SIZE - 1]; + __syncthreads(); } #endif }