vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-10-03 19:45:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge commit '2a74817f93' into concedo_experimental

Browse source 
# Conflicts:
#	ci/run.sh
#	docs/backend/SYCL.md
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	ggml/src/ggml-opencl/kernels/mul_mm_q4_k_f32_l4_lm.cl
#	ggml/src/ggml-opencl/kernels/mul_mm_q5_k_f32_l4_lm.cl
#	ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32_flat.cl
#	ggml/src/ggml-opencl/kernels/mul_mv_q5_k_f32_flat.cl
#	ggml/src/ggml-sycl/common.hpp
#	ggml/src/ggml-sycl/ggml-sycl.cpp
#	ggml/src/ggml-webgpu/ggml-webgpu.cpp
#	tests/test-backend-ops.cpp
#	tests/test-save-load-state.cpp
This commit is contained in:
Concedo 2026-09-02 23:13:14 +08:00
commit 3bf2e14605
17 changed files with 776 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

16
ggml/src/ggml-cuda/ggml-cuda.cu
View file

@ -1807,7 +1807,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
return false;
}
if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) {
if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] > get_mmvq_mmid_max_batch(src0->type, cc)) {
return false;
}
@ -2991,9 +2991,10 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
};
bool is_ok = true;
// exception for topk-moe, as each row is read entirely before writing
if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) {
return true;
// one block reads all logits before it writes, so logits may alias the out nodes
const ggml_tensor * logits_may_alias = nullptr;
if (is_topk_moe && ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK) {
logits_may_alias = cgraph->nodes[node_idx]->src[0];
}
for (int i = 0; i < out_count; ++i) {
@ -3007,7 +3008,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
const ggml_tensor * src = cgraph->nodes[j]->src[src_idx];
if (!src || src->op == GGML_OP_NONE) {
if (!src || src->op == GGML_OP_NONE || src == logits_may_alias) {
continue;
}
@ -5285,6 +5286,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
case GGML_OP_SUM:
return ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_TOP_K:
#if defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
return true;
#else
return op->src[0]->ne[0] <= 1024;
#endif // defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
case GGML_OP_ARGSORT:
#ifndef GGML_CUDA_USE_CUB
return op->src[0]->ne[0] <= 1024;

115
ggml/src/ggml-cuda/mmvq.cu
View file

@ -773,10 +773,10 @@ static __global__ void mul_mat_vec_q(
// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst)
// Block: (warp_size, ncols_dst) - each warp handles one token independently.
// No shared memory reduction needed since each warp works alone.
template <ggml_type type, int c_rows_per_block>
template <ggml_type type, int c_rows_per_block, bool has_fusion = false>
__launch_bounds__(get_mmvq_mmid_max_batch_for_device<type>()*ggml_cuda_get_physical_warp_size(), 1)
static __global__ void mul_mat_vec_q_moe(
const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr,
const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion,
float * dst_ptr,
const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
@ -794,6 +794,29 @@ static __global__ void mul_mat_vec_q_moe(
constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type);
// fuse gate, bias, scales, and glu_op into the up projection
bool use_gate = false;
const void * vgate = nullptr;
const float * x_bias = nullptr;
const float * gate_bias = nullptr;
const float * x_scale = nullptr;
const float * gate_scale = nullptr;
ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU;
float glu_limit = 0.0f;
if constexpr (has_fusion) {
use_gate = fusion.gate != nullptr;
vgate = fusion.gate;
x_bias = (const float *) fusion.x_bias;
gate_bias = (const float *) fusion.gate_bias;
active_glu = fusion.glu_op;
glu_limit = fusion.glu_limit;
if constexpr (type == GGML_TYPE_NVFP4) {
x_scale = (const float *) fusion.x_scale;
gate_scale = (const float *) fusion.gate_scale;
}
}
const uint32_t token_idx = threadIdx.y;
const int row0 = c_rows_per_block*blockIdx.x;
const int blocks_per_row_x = ncols_x / qk;
@ -814,6 +837,7 @@ static __global__ void mul_mat_vec_q_moe(
// partial sum for each thread
float tmp[c_rows_per_block] = {0.0f};
float tmp_gate[c_rows_per_block] = {0.0f};
for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) {
const int kby = kbx * (qk/QK8_1);
@ -822,6 +846,11 @@ static __global__ void mul_mat_vec_q_moe(
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
if constexpr (has_fusion) {
if (use_gate) {
tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs);
}
}
}
}
@ -831,11 +860,63 @@ static __global__ void mul_mat_vec_q_moe(
#pragma unroll
for (int i = 0; i < c_rows_per_block; ++i) {
tmp[i] = warp_reduce_sum<warp_size>(tmp[i]);
if constexpr (has_fusion) {
if (use_gate) {
tmp_gate[i] = warp_reduce_sum<warp_size>(tmp_gate[i]);
}
}
}
// Write results
if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) {
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = tmp[threadIdx.x];
float result = tmp[threadIdx.x];
if constexpr (has_fusion) {
const uint32_t bias_idx = channel_x*stride_channel_dst + row0 + threadIdx.x;
if constexpr (type == GGML_TYPE_NVFP4) {
if (x_scale) {
result *= x_scale[channel_x];
}
}
if (x_bias) {
result += x_bias[bias_idx];
}
if (use_gate) {
float gate_value = tmp_gate[threadIdx.x];
if constexpr (type == GGML_TYPE_NVFP4) {
if (gate_scale) {
gate_value *= gate_scale[channel_x];
}
}
if (gate_bias) {
gate_value += gate_bias[bias_idx];
}
switch (active_glu) {
case GGML_GLU_OP_SWIGLU:
result *= ggml_cuda_op_silu_single(gate_value);
break;
case GGML_GLU_OP_GEGLU:
result *= ggml_cuda_op_gelu_single(gate_value);
break;
case GGML_GLU_OP_SWIGLU_OAI:
result = ggml_cuda_op_swiglu_oai_single(gate_value, result);
break;
case GGML_GLU_OP_SWIGLU_CLAMP:
result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit);
break;
default:
result = result * gate_value;
break;
}
}
}
dst[channel_dst*stride_channel_dst + token_idx*stride_col_dst + row0 + threadIdx.x] = result;
}
if constexpr (!has_fusion) {
GGML_UNUSED_VARS(use_gate, tmp_gate, vgate, x_bias, gate_bias, active_glu, glu_limit, x_scale, gate_scale);
} else if constexpr (type != GGML_TYPE_NVFP4) {
GGML_UNUSED_VARS(x_scale, gate_scale);
}
}
@ -885,7 +966,7 @@ static void mul_mat_vec_q_switch_fusion(
template <ggml_type type>
static void mul_mat_vec_q_moe_launch(
const void * vx, const void * vy, const int32_t * ids, float * dst,
const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst,
const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x,
const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst,
const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst,
@ -898,11 +979,22 @@ static void mul_mat_vec_q_moe_launch(
const dim3 block_dims(warp_size, ncols_dst);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block>, launch_params,
vx, vy, ids, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr ||
fusion.x_scale != nullptr || fusion.gate_scale != nullptr;
if (has_fusion) {
ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, true>, launch_params,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
} else {
ggml_cuda_kernel_launch(mul_mat_vec_q_moe<type, rows_per_block, false>, launch_params,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride);
}
}
template <ggml_type type>
@ -998,7 +1090,7 @@ static void mul_mat_vec_q_switch_ncols_dst(
if (has_ids && ncols_dst > 1) {
// Multi-token MUL_MAT_ID path - dedicated MoE kernel
mul_mat_vec_q_moe_launch<type>(
vx, vy, ids, dst, ncols_x, nchannels_y_fd, nrows_x,
vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x,
stride_row_x, stride_col_y, stride_col_dst,
stride_channel_x, stride_channel_y, stride_channel_dst,
ncols_dst, ids_stride, warp_size, nchannels_dst, stream);
@ -1280,7 +1372,8 @@ void ggml_cuda_mul_mat_vec_q(
ggml_cuda_mm_fusion_args_device fusion_local{};
if (fusion) {
GGML_ASSERT( !ids || dst->ne[2] == 1);
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
GGML_ASSERT( !ids || dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc));
GGML_ASSERT( ids || dst->ne[1] == 1);
// Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is
// non-negligible for some models such as gpt-oss-20b

180
ggml/src/ggml-cuda/top-k.cu
View file

@ -48,6 +48,168 @@ static int next_power_of_2(int x) {
#endif // CUB_TOP_K_AVAILABLE
#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP)
static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) {
const uint32_t bits = __float_as_uint(value);
const uint32_t mask = (uint32_t) (-(int32_t) (bits >> 31)) | 0x80000000U;
return bits ^ mask;
}
struct top_k_radix_state {
uint32_t prefix;
uint32_t prefix_mask;
int rank;
int greater_count;
int equal_count;
};
static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) {
const int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row < nrows) {
states[row] = {0, 0, k, 0, 0};
}
}
template<int BLOCK_SIZE, int RADIX_BITS>
static __global__ void top_k_radix_histogram(
const float * __restrict__ src,
const top_k_radix_state * __restrict__ states,
int * __restrict__ block_histograms,
int ncols,
int blocks_per_row,
int shift) {
constexpr int NBINS = 1 << RADIX_BITS;
const int row = blockIdx.x / blocks_per_row;
const int row_block = blockIdx.x % blocks_per_row;
const int tid = threadIdx.x;
const float * row_src = src + (size_t) row * ncols;
__shared__ int histogram[NBINS];
histogram[tid] = 0;
__syncthreads();
const top_k_radix_state state = states[row];
for (int col = row_block * BLOCK_SIZE + tid;
col < ncols;
col += blocks_per_row * BLOCK_SIZE) {
const uint32_t key = top_k_float_to_ordered(row_src[col]);
if ((key & state.prefix_mask) == state.prefix) {
atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1);
}
}
__syncthreads();
const size_t histogram_offset =
((size_t) row * blocks_per_row + row_block) * NBINS;
block_histograms[histogram_offset + tid] = histogram[tid];
}
template<int BLOCK_SIZE, int RADIX_BITS>
static __global__ void top_k_radix_select(
const int * __restrict__ block_histograms,
top_k_radix_state * __restrict__ states,
int blocks_per_row,
int shift) {
constexpr int NBINS = 1 << RADIX_BITS;
const int row = blockIdx.x;
const int tid = threadIdx.x;
__shared__ int histogram[NBINS];
int count = 0;
for (int row_block = 0; row_block < blocks_per_row; ++row_block) {
const size_t offset = ((size_t) row * blocks_per_row + row_block) * NBINS;
count += block_histograms[offset + tid];
}
histogram[tid] = count;
__syncthreads();
if (tid == 0) {
top_k_radix_state state = states[row];
int bin = NBINS - 1;
while (bin > 0 && histogram[bin] < state.rank) {
state.rank -= histogram[bin--];
}
state.prefix |= (uint32_t) bin << shift;
state.prefix_mask |= (uint32_t) (NBINS - 1) << shift;
states[row] = state;
}
}
static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) {
const int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row < nrows) {
states[row].greater_count = 0;
states[row].equal_count = 0;
}
}
template<int BLOCK_SIZE>
static __global__ void top_k_radix_gather(
const float * __restrict__ src,
int * __restrict__ dst,
top_k_radix_state * __restrict__ states,
int ncols,
int k,
int blocks_per_row) {
const int row = blockIdx.x / blocks_per_row;
const int row_block = blockIdx.x % blocks_per_row;
const int tid = threadIdx.x;
const float * row_src = src + (size_t) row * ncols;
int * row_dst = dst + (size_t) row * k;
top_k_radix_state * state = &states[row];
for (int col = row_block * BLOCK_SIZE + tid;
col < ncols;
col += blocks_per_row * BLOCK_SIZE) {
const uint32_t key = top_k_float_to_ordered(row_src[col]);
if (key > state->prefix) {
const int pos = atomicAdd(&state->greater_count, 1);
row_dst[pos] = col;
} else if (key == state->prefix) {
const int pos = atomicAdd(&state->equal_count, 1);
if (pos < state->rank) {
row_dst[k - state->rank + pos] = col;
}
}
}
}
static void top_k_radix_cuda(
ggml_cuda_pool & pool,
const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) {
constexpr int BLOCK_SIZE = 256;
constexpr int RADIX_BITS = 8;
constexpr int NBINS = 1 << RADIX_BITS;
const int blocks_per_row = std::min((ncols + 1023) / 1024, 64);
ggml_cuda_pool_alloc<top_k_radix_state> states_alloc(pool, nrows);
ggml_cuda_pool_alloc<int> histograms_alloc(pool, (size_t) nrows * blocks_per_row * NBINS);
top_k_radix_state * states = states_alloc.get();
int * histograms = histograms_alloc.get();
top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k);
const dim3 row_grid(blocks_per_row * nrows);
for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS>
<<<row_grid, BLOCK_SIZE, 0, stream>>>(
src, states, histograms, ncols, blocks_per_row, shift);
top_k_radix_select<BLOCK_SIZE, RADIX_BITS>
<<<nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift);
}
top_k_radix_reset_counters
<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows);
top_k_radix_gather<BLOCK_SIZE>
<<<row_grid, BLOCK_SIZE, 0, stream>>>(
src, dst, states, ncols, k, blocks_per_row);
}
#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP)
void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const float * src0_d = (const float *) src0->data;
@ -96,10 +258,18 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
dst_d += k * iter_nrows;
}
#else // GGML_CUDA_USE_CUB
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
cudaMemcpyDeviceToDevice, stream));
#if defined(GGML_USE_HIP)
if (ncols > 1024) {
top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream);
} else {
#endif // defined(GGML_USE_HIP)
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
cudaMemcpyDeviceToDevice, stream));
#if defined(GGML_USE_HIP)
}
#endif // defined(GGML_USE_HIP)
#endif
}

