vulkan: add GATED_DELTA_NET op support (#20334)
* vulkan: add GATED_DELTA_NET op support Implements the fused gated delta net recurrence as a Vulkan compute shader with full support for scalar gate, KDA vector gate, GQA broadcast, multi-token sequences, and permuted (non-contiguous) q/k inputs. Specialization constants select head size (32/64/128) and KDA mode at pipeline creation time. Passes all 13 test-backend-ops cases on AMD Radeon 890M (RADV GFX1150). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * vulkan: optimize GATED_DELTA_NET shader (Phase 1) - vec4 dot products on all inner loops (dp4 hardware intrinsic) - Cache exp(g) in shared memory for KDA path, eliminating ~32K redundant global reads and ~16K redundant exp() calls per token - vec4 fused decay + rank-1 update (3 vec4 ops vs 12 scalar ops) - Add perf benchmark cases for GATED_DELTA_NET to test-backend-ops KDA TG: +5.4% throughput. Non-KDA: no regressions. 13/13 test-backend-ops passing on AMD Radeon 890M (RADV GFX1150). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * vulkan: address review feedback for GATED_DELTA_NET Pipeline array refactor [3][2], A_TYPE/D_TYPE/FLOAT_TYPE shader macros, scale in push constants, supports_op fix, dispatch restructuring. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * vulkan: use FLOAT_TYPE for buffer/shared declarations, align formatting Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * vulkan: add explicit FLOAT_TYPE casts for buffer loads Wrap data_q, data_k, and data_g buffer reads with FLOAT_TYPE() casts to ensure correct behavior across all Vulkan configurations. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * vulkan: fix Q/K broadcast for interleaved head layout Adapt to the interleaved broadcast convention from #20340: head_id / rq1 → head_id % neq1 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Progeny Alpha <ProgenyAlpha@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
ProgenyAlpha
co-authored by
Claude Opus 4.6
Progeny Alpha
Claude Opus 4.6
Progeny Alpha
parent
c3e3f9e533
commit
deee23863b
@@ -0,0 +1,128 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
layout(constant_id = 0) const uint S_V = 128;
|
||||
layout(constant_id = 1) const uint KDA = 0;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform Parameters {
|
||||
uint H;
|
||||
uint n_tokens;
|
||||
uint n_seqs;
|
||||
uint s_off;
|
||||
uint sq1, sq2, sq3;
|
||||
uint sv1, sv2, sv3;
|
||||
uint sb1, sb2, sb3;
|
||||
uint neq1, rq3;
|
||||
float scale;
|
||||
};
|
||||
|
||||
layout(binding = 0) readonly buffer QBuf { FLOAT_TYPE data_q[]; };
|
||||
layout(binding = 1) readonly buffer KBuf { FLOAT_TYPE data_k[]; };
|
||||
layout(binding = 2) readonly buffer VBuf { FLOAT_TYPE data_v[]; };
|
||||
layout(binding = 3) readonly buffer GBuf { FLOAT_TYPE data_g[]; };
|
||||
layout(binding = 4) readonly buffer BetaBuf { FLOAT_TYPE data_beta[]; };
|
||||
layout(binding = 5) readonly buffer StateBuf { FLOAT_TYPE data_state[]; };
|
||||
layout(binding = 6) buffer DstBuf { FLOAT_TYPE data_dst[]; };
|
||||
|
||||
shared FLOAT_TYPE s_k[S_V];
|
||||
shared FLOAT_TYPE s_q[S_V];
|
||||
shared FLOAT_TYPE s_g[S_V]; // KDA only: cached exp(g[i])
|
||||
|
||||
void main() {
|
||||
const uint head_id = gl_WorkGroupID.x;
|
||||
const uint seq_id = gl_WorkGroupID.y;
|
||||
const uint col = gl_LocalInvocationID.x;
|
||||
|
||||
const uint iq1 = head_id % neq1;
|
||||
const uint iq3 = seq_id / rq3;
|
||||
|
||||
const uint state_size = S_V * S_V;
|
||||
const uint state_base = (seq_id * H + head_id) * state_size;
|
||||
|
||||
FLOAT_TYPE state[S_V];
|
||||
[[unroll]] for (uint i = 0; i < S_V; i++) {
|
||||
state[i] = FLOAT_TYPE(data_state[state_base + i * S_V + col]);
|
||||
}
|
||||
|
||||
uint attn_off = (seq_id * n_tokens * H + head_id) * S_V;
|
||||
|
||||
for (uint t = 0; t < n_tokens; t++) {
|
||||
const uint q_off = iq3 * sq3 + t * sq2 + iq1 * sq1;
|
||||
const uint k_off = q_off;
|
||||
const uint v_off = seq_id * sv3 + t * sv2 + head_id * sv1;
|
||||
|
||||
s_q[col] = FLOAT_TYPE(data_q[q_off + col]);
|
||||
s_k[col] = FLOAT_TYPE(data_k[k_off + col]);
|
||||
|
||||
const uint gb_off = seq_id * sb3 + t * sb2 + head_id * sb1;
|
||||
|
||||
if (KDA != 0) {
|
||||
const uint g_base = gb_off * S_V;
|
||||
s_g[col] = exp(FLOAT_TYPE(data_g[g_base + col]));
|
||||
}
|
||||
|
||||
barrier();
|
||||
|
||||
const FLOAT_TYPE v_val = FLOAT_TYPE(data_v[v_off + col]);
|
||||
const FLOAT_TYPE beta_val = FLOAT_TYPE(data_beta[gb_off]);
|
||||
|
||||
if (KDA == 0) {
|
||||
const FLOAT_TYPE g_val = exp(FLOAT_TYPE(data_g[gb_off]));
|
||||
|
||||
FLOAT_TYPE kv_col = 0.0;
|
||||
[[unroll]] for (uint i = 0; i < S_V; i += 4) {
|
||||
kv_col += dot(
|
||||
vec4(state[i], state[i+1], state[i+2], state[i+3]),
|
||||
vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3])
|
||||
);
|
||||
}
|
||||
|
||||
FLOAT_TYPE delta_col = (v_val - g_val * kv_col) * beta_val;
|
||||
|
||||
FLOAT_TYPE attn_col = 0.0;
|
||||
[[unroll]] for (uint i = 0; i < S_V; i += 4) {
|
||||
vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
|
||||
vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
|
||||
sv = g_val * sv + kv * delta_col;
|
||||
state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w;
|
||||
|
||||
attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3]));
|
||||
}
|
||||
|
||||
data_dst[attn_off + col] = attn_col * scale;
|
||||
} else {
|
||||
FLOAT_TYPE kv_col = 0.0;
|
||||
[[unroll]] for (uint i = 0; i < S_V; i += 4) {
|
||||
vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]);
|
||||
vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
|
||||
vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
|
||||
kv_col += dot(gv * sv, kv);
|
||||
}
|
||||
|
||||
FLOAT_TYPE delta_col = (v_val - kv_col) * beta_val;
|
||||
|
||||
FLOAT_TYPE attn_col = 0.0;
|
||||
[[unroll]] for (uint i = 0; i < S_V; i += 4) {
|
||||
vec4 gv = vec4(s_g[i], s_g[i+1], s_g[i+2], s_g[i+3]);
|
||||
vec4 sv = vec4(state[i], state[i+1], state[i+2], state[i+3]);
|
||||
vec4 kv = vec4(s_k[i], s_k[i+1], s_k[i+2], s_k[i+3]);
|
||||
sv = gv * sv + kv * delta_col;
|
||||
state[i] = sv.x; state[i+1] = sv.y; state[i+2] = sv.z; state[i+3] = sv.w;
|
||||
|
||||
attn_col += dot(sv, vec4(s_q[i], s_q[i+1], s_q[i+2], s_q[i+3]));
|
||||
}
|
||||
|
||||
data_dst[attn_off + col] = attn_col * scale;
|
||||
}
|
||||
|
||||
attn_off += S_V * H;
|
||||
barrier();
|
||||
}
|
||||
|
||||
[[unroll]] for (uint i = 0; i < S_V; i++) {
|
||||
data_dst[s_off + state_base + i * S_V + col] = state[i];
|
||||
}
|
||||
}
|
||||
@@ -987,6 +987,8 @@ void process_shaders() {
|
||||
|
||||
string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
|
||||
|
||||
string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}}));
|
||||
|
||||
string_to_spv("opt_step_adamw_f32", "opt_step_adamw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
|
||||
string_to_spv("opt_step_sgd_f32", "opt_step_sgd.comp", merge_maps(base_dict, {{"A_TYPE", "float"}}));
|
||||
|
||||
|
||||
Reference in New Issue
Block a user