Commit 69ee2ecc0 for llama.cpp
commit 69ee2ecc0150ba7d1204a7f7393dfd651569db8c
Author: Johannes Gäßler <johannesg@5d6.de>
Date: Thu Nov 5 11:33:25 2026 +0100
CUDA: linearize MoE data layout for MMQ/MMF (#31047)
* CUDA: linearize MoE data layout for MMQ/MMF
* Apply MMQ tile size optimization more generally
* Fix F16 performance regression
diff --git a/ggml/src/ggml-cuda/mmf.cu b/ggml/src/ggml-cuda/mmf.cu
index 646a5899c..040d390e0 100644
--- a/ggml/src/ggml-cuda/mmf.cu
+++ b/ggml/src/ggml-cuda/mmf.cu
@@ -50,7 +50,8 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
mmf_ids_data * ids_info_ptr = nullptr;
ggml_cuda_pool_alloc<int32_t> ids_src_compact_dev;
ggml_cuda_pool_alloc<int32_t> ids_dst_compact_dev;
- ggml_cuda_pool_alloc<int32_t> expert_bounds_dev;
+
+ ggml_cuda_pool_alloc<ggml_cuda_mm_id_tile_spec> tile_spec_dev;
// For MUL_MAT_ID the memory layout is different than for MUL_MAT:
const int64_t ncols_dst = ids ? ne2 : ne1;
@@ -74,25 +75,28 @@ void ggml_cuda_mul_mat_f(ggml_backend_cuda_context & ctx, const ggml_tensor * sr
const int64_t n_experts = ne02;
const int64_t n_tokens = ne12;
const int64_t ne_get_rows = n_tokens * n_expert_used;
+ const int64_t max_tiles = (ne_get_rows + 16 - 1) / 16 + ne02;
ids_src_compact_dev.alloc(ctx.pool(), ne_get_rows);
ids_dst_compact_dev.alloc(ctx.pool(), ne_get_rows);
- expert_bounds_dev.alloc(ctx.pool(), n_experts + 1);
+ tile_spec_dev.alloc(ctx.pool(), max_tiles);
const int si1 = static_cast<int>(ids_s1);
const int sis1 = static_cast<int>(src1->nb[2] / src1->nb[1]);
GGML_ASSERT(sis1 > 0);
- ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), expert_bounds_dev.get(),
- static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), si1, sis1, /*write_inverse =*/ false, ctx.stream());
+ ggml_cuda_launch_mm_ids_helper(ids_d, ids_src_compact_dev.get(), ids_dst_compact_dev.get(), tile_spec_dev.get(),
+ static_cast<int>(n_experts), static_cast<int>(n_tokens), static_cast<int>(n_expert_used), static_cast<int>(ne11), 16,
+ si1, sis1, /*write_inverse =*/ false, ctx.stream());
CUDA_CHECK(cudaGetLastError());
- ids_info.ids_src_compact = ids_src_compact_dev.get();
- ids_info.ids_dst_compact = ids_dst_compact_dev.get();
- ids_info.expert_bounds_dev = expert_bounds_dev.get();
- ids_info.n_experts = static_cast<int>(n_experts);
- ids_info.sis1 = sis1;
+ ids_info.ids_src_compact = ids_src_compact_dev.get();
+ ids_info.ids_dst_compact = ids_dst_compact_dev.get();
+ ids_info.tile_spec_dev = tile_spec_dev.get();
+ ids_info.n_experts = static_cast<int>(n_experts);
+ ids_info.max_tiles = static_cast<int>(max_tiles);
+ ids_info.sis1 = sis1;
ids_info_ptr = &ids_info;
}
diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh
index 429c8874c..c49f6c47a 100644
--- a/ggml/src/ggml-cuda/mmf.cuh
+++ b/ggml/src/ggml-cuda/mmf.cuh
@@ -1,7 +1,8 @@
#pragma once
-#include "mma.cuh"
#include "common.cuh"
+#include "mma.cuh"
+#include "mmid.cuh"
#include "convert.cuh"
using namespace ggml_cuda_mma;
@@ -36,8 +37,9 @@ static constexpr __device__ int mmf_get_padding() {
struct mmf_ids_data {
const int32_t * ids_src_compact = nullptr;
const int32_t * ids_dst_compact = nullptr;
- const int32_t * expert_bounds_dev = nullptr;
+ const ggml_cuda_mm_id_tile_spec * tile_spec_dev = nullptr;
int n_experts = 0;
+ int max_tiles = 0;
int sis1 = 0;
};
@@ -308,7 +310,7 @@ __launch_bounds__(ggml_cuda_get_physical_warp_size()*nwarps, 1)
static __global__ void mul_mat_f_ids(
const T * __restrict__ x, const float * __restrict__ y,
const int32_t * __restrict__ ids_src_compact, const int32_t * __restrict__ ids_dst_compact,
- const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
+ const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst,
const int ncols, const int ncols_dst_total, const int nchannels_dst, const int stride_row, const int stride_col_y, const int stride_col_dst,
const int channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const int sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
@@ -349,21 +351,20 @@ static __global__ void mul_mat_f_ids(
constexpr int ntA = rows_per_block / tile_A::I;
constexpr int ntB = (cols_per_block + tile_B::I - 1) / tile_B::I;
- const int row0 = blockIdx.x * rows_per_block;
+ const int row0 = blockIdx.x * rows_per_block;
+ const int tile_idx = blockIdx.y;
- const int expert_idx = blockIdx.y;
- const int expert_start = expert_bounds[expert_idx];
- const int expert_end = expert_bounds[expert_idx + 1];
+ const ggml_cuda_mm_id_tile_spec ts = tile_spec[tile_idx];
+
+ const int expert_idx = ts.expert;
+ const int expert_start = ts.col_low;
+ const int expert_end = ts.col_high;
const int ncols_expert = expert_end - expert_start;
- const int tiles_for_expert = (ncols_expert + cols_per_block - 1) / cols_per_block;
- const int tile_idx = blockIdx.z;
- if (tile_idx >= tiles_for_expert) {
+ if (ncols_expert == 0) {
return;
}
- const int col_base = tile_idx * cols_per_block;
-
GGML_UNUSED(channel_ratio);
const int channel_x = expert_idx;
@@ -409,10 +410,9 @@ static __global__ void mul_mat_f_ids(
#pragma unroll
for (int j0 = 0; j0 < tile_B::I; ++j0) {
const int j = j0 + tile_idx_local*tile_B::I;
- const int global_j = col_base + j;
float val = 0.0f;
- if (j < cols_per_block && global_j < ncols_expert) {
- const int src_entry = ids_src_expert[global_j];
+ if (j < cols_per_block && j < ncols_expert) {
+ const int src_entry = ids_src_expert[j];
const uint2 qrm = fast_div_modulo((uint32_t) src_entry, sis1_fd);
const int token = (int) qrm.x;
const int channel = (int) qrm.y;
@@ -462,10 +462,9 @@ static __global__ void mul_mat_f_ids(
#pragma unroll
for (int j0 = 0; j0 < tile_B::I; ++j0) {
const int j = j0 + tile_idx_local*tile_B::I;
- const int global_j = col_base + j;
float2 tmp = make_float2(0.0f, 0.0f);
- if (j < cols_per_block && global_j < ncols_expert) {
- const int src_entry = ids_src_expert[global_j];
+ if (j < cols_per_block && j < ncols_expert) {
+ const int src_entry = ids_src_expert[j];
const uint2 qrm = fast_div_modulo((uint32_t) src_entry, sis1_fd);
const int token = (int) qrm.x;
const int channel = (int) qrm.y;
@@ -564,9 +563,8 @@ static __global__ void mul_mat_f_ids(
}
}
- const int global_j = col_base + j;
- if (j < cols_per_block && global_j < ncols_expert && nchannels_dst > 0) {
- const int dst_entry = ids_dst_expert[global_j];
+ if (j < cols_per_block && j < ncols_expert && nchannels_dst > 0) {
+ const int dst_entry = ids_dst_expert[j];
const uint2 qrm = fast_div_modulo((uint32_t) dst_entry, nch_fd);
const int token = (int) qrm.x;
if (token < ncols_dst_total) {
@@ -580,7 +578,7 @@ static __global__ void mul_mat_f_ids(
}
}
#else
- GGML_UNUSED_VARS(x, y, ids_src_compact, ids_dst_compact, expert_bounds, dst,
+ GGML_UNUSED_VARS(x, y, ids_src_compact, ids_dst_compact, tile_spec, dst,
ncols, ncols_dst_total, nchannels_dst, stride_row, stride_col_y, stride_col_dst,
channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, sis1_fd, nch_fd);
@@ -603,17 +601,16 @@ static inline void mul_mat_f_switch_ids(
// Use the compact-ids kernel only for larger tiles; for small ncols_dst (< 16)
// we prefer the normal mul_mat_f path with has_ids=true.
if (has_ids_data && ncols_dst > 16) {
- const int max_tiles = (int) ((ncols_dst + cols_per_block - 1) / cols_per_block);
- if (max_tiles == 0) {
+ if (ids_data->max_tiles == 0) {
return;
}
- dim3 block_nums_ids(block_nums.x, ids_data->n_experts, max_tiles);
+ dim3 block_nums_ids(block_nums.x, ids_data->max_tiles, 1);
const uint3 sis1_fd = ids_data->sis1 > 0 ? init_fastdiv_values((uint32_t) ids_data->sis1) : make_uint3(0, 0, 1);
const uint3 nch_fd = init_fastdiv_values((uint32_t) nchannels_dst);
mul_mat_f_ids<T, rows_per_block, cols_per_block, nwarps><<<block_nums_ids, block_dims, nbytes_shared_total, stream>>>
- (x, y, ids_data->ids_src_compact, ids_data->ids_dst_compact, ids_data->expert_bounds_dev, dst,
+ (x, y, ids_data->ids_src_compact, ids_data->ids_dst_compact, ids_data->tile_spec_dev, dst,
ncols_x, ncols_dst, nchannels_dst, stride_row, stride_col_y, stride_col_dst,
channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst,
sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst,
diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu
index 12ae14ce9..64de4fb27 100644
--- a/ggml/src/ggml-cuda/mmid.cu
+++ b/ggml/src/ggml-cuda/mmid.cu
@@ -27,12 +27,16 @@ template <> struct mm_ids_pow2<0> { static constexpr int value = 1; };
// Helper function for mul_mat_id, converts ids to a more convenient format.
// ids_src1 describes how to permute the flattened column indices of src1 in order to get a compact src1 tensor sorted by expert.
// ids_dst describes the same mapping but for the dst tensor.
-// The upper and lower bounds for the ith expert in the compact src1 tensor are stored in expert_bounds[i:i+1].
+// The first (n_tokens*n_expert_used + J - 1) / tile_width elements of tile_spec contain the specication for
+// matrix multiplication tiles at token positions which are exactly divided by tile_width.
+// Additionally the following n_expert elements of tile_spec may contain additional tiles that fall between these seams.
template <int n_expert_used_template>
__launch_bounds__(ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mm_ids_helper(
- const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse) {
+ const int32_t * __restrict__ ids,
+ int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+ const int n_tokens, const int n_expert_used_var, const int nchannels_y, const uint3 tile_width,
+ const int si1, const int sis1, const bool write_inverse) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
const int n_expert_used = n_expert_used_template == 0 ? n_expert_used_var : n_expert_used_template;
const int expert = blockIdx.x;
@@ -118,19 +122,29 @@ static __global__ void mm_ids_helper(
return;
}
- expert_bounds[expert] = nex_prev;
+ int col_low = nex_prev;
+ int col_high = nex_prev + min(int(tile_width.z), it_compact);
- if (expert < static_cast<int>(gridDim.x) - 1) {
- return;
+ // The while loop covers all tiles at the seams containing the column with % tile_width == 0.
+ while (col_high > fastdiv(col_low + tile_width.z - 1, tile_width) * tile_width.z) {
+ tile_spec[fastdiv(col_low + tile_width.z - 1, tile_width)]
+ = {expert, col_low, col_high, /*padding =*/ 0};
+
+ col_low = col_high;
+ col_high = min(col_high + int(tile_width.z), nex_prev + it_compact);
}
- expert_bounds[gridDim.x] = nex_prev + it_compact;
+ // Finally, at most one extra tile_spec between the seams per expert.
+ tile_spec[fastdiv(n_tokens*n_expert_used + tile_width.z - 1, tile_width) + expert]
+ = {expert, col_low, col_high, /*padding =*/ 0};
}
template <int n_expert_used_template>
static void launch_mm_ids_helper(
- const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+ const int32_t * __restrict__ ids,
+ int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+ const int n_experts, const int n_tokens, const int n_expert_used_var, const int nchannels_y, const int tile_width,
+ const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
GGML_ASSERT(n_tokens < (1 << 22) && "too few bits in mm_ids_helper_store");
GGML_ASSERT(n_expert_used_var < (1 << 10) && "too few bits in mm_ids_helper_store");
@@ -139,41 +153,53 @@ static void launch_mm_ids_helper(
const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
CUDA_SET_SHARED_MEMORY_LIMIT(mm_ids_helper<n_expert_used_template>, smpbo);
+ const uint3 tile_width_fd = init_fastdiv_values(tile_width);
+
const dim3 num_blocks(n_experts, 1, 1);
const dim3 block_size(warp_size, 1, 1);
const size_t nbytes_shared = n_tokens*sizeof(mm_ids_helper_store);
GGML_ASSERT(nbytes_shared <= smpbo);
mm_ids_helper<n_expert_used_template><<<num_blocks, block_size, nbytes_shared, stream>>>
- (ids, ids_src1, ids_dst, expert_bounds, n_tokens, n_expert_used_var, nchannels_y, si1, sis1, write_inverse);
+ (ids, ids_src1, ids_dst, tile_spec, n_tokens, n_expert_used_var, nchannels_y, tile_width_fd, si1, sis1, write_inverse);
}
void ggml_cuda_launch_mm_ids_helper(
- const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds,
- const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
+ const int32_t * __restrict__ ids,
+ int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+ const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int tile_width,
+ const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) {
switch (n_expert_used) {
case 2:
- launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 4:
- launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper< 4>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 6:
- launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper< 6>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 8:
- launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper< 8>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 10:
- launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper<10>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 16:
- launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper<16>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
case 32:
- launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper<32>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
default:
- launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream);
+ launch_mm_ids_helper< 0>(ids, ids_src1, ids_dst, tile_spec,
+ n_experts, n_tokens, n_expert_used, nchannels_y, tile_width, si1, sis1, write_inverse, stream);
break;
}
}
diff --git a/ggml/src/ggml-cuda/mmid.cuh b/ggml/src/ggml-cuda/mmid.cuh
index 74c2db433..dc878a71c 100644
--- a/ggml/src/ggml-cuda/mmid.cuh
+++ b/ggml/src/ggml-cuda/mmid.cuh
@@ -1,5 +1,16 @@
#pragma once
+struct ggml_cuda_mm_id_tile_spec {
+ int32_t expert;
+ int32_t col_low;
+ int32_t col_high;
+
+ int32_t padding;
+};
+static_assert(sizeof(ggml_cuda_mm_id_tile_spec) == 16, "bad ggml_cuda_mm_id_tile_spec size");
+
void ggml_cuda_launch_mm_ids_helper(
- const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds,
- int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream);
+ const int32_t * __restrict__ ids,
+ int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec,
+ const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int tile_width,
+ const int si1, const int sis1, const bool write_inverse, cudaStream_t stream);
diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu
index 3d4b01d93..367657508 100644
--- a/ggml/src/ggml-cuda/mmq.cu
+++ b/ggml/src/ggml-cuda/mmq.cu
@@ -247,10 +247,7 @@ void ggml_cuda_mul_mat_q(
ncols_opt = ne12;
// Each expert only sees ne12*n_expert_used/ne02 tokens on average.
- // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
- if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
- ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
- }
+ ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
}
J_best = ggml_cuda_mmq_get_J_best(src0->type, ncols_opt, fallback, cc, smpbo, prec_src1, &nthreads_best);
@@ -264,6 +261,8 @@ void ggml_cuda_mul_mat_q(
* src1_load_chunk_size;
if (!ids) {
+ const int ntj = (ne1 + J_best - 1) / J_best;
+
const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block + src1_q8_1_padding;
ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
@@ -301,7 +300,7 @@ void ggml_cuda_mul_mat_q(
ne00, ne01, ne1, s01, ne11, s1,
ne02, ne12, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- ne1, J_best, fusion_device};
+ ntj, J_best, fusion_device};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
return;
}
@@ -316,7 +315,9 @@ void ggml_cuda_mul_mat_q(
ggml_cuda_pool_alloc<int32_t> ids_src1(ctx.pool(), ne_get_rows);
ggml_cuda_pool_alloc<int32_t> ids_dst(ctx.pool(), ne_get_rows + J_best-1); // Needs to be padded for unconditional memory access.
- ggml_cuda_pool_alloc<int32_t> expert_bounds(ctx.pool(), ne02 + 1);
+
+ const int ntj = (ne_get_rows + J_best - 1) / J_best + ne02;
+ ggml_cuda_pool_alloc<ggml_cuda_mm_id_tile_spec> tile_spec(ctx.pool(), ntj);
// gate/up activations are broadcast across experts (ne11 == 1): quantize each token once and
// scatter to its slots. ids_src1 then holds the inverse map (token slot -> compact row).
@@ -327,8 +328,8 @@ void ggml_cuda_mul_mat_q(
const int si1 = ids->nb[1] / ggml_element_size(ids);
const int sis1 = nb12 / nb11;
- ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(),
- ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
+ ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), tile_spec.get(),
+ ne02, ne12, n_expert_used, ne11, J_best, si1, sis1, /*write_inverse =*/ dedup_bcast, stream);
CUDA_CHECK(cudaGetLastError());
}
@@ -375,12 +376,12 @@ void ggml_cuda_mul_mat_q(
// Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
const mmq_args args = {
- src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
+ src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), tile_spec.get(), dst_d,
src1_scale.ptr,
ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
- ne02, ne02, s02, s12, s2,
+ ne02, 1, s02, s12, s2,
ne03, ne13, s03, s13, s3,
- ne12, J_best, fusion_device};
+ ntj, J_best, fusion_device};
ggml_cuda_mul_mat_q_switch_type(ctx, args, stream, prec_src1);
}
diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh
index 474b19dfe..4d67555ef 100644
--- a/ggml/src/ggml-cuda/mmq.cuh
+++ b/ggml/src/ggml-cuda/mmq.cuh
@@ -1,6 +1,7 @@
#pragma once
#include "common.cuh"
+#include "mmid.cuh"
#include "unary.cuh"
#include <climits>
@@ -14,8 +15,8 @@
typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, const char * __restrict__ x_gate, int * x_tile,
const int kbx0, const int i_max, const int stride);
typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
-typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
- float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ ids_dst,
+ float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale);
enum mmq_q8_1_ds_layout {
@@ -462,22 +463,21 @@ static __device__ __forceinline__ const block_t * ggml_cuda_mmq_x_row(
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
- const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale) {
GGML_UNUSED_VARS(glu_op, glu_limit, x_scale, gate_scale);
-
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
static_assert(I % warp_size == 0, "DP4A MMQ tile height must be a multiple of the physical warp size");
- const bool y_scale_used = y_scale != nullptr;
+ const bool y_scale_used = type == GGML_TYPE_NVFP4 && y_scale != nullptr;
#pragma unroll
for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
- if (j > j_max) {
+ if (j >= j_sup) {
return;
}
@@ -489,15 +489,10 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
continue;
}
- if constexpr (type == GGML_TYPE_NVFP4) {
- if (y_scale_used) {
- dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
- } else {
- dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
- }
+ if (y_scale_used) {
+ dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
} else {
dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
- GGML_UNUSED(y_scale_used);
}
}
}
@@ -508,7 +503,7 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
template<ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fusion>
static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
- const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max,
+ const float * __restrict__ y_scale, const int stride, const int i_max, const int j_sup,
const ggml_glu_op glu_op, const float glu_limit, const float * __restrict__ x_scale, const float * __restrict__ gate_scale) {
#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -523,7 +518,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);
- const bool y_scale_used = y_scale != nullptr;
+ const bool y_scale_used = type == GGML_TYPE_NVFP4 && y_scale != nullptr;
// Fused NVFP4: load the weight scales once, 1 without scales.
constexpr bool fused_nvfp4 = has_fusion && type == GGML_TYPE_NVFP4;
@@ -538,7 +533,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
for (int l = 0; l < (has_fusion ? tile_C::ne/2 : tile_C::ne); ++l) {
const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l);
- if (j > j_max) {
+ if (j >= j_sup) {
continue;
}
@@ -548,7 +543,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
continue;
}
- if constexpr (has_fusion) {
+ if (has_fusion) {
const int sum_index = (j0/tile_C::J + n)*tile_C::ne + l;
float up = sum[sum_index];
float gate = sum[sum_index + tile_C::ne/2];
@@ -561,15 +556,10 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
gate *= scale_gate;
}
dst[ids_dst[j]*stride + i] = ggml_cuda_op_glu_single(glu_op, gate, up, glu_limit);
- } else if constexpr (type == GGML_TYPE_NVFP4) {
- if (y_scale_used) {
- dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
- } else {
- dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
- }
+ } else if (y_scale_used) {
+ dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
} else {
dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
- GGML_UNUSED(y_scale_used);
}
}
}
@@ -937,7 +927,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
const float * __restrict__ x_scale, const float * __restrict__ gate_scale,
const int stride_row_x, const int ncols_y, const int stride_col_dst,
- const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
+ const int tile_x_i_max, const int tile_y_j_sup, const int kb0_start, const int kb0_stop) {
static_assert(!(fixup && has_fusion), "the fused gate/up path does not split K");
static_assert(!(fallback && has_fusion), "the fused gate/up path needs nrows_x % 128 == 0");
@@ -970,7 +960,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int);
for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) {
- load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x);
+ load_tiles(x, x_gate, tile_x, offset_x + kb0, tile_x_i_max, stride_row_x);
{
const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz;
#pragma unroll
@@ -1007,7 +997,7 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
if (fixup) {
write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J, glu_op, glu_limit, x_scale, gate_scale);
} else {
- write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j, glu_op, glu_limit, x_scale, gate_scale);
+ write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_i_max, tile_y_j_sup, glu_op, glu_limit, x_scale, gate_scale);
}
}
@@ -1018,13 +1008,13 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1, bool has_fu
__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1), ggml_cuda_mmq_get_occupancy(type, J, fallback, prec_src1))
static __global__ void mul_mat_q(
const char * __restrict__ x, const char * __restrict__ x_gate, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
- const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+ const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst, float * __restrict__ tmp_fixup,
const float * __restrict__ y_scale, const ggml_glu_op glu_op, const float glu_limit,
const float * __restrict__ x_scale, const float * __restrict__ gate_scale,
const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
- const uint3 ntx) {
+ const uint3 ntj) {
#ifdef TURING_MMA_AVAILABLE
constexpr bool fusion_available = true;
@@ -1049,7 +1039,7 @@ static __global__ void mul_mat_q(
constexpr int qk = ggml_cuda_type_traits<type>::qk;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
- const uint32_t nty = (nrows_x + I - 1) / I; // Number of tiles y
+ const uint32_t nti = (nrows_x + I - 1) / I; // Number of tiles in I direction
// Initialize the ids for writing back data with just the index.
// For regular matrix multiplications this is never changed.
@@ -1074,35 +1064,28 @@ static __global__ void mul_mat_q(
const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y);
const int wt = tmp2.x;
- const int zt = tmp2.y;
+ int zt = tmp2.y;
const int jt = blockIdx.x;
const int it = blockIdx.y;
// Defaults for regular matrix multiplication:
- int col_low = 0;
- int col_high = ncols_dst;
- int col_diff = ncols_dst;
+ int col_low = jt*J;
+ int col_high = min(col_low + J, ncols_dst);
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
- int offset_y_scale;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
- } else {
- GGML_UNUSED(offset_y_scale);
- }
if (ids_dst) {
- col_low = expert_bounds[zt + 0];
- col_high = expert_bounds[zt + 1];
- col_diff = col_high - col_low;
-
- offset_y = 0;
- offset_dst = 0;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = 0;
- }
+ const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+ zt = tsj.expert;
+ col_low = tsj.col_low;
+ col_high = tsj.col_high;
- if (jt*J >= col_diff) {
+ offset_y = 0;
+ offset_y_scale = 0;
+ offset_dst = 0;
+
+ if (col_low == col_high) {
return;
}
@@ -1115,21 +1098,18 @@ static __global__ void mul_mat_q(
break;
}
- ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+ ids_dst_shared[j] = ids_dst[col_low + j];
}
__syncthreads();
}
- offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
- offset_dst += it*I_dst;
- const float * y_scale_tile = nullptr;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale += col_low + jt*J;
- y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
- }
+ offset_y += col_low * (sizeof(block_q8_1_mmq)/sizeof(int));
+ offset_y_scale += col_low;
+ offset_dst += it*I_dst;
- const int tile_x_max_i = nrows_x - it*I_dst - 1;
- const int tile_y_max_j = col_diff - jt*J - 1;
+ const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
+
+ const int tile_x_i_max = nrows_x - it*I_dst - 1;
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I_dst*stride_row_x;
@@ -1141,7 +1121,7 @@ static __global__ void mul_mat_q(
mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, has_fusion>
(x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
glu_op, glu_limit, x_scale_tile, gate_scale_tile, stride_row_x, ncols_y, stride_col_dst,
- tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
+ tile_x_i_max, col_high - col_low, 0, blocks_per_ne00.z);
return;
}
@@ -1149,8 +1129,8 @@ static __global__ void mul_mat_q(
constexpr int blocks_per_iter = ITER_K / qk;
// kbc == k block continuous, current index in continuous ijk space.
- int kbc = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
- int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+ int kbc = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
+ int kbc_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter;
kbc_stop -= fastmodulo(kbc_stop, blocks_per_ne00) % blocks_per_iter;
@@ -1160,41 +1140,34 @@ static __global__ void mul_mat_q(
int kb0_stop = min(blocks_per_ne00.z, uint32_t(kb0_start + kbc_stop - kbc));
while (kbc < kbc_stop && kb0_stop == int(blocks_per_ne00.z)) {
int tmp = fastdiv(kbc, blocks_per_ne00);
- uint2 tmp2 = fast_div_modulo(tmp, ntx);
+ uint2 tmp2 = fast_div_modulo(tmp, ntj);
const int jt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nchannels_y);
- const int zt = tmp2.y;
+ int zt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nsamples_y);
const int wt = tmp2.y;
const int it = tmp2.x;
// Defaults for regular matrix multiplication:
- int col_low = 0;
- int col_high = ncols_dst;
- int col_diff = ncols_dst;
+ int col_low = jt*J;
+ int col_high = min(col_low + J, ncols_dst);
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
- int offset_y_scale;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
- } else {
- GGML_UNUSED(offset_y_scale);
- }
if (ids_dst) {
- col_low = expert_bounds[zt + 0];
- col_high = expert_bounds[zt + 1];
- col_diff = col_high - col_low;
-
- offset_y = 0;
- offset_dst = 0;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = 0;
- }
+ const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+ zt = tsj.expert;
+ col_low = tsj.col_low;
+ col_high = tsj.col_high;
+
+ offset_y = 0;
+ offset_y_scale = 0;
+ offset_dst = 0;
- if (jt*J >= col_diff) {
+ if (col_low == col_high) {
kbc += blocks_per_ne00.z;
kbc -= fastmodulo(kbc, blocks_per_ne00);
@@ -1213,21 +1186,18 @@ static __global__ void mul_mat_q(
break;
}
- ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+ ids_dst_shared[j] = ids_dst[col_low + j];
}
__syncthreads();
}
- offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
- offset_dst += it*I;
- const float * y_scale_tile = nullptr;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale += col_low + jt * J;
- y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
- }
+ offset_y += col_low * (sizeof(block_q8_1_mmq) / sizeof(int));
+ offset_y_scale += col_low;
+ offset_dst += it*I;
+
+ const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
- const int tile_x_max_i = nrows_x - it*I - 1;
- const int tile_y_max_j = col_diff - jt*J - 1;
+ const int tile_x_i_max = nrows_x - it*I - 1;
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
@@ -1235,7 +1205,7 @@ static __global__ void mul_mat_q(
mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
(x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
- tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+ tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);
kbc += blocks_per_ne00.z;
kbc -= fastmodulo(kbc, blocks_per_ne00);
@@ -1249,41 +1219,34 @@ static __global__ void mul_mat_q(
}
int tmp = fastdiv(kbc, blocks_per_ne00);
- uint2 tmp2 = fast_div_modulo(tmp, ntx);
+ uint2 tmp2 = fast_div_modulo(tmp, ntj);
const int jt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nchannels_y);
- const int zt = tmp2.y;
+ int zt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nsamples_y);
const int wt = tmp2.y;
const int it = tmp2.x;
// Defaults for regular matrix multiplication:
- int col_low = 0;
- int col_high = ncols_dst;
- int col_diff = ncols_dst;
+ int col_low = jt*J;
+ int col_high = min(col_low + J, ncols_dst);
int offset_y = wt*stride_sample_y + zt*stride_channel_y;
+ int offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
- int offset_y_scale;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
- } else {
- GGML_UNUSED(offset_y_scale);
- }
if (ids_dst) {
- col_low = expert_bounds[zt + 0];
- col_high = expert_bounds[zt + 1];
- col_diff = col_high - col_low;
+ const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+ zt = tsj.expert;
+ col_low = tsj.col_low;
+ col_high = tsj.col_high;
- offset_y = 0;
- offset_dst = 0;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale = 0;
- }
+ offset_y = 0;
+ offset_y_scale = 0;
+ offset_dst = 0;
- if (jt*J >= col_diff) {
+ if (col_low == col_high) {
return;
}
@@ -1302,16 +1265,13 @@ static __global__ void mul_mat_q(
__syncthreads();
}
- offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
- offset_dst += it*I;
- const float * y_scale_tile = nullptr;
- if constexpr (type == GGML_TYPE_NVFP4) {
- offset_y_scale += col_low + jt * J;
- y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
- }
+ offset_y += col_low * (sizeof(block_q8_1_mmq) / sizeof(int));
+ offset_y_scale += col_low;
+ offset_dst += it*I;
+
+ const float * y_scale_tile = type == GGML_TYPE_NVFP4 && y_scale ? y_scale + offset_y_scale : nullptr;
- const int tile_x_max_i = nrows_x - it*I - 1;
- const int tile_y_max_j = col_diff - jt*J - 1;
+ const int tile_x_i_max = nrows_x - it*I - 1;
const int offset_x = fastdiv(wt, sample_ratio)*stride_sample_x + fastdiv(zt, channel_ratio)*stride_channel_x + it*I*stride_row_x;
@@ -1319,16 +1279,16 @@ static __global__ void mul_mat_q(
mul_mat_q_process_tile<type, J, fallback, fixup, prec_src1, false>
(x, x_gate, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
glu_op, glu_limit, x_scale, gate_scale, stride_row_x, ncols_y, stride_col_dst,
- tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
+ tile_x_i_max, col_high - col_low, kb0_start, kb0_stop);
}
template <ggml_type type, int J, bool fallback, ggml_prec prec_src1>
__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1)/2, 1)
static __global__ void mul_mat_q_stream_k_fixup(
- const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ expert_bounds, float * __restrict__ dst,
+ const int32_t * __restrict__ ids_dst, const ggml_cuda_mm_id_tile_spec * __restrict__ tile_spec, float * __restrict__ dst,
float * __restrict__ tmp_last_tile, const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst,
const int stride_col_dst, const uint3 nchannels_y, const int stride_channel_dst, const uint3 nsamples_y,
- const int stride_sample_dst, const uint3 ntx) {
+ const int stride_sample_dst, const uint3 ntj) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nwarps = (ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / 2) / warp_size;
constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
@@ -1339,13 +1299,13 @@ static __global__ void mul_mat_q_stream_k_fixup(
float sum[J / nwarps] = {0.0f};
const int i = blockIdx.y*warp_size + threadIdx.x;
- const int nty = (nrows_x + I - 1) / I;
+ const int nti = (nrows_x + I - 1) / I;
const int bidx0 = blockIdx.x;
// kbc == k block continuous, current index in continuous ijk space.
- int kbc0 = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
- int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+ int kbc0 = int64_t(blockIdx.x) *(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
+ int kbc0_stop = int64_t(blockIdx.x + 1)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
kbc0 -= fastmodulo(kbc0, blocks_per_ne00) % blocks_per_iter;
kbc0_stop -= fastmodulo(kbc0_stop, blocks_per_ne00) % blocks_per_iter;
@@ -1364,7 +1324,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
int bidx = bidx0 - 1;
int kbc_stop = kbc0;
while(true) {
- int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntx.z*nty*blocks_per_ne00.z) / gridDim.x;
+ int kbc = int64_t(bidx)*(nsamples_y.z*nchannels_y.z*ntj.z*nti*blocks_per_ne00.z) / gridDim.x;
kbc -= fastmodulo(kbc, blocks_per_ne00) % blocks_per_iter;
if (kbc == kbc_stop) { // Did not have any data.
@@ -1396,11 +1356,11 @@ static __global__ void mul_mat_q_stream_k_fixup(
}
int tmp = fastdiv(kbc0, blocks_per_ne00);
- uint2 tmp2 = fast_div_modulo(tmp, ntx);
+ uint2 tmp2 = fast_div_modulo(tmp, ntj);
const int jt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nchannels_y);
- const int zt = tmp2.y;
+ int zt = tmp2.y;
tmp = tmp2.x;
tmp2 = fast_div_modulo(tmp, nsamples_y);
const int wt = tmp2.y;
@@ -1411,7 +1371,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
dst += offset_dst;
const int i_max = nrows_x - it*I - 1;
- const int j_max = ncols_dst - jt*J - 1;
+ const int j_sup = ncols_dst - jt*J;
if (fallback && i > i_max) {
return;
}
@@ -1420,7 +1380,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
- if (j > j_max) {
+ if (j >= j_sup) {
return;
}
@@ -1430,12 +1390,14 @@ static __global__ void mul_mat_q_stream_k_fixup(
}
__shared__ int ids_dst_shared[J];
- const int col_low = expert_bounds[zt + 0];
- const int col_high = expert_bounds[zt + 1];
- const int col_diff = col_high - col_low;
+ const ggml_cuda_mm_id_tile_spec tsj = tile_spec[jt];
+ zt = tsj.expert;
+ const int col_low = tsj.col_low;
+ const int col_high = tsj.col_high;
+ const int j_sup = col_high - col_low;
for (int j = threadIdx.y*warp_size + threadIdx.x; j < J; j += nwarps*warp_size) {
- ids_dst_shared[j] = ids_dst[col_low + jt*J + j];
+ ids_dst_shared[j] = ids_dst[col_low + j];
}
__syncthreads();
@@ -1443,7 +1405,6 @@ static __global__ void mul_mat_q_stream_k_fixup(
dst += offset_dst;
const int i_max = nrows_x - it*I - 1;
- const int j_max = col_diff - jt*J - 1;
if (fallback && i > i_max) {
return;
}
@@ -1452,7 +1413,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
for (int j0 = 0; j0 < J; j0 += nwarps) {
const int j = j0 + threadIdx.y;
- if (j > j_max) {
+ if (j >= j_sup) {
return;
}
@@ -1461,12 +1422,12 @@ static __global__ void mul_mat_q_stream_k_fixup(
}
struct mmq_args {
- const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
- const float * y_scale;
+ const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst;
+ const ggml_cuda_mm_id_tile_spec * tile_spec; float * dst; const float * y_scale;
int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
- int64_t ncols_max;
+ int64_t ntj;
int J_best; // Tile width in ne11(dense)/ne12(MoE) direction to use for optimal performance.
ggml_cuda_mm_fusion_args_device fusion = {}; // gate != nullptr selects the fused gate/up/GLU path
};
@@ -1498,39 +1459,38 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
}
const int I_dst = has_fusion ? config.I/2 : config.I; // Rows of dst per tile, see mul_mat_q.
- const int nty = (args.nrows_x + I_dst - 1) / I_dst;
- const int ntx = (args.ncols_max + config.J - 1) / config.J;
- const int ntzw = args.nchannels_y * args.nsamples_y;
+ const int nti = (args.nrows_x + I_dst - 1) / I_dst;
+ const int ntzw = args.nchannels_y * args.nsamples_y;
// Column tiles innermost so that the concurrently running blocks share the same tile of x and read it from L2.
- const dim3 block_nums_xy_tiling(ntx, nty, ntzw);
+ const dim3 block_nums_xy_tiling(args.ntj, nti, ntzw);
- GGML_ASSERT(args.nchannels_y % args.nchannels_x == 0);
- GGML_ASSERT(args.nsamples_y % args.nsamples_x == 0);
- const int channel_ratio = args.nchannels_y / args.nchannels_x;
- const int sample_ratio = args.nsamples_y / args.nsamples_x;
+ GGML_ASSERT(args.ids_dst || args.nchannels_y % args.nchannels_x == 0);
+ GGML_ASSERT( args.nsamples_y % args.nsamples_x == 0);
+ const int channel_ratio = args.ids_dst ? 1 : args.nchannels_y / args.nchannels_x;
+ const int sample_ratio = args.nsamples_y / args.nsamples_x;
const uint3 blocks_per_ne00_fd = init_fastdiv_values(args.ncols_x / ggml_cuda_type_traits<type>::qk);
- const uint3 ntx_fd = init_fastdiv_values(ntx);
+ const uint3 ntj_fd = init_fastdiv_values(args.ntj);
const uint3 nchannels_y_fd = init_fastdiv_values(args.nchannels_y);
const uint3 nsamples_y_fd = init_fastdiv_values(args.nsamples_y);
const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio);
const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio);
if (has_fusion || !config.stream_k) {
- GGML_ASSERT(nty <= 65535); // gridDim.y limit
+ GGML_ASSERT(nti <= 65535); // gridDim.y limit
mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
- (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
- (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
+ (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.tile_spec, args.dst, nullptr,
+ args.y_scale, args.fusion.glu_op, args.fusion.glu_limit, (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
+ ntj_fd);
return;
}
// For the stream-k kernel it is possible to run it with tiling by setting the number of CUDA blocks equal to the number of tiles.
// This is worthwhile if the efficiency of tiling is high and skipping the fixup kernel is more important.
- const int ntiles_dst = ntx * nty * ntzw;
+ const int ntiles_dst = args.ntj * nti * ntzw;
const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);
@@ -1551,12 +1511,12 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
mul_mat_q<type, J, fallback, prec_src1, has_fusion><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
- (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale, args.fusion.glu_op, args.fusion.glu_limit,
- (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
+ (args.x, (const char *) args.fusion.gate, args.y, args.ids_dst, args.tile_spec, args.dst, tmp_fixup.ptr,
+ args.y_scale, args.fusion.glu_op, args.fusion.glu_limit, (const float *) args.fusion.x_scale, (const float *) args.fusion.gate_scale,
blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
- ntx_fd);
+ ntj_fd);
if (!fixup_needed) {
return;
@@ -1564,9 +1524,9 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
CUDA_CHECK(cudaGetLastError());
mul_mat_q_stream_k_fixup<type, J, fallback, prec_src1><<<block_nums_fixup, block_dims_fixup, 0, stream>>>
- (args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
+ (args.ids_dst, args.tile_spec, args.dst, tmp_fixup.ptr, blocks_per_ne00_fd, args.nrows_x, args.ncols_dst,
args.nrows_dst, nchannels_y_fd, args.stride_channel_dst, nsamples_y_fd, args.stride_sample_dst,
- ntx_fd);
+ ntj_fd);
}
template <ggml_type type, bool fallback, ggml_prec prec_src1, bool has_fusion>