24
ggml/src/ggml-cuda/topk-moe.cu
View file

@ -88,15 +88,16 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models
*/
template <int n_experts, bool has_bias>
__launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits,
float * weights,
int32_t * ids,
float * bias,
const int n_rows,
const int n_expert_used,
const float clamp_val,
const float scale_val,
const topk_moe_config config) {
__launch_bounds__(TOPK_MOE_ROWS_PER_BLOCK * WARP_SIZE, 1)
__global__ void topk_moe_cuda(const float * logits,
float * weights,
int32_t * ids,
float * bias,
const int n_rows,
const int n_expert_used,
const float clamp_val,
const float scale_val,
const topk_moe_config config) {
const int row = blockIdx.x * blockDim.y + threadIdx.y;
if (row >= n_rows) {
return;
@ -123,6 +124,9 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
wt[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? logits[expert] : -INFINITY;
}
// Weights and IDs can alias logits, so wait until every row in the block reads its logits.
__syncthreads();
if (!config.delayed_softmax) {
if (config.use_sigmoid) {
sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
@ -282,7 +286,7 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
const topk_moe_config config) {
GGML_ASSERT(!(config.with_norm && config.delayed_softmax) &&
"delayed softmax is not supported with weight normalization");
const int rows_per_block = 4;
const int rows_per_block = TOPK_MOE_ROWS_PER_BLOCK;
dim3 grid_dims((n_rows + rows_per_block - 1) / rows_per_block, 1, 1);
dim3 block_dims(WARP_SIZE, rows_per_block, 1);
cudaStream_t stream = ctx.stream();

3
ggml/src/ggml-cuda/topk-moe.cuh
View file

@ -3,6 +3,9 @@
#include <initializer_list>
// Rows that one CUDA block handles.
#define TOPK_MOE_ROWS_PER_BLOCK 8
struct ggml_cuda_topk_moe_args {
bool sigmoid{};
bool sqrt_softplus{};

19
ggml/src/ggml-metal/ggml-metal-device.cpp
View file

@ -1336,7 +1336,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_
return res;
}
// note: reuse the argsort kernel for top_k
// note: reuse the argsort kernel for the bitonic top_k fallback
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_TOP_K);
@ -1364,6 +1364,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_TOP_K);
char base[256];
char name[256];
snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type));
snprintf(name, 256, "%s", base);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
}
return res;
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_TOP_K);

