Skip to content

perf(cuda): add in-place GDN state write #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dusterbloom wants to merge 1 commit into
base: luce-dflash
Choose a base branch
from
Open

perf(cuda): add in-place GDN state write #24

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2547,6 +2547,15 @@ extern "C" {
struct ggml_tensor * beta,
struct ggml_tensor * state);

GGML_API struct ggml_tensor * ggml_gated_delta_net_inplace(
struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * v,
struct ggml_tensor * g,
struct ggml_tensor * beta,
struct ggml_tensor * state);

GGML_API void ggml_gated_delta_net_set_skip_intermediate(
struct ggml_tensor * tensor,
bool skip_intermediate);
Expand Down
34 changes: 20 additions & 14 deletions ggml/src/ggml-cuda/gated_delta_net.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ gated_delta_net_cuda(const float * q,
const float * beta,
const float * curr_state,
float * dst,
float * state_out,
const int * parent_ids, // TREE_MODE only; else ignored
InterT * persist_inter, // optional external buffer for per-token intermediates
bool skip_intermediate,
Expand Down Expand Up @@ -81,7 +82,7 @@ gated_delta_net_cuda(const float * q,
const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
const int64_t final_state_elems = S_v * S_v * H * n_seqs;
float * attn_data = dst;
float * state = dst + attn_score_elems;
float * state = state_out;
// intermediate_states region: one S_v*S_v*H*n_seqs state per token. Written
// inside the token loop below (one state per `t`) to enable spec-decode
// rollback without a replay forward pass. See ggml.c::ggml_gated_delta_net.
Expand Down Expand Up @@ -275,6 +276,7 @@ gated_delta_net_cuda_grouped_cols(const float * q,
const float * beta,
const float * curr_state,
float * dst,
float * state_out,
InterT * persist_inter,
bool skip_intermediate,
int64_t H,
Expand Down Expand Up @@ -318,7 +320,7 @@ gated_delta_net_cuda_grouped_cols(const float * q,
const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
const int64_t final_state_elems = S_v * S_v * H * n_seqs;
float * attn_data = dst;
float * state = dst + attn_score_elems;
float * state = state_out;
InterT * inter_states = persist_inter
? persist_inter
: (InterT *)(dst + attn_score_elems + final_state_elems);
Expand Down Expand Up @@ -452,6 +454,7 @@ 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,
float * dst_d,
float * state_out_d,
const int * parent_ids_d,
InterT * persist_inter_d,
int64_t S_v, int64_t H, int64_t n_tokens, int64_t n_seqs,
Expand All @@ -475,19 +478,19 @@ static void launch_gated_delta_net(
switch (S_v) {
case 16:
gated_delta_net_cuda<16, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
case 32:
gated_delta_net_cuda<32, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
case 64: {
gated_delta_net_cuda<64, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
Expand All @@ -505,32 +508,32 @@ static void launch_gated_delta_net(
dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp));
dim3 grouped_block_dims(32, column_groups_per_block, 1);
gated_delta_net_cuda_grouped_cols<128, cols, width, 32, InterT><<<grouped_grid_dims, grouped_block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
} else if (warp_size == 64) {
constexpr int groups_per_warp = 64 / width;
dim3 grouped_grid_dims(H, n_seqs, (groups + column_groups_per_block * groups_per_warp - 1) / (column_groups_per_block * groups_per_warp));
dim3 grouped_block_dims(64, column_groups_per_block, 1);
gated_delta_net_cuda_grouped_cols<128, cols, width, 64, InterT><<<grouped_grid_dims, grouped_block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
} else {
gated_delta_net_cuda<128, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
}
} else {
gated_delta_net_cuda<128, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
}
} else {
gated_delta_net_cuda<128, KDA, TREE_MODE, InterT><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_inter_d, skip_intermediate, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
}
Expand Down Expand Up @@ -586,6 +589,9 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor *

const float * s_d = (const float *) src_state->data;
float * dst_d = (float *) dst->data;
const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
const bool inplace_state = ggml_get_op_params_i32(dst, 1) != 0;
float * state_out_d = inplace_state ? (float *) src_state->data : dst_d + attn_score_elems;
const int * parent_ids_d = src_parent
? (const int *) src_parent->data
: nullptr;
Expand Down Expand Up @@ -640,24 +646,24 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor *
if (kda) { \
if (tree_mode) { \
launch_gated_delta_net<true, true, INTER_T>( \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_typed, \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_typed, \
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \
sb1, sb2, sb3, neqk1, rq3, skip_intermediate, scale, stream); \
} else { \
launch_gated_delta_net<true, false, INTER_T>( \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, nullptr, persist_typed, \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, nullptr, persist_typed, \
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \
sb1, sb2, sb3, neqk1, rq3, skip_intermediate, scale, stream); \
} \
} else { \
if (tree_mode) { \
launch_gated_delta_net<false, true, INTER_T>( \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, parent_ids_d, persist_typed, \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, parent_ids_d, persist_typed, \
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \
sb1, sb2, sb3, neqk1, rq3, skip_intermediate, scale, stream); \
} else { \
launch_gated_delta_net<false, false, INTER_T>( \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, nullptr, persist_typed, \
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_out_d, nullptr, persist_typed, \
S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, \
sb1, sb2, sb3, neqk1, rq3, skip_intermediate, scale, stream); \
} \
Expand Down
14 changes: 14 additions & 0 deletions ggml/src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6262,6 +6262,7 @@ struct ggml_tensor * ggml_gated_delta_net(
// roll back SSM state to the accepted prefix without a full replay forward pass.
const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs + S_v * n_tokens * n_seqs, 1, 1 };
struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
ggml_set_op_params_i32(result, 1, 0);

result->op = GGML_OP_GATED_DELTA_NET;
result->src[0] = q;
Expand All @@ -6274,6 +6275,19 @@ struct ggml_tensor * ggml_gated_delta_net(
return result;
}

struct ggml_tensor * ggml_gated_delta_net_inplace(
struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * v,
struct ggml_tensor * g,
struct ggml_tensor * beta,
struct ggml_tensor * state) {
struct ggml_tensor * result = ggml_gated_delta_net(ctx, q, k, v, g, beta, state);
ggml_set_op_params_i32(result, 1, 1);
return result;
}

void ggml_gated_delta_net_set_skip_intermediate(
struct ggml_tensor * tensor,
bool skip_intermediate) {
Expand Down
Loading