vulkan: add DeepSeek-V4 hyper-connection fused ops (DSV4_HC_COMB/PRE/POST) (#26578)
* vulkan: add DeepSeek-V4 hyper-connection fused ops (DSV4_HC_COMB/PRE/POST) CUDA has these ops from the DeepSeek-V4 merge and Metal gained them in PR 26459. Vulkan was the last major backend running the unfused primitive chain. On DeepSeek-V4-Flash the unfused Sinkhorn comb chain alone takes about 32% of decode op time on gfx1151 (Strix Halo), spread over roughly 16k dispatches per token. dsv4_hc_comb runs the full 20-iteration Sinkhorn in registers. A token's 4x4 comb matrix lives in 16 consecutive subgroup lanes, with idst in bits 0-1 and isrc in bits 2-3 to match the CPU reference layout, so subgroupShuffleXor by 1|2 reduces rows and by 4|8 reduces columns. One dispatch replaces about 137 strictly ordered node executions per site. The shuffle masks never cross a 16-lane boundary, so a subgroup of size 64 packs 4 independent tokens. dsv4_hc_pre and dsv4_hc_post handle the elementwise stream collapse and fan-out, with per-token coefficients staged in shared memory. GGML_VK_DISABLE_DSV4_HC disables all three ops. The _COMB, _PRE and _POST variants gate each op independently so a single kernel can be bisected against the unfused graph. Adds eval cases at the production n_iter=20 across batch sizes that cross subgroup and workgroup boundaries. * vulkan: dsv4 hc review fixes Drop the per-op env-var disables and device flags, the stride divisibility check (ggml guarantees it) and the workgroup-count fallback in supports_op. Trim the comb shader comments to the lane layout. --------- Co-authored-by: Kevin Hopper <no-reply@maestro.press>
This commit is contained in:
Kevin Hopper
co-authored by
Kevin Hopper
Kevin Hopper
parent
0c963452ea
commit
7a333e7240
@@ -1110,6 +1110,9 @@ struct vk_device_struct {
|
||||
vk_pipeline pipeline_cumsum_multipass2_f32;
|
||||
vk_pipeline pipeline_argmax_f32;
|
||||
vk_pipeline pipeline_count_equal_i32;
|
||||
vk_pipeline pipeline_dsv4_hc_comb_f32;
|
||||
vk_pipeline pipeline_dsv4_hc_pre_f32;
|
||||
vk_pipeline pipeline_dsv4_hc_post_f32;
|
||||
std::map<vk_solve_tri_pipeline_state, vk_pipeline> pipeline_solve_tri_f32;
|
||||
vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16;
|
||||
vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16;
|
||||
@@ -1467,6 +1470,53 @@ struct vk_op_fwht_push_constants {
|
||||
float scale;
|
||||
};
|
||||
|
||||
struct vk_op_dsv4_hc_comb_push_constants {
|
||||
uint32_t n_tokens;
|
||||
|
||||
uint32_t nbm0; uint32_t nbm1;
|
||||
uint32_t nbs0;
|
||||
uint32_t nbb0;
|
||||
uint32_t nbd0; uint32_t nbd1; uint32_t nbd2;
|
||||
|
||||
uint32_t m_offset;
|
||||
uint32_t s_offset;
|
||||
uint32_t b_offset;
|
||||
uint32_t d_offset;
|
||||
|
||||
float eps;
|
||||
uint32_t n_iter;
|
||||
};
|
||||
|
||||
struct vk_op_dsv4_hc_pre_push_constants {
|
||||
uint32_t n_embd;
|
||||
uint32_t n_tokens;
|
||||
|
||||
uint32_t nbx0; uint32_t nbx1; uint32_t nbx2;
|
||||
uint32_t nbw0; uint32_t nbw1;
|
||||
uint32_t nbd0; uint32_t nbd1;
|
||||
|
||||
uint32_t x_offset;
|
||||
uint32_t w_offset;
|
||||
uint32_t d_offset;
|
||||
};
|
||||
|
||||
struct vk_op_dsv4_hc_post_push_constants {
|
||||
uint32_t n_embd;
|
||||
uint32_t n_tokens;
|
||||
|
||||
uint32_t nbx0; uint32_t nbx1;
|
||||
uint32_t nbr0; uint32_t nbr1; uint32_t nbr2;
|
||||
uint32_t nbp0; uint32_t nbp1;
|
||||
uint32_t nbc0; uint32_t nbc1; uint32_t nbc2;
|
||||
uint32_t nbd0; uint32_t nbd1; uint32_t nbd2;
|
||||
|
||||
uint32_t x_offset;
|
||||
uint32_t r_offset;
|
||||
uint32_t p_offset;
|
||||
uint32_t c_offset;
|
||||
uint32_t d_offset;
|
||||
};
|
||||
|
||||
struct vk_op_count_experts_push_constants {
|
||||
uint32_t ne00;
|
||||
uint32_t ne01;
|
||||
@@ -2631,6 +2681,32 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk
|
||||
GGML_UNUSED(src3);
|
||||
}
|
||||
|
||||
template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_comb_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
|
||||
p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
|
||||
p.s_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
|
||||
p.b_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type);
|
||||
p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
|
||||
|
||||
GGML_UNUSED(src3);
|
||||
}
|
||||
|
||||
template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_pre_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
|
||||
p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
|
||||
p.w_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
|
||||
p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
|
||||
|
||||
GGML_UNUSED(src2);
|
||||
GGML_UNUSED(src3);
|
||||
}
|
||||
|
||||
template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_post_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
|
||||
p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
|
||||
p.r_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
|
||||
p.p_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type);
|
||||
p.c_offset = get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type);
|
||||
p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
|
||||
}
|
||||
|
||||
struct ggml_backend_vk_buffer_context {
|
||||
vk_device_ref device;
|
||||
vk_buffer dev_buffer;
|
||||
@@ -5977,6 +6053,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true);
|
||||
}
|
||||
|
||||
// comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it
|
||||
// needs at least 16 lanes, pinned to a known size.
|
||||
if (device->subgroup_basic && device->subgroup_shuffle && device->subgroup_require_full_support && device->subgroup_size >= 16) {
|
||||
const uint32_t tokens_per_workgroup = 4 * (device->subgroup_size / 16);
|
||||
ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, sizeof(vk_op_dsv4_hc_comb_push_constants), {tokens_per_workgroup, 1, 1}, { device->subgroup_size }, 1, true, true, device->subgroup_size);
|
||||
}
|
||||
|
||||
ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, { 256 }, 1);
|
||||
ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256 }, 1);
|
||||
|
||||
for (auto &s : device->pipeline_solve_tri_f32) {
|
||||
const vk_solve_tri_pipeline_state &state = s.first;
|
||||
|
||||
@@ -10204,6 +10290,98 @@ static void ggml_vk_fwht(ggml_backend_vk_context * ctx, vk_context& subctx, cons
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, dst_buf }, pc, { workgroups_x, 1, 1 });
|
||||
}
|
||||
|
||||
static uint32_t ggml_vk_nb_elem(const ggml_tensor * t, int i) {
|
||||
return (uint32_t)(t->nb[i] / ggml_type_size(t->type));
|
||||
}
|
||||
|
||||
static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst) {
|
||||
VK_LOG_DEBUG("ggml_vk_dsv4_hc_comb(" << mixes << ", " << scale << ", " << base << ", " << dst << ")");
|
||||
|
||||
vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32;
|
||||
GGML_ASSERT(pipeline != nullptr);
|
||||
|
||||
const uint32_t n_tokens = (uint32_t)mixes->ne[1];
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
|
||||
|
||||
const vk_subbuffer mixes_buf = ggml_vk_tensor_subbuffer(ctx, mixes, true);
|
||||
const vk_subbuffer scale_buf = ggml_vk_tensor_subbuffer(ctx, scale, true);
|
||||
const vk_subbuffer base_buf = ggml_vk_tensor_subbuffer(ctx, base, true);
|
||||
const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
|
||||
|
||||
vk_op_dsv4_hc_comb_push_constants pc = {
|
||||
n_tokens,
|
||||
ggml_vk_nb_elem(mixes, 0), ggml_vk_nb_elem(mixes, 1),
|
||||
ggml_vk_nb_elem(scale, 0),
|
||||
ggml_vk_nb_elem(base, 0),
|
||||
ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2),
|
||||
0, 0, 0, 0,
|
||||
ggml_get_op_params_f32(dst, 0),
|
||||
(uint32_t)ggml_get_op_params_i32(dst, 1),
|
||||
};
|
||||
init_pushconst_tensor_offsets(ctx, pc, mixes, scale, base, nullptr, dst);
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { mixes_buf, scale_buf, base_buf, dst_buf }, pc, { n_tokens, 1, 1 });
|
||||
}
|
||||
|
||||
static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst) {
|
||||
VK_LOG_DEBUG("ggml_vk_dsv4_hc_pre(" << x << ", " << weights << ", " << dst << ")");
|
||||
|
||||
vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32;
|
||||
GGML_ASSERT(pipeline != nullptr);
|
||||
|
||||
const uint32_t n_embd = (uint32_t)x->ne[0];
|
||||
const uint32_t n_tokens = (uint32_t)x->ne[2];
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
|
||||
|
||||
const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true);
|
||||
const vk_subbuffer w_buf = ggml_vk_tensor_subbuffer(ctx, weights, true);
|
||||
const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
|
||||
|
||||
vk_op_dsv4_hc_pre_push_constants pc = {
|
||||
n_embd, n_tokens,
|
||||
ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), ggml_vk_nb_elem(x, 2),
|
||||
ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1),
|
||||
ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1),
|
||||
0, 0, 0,
|
||||
};
|
||||
init_pushconst_tensor_offsets(ctx, pc, x, weights, nullptr, nullptr, dst);
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, w_buf, d_buf }, pc, { n_embd, n_tokens, 1 });
|
||||
}
|
||||
|
||||
static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst) {
|
||||
VK_LOG_DEBUG("ggml_vk_dsv4_hc_post(" << x << ", " << residual << ", " << post << ", " << comb << ", " << dst << ")");
|
||||
|
||||
vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32;
|
||||
GGML_ASSERT(pipeline != nullptr);
|
||||
|
||||
const uint32_t n_embd = (uint32_t)x->ne[0];
|
||||
const uint32_t n_tokens = (uint32_t)x->ne[1];
|
||||
|
||||
ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
|
||||
|
||||
const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true);
|
||||
const vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, residual, true);
|
||||
const vk_subbuffer p_buf = ggml_vk_tensor_subbuffer(ctx, post, true);
|
||||
const vk_subbuffer c_buf = ggml_vk_tensor_subbuffer(ctx, comb, true);
|
||||
const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
|
||||
|
||||
vk_op_dsv4_hc_post_push_constants pc = {
|
||||
n_embd, n_tokens,
|
||||
ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1),
|
||||
ggml_vk_nb_elem(residual, 0), ggml_vk_nb_elem(residual, 1), ggml_vk_nb_elem(residual, 2),
|
||||
ggml_vk_nb_elem(post, 0), ggml_vk_nb_elem(post, 1),
|
||||
ggml_vk_nb_elem(comb, 0), ggml_vk_nb_elem(comb, 1), ggml_vk_nb_elem(comb, 2),
|
||||
ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2),
|
||||
0, 0, 0, 0, 0,
|
||||
};
|
||||
init_pushconst_tensor_offsets(ctx, pc, x, residual, post, comb, dst);
|
||||
|
||||
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 });
|
||||
}
|
||||
|
||||
static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) {
|
||||
ggml_tensor * dst = cgraph->nodes[node_idx];
|
||||
ggml_tensor * src0 = dst->src[0];
|
||||
@@ -16222,6 +16400,18 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
|
||||
case GGML_OP_CUMSUM:
|
||||
ggml_vk_cumsum(ctx, compute_ctx, src0, node);
|
||||
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
ggml_vk_dsv4_hc_comb(ctx, compute_ctx, src0, src1, src2, node);
|
||||
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
ggml_vk_dsv4_hc_pre(ctx, compute_ctx, src0, src1, node);
|
||||
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
ggml_vk_dsv4_hc_post(ctx, compute_ctx, src0, src1, src2, src3, node);
|
||||
|
||||
break;
|
||||
case GGML_OP_MEAN:
|
||||
ggml_vk_mean(ctx, compute_ctx, src0, node);
|
||||
@@ -19289,6 +19479,31 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
|
||||
}
|
||||
return false;
|
||||
}
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
if (op->type != GGML_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
for (uint32_t i = 0; i < GGML_MAX_SRC; ++i) {
|
||||
if (op->src[i] && op->src[i]->type != GGML_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
// hc is hardcoded to 4 in the shaders. ggml only constrains it
|
||||
// to 4 for COMB, so PRE/POST have to be checked here.
|
||||
if (op->op == GGML_OP_DSV4_HC_PRE && op->src[0]->ne[1] != 4) {
|
||||
return false;
|
||||
}
|
||||
if (op->op == GGML_OP_DSV4_HC_POST && op->src[1]->ne[1] != 4) {
|
||||
return false;
|
||||
}
|
||||
if (op->op == GGML_OP_DSV4_HC_COMB) {
|
||||
return device->pipeline_dsv4_hc_comb_f32 != nullptr;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
case GGML_OP_SOLVE_TRI:
|
||||
{
|
||||
if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32) {
|
||||
@@ -20277,6 +20492,13 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
|
||||
tensor_clone = ggml_sum_rows(ggml_ctx, src_clone[0]);
|
||||
} else if (tensor->op == GGML_OP_CUMSUM) {
|
||||
tensor_clone = ggml_cumsum(ggml_ctx, src_clone[0]);
|
||||
} else if (tensor->op == GGML_OP_DSV4_HC_COMB) {
|
||||
tensor_clone = ggml_dsv4_hc_comb(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
|
||||
ggml_get_op_params_f32(tensor, 0), ggml_get_op_params_i32(tensor, 1));
|
||||
} else if (tensor->op == GGML_OP_DSV4_HC_PRE) {
|
||||
tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]);
|
||||
} else if (tensor->op == GGML_OP_DSV4_HC_POST) {
|
||||
tensor_clone = ggml_dsv4_hc_post(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]);
|
||||
} else if (tensor->op == GGML_OP_MEAN) {
|
||||
tensor_clone = ggml_mean(ggml_ctx, src_clone[0]);
|
||||
} else if (tensor->op == GGML_OP_ARGMAX) {
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#extension GL_KHR_shader_subgroup_basic : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle : require
|
||||
|
||||
// 16 lanes per token, indexed idst + hc*isrc: idst in bits 0..1, isrc in bits 2..3,
|
||||
// so subgroupShuffleXor by 1|2 reduces a row and by 4|8 a column.
|
||||
|
||||
layout(constant_id = 0) const uint SUBGROUP_SIZE = 32;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter
|
||||
{
|
||||
uint n_tokens;
|
||||
|
||||
uint nbm0; uint nbm1; // mixes
|
||||
uint nbs0; // scale
|
||||
uint nbb0; // base
|
||||
uint nbd0; uint nbd1; uint nbd2; // dst
|
||||
|
||||
uint m_offset;
|
||||
uint s_offset;
|
||||
uint b_offset;
|
||||
uint d_offset;
|
||||
|
||||
float eps;
|
||||
uint n_iter;
|
||||
};
|
||||
|
||||
layout(binding = 0, std430) readonly buffer M { float data_m[]; };
|
||||
layout(binding = 1, std430) readonly buffer S { float data_s[]; };
|
||||
layout(binding = 2, std430) readonly buffer B { float data_b[]; };
|
||||
layout(binding = 3, std430) writeonly buffer D { float data_d[]; };
|
||||
|
||||
const uint hc = 4;
|
||||
const uint comb_offset = 2 * hc;
|
||||
|
||||
const uint TOKENS_PER_SUBGROUP = SUBGROUP_SIZE / 16;
|
||||
|
||||
void main() {
|
||||
const uint lane = gl_SubgroupInvocationID;
|
||||
const uint blk = lane >> 4; // which 16-lane block, i.e. which token
|
||||
const uint idx = lane & 15; // idst + hc*isrc
|
||||
|
||||
const uint sg = gl_WorkGroupID.x * gl_WorkGroupSize.y + gl_SubgroupID;
|
||||
const uint it = sg * TOKENS_PER_SUBGROUP + blk;
|
||||
|
||||
// no early return, the shuffles need every lane; out-of-range blocks compute a discarded value
|
||||
const bool in_range = it < n_tokens;
|
||||
|
||||
const float scale_comb = data_s[s_offset + 2 * nbs0];
|
||||
|
||||
float v = 0.0f;
|
||||
if (in_range) {
|
||||
v = data_m[m_offset + (comb_offset + idx) * nbm0 + it * nbm1] * scale_comb
|
||||
+ data_b[b_offset + (comb_offset + idx) * nbb0];
|
||||
}
|
||||
|
||||
// Softmax across destinations: the four lanes sharing an isrc.
|
||||
float vmax = max(v, subgroupShuffleXor(v, 1));
|
||||
vmax = max(vmax, subgroupShuffleXor(vmax, 2));
|
||||
v = exp(v - vmax);
|
||||
|
||||
float sum = v + subgroupShuffleXor(v, 1);
|
||||
sum += subgroupShuffleXor(sum, 2);
|
||||
v = v / sum + eps;
|
||||
|
||||
// Normalize columns: equal destination indices are four lanes apart.
|
||||
sum = v + subgroupShuffleXor(v, 4);
|
||||
sum += subgroupShuffleXor(sum, 8);
|
||||
v /= sum + eps;
|
||||
|
||||
for (uint i = 1; i < n_iter; ++i) {
|
||||
sum = v + subgroupShuffleXor(v, 1);
|
||||
sum += subgroupShuffleXor(sum, 2);
|
||||
v /= sum + eps;
|
||||
|
||||
sum = v + subgroupShuffleXor(v, 4);
|
||||
sum += subgroupShuffleXor(sum, 8);
|
||||
v /= sum + eps;
|
||||
}
|
||||
|
||||
if (in_range) {
|
||||
const uint idst = idx & 3;
|
||||
const uint isrc = idx >> 2;
|
||||
data_d[d_offset + idst * nbd0 + isrc * nbd1 + it * nbd2] = v;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
// Fan one stream back out to hc streams and add the combination-weighted
|
||||
// residuals:
|
||||
//
|
||||
// dst[i0, idst, it] = x[i0, it]*post[idst, it]
|
||||
// + sum_isrc residual[i0, isrc, it]*comb[idst, isrc, it]
|
||||
|
||||
layout(constant_id = 0) const uint BLOCK_SIZE = 256;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter
|
||||
{
|
||||
uint n_embd;
|
||||
uint n_tokens;
|
||||
|
||||
uint nbx0; uint nbx1; // x
|
||||
uint nbr0; uint nbr1; uint nbr2; // residual
|
||||
uint nbp0; uint nbp1; // post
|
||||
uint nbc0; uint nbc1; uint nbc2; // comb
|
||||
uint nbd0; uint nbd1; uint nbd2; // dst
|
||||
|
||||
uint x_offset;
|
||||
uint r_offset;
|
||||
uint p_offset;
|
||||
uint c_offset;
|
||||
uint d_offset;
|
||||
};
|
||||
|
||||
layout(binding = 0, std430) readonly buffer X { float data_x[]; };
|
||||
layout(binding = 1, std430) readonly buffer R { float data_r[]; };
|
||||
layout(binding = 2, std430) readonly buffer P { float data_p[]; };
|
||||
layout(binding = 3, std430) readonly buffer C { float data_c[]; };
|
||||
layout(binding = 4, std430) writeonly buffer D { float data_d[]; };
|
||||
|
||||
const uint hc = 4;
|
||||
|
||||
shared float post_s[hc];
|
||||
shared float comb_s[hc * hc];
|
||||
|
||||
void main() {
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
const uint it = gl_WorkGroupID.y;
|
||||
|
||||
if (tid < hc) {
|
||||
post_s[tid] = data_p[p_offset + tid * nbp0 + it * nbp1];
|
||||
}
|
||||
if (tid < hc * hc) {
|
||||
const uint idst = tid & 3;
|
||||
const uint isrc = tid >> 2;
|
||||
comb_s[tid] = data_c[c_offset + idst * nbc0 + isrc * nbc1 + it * nbc2];
|
||||
}
|
||||
barrier();
|
||||
|
||||
// After the barrier, so every invocation reaches it.
|
||||
const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid;
|
||||
if (i0 >= n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float xv = data_x[x_offset + i0 * nbx0 + it * nbx1];
|
||||
|
||||
const uint rb = r_offset + i0 * nbr0 + it * nbr2;
|
||||
|
||||
float r[hc];
|
||||
[[unroll]]
|
||||
for (uint isrc = 0; isrc < hc; ++isrc) {
|
||||
r[isrc] = data_r[rb + isrc * nbr1];
|
||||
}
|
||||
|
||||
[[unroll]]
|
||||
for (uint idst = 0; idst < hc; ++idst) {
|
||||
float result = xv * post_s[idst];
|
||||
[[unroll]]
|
||||
for (uint isrc = 0; isrc < hc; ++isrc) {
|
||||
result = fma(r[isrc], comb_s[idst + hc * isrc], result);
|
||||
}
|
||||
data_d[d_offset + i0 * nbd0 + idst * nbd1 + it * nbd2] = result;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
#version 450
|
||||
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
|
||||
// Collapse the hc residual streams of a token into one, weighted per stream:
|
||||
//
|
||||
// dst[i0, it] = sum_ih x[i0, ih, it] * weights[ih, it]
|
||||
|
||||
layout(constant_id = 0) const uint BLOCK_SIZE = 256;
|
||||
|
||||
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout(push_constant) uniform parameter
|
||||
{
|
||||
uint n_embd;
|
||||
uint n_tokens;
|
||||
|
||||
uint nbx0; uint nbx1; uint nbx2; // x
|
||||
uint nbw0; uint nbw1; // weights
|
||||
uint nbd0; uint nbd1; // dst
|
||||
|
||||
uint x_offset;
|
||||
uint w_offset;
|
||||
uint d_offset;
|
||||
};
|
||||
|
||||
layout(binding = 0, std430) readonly buffer X { float data_x[]; };
|
||||
layout(binding = 1, std430) readonly buffer W { float data_w[]; };
|
||||
layout(binding = 2, std430) writeonly buffer D { float data_d[]; };
|
||||
|
||||
const uint hc = 4;
|
||||
|
||||
shared float w[hc];
|
||||
|
||||
void main() {
|
||||
const uint tid = gl_LocalInvocationID.x;
|
||||
const uint it = gl_WorkGroupID.y;
|
||||
|
||||
if (tid < hc) {
|
||||
w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1];
|
||||
}
|
||||
barrier();
|
||||
|
||||
// After the barrier, so every invocation reaches it.
|
||||
const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid;
|
||||
if (i0 >= n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
const uint xb = x_offset + i0 * nbx0 + it * nbx2;
|
||||
|
||||
float result = 0.0f;
|
||||
[[unroll]]
|
||||
for (uint ih = 0; ih < hc; ++ih) {
|
||||
result = fma(data_x[xb + ih * nbx1], w[ih], result);
|
||||
}
|
||||
|
||||
data_d[d_offset + i0 * nbd0 + it * nbd1] = result;
|
||||
}
|
||||
@@ -1042,6 +1042,9 @@ void process_shaders() {
|
||||
string_to_spv("fwht_f32", "fwht.comp", {});
|
||||
string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}});
|
||||
string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}}));
|
||||
string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {});
|
||||
string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {});
|
||||
string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {});
|
||||
string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
|
||||
string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
|
||||
string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
|
||||
|
||||
@@ -8807,6 +8807,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 4));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(257, 8));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 20));
|
||||
// production n_iter (DeepSeek-V4 uses 20) across batch sizes that cross
|
||||
// subgroup and workgroup boundaries; 1 = single-token decode
|
||||
for (int64_t n_tokens : {1, 256, 336, 512, 513, 1024, 2048}) {
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(n_tokens, 20));
|
||||
}
|
||||
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(31, 17));
|
||||
|
||||
Reference in New Issue
Block a user