1
ggml/src/ggml-metal/ggml-metal-device.h
View file

@ -145,6 +145,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse );
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op);

11
ggml/src/ggml-metal/ggml-metal-impl.h
View file

@ -1189,6 +1189,17 @@ typedef struct {
int32_t len;
} ggml_metal_kargs_argsort_merge;
typedef struct {
int32_t ne00; // number of columns (elements per row)
int32_t ne01; // rows
int32_t ne02;
int32_t ne03;
uint64_t nb01; // row stride in src0
uint64_t nb02;
uint64_t nb03;
int32_t top_k; // k
} ggml_metal_kargs_top_k;
typedef struct {
int32_t nrows;
} ggml_metal_kargs_fwht;

72
ggml/src/ggml-metal/ggml-metal-ops.cpp
View file

@ -5091,7 +5091,9 @@ int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) {
return 1;
}
int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
// bitonic-sort + merge fallback: efficient when k is small and there are few rows,
// where the single-workgroup-per-row radix-select cannot reach enough parallelism
static void ggml_metal_op_top_k_bitonic(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
ggml_metal_library_t lib = ctx->lib;
@ -5199,6 +5201,74 @@ int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
len <<= 1;
}
}
// radix-select: one workgroup per row. Maps each float to an order-preserving unsigned
// key, finds the k-th largest via 4 radix-8 histogram passes, then compacts the top-k
// indices. Fast for large k and/or many rows.
static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
auto pipeline = ggml_metal_library_get_pipeline_top_k_radix(lib, op);
// one workgroup per row; radix-select the k-th largest value
const int nth = std::min(1024, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
ggml_metal_kargs_top_k args = {
/*.ne00 =*/ ne00,
/*.ne01 =*/ ne01,
/*.ne02 =*/ ne02,
/*.ne03 =*/ ne03,
/*.nb01 =*/ nb01,
/*.nb02 =*/ nb02,
/*.nb03 =*/ nb03,
/*.top_k =*/ (int32_t) op->ne[0],
};
// shared memory: 256-entry histogram + bucket/above scalars + output counter
const size_t smem_histo = GGML_PAD(256*sizeof(uint32_t), 16);
const size_t smem_bucket = GGML_PAD( sizeof(uint32_t), 16);
const size_t smem_above = GGML_PAD( sizeof(uint32_t), 16);
const size_t smem_out = GGML_PAD( sizeof(uint32_t), 16);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_histo, 0);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_bucket, 1);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_above, 2);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_out, 3);
ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1);
}
int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
// radix-select has a fixed single-workgroup-per-row cost (~50-60us) that is only
// amortized for long rows, many rows, or a large k; otherwise the bitonic path wins
const int ncols = op->src[0]->ne[0];
const int k = op->ne[0];
const int nrows = ggml_nrows(op->src[0]);
const bool use_radix =
ncols > 2048 && (k > 64 || (nrows > 4 && ncols >= 8192));
if (use_radix) {
ggml_metal_op_top_k_radix(ctx, idx);
} else {
ggml_metal_op_top_k_bitonic(ctx, idx);
}
return 1;
}

