mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-09-07 21:41:19 +02:00
ggml-metal: add chunked SSD MMA for Mamba-2 prefill optimization (#26647)
* metal: WIP chunked SSD SSM_SCAN kernels for multi-token prefill * metal: drop scalar SSD path; MMA + sequential tail * drop WIP ssm scan test noise * remove state_from_dst and rename CS and NSG constants * remove unrelated added whitespace padding * added clarity to mma_tokens calculation * added clarity to use_mma bool checks * added comments to metal ssd op constants for clarity * reserve K tokens for sequential kernel rollback snapshots * reset concurrency between mma and seq tail * remove print args no longer used * fixed comment to no longer point to specific line * add FC_SSM_SCAN so seq path skips token offlset unless it's mma tail * added changes to new ssm.metal for rebase after ggml-metal.metal refactor * specialize ssm_scan tail with a template instead of a function constant --------- Co-authored-by: dpantaleoni <dominikpantaleoni@gmail.com> Co-authored-by: forforever73 <690105611@qq.com>
This commit is contained in:
Dominik Pantaleoni
committed by
GitHub
GitHub
parent
5d5cb4c3a4
commit
11cd988428
@@ -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];
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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<bool TAIL>
|
||||
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<false>) kernel_ssm_scan_t;
|
||||
|
||||
template [[host_name("kernel_ssm_scan_f32")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl<false>;
|
||||
template [[host_name("kernel_ssm_scan_f32_tail")]] kernel kernel_ssm_scan_t kernel_ssm_scan_impl<true>;
|
||||
|
||||
// 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<float, 8>(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<float, 8>(0.0f);
|
||||
simdgroup_float8x8 y_inter = make_filled_simdgroup_matrix<float, 8>(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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<std::unique_ptr<test_case>> 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));
|
||||
|
||||
Reference in New Issue
Block a user