Compare commits

..
4 Commits
Author SHA1 Message Date
b387ddfd84 vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize (#27812)
* vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize

is_src_of doesn't treat two views of one tensor as dependent, so the optimizer reorders nodes across aliased reads and writes. 

Result: silently wrong tokens under greedy decoding, different output on every server start, and invalid speculative-decoding acceptance, with nothing logged.

Hits Qwen3.8's recurrent state (and any model with view-aliased state) on AMD and NVIDIA Vulkan.  CUDA is clean. 

Compare view_src bases on both sides.

Fixes #27805

* vulkan: don't treat view/no-op nodes as aliasing dependencies

Nodes whose op is NONE, RESHAPE, TRANSPOSE, VIEW or PERMUTE execute nothing, so aliasing through them is not a real dependency. The previous base comparison matched them anyway, which only costs the optimizer reordering freedom.

Co-authored-by: Jeff Bolz <jbolz@nvidia.com>

* vulkan: make the lambda parameter const and capture is_empty in is_src_of

Code will not compile without these changes.  
is_src_of has an empty capture list, so is_empty was not visible inside it, and is_empty took a non-const pointer, while is_src_of receives const ones. Other call sites pass non-const pointers, which still convert as usual.

---------

Co-authored-by: Jeff Bolz <jbolz@nvidia.com>
2026-08-28 19:12:33 +02:00
a43c3986b4 ggml : fix conv_transpose_2d for multiple batches (#26132)
* ggml : fix conv_transpose_2d for multiple batches

ggml_compute_forward_conv_transpose_2d_impl only computed the first
batch (ne[3] of the destination); every batch after the first was left
as zero. Both the src1 permutation and the main compute loop now iterate
over the batch dimension, and the work buffer size in ggml_graph_plan is
scaled by the src1 batch count so the extra permuted batches fit. A
multi-batch test case is added to test-backend-ops.

Fixes ggml-org/ggml#1448

* metal : fix conv_transpose_2d for multiple batches

The kernel only computed batch 0 of the input (src1->ne[3]); every
output batch after the first was left as zero, so multi-batch
conv_transpose_2d results diverged from the CPU reference.

The grid now covers all batches (OW x OH x OC x N), the kernel decodes
the batch from the grid z coordinate and offsets both the input and
destination indices accordingly. nb3 is passed in the kernel args.

Assisted-by: pi:llama.cpp/Qwen3.8-27B

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
2026-08-28 20:09:08 +03:00
90c26fcd4b Vulkan: add hoisting support for row IDs and expert count in shaders (#26686)
* vulkan: add hoisting support for row IDs and expert count in shaders

* use hoisted row ids in coopmat2

* vulkan: address review feedback on count_experts
- use vk_op_count_experts_push_constants instead of a raw uint vector
- apply the fastdiv trick to the ne00 div/mod in count_experts
- compute the per-expert offsets with subgroupExclusiveAdd when the
  device supports it, keeping the serial path as fallback
- document the data_d layout and the hoisted_row_id_words bound
- drop a leftover debug print in ggml_vk_matmul_id

* vulkan: use init_pushconst_fastdiv for count_experts push constants

* vulkan: refine comments for row ID hoisting and data layout in count_experts shader

* Whitespace

---------

Co-authored-by: Jeff Bolz <jbolz@nvidia.com>
2026-08-28 16:52:49 +02:00
Georgi GerganovandGitHub 8663224818 context : disable non-fused GDN and LID ops (#27877) 2026-08-28 16:34:26 +03:00
14 changed files with 275 additions and 81 deletions
+2 -1
View File
@@ -2936,12 +2936,13 @@ struct ggml_cplan ggml_graph_plan(
const int64_t ne10 = node->src[1]->ne[0]; // W
const int64_t ne11 = node->src[1]->ne[1]; // H
const int64_t ne12 = node->src[1]->ne[2]; // Channels In
const int64_t ne13 = node->src[1]->ne[3]; // Batch
GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32);
GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32);
cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03;
cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12;
cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13;
} break;
case GGML_OP_TOP_K:
+32 -26
View File
@@ -7267,18 +7267,21 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
}
}
// permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh)
// permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches
{
kernel_t * const wdata = (kernel_t *) params->wdata + nk;
for (int i12 = 0; i12 < ne12; i12++) {
for (int i11 = 0; i11 < ne11; i11++) {
const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11);
kernel_t * dst_data = wdata + i11*ne10*ne12;
for (int i10 = 0; i10 < ne10; i10++) {
if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
} else {
dst_data[i10*ne12 + i12] = src[i10];
for (int i13 = 0; i13 < ne13; i13++) {
kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12;
for (int i12 = 0; i12 < ne12; i12++) {
for (int i11 = 0; i11 < ne11; i11++) {
const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
kernel_t * dst_data = wdata_b + i11*ne10*ne12;
for (int i10 = 0; i10 < ne10; i10++) {
if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]);
} else {
dst_data[i10*ne12 + i12] = src[i10];
}
}
}
}
@@ -7305,24 +7308,27 @@ static void ggml_compute_forward_conv_transpose_2d_impl(
kernel_t * const wdata_src = wdata + nk;
for (int i2 = ip0; i2 < ip1; i2++) { // Cout
float * dst_data = (float *)((char *) dst->data + i2*nb2);
kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03;
for (int i11 = 0; i11 < ne11; i11++) {
for (int i10 = 0; i10 < ne10; i10++) {
const int i1n = i11*ne10*ne12 + i10*ne12;
for (int i01 = 0; i01 < ne01; i01++) {
for (int i00 = 0; i00 < ne00; i00++) {
float v = 0;
if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
ggml_vec_dot_f16(ne03, &v, 0,
wdata_src + i1n, 0,
wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
} else {
ggml_vec_dot_f32(ne03, &v, 0,
wdata_src + i1n, 0,
wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
for (int i3 = 0; i3 < ne3; i3++) { // batch
float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2);
kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12;
for (int i11 = 0; i11 < ne11; i11++) {
for (int i10 = 0; i10 < ne10; i10++) {
const int i1n = i11*ne10*ne12 + i10*ne12;
for (int i01 = 0; i01 < ne01; i01++) {
for (int i00 = 0; i00 < ne00; i00++) {
float v = 0;
if constexpr (std::is_same_v<kernel_t, ggml_fp16_t>) {
ggml_vec_dot_f16(ne03, &v, 0,
wdata_src_b + i1n, 0,
wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
} else {
ggml_vec_dot_f32(ne03, &v, 0,
wdata_src_b + i1n, 0,
wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1);
}
dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
}
dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v;
}
}
}
+1
View File
@@ -660,6 +660,7 @@ typedef struct {
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ggml_metal_kargs_conv_transpose_2d;
typedef struct {
+3 -1
View File
@@ -4645,6 +4645,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
const int32_t OW = op->ne[0];
const int32_t OH = op->ne[1];
const int32_t OC = op->ne[2];
const int32_t N = op->src[1]->ne[3];
ggml_metal_kargs_conv_transpose_2d args = {
/*.IC =*/ IC,
@@ -4657,6 +4658,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
/*.nb0 =*/ nb0,
/*.nb1 =*/ nb1,
/*.nb2 =*/ nb2,
/*.nb3 =*/ nb3,
};
auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op);
@@ -4671,7 +4673,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1);
ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1);
return 1;
}
+4 -3
View File
@@ -366,7 +366,8 @@ kernel void kernel_conv_transpose_2d(
const int64_t out_x = tgpig[0];
const int64_t out_y = tgpig[1];
const int64_t out_c = tgpig[2];
const int64_t batch = tgpig[2] / args.OC;
const int64_t out_c = tgpig[2] % args.OC;
const int64_t kw = tpitg[0];
const int64_t kh = tpitg[1];
@@ -390,7 +391,7 @@ kernel void kernel_conv_transpose_2d(
if (in_x >= args.IW) continue;
const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x;
const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x;
const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw;
v += (float)src0[kernel_idx] * src1[input_idx];
@@ -408,7 +409,7 @@ kernel void kernel_conv_transpose_2d(
total += shared_sum[i];
}
device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2);
device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0);
dst_ptr[0] = total;
}
}
+54 -16
View File
@@ -1348,6 +1348,8 @@ struct vk_mat_mat_id_push_constants {
uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d;
uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11;
uint32_t padded_N;
uint32_t n_experts;
uint32_t hoist_row_ids;
};
struct vk_mat_vec_id_push_constants {
uint32_t ncols;
@@ -1428,6 +1430,10 @@ struct vk_op_count_experts_push_constants {
uint32_t nb00;
uint32_t nb01;
uint32_t a_offset;
uint32_t n_experts;
uint32_t hoist_row_ids;
uint32_t ne00mp;
uint32_t ne00L;
};
struct vk_op_glu_push_constants {
@@ -1606,6 +1612,10 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) {
init_fastdiv_values(p.ne20, p.ne2_0mp, p.ne2_0L);
}
template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) {
init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L);
}
struct vk_op_binary_push_constants {
uint32_t ne;
uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03;
@@ -5839,7 +5849,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1);
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);
if (device->subgroup_arithmetic && device->subgroup_require_full_support) {
ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true);
} else {
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);
}
for (auto &s : device->pipeline_solve_tri_f32) {
const vk_solve_tri_pipeline_state &state = s.first;
@@ -8970,13 +8984,13 @@ static void ggml_vk_matmul_id(
uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11,
uint32_t padded_n) {
uint32_t padded_n, bool hoist_row_ids) {
VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " <<
"m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " <<
"batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " <<
"n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")");
const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d,
nei0, nei1, nbi1, ne11, padded_n };
nei0, nei1, nbi1, ne11, padded_n, n_as, uint32_t(hoist_row_ids) };
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as });
}
@@ -10162,6 +10176,12 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
// const uint64_t ne23 = dst->ne[3];
const uint64_t n_as = ne02;
// n_as counts, n_as offsets, one total, then one packed row id per (expert, token).
// Hoisting requires 16-bit indices for the packing and a table that fits one binding.
const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1;
const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff &&
hoisted_row_id_words * sizeof(uint32_t) <=
ctx->device->properties.limits.maxStorageBufferRange;
ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context;
ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context;
@@ -10302,7 +10322,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
}
vk_pipeline count_experts = ctx->device->pipeline_count_experts;
uint32_t expert_count_size = sizeof(uint32_t) * n_as;
const size_t expert_data_size = sizeof(uint32_t) *
(hoist_row_ids ? hoisted_row_id_words : n_as);
{
if (
@@ -10318,8 +10339,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
ctx->prealloc_size_y = y_sz;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_size_split_k < expert_count_size) {
ctx->prealloc_size_split_k = expert_count_size;
if (ctx->prealloc_size_split_k < expert_data_size) {
ctx->prealloc_size_split_k = expert_data_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
@@ -10385,18 +10406,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
}
}
// Count how many times each expert is used
vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size };
if (ctx->prealloc_split_k_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}
{
const std::vector<uint32_t> pc = { (uint32_t)nei0,
vk_op_count_experts_push_constants pc = { (uint32_t)nei0,
(uint32_t)nei1,
(uint32_t)(nbi0 / ggml_type_size(ids->type)),
(uint32_t)(nbi1 / ggml_type_size(ids->type)),
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) };
(uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)),
(uint32_t)n_as,
uint32_t(hoist_row_ids),
0, 0 };
init_pushconst_fastdiv(pc);
ggml_vk_dispatch_pipeline(ctx, subctx, count_experts,
{ vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1});
{ vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc,
{ hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1});
}
if (x_non_contig) {
@@ -10465,7 +10491,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
{ d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf,
ne01, ne21, ne10, ne10, stride_b_y, ne01,
stride_batch_x, stride_batch_y, ne20*ne21,
n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n
n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n, hoist_row_ids
); // NOLINT
if (x_non_contig || qx_needs_dequant) {
@@ -17774,20 +17800,32 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
return;
}
auto const &is_empty = [](ggml_tensor * node) -> bool {
auto const &is_empty = [](const ggml_tensor * node) -> bool {
return node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE;
};
auto const &is_src_of = [](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
auto const &is_src_of = [&is_empty](const ggml_tensor *dst, const ggml_tensor *src) -> bool {
auto const &base = [](const ggml_tensor * tensor) {
return tensor->view_src ? tensor->view_src : tensor;
};
for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
if (dst->src[s] == src) {
return true;
}
if (is_empty(dst) || is_empty(src)) {
continue;
}
// A source view of dst may read storage written through a different view by src.
if (dst->src[s] && base(dst->src[s]) == base(src)) {
return true;
}
// Moving dst forward may overwrite storage still read through a view by src.
if (src->src[s] && base(dst) == base(src->src[s])) {
return true;
}
}
// implicit dependency if they view the same tensor
const ggml_tensor *dst2 = dst->view_src ? dst->view_src : dst;
const ggml_tensor *src2 = src->view_src ? src->view_src : src;
if (dst2 == src2) {
if (base(dst) == base(src)) {
return true;
}
return false;
@@ -2,6 +2,11 @@
#extension GL_EXT_control_flow_attributes : enable
#ifdef USE_SUBGROUPS
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#endif
#include "types.glsl"
layout (push_constant) uniform parameter
@@ -11,6 +16,10 @@ layout (push_constant) uniform parameter
uint32_t nb00;
uint32_t nb01;
uint32_t a_offset;
uint32_t n_experts;
uint32_t hoist_row_ids;
uint32_t ne00mp;
uint32_t ne00L;
} p;
#define BLOCK_SIZE 256
@@ -21,16 +30,98 @@ layout (binding = 0) readonly buffer A {uint data_a[];};
layout (binding = 1) writeonly buffer D {uint data_d[];};
shared uint vals[BLOCK_SIZE];
shared uint offsets[BLOCK_SIZE];
shared uint cursors[BLOCK_SIZE];
// see init_fastdiv_values in ggml-vulkan.cpp
uint fastdiv(uint n, uint mp, uint L) {
uint msbs, lsbs;
// msbs = mulhi(n, mp)
umulExtended(n, mp, msbs, lsbs);
return (msbs + n) >> L;
}
// data_d layout when p.hoist_row_ids is set:
// [0, n_experts) per-expert row count
// [n_experts, 2*n_experts) per-expert start offset into the row id region
// [2*n_experts] total row count
// [2*n_experts + 1, ) row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff)
// Otherwise only data_d[expert_id] is written, holding that expert's row count.
void main() {
const uint expert_id = gl_WorkGroupID.x;
const uint num_elements = p.ne00 * p.ne01;
const uint tid = gl_LocalInvocationID.x;
if (p.hoist_row_ids != 0) {
if (tid < p.n_experts) {
vals[tid] = 0;
}
barrier();
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
if (expert < p.n_experts) {
atomicAdd(vals[expert], 1);
}
}
barrier();
#ifdef USE_SUBGROUPS
if (gl_SubgroupID == 0) {
// pad the trip count so the subgroup ops stay in uniform control flow
const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1);
uint base = 0;
for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) {
const bool in_range = expert < p.n_experts;
const uint count = in_range ? vals[expert] : 0;
const uint offset = base + subgroupExclusiveAdd(count);
if (in_range) {
data_d[expert] = count;
data_d[p.n_experts + expert] = offset;
offsets[expert] = offset;
cursors[expert] = 0;
}
base += subgroupAdd(count);
}
if (subgroupElect()) {
data_d[2 * p.n_experts] = base;
}
}
#else
if (tid == 0) {
uint offset = 0;
for (uint expert = 0; expert < p.n_experts; ++expert) {
const uint count = vals[expert];
data_d[expert] = count;
data_d[p.n_experts + expert] = offset;
offsets[expert] = offset;
cursors[expert] = 0;
offset += count;
}
data_d[2 * p.n_experts] = offset;
}
#endif
barrier();
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
if (expert < p.n_experts) {
const uint row = atomicAdd(cursors[expert], 1);
const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu);
data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id;
}
}
return;
}
uint count = 0;
for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) {
const uint i01 = idx / p.ne00;
const uint i00 = idx % p.ne00;
const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L);
const uint i00 = idx - i01 * p.ne00;
const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00];
count += uint(a == expert_id);
+21 -14
View File
@@ -88,6 +88,9 @@ layout (push_constant) uniform parameter
uint nei1;
uint nbi1;
uint ne11;
uint padded_N;
uint n_experts;
uint hoist_row_ids;
#else
uint base_work_group_z;
uint num_batches;
@@ -214,27 +217,31 @@ void main() {
const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK;
#ifdef MUL_MAT_ID
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#else
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
}
_ne1++;
}
_ne1++;
}
}
}
barrier();
barrier();
#endif
}
// Workgroup has no work
if (ic * BN >= _ne1) return;
@@ -67,6 +67,10 @@ layout (push_constant) uniform parameter
#endif
// N dimension for the B matrix can be >= p.N
uint padded_N;
#ifdef MUL_MAT_ID
uint n_experts;
uint hoist_row_ids;
#endif
} p;
@@ -225,6 +229,23 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
}
barrier();
}
void load_row_ids_hoisted(uint expert_idx, uint ic) {
_ne1 = uint(data_expert_count[expert_idx]);
const uint tile_begin = ic * BN;
const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
const uint ii0 = packed_row_id & 0xffffu;
const uint ii1 = packed_row_id >> 16;
row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0);
}
barrier();
}
#endif
void main() {
@@ -266,7 +287,9 @@ void main() {
const uint ik = gl_WorkGroupID.x / blocks_m;
#ifdef MUL_MAT_ID
if (bitCount(p.nei0) == 1) {
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
@@ -71,4 +71,19 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) {
barrier();
}
#endif // MUL_MAT_ID_USE_SUBGROUPS
void load_row_ids_hoisted(uint expert_idx, uint ic) {
_ne1 = uint(data_expert_count[expert_idx]);
const uint tile_begin = ic * BN;
const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0;
const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]);
const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin;
for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) {
const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]);
row_ids[i] = u16vec2(packed_row_id & 0xffffu, packed_row_id >> 16);
}
barrier();
}
#endif // MUL_MAT_ID
@@ -56,6 +56,9 @@ layout (push_constant) uniform parameter
uint nei1;
uint nbi1;
uint ne11;
uint padded_N;
uint n_experts;
uint hoist_row_ids;
#else
uint base_work_group_z;
uint num_batches;
@@ -157,27 +160,31 @@ void main() {
const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK;
#ifdef MUL_MAT_ID
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
if (p.hoist_row_ids != 0) {
load_row_ids_hoisted(expert_idx, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#ifdef MUL_MAT_ID_USE_SUBGROUPS
if (bitCount(p.nei0) == 1) {
load_row_ids(expert_idx, true, ic);
} else {
load_row_ids(expert_idx, false, ic);
}
#else
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
_ne1 = 0;
for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) {
for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) {
if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) {
if (_ne1 >= ic * BN) {
row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1);
}
_ne1++;
}
_ne1++;
}
}
}
barrier();
barrier();
#endif
}
// Workgroup has no work
if (ic * BN >= _ne1) return;
@@ -1039,6 +1039,7 @@ void process_shaders() {
string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("count_experts", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}));
string_to_spv("count_experts_subgroup", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}, {"USE_SUBGROUPS", "1"}}));
for (std::string dim_str : {"", "_3d"}) {
for (bool bda : {false, true}) {
+3 -3
View File
@@ -231,10 +231,10 @@ llama_context::llama_context(
cparams.fused_gdn_ar = true;
cparams.fused_gdn_ch = true;
cparams.auto_fgdn = true;
cparams.auto_fgdn = false;
cparams.fused_lid = true;
cparams.auto_flid = true;
cparams.fused_lid = true;
cparams.auto_flid = false;
cparams.fused_dsv4_hc_pre = true;
cparams.fused_dsv4_hc_comb = true;
+1
View File
@@ -8785,6 +8785,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_conv_transpose_2d({3, 2, 3, 1}, {2, 2, 1, 3}, 1, kernel_type));
test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 1}, {3, 3, 1, 9}, 2, kernel_type));
test_cases.emplace_back(new test_conv_transpose_2d({129, 63, 35, 1}, {3, 3, 48, 35}, 1, kernel_type));
test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 2}, {3, 3, 1, 9}, 2, kernel_type)); // for multiple batches
}
test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 500, 1, 1}));