190
ggml/src/ggml-metal/ggml-metal-tuning.cpp
View file

@ -68,6 +68,196 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) {
// sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics.
// ref: https://github.com/ggml-org/llama.cpp/pull/27824
constexpr fa_vec_entry_t fa_vec_tuned_table[] = {
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 2, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 3 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 1 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 3 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 32, 32, 3, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 1, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 2, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 1, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 96, 96, 3, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 2, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 192, 128, 3, 4 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 192, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 512, 512, 3, 3 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 2, 2 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 512, 512, 3, 2 }, { 1, 1 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } },
{ { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } },
{ { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } },

105
ggml/src/ggml-metal/kernels/argsort.metal
View file

@ -230,3 +230,108 @@ kernel void kernel_argsort_merge_f32_i32(
template [[host_name("kernel_argsort_merge_f32_i32_asc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_ASC>;
template [[host_name("kernel_argsort_merge_f32_i32_desc")]] kernel argsort_merge_t kernel_argsort_merge_f32_i32<GGML_SORT_ORDER_DESC>;
static inline uint ggml_top_k_f2ui(float x) {
uint y = as_type<uint>(x);
if ((y & 0x80000000u) != 0u) {
y ^= 0xFFFFFFFFu; // negative floats: flip all bits
} else {
y |= 0x80000000u; // positive floats: set the sign bit
}
return y;
}
kernel void kernel_top_k_f32_i32(
constant ggml_metal_kargs_top_k & args,
device const char * src0,
device int32_t * dst,
threadgroup atomic_uint * histo [[threadgroup(0)]],
threadgroup uint * sh_bucket [[threadgroup(1)]],
threadgroup uint * sh_above [[threadgroup(2)]],
threadgroup atomic_uint * out_count [[threadgroup(3)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
const uint ncols = args.ne00;
const uint top_k = args.top_k;
const uint i01 = tgpig[0];
const uint i02 = tgpig[1];
const uint i03 = tgpig[2];
device const float * src0_row = (device const float *) (src0 + args.nb01*i01 + args.nb02*i02 + args.nb03*i03);
device int32_t * dst_row = dst + top_k*(i01 + args.ne01*i02 + args.ne01*args.ne02*i03);
const uint tid = tpitg.x;
const uint ntg_x = ntg.x;
uint prefix = 0; // fixed high bits of the threshold key
uint desired = top_k; // count still needed from the candidate range
for (int shift = 24; shift >= 0; shift -= 8) {
for (uint i = tid; i < 256; i += ntg_x) {
atomic_store_explicit(&histo[i], 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint hi_mask = (shift + 8 >= 32) ? 0u : (0xFFFFFFFFu << uint(shift + 8));
const uint prefix_hi = prefix & hi_mask;
for (uint i = tid; i < ncols; i += ntg_x) {
const uint key = ggml_top_k_f2ui(src0_row[i]);
if ((key & hi_mask) == prefix_hi) {
atomic_fetch_add_explicit(&histo[(key >> uint(shift)) & 0xFFu], 1u, memory_order_relaxed);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// top-down scan for the bucket holding the k-th value
if (tid == 0) {
uint acc = 0;
uint b = 0;
for (int bb = 255; bb >= 0; --bb) {
const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed);
if (acc + c >= desired) {
b = uint(bb);
break;
}
acc += c;
}
*sh_bucket = b;
*sh_above = acc;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
prefix |= *sh_bucket << uint(shift);
desired -= *sh_above;
// ensure every thread has consumed sh_bucket/sh_above before the next pass
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0) {
atomic_store_explicit(out_count, 0u, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// emit everything above the threshold, then fill the rest from ties
const uint threshold = prefix;
for (uint i = tid; i < ncols; i += ntg_x) {
if (ggml_top_k_f2ui(src0_row[i]) > threshold) {
const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
dst_row[pos] = (int32_t) i;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = tid; i < ncols; i += ntg_x) {
if (ggml_top_k_f2ui(src0_row[i]) == threshold) {
const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
if (pos < top_k) {
dst_row[pos] = (int32_t) i;
}
}
}
}

21
src/llama-context.cpp
View file

@ -2917,17 +2917,28 @@ public:
}
if (mbuf_cur.n_tensors == mbuf.n_tensors) {
// same chunking: copy 1:1 by index
// an equal tensor count does not imply the same chunking, e.g. save ranges [2,1] vs restore runs [1,2]
bool same_chunking = true;
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i]));
ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
if (ggml_nbytes(mbuf_cur.cpy[i]) != ggml_nbytes(mbuf.org[i])) {
same_chunking = false;
break;
}
}
if (same_chunking) {
// same chunking: copy 1:1 by index
for (size_t i = 0; i < mbuf_cur.org.size(); ++i) {
ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]);
}
continue;
}
continue;
}
// different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org)
// with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk
// it differently, so copy across tensor boundaries rather than 1:1 by index.
// it differently (even with an equal number of tensors), so copy across tensor boundaries rather than
// 1:1 by index.
const size_t total = mbuf_cur.total_size;
ggml_init_params params_scratch = {

64
src/llama-kv-cache.cpp
View file

@ -2538,6 +2538,24 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo) {
auto & cells = v_cells[strm];
// batch the scatter reads per contiguous run of destination indices
// from inclusive, to exclusive - same convention as cell_ranges_t
// contiguous cells yield a single run covering the whole block
struct cell_run { uint32_t from; uint32_t to; };
std::vector<cell_run> runs;
if (cell_count > 0) {
const auto & idxs = sinfo.idxs[0];
uint32_t i0 = 0;
while (i0 < cell_count) {
uint32_t i1 = i0 + 1;
while (i1 < cell_count && idxs[i1] == idxs[i1 - 1] + 1) {
++i1;
}
runs.push_back({idxs[i0], idxs[i1 - 1] + 1});
i0 = i1;
}
}
uint32_t v_trans;
uint32_t n_layer;
@ -2585,17 +2603,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells, single memcpy
io.read_tensor(k, sinfo.head() * k_size_row, cell_count * k_size_row);
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = sinfo.idxs[0][i] * k_size_row;
io.read_tensor(k, dst_offset, k_size_row);
}
}
for (const auto & r : runs) {
io.read_tensor(k, (size_t) r.from * k_size_row, (size_t) (r.to - r.from) * k_size_row);
}
}
@ -2628,17 +2637,8 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells, single memcpy
io.read_tensor(v, sinfo.head() * v_size_row, cell_count * v_size_row);
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = sinfo.idxs[0][i] * v_size_row;
io.read_tensor(v, dst_offset, v_size_row);
}
}
for (const auto & r : runs) {
io.read_tensor(v, (size_t) r.from * v_size_row, (size_t) (r.to - r.from) * v_size_row);
}
}
} else {
@ -2679,22 +2679,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32
return false;
}
if (cell_count) {
if (sinfo.is_contiguous()) {
// Fast path: contiguous cells
const uint32_t h = sinfo.head();
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
const size_t dst_offset = (h + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, cell_count * v_size_el);
}
} else {
// Slow path: scatter to non-contiguous positions
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
for (uint32_t i = 0; i < cell_count; ++i) {
const size_t dst_offset = (sinfo.idxs[0][i] + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, v_size_el);
}
}
for (uint32_t j = 0; j < n_embd_v_gqa; ++j) {
for (const auto & r : runs) {
const size_t dst_offset = ((size_t) r.from + j * cells.size()) * v_size_el;
io.read_tensor(v, dst_offset, (size_t) (r.to - r.from) * v_size_el);
}
}
}

