const bool kda = (neg0 == S_v);
- // scratch layout per thread: [delta(S_v)]
- const int64_t scratch_per_thread = S_v;
+ // state is 3D (S_v*S_v*H, K, n_seqs); K is the snapshot slot count.
+ const int64_t K = src_state->ne[1];
+ GGML_ASSERT(K >= 1);
+ // per-seq stride in floats (slot 0 of seq s lives at state + s * seq_stride)
+ const int64_t state_seq_stride = src_state->nb[2] / sizeof(float);
+
+ const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0);
const int ith = params->ith;
- float * delta = (float *)params->wdata + ith * scratch_per_thread + CACHE_LINE_SIZE_F32;
+ float * delta = (float *)params->wdata + ith * per_thread + CACHE_LINE_SIZE_F32;
+ float * state_work = K > 1 ? (delta + S_v) : nullptr;
// output layout: [attn_scores | new_states]
- // attn_scores: S_v * H * n_tokens * n_seqs floats
- // new_states: S_v * S_v * H * n_seqs floats
- const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
+ // attn_scores: S_v * H * n_tokens * n_seqs floats
+ // new_states: S_v * S_v * H * n_seqs * K floats (K snapshot slots; last min(n_tokens, K))
+ const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
+ const int64_t state_size_per_snap = S_v * S_v * H * n_seqs;
float * attn_out_base = (float *)dst->data;
float * state_out_base = (float *)dst->data + attn_score_elems;
+ // snapshot slot mapping: target_slot = t - shift. When n_tokens < K only the last
+ // n_tokens slots are written; earlier slots are left untouched (caller-owned).
+ const int64_t shift = n_tokens - K;
+
const float * state_in_base = (const float *)src_state->data;
//const int64_t rq1 = nev1 / neq1;
const int64_t iq3 = iv3 / rq3;
const int64_t ik3 = iv3 / rk3;
- float * s_out = state_out_base + (iv3 * H + iv1) * S_v * S_v;
+ // For K=1, write directly to the single output slot to avoid an extra memcpy at the end.
+ // For K>1, work in scratch and copy out per-token when the slot is in range.
+ float * s_out = (K > 1)
+ ? state_work
+ : state_out_base + (iv3 * H + iv1) * S_v * S_v;
- // copy input state into output buffer and operate in-place
- const float * s_in = state_in_base + (iv3 * H + iv1) * S_v * S_v;
+ // copy input state into the working buffer and operate in-place
+ // state layout (D, K, n_seqs): slot 0 of seq iv3 starts at iv3 * state_seq_stride.
+ const float * s_in = state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v;
memcpy(s_out, s_in, S_v * S_v * sizeof(float));
// attn output pointer for first token of this (head, seq)
}
attn_data += S_v * H; // advance to next token
+
+ if (K > 1) {
+ const int64_t target_slot = t - shift;
+ if (target_slot >= 0 && target_slot < K) {
+ float * curr_state_o = state_out_base + target_slot * state_size_per_snap +
+ (iv3 * H + iv1) * S_v * S_v;
+ memcpy(curr_state_o, s_out, S_v * S_v * sizeof(float));
+ }
+ }
}
}
}
#include "gated_delta_net.cuh"
-template <int S_v, bool KDA>
+template <int S_v, bool KDA, bool keep_rs_t>
__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2)
gated_delta_net_cuda(const float * q,
const float * k,
int64_t sb3,
const uint3 neqk1_magic,
const uint3 rq3_magic,
- float scale) {
+ float scale,
+ int K) {
const uint32_t h_idx = blockIdx.x;
const uint32_t sequence = blockIdx.y;
// each warp owns one column, using warp-level primitives to reduce across rows
float * attn_data = dst;
float * state = dst + attn_score_elems;
- const int64_t state_offset = (sequence * H + h_idx) * S_v * S_v;
- state += state_offset;
- curr_state += state_offset + col * S_v;
+ // input state layout (D, K, n_seqs) — seq stride is K * D = K * H * S_v * S_v.
+ // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
+ const int64_t state_in_offset = sequence * K * H * S_v * S_v + h_idx * S_v * S_v;
+ const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v;
+ const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
+ state += state_out_offset;
+ curr_state += state_in_offset + col * S_v;
attn_data += (sequence * n_tokens * H + h_idx) * S_v;
constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
s_shard[r] = curr_state[i];
}
+ // slot mapping: target_slot = t - shift. When n_tokens < K only the last n_tokens slots
+ // are written; earlier slots are left untouched (caller-owned).
+ const int shift = (int) n_tokens - K;
+
for (int t = 0; t < n_tokens; t++) {
const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1;
const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1;
}
attn_data += S_v * H;
+
+ if constexpr (keep_rs_t) {
+ const int target_slot = t - shift;
+ if (target_slot >= 0 && target_slot < K) {
+ float * curr_state = (dst + attn_score_elems) + target_slot * state_size_per_token + state_out_offset;
+#pragma unroll
+ for (int r = 0; r < rows_per_lane; r++) {
+ const int i = r * warp_size + lane;
+ curr_state[col * S_v + i] = s_shard[r];
+ }
+ }
+ }
}
- // Write state back to global memory (transposed layout)
+ if constexpr (!keep_rs_t) {
#pragma unroll
- for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
- state[col * S_v + i] = s_shard[r];
+ for (int r = 0; r < rows_per_lane; r++) {
+ const int i = r * warp_size + lane;
+ state[col * S_v + i] = s_shard[r];
+ }
}
}
-template <bool KDA>
+template <bool KDA, bool keep_rs_t>
static void launch_gated_delta_net(
const float * q_d, const float * k_d, const float * v_d,
const float * g_d, const float * b_d, const float * s_d,
int64_t sv1, int64_t sv2, int64_t sv3,
int64_t sb1, int64_t sb2, int64_t sb3,
int64_t neqk1, int64_t rq3,
- float scale, cudaStream_t stream) {
+ float scale, int K, cudaStream_t stream) {
//TODO: Add chunked kernel for even faster pre-fill
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
const int num_warps = 4;
switch (S_v) {
case 16:
- gated_delta_net_cuda<16, KDA><<<grid_dims, block_dims, 0, stream>>>(
+ gated_delta_net_cuda<16, KDA, keep_rs_t><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
break;
case 32:
- gated_delta_net_cuda<32, KDA><<<grid_dims, block_dims, 0, stream>>>(
+ gated_delta_net_cuda<32, KDA, keep_rs_t><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
break;
case 64: {
- gated_delta_net_cuda<64, KDA><<<grid_dims, block_dims, 0, stream>>>(
+ gated_delta_net_cuda<64, KDA, keep_rs_t><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
break;
}
case 128: {
- gated_delta_net_cuda<128, KDA><<<grid_dims, block_dims, 0, stream>>>(
+ gated_delta_net_cuda<128, KDA, keep_rs_t><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
+ sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
break;
}
default:
cudaStream_t stream = ctx.stream();
+ // state is 3D (S_v*S_v*H, K, n_seqs); K is the snapshot slot count.
+ const int K = (int) src_state->ne[1];
+ const bool keep_rs = K > 1;
+
if (kda) {
- launch_gated_delta_net<true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
- S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1, rq3, scale, stream);
+ if (keep_rs) {
+ launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
+ sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+ } else {
+ launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
+ sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+ }
} else {
- launch_gated_delta_net<false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
- S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
- sb1, sb2, sb3, neqk1, rq3, scale, stream);
+ if (keep_rs) {
+ launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
+ sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+ } else {
+ launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+ S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
+ sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+ }
}
}
constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]];
constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]];
+constant short FC_gated_delta_net_K [[function_constant(FC_GATED_DELTA_NET + 2)]];
#if 1
template<short NSG>
uint3 ntg[[threads_per_threadgroup]]) {
#define S_v FC_gated_delta_net_ne20
#define G FC_gated_delta_net_ne30
+#define K FC_gated_delta_net_K
const uint tx = tpitg.x;
const uint ty = tpitg.y;
- const uint i23 = tgpig.z; // B
- const uint i21 = tgpig.y; // H
- const uint i20 = tgpig.x*NSG + ty;
+ const uint i23 = tgpig.z; // B (n_seqs)
+ const uint i21 = tgpig.y; // H (head)
+ const uint i20 = tgpig.x*NSG + ty; // row within S_v
const uint i01 = i21 % args.ne01;
const uint i11 = i21 % args.ne11;
const float scale = 1.0f / sqrt((float)S_v);
+ // input state layout (D, K, n_seqs): per-seq stride is K*H*D; we read slot 0.
// state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous
- device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
+ const uint state_in_base = (i23*K*args.ne21 + i21)*S_v*S_v + i20*S_v;
+ device const float * s_ptr = (device const float *) (s) + state_in_base;
float ls[NSG];
device const float * b_ptr = (device const float *) (b) + (i23*args.ne22*args.ne21 + i21);
device const float * g_ptr = (device const float *) (g) + (i23*args.ne22*args.ne21 + i21)*G;
+ // snapshot slot mapping: target_slot = t - shift. When n_tokens < K, only the last
+ // n_tokens slots are written; earlier slots are left untouched (caller-owned).
+ const int shift = (int)args.ne22 - (int)K;
+
+ // output state base offset: after attention scores
+ const uint attn_size = args.ne22 * args.ne21 * S_v * args.ne23;
+ // output state per-slot size: S_v * S_v * H * n_seqs
+ const uint state_size_per_snap = S_v * S_v * args.ne21 * args.ne23;
+ // per-(seq,head) offset within a slot
+ const uint state_out_base = (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
+
for (short t = 0; t < args.ne22; t++) {
float s_k = 0.0f;
b_ptr += args.ne21;
g_ptr += args.ne21*G;
- }
- device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20*S_v;
+ if (K > 1u) {
+ const int target_slot = (int)t - shift;
+ if (target_slot >= 0 && target_slot < (int)K) {
+ device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base;
+ FOR_UNROLL (short j = 0; j < NSG; j++) {
+ const short is = tx*NSG + j;
+ dst_state[is] = ls[j];
+ }
+ }
+ }
+ }
- FOR_UNROLL (short j = 0; j < NSG; j++) {
- const short is = tx*NSG + j;
- dst_state[is] = ls[j];
+ if (K == 1u) {
+ device float * dst_state = (device float *) (dst) + attn_size + state_out_base;
+ FOR_UNROLL (short j = 0; j < NSG; j++) {
+ const short is = tx*NSG + j;
+ dst_state[is] = ls[j];
+ }
}
#undef S_v
#undef G
+#undef K
}
typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t;
overview / pkg / ggml / sources / whisper.cpp / commitdiff