diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 4036cc21d..a82caa5e4 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -572,7 +572,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op) { +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op, bool tail) { GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); char base[256]; @@ -580,7 +580,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me const int nsg = (ne00 + 31)/32; - snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type)); + snprintf(base, 256, "kernel_ssm_scan_%s%s", ggml_type_name(op->src[0]->type), tail ? "_tail" : ""); snprintf(name, 256, "%s_nsg=%d", base, nsg); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); @@ -598,6 +598,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_ssm_scan_ssd_mma_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + // acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup + res.smem = (3*OP_SSM_SCAN_SSD_CS + + OP_SSM_SCAN_SSD_CS*OP_SSM_SCAN_SSD_HD + + OP_SSM_SCAN_SSD_NSG*8*OP_SSM_SCAN_SSD_CS + + OP_SSM_SCAN_SSD_NSG*2*8*8)*sizeof(float); + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) { char base[256]; char name[256]; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 6c39428c7..003b688db 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -129,7 +129,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tail); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index f0b779979..9becf0479 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -158,6 +158,10 @@ #define OP_SUM_ROWS_NUM_SUM_ROWS 10 #define OP_SUM_ROWS_NUM_MEAN 11 +#define OP_SSM_SCAN_SSD_CS 64 // Metal-specific; Chunk Size; 64 is largest multiple of 8 (simdgroup tile) fitting into 32 KiB Metal threadgroup mem limit (~26.75 KiB shared mem; see smem layout comment in kernel_ssm_scan_ssd_mma_f32) +#define OP_SSM_SCAN_SSD_HD 64 // Metal-specific; Head Dim the MMA kernel is specialized for (Mamba-2); use_mma gates on d_inner == this +#define OP_SSM_SCAN_SSD_NSG 4 // Metal-specific; Number of SimdGroups per threadgroup; NSG*32 == threads dispatched per threadgroup + // kernel argument structs // // - element counters (e.g. ne00) typically use int32_t to reduce register usage @@ -893,6 +897,8 @@ typedef struct { int64_t n_head; int64_t n_group; int64_t n_seq_tokens; + int64_t n_seq_tokens_total; + int64_t token_offset; int64_t n_seqs; int64_t K; uint64_t s_off; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 1c3bb936b..75de0f6dd 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1677,6 +1677,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; + const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); @@ -1722,6 +1723,8 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { /*.n_head =*/ n_head, /*.n_group =*/ n_group, /*.n_seq_tokens =*/ n_seq_tokens, + /*.n_seq_tokens_total =*/ n_seq_tokens, + /*.token_offset =*/ 0, /*.n_seqs =*/ n_seqs, /*.K =*/ K, /*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float), @@ -1751,26 +1754,53 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { /*.nb0 =*/ nb0, }; - auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op); + constexpr int64_t CHUNK = OP_SSM_SCAN_SSD_CS; - GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + const int64_t snap_reserve = K > 1 ? K : 0; // tokens reserved for sequential kernel rollback snapshots + const int64_t mma_tokens = ((n_seq_tokens - snap_reserve) / CHUNK) * CHUNK; // largest multiple of CHUNK that leaves snap_reserve for the tail + const bool use_mma = + mma_tokens > 0 && + ne30 == 1 && // checks that A tensor is set to scalar decay per head (A shape {1, n_head}) + props_dev->has_simdgroup_mm && // hardware check for M1 or newer + d_state % 8 == 0 && // d_state must be multiple of 8 to align with simdgroup_float 8x8 tiles + d_inner == OP_SSM_SCAN_SSD_HD; // mma kernel is specialized for the Mamba-2 head dim; this checks it - const size_t smem = pipeline.smem; + const auto dispatch = [&](ggml_metal_pipeline_with_params pipeline, int64_t nth, int64_t n_tg_x) { + GGML_ASSERT(nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + GGML_ASSERT(pipeline.smem <= props_dev->max_theadgroup_memory_size); - ggml_metal_encoder_set_pipeline(enc, pipeline); - ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8); + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8); + ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); + ggml_metal_encoder_dispatch_threadgroups(enc, n_tg_x, n_head, n_seqs, nth, 1, 1); + }; - ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + if (!use_mma) { + dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, false), d_state, d_inner); + return 1; + } - ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1); + args.n_seq_tokens = mma_tokens; + dispatch( + ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(lib, op), + OP_SSM_SCAN_SSD_NSG*32, + 1); + + if (mma_tokens < n_seq_tokens) { + ggml_metal_op_concurrency_reset(ctx); + + args.n_seq_tokens = n_seq_tokens - mma_tokens; + args.token_offset = mma_tokens; + dispatch(ggml_metal_library_get_pipeline_ssm_scan(lib, op, true), d_state, d_inner); + } return 1; } diff --git a/ggml/src/ggml-metal/kernels/ssm.metal b/ggml/src/ggml-metal/kernels/ssm.metal index be065c4ff..d3118a831 100644 --- a/ggml/src/ggml-metal/kernels/ssm.metal +++ b/ggml/src/ggml-metal/kernels/ssm.metal @@ -159,7 +159,9 @@ kernel void kernel_ssm_conv_f32_f32_batched_4( // ref: ggml.c:ggml_compute_forward_ssm_scan_f32, Mamba-2 part // Optimized version: reduces redundant memory loads by having one thread load shared values -kernel void kernel_ssm_scan_f32( +// TAIL == false is the whole-sequence / decode path: token_offset folds away at compile time. +template +kernel void kernel_ssm_scan_impl( constant ggml_metal_kargs_ssm_scan & args, device const void * src0, device const void * src1, @@ -200,13 +202,17 @@ kernel void kernel_ssm_scan_f32( const int32_t n_t = args.n_seq_tokens; const int32_t n_s = args.n_seqs; const int32_t K = args.K; + const int32_t n_t_total = TAIL ? args.n_seq_tokens_total : n_t; + const int32_t t_off = TAIL ? args.token_offset : 0; const int32_t s_off = args.s_off; device const int32_t * ids = (device const int32_t *) src6; - device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03); device float * s_buff = (device float *) ((device char *) dst + ir*args.nb02 + i3*args.nb03 + s_off); + device const float * s0_buff = t_off != 0 ? + s_buff : + (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03); const int32_t i = i0 + i1*nc; const int32_t g = ir / (nh / ng); // repeat_interleave @@ -218,12 +224,12 @@ kernel void kernel_ssm_scan_f32( const float A0 = A[i0%args.ne30]; - device const float * x = (device const float *)((device const char *) src1 + i1*args.nb10 + ir*args.nb11 + i3*args.nb13); // {dim, nh, nt, ns} - device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20 + i3*args.nb22); // {nh, nt, ns} - device const float * B = (device const float *)((device const char *) src4 + g*args.nb41 + i3*args.nb43); // {d_state, ng, nt, ns} - device const float * C = (device const float *)((device const char *) src5 + g*args.nb51 + i3*args.nb53); // {d_state, ng, nt, ns} + device const float * x = (device const float *)((device const char *) src1 + i1*args.nb10 + ir*args.nb11 + t_off*args.nb12 + i3*args.nb13); // {dim, nh, nt, ns} + device const float * dt = (device const float *)((device const char *) src2 + ir*args.nb20 + t_off*args.nb21 + i3*args.nb22); // {nh, nt, ns} + device const float * B = (device const float *)((device const char *) src4 + g*args.nb41 + t_off*args.nb42 + i3*args.nb43); // {d_state, ng, nt, ns} + device const float * C = (device const float *)((device const char *) src5 + g*args.nb51 + t_off*args.nb52 + i3*args.nb53); // {d_state, ng, nt, ns} - device float * y = dst + (i1 + ir*(nr) + i3*(n_t*nh*nr)); // {dim, nh, nt, ns} + device float * y = dst + (i1 + ir*nr + t_off*nh*nr + i3*(n_t_total*nh*nr)); // {dim, nh, nt, ns} for (int i2 = 0; i2 < n_t; i2 += sgptg) { threadgroup_barrier(mem_flags::mem_threadgroup); @@ -285,3 +291,183 @@ kernel void kernel_ssm_scan_f32( s_buff[i] = s; } + +typedef decltype(kernel_ssm_scan_impl) kernel_ssm_scan_t; + +template [[host_name("kernel_ssm_scan_f32")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl; +template [[host_name("kernel_ssm_scan_f32_tail")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl; + +// Chunked SSD SSM scan via Metal simdgroup MMatrix Multiply-Accumulate (simdgroup_float8x8) fast path. +// One threadgroup per (head, sequence) and tokens are processed in chunks. +// C*B^T computed in each chunk one time and reused across the head_dim channel tiles. +kernel void kernel_ssm_scan_ssd_mma_f32( + constant ggml_metal_kargs_ssm_scan & args, + device const void * src0, + device const void * src1, + device const void * src2, + device const void * src3, + device const void * src4, + device const void * src5, + device const void * src6, + device float * dst, + threadgroup float * shared [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]]) { + constexpr short CS = OP_SSM_SCAN_SSD_CS; + constexpr short TC = 8; // Tile Count of each edge in a simdgroup 8x8 tile + constexpr short HD = OP_SSM_SCAN_SSD_HD; + constexpr short NSG = OP_SSM_SCAN_SSD_NSG; + + // acs/exp(acs)/state-decay vectors, dtX[CS][HD], four private SAM row tiles [8][CS], + // and two 8x8 scratch tiles per simdgroup. Total: 26.75 KiB. + threadgroup float * shared_acs = shared; + threadgroup float * shared_exp_acs = shared + CS; + threadgroup float * shared_state_decay = shared + 2*CS; + threadgroup float * shared_dtx = shared + 3*CS; + threadgroup float * shared_sam = shared + 3*CS + CS*HD; + threadgroup float * sam_rows = shared_sam + sgitg*TC*CS; + threadgroup float * shared_tile = shared_sam + NSG*TC*CS; + threadgroup float * tile0 = shared_tile + sgitg*2*TC*TC; + threadgroup float * tile1 = tile0 + TC*TC; + + const int32_t ir = tgpig.y; // current head + const int32_t i3 = tgpig.z; // current seq + + const int32_t nc = args.d_state; + const int32_t nr = args.d_inner; + const int32_t nh = args.n_head; + const int32_t ng = args.n_group; + const int32_t n_t = args.n_seq_tokens; + const int32_t n_t_total = args.n_seq_tokens_total; + const int32_t g = ir / (nh / ng); + + device const int32_t * ids = (device const int32_t *) src6; + + device const float * s0_buff = (device const float *) ((device const char *) src0 + ir*args.nb02 + ids[i3]*args.nb03); + device float * s_buff = (device float *) ((device char *) dst + ir*args.nb02 + i3*args.nb03 + args.s_off); + + device const float * A = (device const float *) ((device const char *) src3 + ir*args.nb31); + device const float * x = (device const float *) ((device const char *) src1 + ir*args.nb11 + i3*args.nb13); + device const float * dt = (device const float *) ((device const char *) src2 + ir*args.nb20 + i3*args.nb22); + device const float * B = (device const float *) ((device const char *) src4 + g*args.nb41 + i3*args.nb43); + device const float * C = (device const float *) ((device const char *) src5 + g*args.nb51 + i3*args.nb53); + + device float * y = dst + (ir*nr + i3*(n_t_total*nh*nr)); + + for (int32_t t0 = 0; t0 < n_t; t0 += CS) { + for (int32_t idx = tiitg; idx < CS*HD; idx += NSG*N_SIMDWIDTH) { + const int32_t t = idx / HD; + const int32_t c = idx % HD; + const float dt0 = dt[(t0 + t) * (int32_t) args.ns21]; + const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0; + shared_dtx[idx] = x[(t0 + t) * (int32_t) args.ns12 + c] * dtsp; + } + if (tiitg < CS) { + const float dt0 = dt[(t0 + tiitg) * (int32_t) args.ns21]; + const float dtsp = dt0 <= 20.0f ? log(1.0f + exp(dt0)) : dt0; + shared_acs[tiitg] = dtsp * A[0]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tiitg == 0) { + float acc = 0.0f; + for (short t = 0; t < CS; ++t) { + acc += shared_acs[t]; + shared_acs[t] = acc; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg < CS) { + shared_exp_acs[tiitg] = exp(shared_acs[tiitg]); + shared_state_decay[tiitg] = exp(shared_acs[CS - 1] - shared_acs[tiitg]); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + device const float * state = t0 == 0 ? s0_buff : s_buff; + + // Build one 8x64 row tile of SAM per simdgroup, then reuse it across every channel tile. + for (short ib = sgitg; ib < CS/TC; ib += NSG) { + for (short jb = 0; jb <= ib; ++jb) { + simdgroup_float8x8 cb = make_filled_simdgroup_matrix(0.0f); + + for (int32_t k0 = 0; k0 < nc; k0 += TC) { + simdgroup_float8x8 mc; + simdgroup_float8x8 mb; + simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52); + simdgroup_load(mb, B + (t0 + jb*TC)*(int32_t) args.ns42 + k0, args.ns42, 0, true); + simdgroup_multiply_accumulate(cb, mc, mb, cb); + } + + threadgroup float * sam = sam_rows + jb*TC; + simdgroup_store(cb, sam, CS); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) { + const short ri = e / TC; + const short rj = e % TC; + const short i = ib*TC + ri; + const short j = jb*TC + rj; + sam[ri*CS + rj] = j <= i ? + sam[ri*CS + rj] * exp(shared_acs[i] - shared_acs[j]) : 0.0f; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + + for (short ch = 0; ch < HD/TC; ++ch) { + simdgroup_float8x8 y_diag = make_filled_simdgroup_matrix(0.0f); + simdgroup_float8x8 y_inter = make_filled_simdgroup_matrix(0.0f); + + for (short jb = 0; jb <= ib; ++jb) { + simdgroup_float8x8 sam; + simdgroup_float8x8 mdtx; + simdgroup_load(sam, sam_rows + jb*TC, CS); + simdgroup_load(mdtx, shared_dtx + jb*TC*HD + ch*TC, HD); + simdgroup_multiply_accumulate(y_diag, sam, mdtx, y_diag); + } + + for (int32_t k0 = 0; k0 < nc; k0 += TC) { + simdgroup_float8x8 mc; + simdgroup_float8x8 ms; + simdgroup_load(mc, C + (t0 + ib*TC)*(int32_t) args.ns52 + k0, args.ns52); + simdgroup_load(ms, state + ch*TC*nc + k0, nc, 0, true); + simdgroup_multiply_accumulate(y_inter, mc, ms, y_inter); + } + + simdgroup_store(y_diag, tile0, TC); + simdgroup_store(y_inter, tile1, TC); + simdgroup_barrier(mem_flags::mem_threadgroup); + for (short e = tiisg; e < TC*TC; e += N_SIMDWIDTH) { + const short ri = e / TC; + const short ci = e % TC; + const int32_t token = t0 + ib*TC + ri; + y[token*nh*nr + ch*TC + ci] = + tile0[e] + shared_exp_acs[ib*TC + ri] * tile1[e]; + } + simdgroup_barrier(mem_flags::mem_threadgroup); + } + } + + // All simdgroups must finish reading s_buff before any thread overwrites it. + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + + // Keep the carried-state reduction in token order. Reassociating this particular product + // with MMA compounds rounding differences at every chunk boundary; CB, y_diag, and C*S + // remain on the matrix unit. + const float chunk_decay = exp(shared_acs[CS - 1]); + for (int32_t idx = tiitg; idx < nc*HD; idx += NSG*N_SIMDWIDTH) { + const int32_t ci = idx / nc; + const int32_t si = idx % nc; + float state_c = 0.0f; + for (short t = 0; t < CS; ++t) { + state_c += shared_state_decay[t] * + B[(t0 + t)*(int32_t) args.ns42 + si] * + shared_dtx[t*HD + ci]; + } + s_buff[idx] = chunk_decay * state[idx] + state_c; + } + + // All state tiles must be visible before the next chunk consumes s_buff as S_prev. + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + } +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 99014df09..53e93a144 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4120,9 +4120,10 @@ struct test_ssm_scan : public test_case { const int64_t n_seqs; const bool xbc_overlap; const int64_t K; + const bool weak_decay; std::string vars() override { - return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K); + return VARS_TO_STR10(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K, weak_decay); } test_ssm_scan(ggml_type type = GGML_TYPE_F32, @@ -4133,8 +4134,9 @@ struct test_ssm_scan : public test_case { int64_t n_seq_tokens = 32, int64_t n_seqs = 32, bool xbc_overlap = false, - int64_t K = 1) - : 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), K(K) {} + int64_t K = 1, + bool weak_decay = 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), K(K), weak_decay(weak_decay) {} double max_nmse_err() override { // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32. @@ -4187,7 +4189,7 @@ struct test_ssm_scan : public test_case { continue; } 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); + init_tensor_uniform(t, weak_decay ? -0.02f : -1.0f, weak_decay ? -0.005f : -0.5f); } else { init_tensor_uniform(t); } @@ -9111,6 +9113,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 64, 4)); // Metal SSD one chunk MMA only, no seq tail + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 65, 2)); // SSD one chunk + 1-token sequential tail + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2)); // SSD multi-chunk, no tail (exercises the chunk-to-chunk state handoff) + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 128, 2, false, /*K=*/1, /*weak_decay=*/true)); // SSD multi-chunk, carried state not numerically negligible 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));