25
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte
View file

@ -7,7 +7,12 @@
ChatMessageSystem,
ChatMessageUser
} from '$lib/components/app/chat';
import { REASONING_TAGS, ROUTES, SYSTEM_MESSAGE_PLACEHOLDER } from '$lib/constants';
import {
AGENTIC_TEXT_COPY_SEPARATOR,
REASONING_TAGS,
ROUTES,
SYSTEM_MESSAGE_PLACEHOLDER
} from '$lib/constants';
import { setChatMessageActionsContext, setChatMessageEditContext } from '$lib/contexts';
import { AgenticSectionType, AttachmentType, MessageRole } from '$lib/enums';
import { DatabaseService } from '$lib/services/database.service';
@ -237,6 +242,24 @@
}
function handleCopy() {
// Agentic sessions render as a single entry anchored on the first assistant
// turn, whose own content is typically just the first tool call. Copy the
// text sections of the whole session so the clipboard matches the visible
// response instead of the anchor turn.
if (message.role === MessageRole.ASSISTANT) {
const sections = deriveAgenticSections(message, toolMessages, [], false);
const text = sections
.filter((section) => section.type === AgenticSectionType.TEXT)
.map((section) => section.content)
.join(AGENTIC_TEXT_COPY_SEPARATOR);
if (text) {
chatActions.copy(message, text);
return;
}
}
chatActions.copy(message);
}

4
tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte
View file

@ -29,10 +29,10 @@
refreshAllMessages();
},
copy: async (message: DatabaseMessage) => {
copy: async (message: DatabaseMessage, contentOverride?: string) => {
const asPlainText = Boolean(currentConfig.copyTextAttachmentsAsPlainText);
const clipboardContent = formatMessageForClipboard(
message.content,
contentOverride ?? message.content,
message.extra,
asPlainText
);

4
tools/ui/src/lib/constants/agentic.constants.ts
View file

@ -20,6 +20,10 @@ export const SEARCH_SUMMARY = {
// wraps mid-paragraph.
export const RESULT_STAT_SEPARATOR = ' - ';
// Separator between the assistant text sections of a grouped agentic
// session when they are joined for the clipboard.
export const AGENTIC_TEXT_COPY_SEPARATOR = '\n\n';
export const DEFAULT_AGENTIC_CONFIG: AgenticConfig = {
enabled: true,
maxTurns: 100

2
tools/ui/src/lib/types/chat.d.ts vendored
View file

@ -249,7 +249,7 @@ export interface ChatMessageDeletionInfo {
* refresh + user-action notification), passed to each ChatMessage as a prop.
*/
export interface ChatMessageActions {
copy: (message: DatabaseMessage) => void;
copy: (message: DatabaseMessage, contentOverride?: string) => void;
delete: (message: DatabaseMessage) => void;
navigateToSibling: (siblingId: string) => void;
editWithBranching: (