hugo/DS4Server
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

Add GLM 5.3 Flash support

Browse Source
This commit is contained in:
Georg Bauer
2026-09-01 19:06:20 +02:00
parent 9a33c61ea6
commit 46d6a976a5
31 changed files with 7806 additions and 447 deletions
Download Patch File Download Diff File Expand all files Collapse all files

274
metal/dsv4_misc.metal
View File

@@ -195,6 +195,17 @@ struct ds4_metal_args_glm_store_indexer_k {
float pad1;
};
struct ds4_metal_args_glm53_indexer_pool_update {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t head_dim;
uint32_t pool_size;
uint32_t cache_f16;
float eps;
uint32_t pad0;
};
struct ds4_metal_args_glm_attention_full {
uint32_t pos0;
uint32_t n_tokens;
@@ -221,6 +232,15 @@ struct ds4_metal_args_glm_fill_selected_range_batch {
uint32_t pad_row;
};
struct ds4_metal_args_glm53_expand_pool_selection {
uint32_t n_tokens;
uint32_t pos0;
uint32_t selected_pools;
uint32_t index_topk;
uint32_t pool_size;
uint32_t output_width;
};
struct ds4_metal_args_glm_indexer_rope_tail {
uint32_t n_tokens;
uint32_t n_head;
@@ -252,6 +272,8 @@ struct ds4_metal_args_glm_indexer_scores_batch {
uint32_t head_dim;
uint32_t pos0;
uint32_t cache_f16;
uint32_t row_group_size;
uint32_t pad0;
uint64_t q_token_stride;
uint64_t q_head_stride;
uint64_t weights_token_stride;
@@ -259,6 +281,13 @@ struct ds4_metal_args_glm_indexer_scores_batch {
float scale;
};
static inline uint glm_indexer_batch_visible_rows(
constant ds4_metal_args_glm_indexer_scores_batch &args,
uint token) {
const uint group = max(args.row_group_size, 1u);
return min((args.pos0 + token + 1u) / group, args.n_rows);
}
struct ds4_metal_args_glm_qk_lowrank {
uint32_t n_head;
uint32_t kv_lora_dim;
@@ -933,6 +962,117 @@ kernel void kernel_glm_store_indexer_k(
}
}
static inline float glm53_pool_bf16_to_f32(ushort value) {
return as_type<float>((uint)value << 16);
}
kernel void kernel_glm53_indexer_pool_update(
constant ds4_metal_args_glm53_indexer_pool_update &args,
device const char *raw_k,
device const char *gate,
device const float *norm_weight,
device const float *norm_bias,
device const ushort *ape,
device char *pool_cache,
device float *tail_k,
device float *tail_gate,
threadgroup float *shared [[threadgroup(0)]],
uint tid [[thread_index_in_threadgroup]],
uint3 tgpig [[threadgroup_position_in_grid]]) {
if (args.head_dim == 0u || args.pool_size == 0u ||
tid >= args.head_dim || args.n_tokens == 0u) return;
const uint pool = args.pos0 / args.pool_size + tgpig.x;
const uint pool_start = pool * args.pool_size;
const uint input_end = args.pos0 + args.n_tokens;
if (pool_start >= input_end || pool_start + args.pool_size <= args.pos0) return;
threadgroup float *rows = shared;
threadgroup float *mean = rows + args.pool_size * args.head_dim;
threadgroup float *inv = mean + args.pool_size;
const bool complete = pool_start + args.pool_size <= input_end;
for (uint r = 0; r < args.pool_size; r++) {
const uint pos = pool_start + r;
float k_value = 0.0f;
float gate_value = 0.0f;
if (pos >= args.pos0 && pos < input_end) {
const uint src_row = pos - args.pos0;
k_value = ((device const float *)raw_k)[
(uint64_t)src_row * args.head_dim + tid];
gate_value = ((device const float *)gate)[
(uint64_t)src_row * args.head_dim + tid];
if (!complete) {
tail_k[(uint64_t)r * args.head_dim + tid] = k_value;
tail_gate[(uint64_t)r * args.head_dim + tid] = gate_value;
}
} else {
k_value = tail_k[(uint64_t)r * args.head_dim + tid];
gate_value = tail_gate[(uint64_t)r * args.head_dim + tid];
}
rows[(uint64_t)r * args.head_dim + tid] = k_value;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (!complete || pool >= (args.cache_cap + args.pool_size - 1u) / args.pool_size) {
return;
}
if (tid < args.pool_size) {
const uint r = tid;
float sum = 0.0f;
for (uint d = 0; d < args.head_dim; d++) {
sum += rows[(uint64_t)r * args.head_dim + d];
}
const float m = sum / (float)args.head_dim;
float ss = 0.0f;
for (uint d = 0; d < args.head_dim; d++) {
const float delta = rows[(uint64_t)r * args.head_dim + d] - m;
ss += delta * delta;
}
mean[r] = m;
inv[r] = rsqrt(ss / (float)args.head_dim + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float max_logit = -INFINITY;
float logits[4];
for (uint r = 0; r < args.pool_size; r++) {
const uint pos = pool_start + r;
float gate_value;
if (pos >= args.pos0) {
const uint src_row = pos - args.pos0;
gate_value = ((device const float *)gate)[
(uint64_t)src_row * args.head_dim + tid];
} else {
gate_value = tail_gate[(uint64_t)r * args.head_dim + tid];
}
logits[r] = gate_value +
glm53_pool_bf16_to_f32(ape[(uint64_t)r * args.head_dim + tid]);
max_logit = max(max_logit, logits[r]);
}
float denom = 0.0f;
for (uint r = 0; r < args.pool_size; r++) {
logits[r] = exp(logits[r] - max_logit);
denom += logits[r];
}
float pooled = 0.0f;
for (uint r = 0; r < args.pool_size; r++) {
const float normalized =
(rows[(uint64_t)r * args.head_dim + tid] - mean[r]) * inv[r] *
norm_weight[tid] + norm_bias[tid];
pooled += (logits[r] / denom) * normalized;
}
const uint64_t dst_index = (uint64_t)pool * args.head_dim + tid;
if (args.cache_f16 != 0u) {
((device half *)pool_cache)[dst_index] = (half)pooled;
} else {
((device float *)pool_cache)[dst_index] = pooled;
}
}
static inline void glm_dense_cache_store_f32_or_f16(
device char *base,
uint64_t index,
@@ -1011,7 +1151,8 @@ kernel void kernel_glm_build_kv_cache(
corr_dims);
}
const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope;
const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1096,7 +1237,8 @@ kernel void kernel_glm_build_kv_cache_decode_group4(
corr_dims);
}
const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope;
const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += 512u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1179,7 +1321,8 @@ kernel void kernel_glm_build_kv_cache_flash(
corr_dims);
}
const float theta_base = (float)pos;
const float inv_ndims = -1.0f / (float)args.qk_rope;
const float inv_ndims = args.qk_rope != 0u ?
-1.0f / (float)args.qk_rope : 0.0f;
for (uint r = tid * 2u; r < args.qk_rope; r += nth * 2u) {
#ifdef DS4_METAL_ROPE_EXP2_LOG2
const float theta = theta_base * exp2(inv_ndims * (float)r * log2(args.freq_base));
@@ -1400,6 +1543,35 @@ kernel void kernel_glm_fill_selected_range_batch(
selected[gid] = slot < visible ? slot : args.pad_row;
}
kernel void kernel_glm53_expand_pool_selection(
constant ds4_metal_args_glm53_expand_pool_selection &args,
device const uint32_t *pool_selected,
device uint32_t *raw_selected,
uint gid [[thread_position_in_grid]]) {
const uint total = args.n_tokens * args.output_width;
if (gid >= total || args.output_width == 0u || args.pool_size == 0u) return;
const uint token = gid / args.output_width;
const uint slot = gid - token * args.output_width;
uint value = 0xffffffffu;
if (slot < args.index_topk) {
const uint pool_slot = slot / args.pool_size;
if (pool_slot < args.selected_pools) {
const uint pool = pool_selected[
(uint64_t)token * args.selected_pools + pool_slot];
value = pool * args.pool_size + slot % args.pool_size;
}
} else {
const uint tail_slot = slot - args.index_topk;
const uint visible = args.pos0 + token + 1u;
const uint tail_count = visible % args.pool_size;
if (tail_slot < tail_count) {
value = visible - tail_count + tail_slot;
}
}
raw_selected[gid] = value;
}
kernel void kernel_glm_indexer_rope_tail_f32(
constant ds4_metal_args_glm_indexer_rope_tail & args,
device char *x,
@@ -1849,7 +2021,7 @@ kernel void kernel_glm_indexer_scores_batch(
device float *dst = (device float *)(scores +
(uint64_t)token * args.score_token_stride) + row;
const uint visible = min(args.pos0 + token + 1u, args.n_rows);
const uint visible = glm_indexer_batch_visible_rows(args, token);
if (row >= visible) {
if (tid == 0) *dst = -INFINITY;
return;
@@ -1910,7 +2082,7 @@ kernel void kernel_glm_indexer_scores_tiled_f32(
const uint last_token = min(token_base + TM, args.n_tokens);
const uint max_visible = last_token > token_base ?
min(args.pos0 + last_token, args.n_rows) : 0u;
glm_indexer_batch_visible_rows(args, last_token - 1u) : 0u;
if (row_base >= max_visible) {
for (uint i = tid; i < TM*TN; i += 128) {
@@ -2005,13 +2177,13 @@ kernel void kernel_glm_indexer_scores_tiled_f32(
}
if (token0 < args.n_tokens && row0 < args.n_rows) {
const uint visible = min(args.pos0 + token0 + 1u, args.n_rows);
const uint visible = glm_indexer_batch_visible_rows(args, token0);
device float *dst = (device float *)(scores +
(uint64_t)token0 * args.score_token_stride) + row0;
*dst = row0 < visible ? acc0 : -INFINITY;
}
if (token1 < args.n_tokens && row1 < args.n_rows) {
const uint visible = min(args.pos0 + token1 + 1u, args.n_rows);
const uint visible = glm_indexer_batch_visible_rows(args, token1);
device float *dst = (device float *)(scores +
(uint64_t)token1 * args.score_token_stride) + row1;
*dst = row1 < visible ? acc1 : -INFINITY;
@@ -2043,7 +2215,7 @@ kernel void kernel_glm_indexer_scores_tiled(
const uint last_token = min(token_base + TM, args.n_tokens);
const uint max_visible = last_token > token_base ?
min(args.pos0 + last_token, args.n_rows) : 0u;
glm_indexer_batch_visible_rows(args, last_token - 1u) : 0u;
if (row_base >= max_visible) {
for (uint i = tid; i < TM*TN; i += 128) {
@@ -2138,13 +2310,13 @@ kernel void kernel_glm_indexer_scores_tiled(
}
if (token0 < args.n_tokens && row0 < args.n_rows) {
const uint visible = min(args.pos0 + token0 + 1u, args.n_rows);
const uint visible = glm_indexer_batch_visible_rows(args, token0);
device float *dst = (device float *)(scores +
(uint64_t)token0 * args.score_token_stride) + row0;
*dst = row0 < visible ? acc0 : -INFINITY;
}
if (token1 < args.n_tokens && row1 < args.n_rows) {
const uint visible = min(args.pos0 + token1 + 1u, args.n_rows);
const uint visible = glm_indexer_batch_visible_rows(args, token1);
device float *dst = (device float *)(scores +
(uint64_t)token1 * args.score_token_stride) + row1;
*dst = row1 < visible ? acc1 : -INFINITY;
@@ -2241,7 +2413,6 @@ kernel void kernel_glm_qk_lowrank_q8_0_glm52_sg(
ushort3 ntg_u [[threads_per_threadgroup]],
ushort tiisg [[thread_index_in_simdgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
constexpr uint n_head = 64u;
constexpr uint kv_lora_dim = 512u;
constexpr uint qk_nope = 192u;
constexpr uint qk_dim = 256u;
@@ -2249,8 +2420,8 @@ kernel void kernel_glm_qk_lowrank_q8_0_glm52_sg(
const uint head = tgpig.x;
const uint wt = args.weight_type;
if (head >= n_head ||
args.n_head != n_head ||
if (head >= args.n_head ||
(args.n_head != 32u && args.n_head != 64u) ||
args.kv_lora_dim != kv_lora_dim ||
args.qk_nope != qk_nope ||
args.qk_dim != qk_dim ||
@@ -3624,7 +3795,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl(
args.n_selected == 0u ||
args.cache_f16 == 0u ||
args.kv_lora_dim != 512u ||
args.qk_rope != 64u) {
(args.qk_rope != 0u && args.qk_rope != 64u)) {
return;
}
@@ -3668,7 +3839,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_impl(
}
float corr_dims[2] = {0.0f, 0.0f};
if (args.ext_factor != 0.0f) {
if (args.qk_rope != 0u && args.ext_factor != 0.0f) {
glm_rope_yarn_corr_dims((int)args.qk_rope,
(int)args.n_ctx_orig,
args.freq_base,
@@ -3822,9 +3993,8 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint head = tgpig.x * group_heads + head_in_group + args.head_base;
if (token >= args.n_tokens ||
args.n_selected == 0u ||
args.cache_f16 == 0u ||
args.kv_lora_dim != 512u ||
args.qk_rope != 64u) {
(args.qk_rope != 0u && args.qk_rope != 64u)) {
return;
}
@@ -3869,7 +4039,7 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
}
float corr_dims[2] = {0.0f, 0.0f};
if (args.ext_factor != 0.0f) {
if (args.qk_rope != 0u && args.ext_factor != 0.0f) {
glm_rope_yarn_corr_dims((int)args.qk_rope,
(int)args.n_ctx_orig,
args.freq_base,
@@ -3891,10 +4061,17 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint rr = off / kv_vecs;
const uint vv = off - rr * kv_vecs;
const uint row = base + rr;
device const half4 *src =
(device const half4 *)((device const half *)kv_lora_cache +
(uint64_t)row * args.kv_lora_dim);
kv_shared[off] = src[vv];
if (args.cache_f16 != 0u) {
device const half4 *src =
(device const half4 *)((device const half *)kv_lora_cache +
(uint64_t)row * args.kv_lora_dim);
kv_shared[off] = src[vv];
} else {
device const float4 *src =
(device const float4 *)((device const float *)kv_lora_cache +
(uint64_t)row * args.kv_lora_dim);
kv_shared[off] = (half4)src[vv];
}
}
for (uint off = tid; off < rows * rope_vecs; off += 256u) {
const uint rr = off / rope_vecs;
@@ -3903,29 +4080,31 @@ kernel void kernel_glm_attention_indexed_batch_lora_group8_vec_causal_impl(
const uint row = base + rr;
const uint64_t rope_base = (uint64_t)row * args.qk_rope;
const float2 y0 =
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache,
rope_base,
r,
row,
args.qk_rope,
args.freq_base,
args.freq_scale,
args.ext_factor,
args.attn_factor,
corr_dims[0],
corr_dims[1]);
glm_cache_load_rotated_rope_pair(k_rope_cache,
rope_base,
r,
row,
args.qk_rope,
args.cache_f16,
args.freq_base,
args.freq_scale,
args.ext_factor,
args.attn_factor,
corr_dims[0],
corr_dims[1]);
const float2 y1 =
glm_cache_load_rotated_rope_pair_f16_only(k_rope_cache,
rope_base,
r + 2u,
row,
args.qk_rope,
args.freq_base,
args.freq_scale,
args.ext_factor,
args.attn_factor,
corr_dims[0],
corr_dims[1]);
glm_cache_load_rotated_rope_pair(k_rope_cache,
rope_base,
r + 2u,
row,
args.qk_rope,
args.cache_f16,
args.freq_base,
args.freq_scale,
args.ext_factor,
args.attn_factor,
corr_dims[0],
corr_dims[1]);
rope_shared[off] = float4(y0.x, y0.y, y1.x, y1.y);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -4586,14 +4765,15 @@ kernel void kernel_glm_router_select_one(
threadgroup float *scratch [[threadgroup(0)]],
uint token [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]]) {
const uint sort_width = args.n_expert > 256u ? 512u : 256u;
threadgroup float *sel_scores = scratch;
threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + 256);
threadgroup int32_t *idx = (threadgroup int32_t *)(scratch + sort_width);
device const float *token_logits = logits + (uint64_t)token * args.n_expert;
device int32_t *token_selected = selected + (uint64_t)token * args.n_expert_used;
device float *token_weights = weights + (uint64_t)token * args.n_expert_used;
device float *token_probs = probs + (uint64_t)token * args.n_expert;
const uint n_expert = min(args.n_expert, 256u);
const uint n_expert = min(args.n_expert, 512u);
const bool active = tid < n_expert;
const float p = active ? ds4_glm_router_sigmoid(token_logits[tid]) : 0.0f;
if (active) token_probs[tid] = p;
@@ -4601,7 +4781,7 @@ kernel void kernel_glm_router_select_one(
idx[tid] = (int32_t)tid;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k = 2; k <= 256; k <<= 1) {
for (uint k = 2; k <= sort_width; k <<= 1) {
for (uint j = k >> 1; j > 0; j >>= 1) {
const uint other = tid ^ j;
if (other > tid) {

143
metal/glm53_bf16.metal Normal file
View File

@@ -0,0 +1,143 @@
// BF16 model-weight kernels used by GLM-5.3 Flash.
static inline float glm53_bf16_to_f32(ushort value) {
return as_type<float>((uint)value << 16);
}
struct glm53_bf16_matmul_args {
uint in_dim;
uint out_dim;
uint n_rows;
};
kernel void kernel_glm53_embedding_bf16(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const int *tokens,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint row = gid.y;
if (d >= args.in_dim || row >= args.n_rows) return;
const int token = tokens[row];
out[(ulong)row * args.in_dim + d] =
token >= 0 && (uint)token < args.out_dim
? glm53_bf16_to_f32(weights[(ulong)(uint)token * args.in_dim + d])
: 0.0f;
}
static inline void glm53_mul_mv_bf16_f32_row(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const float *x,
device float *out,
uint2 tgpig,
ushort lane,
ushort sg,
ushort nsg) {
const uint out_row = tgpig.x * (uint)nsg + sg;
const uint token = tgpig.y;
if (out_row >= args.out_dim || token >= args.n_rows) return;
device const ushort *w = weights + (ulong)out_row * args.in_dim;
device const float *xr = x + (ulong)token * args.in_dim;
float sum = 0.0f;
uint k = lane;
for (; k + 224u < args.in_dim; k += 256u) {
const ushort w0 = w[k];
const ushort w1 = w[k + 32u];
const ushort w2 = w[k + 64u];
const ushort w3 = w[k + 96u];
const ushort w4 = w[k + 128u];
const ushort w5 = w[k + 160u];
const ushort w6 = w[k + 192u];
const ushort w7 = w[k + 224u];
const float x0 = xr[k];
const float x1 = xr[k + 32u];
const float x2 = xr[k + 64u];
const float x3 = xr[k + 96u];
const float x4 = xr[k + 128u];
const float x5 = xr[k + 160u];
const float x6 = xr[k + 192u];
const float x7 = xr[k + 224u];
sum = fma(glm53_bf16_to_f32(w0), x0, sum);
sum = fma(glm53_bf16_to_f32(w1), x1, sum);
sum = fma(glm53_bf16_to_f32(w2), x2, sum);
sum = fma(glm53_bf16_to_f32(w3), x3, sum);
sum = fma(glm53_bf16_to_f32(w4), x4, sum);
sum = fma(glm53_bf16_to_f32(w5), x5, sum);
sum = fma(glm53_bf16_to_f32(w6), x6, sum);
sum = fma(glm53_bf16_to_f32(w7), x7, sum);
}
for (; k < args.in_dim; k += 32u) {
sum = fma(glm53_bf16_to_f32(w[k]), xr[k], sum);
}
sum = simd_sum(sum);
if (lane == 0u) out[(ulong)token * args.out_dim + out_row] = sum;
}
/* One simdgroup owns one output row. Eight independent loads expose enough
* memory-level parallelism for decode without changing the reduction tree. */
kernel void kernel_glm53_mul_mv_bf16_f32(
constant glm53_bf16_matmul_args &args,
device const ushort *weights,
device const float *x,
device float *out,
uint2 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
glm53_mul_mv_bf16_f32_row(args, weights, x, out,
tgpig, lane, sg, nsg);
}
kernel void kernel_glm53_mul_mv_bf16_f32_qkv(
constant glm53_bf16_matmul_args &args,
device const ushort *weights_q,
device const ushort *weights_k,
device const ushort *weights_v,
device const float *x,
device float *out_q,
device float *out_k,
device float *out_v,
uint3 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
device const ushort *weights = tgpig.z == 0u ? weights_q :
(tgpig.z == 1u ? weights_k : weights_v);
device float *out = tgpig.z == 0u ? out_q :
(tgpig.z == 1u ? out_k : out_v);
glm53_mul_mv_bf16_f32_row(args, weights, x, out,
tgpig.xy, lane, sg, nsg);
}
struct glm53_bf16_block16 {
ushort v[16];
};
template <typename type4x4>
void glm53_dequantize_bf16(
device const glm53_bf16_block16 *src,
short il,
thread type4x4 &reg) {
(void)il;
float4x4 values;
for (short i = 0; i < 16; i++) {
values[i / 4][i % 4] = glm53_bf16_to_f32(src->v[i]);
}
reg = (type4x4)values;
}
typedef decltype(kernel_mul_mm<
half, half4x4, simdgroup_half8x8,
half, half2x4, simdgroup_half8x8,
glm53_bf16_block16, 1, glm53_dequantize_bf16,
float, float4x4, float, float2x4>) glm53_mul_mm_bf16_t;
template [[host_name("kernel_glm53_mul_mm_bf16_f32")]]
kernel glm53_mul_mm_bf16_t kernel_mul_mm<
half, half4x4, simdgroup_half8x8,
half, half2x4, simdgroup_half8x8,
glm53_bf16_block16, 1, glm53_dequantize_bf16,
half, half4x4, float, float2x4>;

314
metal/glm53_kda.metal Normal file
View File

@@ -0,0 +1,314 @@
// Kimi Delta Attention kernels, adapted from the kimi-k3 branch.
struct glm53_kda_args {
uint n_heads;
uint n_rows;
float lower_bound;
float norm_eps;
};
/*
* One threadgroup owns one (sequence, head). Four simdgroups update four
* value rows concurrently; every lane owns four adjacent key columns.
*/
kernel void kernel_glm53_kda_decode(
constant glm53_kda_args &args,
device const float *q_in,
device const float *k_in,
device const float *v_in,
device const float *raw_gate,
device const float *raw_beta,
device const float *output_gate,
device const float *q_conv,
device const float *k_conv,
device const float *v_conv,
device const float *a_log,
device const float *dt_bias,
device const float *output_norm,
device float *conv_state,
device float *state,
device float *out,
threadgroup float *scratch [[threadgroup(0)]],
uint2 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
constexpr uint HISTORY = 3u;
const uint row = tgpig.x;
const uint head = tgpig.y;
if (row >= args.n_rows || head >= args.n_heads) return;
threadgroup float *sq = scratch;
threadgroup float *sk = sq + D;
threadgroup float *sd = sk + D;
threadgroup float *sv = sd + D;
threadgroup float *so = sv + D;
threadgroup float *reduce_q = so + D;
threadgroup float *reduce_k = reduce_q + 4u;
threadgroup float *reduce_o = reduce_k + 4u;
threadgroup float *beta_shared = reduce_o + 4u;
const uint projection = args.n_heads * D;
const uint channel = head * D + tid;
const ulong input_base = (ulong)row * projection + head * D;
const ulong conv_row_stride = 3ul * HISTORY * projection;
if (tid < D) {
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
device float *q_state = conv_state +
(ulong)row * conv_row_stride;
device float *k_state = q_state + HISTORY * projection;
device float *v_state = k_state + HISTORY * projection;
for (uint w = 0; w < HISTORY; w++) {
q_acc = fma(q_state[(ulong)w * projection + channel],
q_conv[(ulong)channel * 4u + w], q_acc);
k_acc = fma(k_state[(ulong)w * projection + channel],
k_conv[(ulong)channel * 4u + w], k_acc);
v_acc = fma(v_state[(ulong)w * projection + channel],
v_conv[(ulong)channel * 4u + w], v_acc);
}
const float q_new = q_in[input_base + tid];
const float k_new = k_in[input_base + tid];
const float v_new = v_in[input_base + tid];
q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc);
k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc);
v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ul * projection + channel];
q_state[2ul * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ul * projection + channel];
k_state[2ul * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ul * projection + channel];
v_state[2ul * projection + channel] = v_new;
sq[tid] = q_acc / (1.0f + exp(-q_acc));
sk[tid] = k_acc / (1.0f + exp(-k_acc));
sv[tid] = v_acc / (1.0f + exp(-v_acc));
const float gate = raw_gate[input_base + tid] + dt_bias[channel];
sd[tid] = exp(args.lower_bound *
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
}
if (tid == 0u) {
beta_shared[0] =
1.0f / (1.0f + exp(-raw_beta[(ulong)row * args.n_heads + head]));
}
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float q_sumsq = sq[tid] * sq[tid];
float k_sumsq = sk[tid] * sk[tid];
q_sumsq = simd_sum(q_sumsq);
k_sumsq = simd_sum(k_sumsq);
if (lane == 0u) {
reduce_q[sg] = q_sumsq;
reduce_k[sg] = k_sumsq;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = simd_sum(q_total);
k_total = simd_sum(k_total);
const float q_scale = rsqrt(q_total + 1.0e-6f) * 0x1.6a09e6p-4f;
const float k_scale = rsqrt(k_total + 1.0e-6f);
if (tid < D) {
sq[tid] *= q_scale;
sk[tid] *= k_scale;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint k0 = lane * 4u;
const float4 q4 = *((threadgroup float4 *)(sq + k0));
const float4 k4 = *((threadgroup float4 *)(sk + k0));
const float4 decay4 = *((threadgroup float4 *)(sd + k0));
const ulong state_head =
((ulong)row * args.n_heads + head) * D * D;
for (uint value = sg; value < D; value += 4u) {
device float4 *hptr =
(device float4 *)(state + state_head + (ulong)value * D + k0);
float4 h = *hptr * decay4;
float hk = dot(h, k4);
hk = simd_sum(hk);
const float delta_v = (sv[value] - hk) * beta_shared[0];
h = fma(k4, float4(delta_v), h);
*hptr = h;
float hq = simd_sum(dot(h, q4));
if (lane == 0u) so[value] = hq;
}
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float o_sumsq = so[tid] * so[tid];
o_sumsq = simd_sum(o_sumsq);
if (lane == 0u) reduce_o[sg] = o_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
float o_total = lane < 4u ? reduce_o[lane] : 0.0f;
o_total = simd_sum(o_total);
const float o_scale = rsqrt(o_total / (float)D + args.norm_eps);
if (tid < D) {
const ulong index = input_base + tid;
const float gate =
1.0f / (1.0f + exp(-output_gate[index]));
out[index] = so[tid] * o_scale * output_norm[tid] * gate;
}
}
kernel void kernel_glm53_kda_prefill_prepare(
constant glm53_kda_args &args,
device float *q,
device float *k,
device float *v,
device float *raw_gate,
device const float *q_conv,
device const float *k_conv,
device const float *v_conv,
device const float *a_log,
device const float *dt_bias,
device float *conv_state,
threadgroup float *scratch [[threadgroup(0)]],
uint head [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
constexpr uint HISTORY = 3u;
if (head >= args.n_heads) return;
threadgroup float *sq = scratch;
threadgroup float *sk = sq + D;
threadgroup float *reduce_q = sk + D;
threadgroup float *reduce_k = reduce_q + 4u;
const uint projection = args.n_heads * D;
const uint channel = head * D + tid;
device float *q_state = conv_state;
device float *k_state = q_state + HISTORY * projection;
device float *v_state = k_state + HISTORY * projection;
for (uint token = 0; token < args.n_rows; token++) {
const ulong index = (ulong)token * projection + channel;
float q_acc = 0.0f;
float k_acc = 0.0f;
float v_acc = 0.0f;
for (uint w = 0; w < HISTORY; w++) {
q_acc = fma(q_state[(ulong)w * projection + channel],
q_conv[(ulong)channel * 4u + w], q_acc);
k_acc = fma(k_state[(ulong)w * projection + channel],
k_conv[(ulong)channel * 4u + w], k_acc);
v_acc = fma(v_state[(ulong)w * projection + channel],
v_conv[(ulong)channel * 4u + w], v_acc);
}
const float q_new = q[index];
const float k_new = k[index];
const float v_new = v[index];
q_acc = fma(q_new, q_conv[(ulong)channel * 4u + 3u], q_acc);
k_acc = fma(k_new, k_conv[(ulong)channel * 4u + 3u], k_acc);
v_acc = fma(v_new, v_conv[(ulong)channel * 4u + 3u], v_acc);
q_state[channel] = q_state[projection + channel];
q_state[projection + channel] = q_state[2ul * projection + channel];
q_state[2ul * projection + channel] = q_new;
k_state[channel] = k_state[projection + channel];
k_state[projection + channel] = k_state[2ul * projection + channel];
k_state[2ul * projection + channel] = k_new;
v_state[channel] = v_state[projection + channel];
v_state[projection + channel] = v_state[2ul * projection + channel];
v_state[2ul * projection + channel] = v_new;
sq[tid] = q_acc / (1.0f + exp(-q_acc));
sk[tid] = k_acc / (1.0f + exp(-k_acc));
v[index] = v_acc / (1.0f + exp(-v_acc));
const float gate = raw_gate[index] + dt_bias[channel];
raw_gate[index] = exp(args.lower_bound *
(1.0f / (1.0f + exp(-exp(a_log[head]) * gate))));
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
float q_sumsq = simd_sum(sq[tid] * sq[tid]);
float k_sumsq = simd_sum(sk[tid] * sk[tid]);
if (lane == 0u) {
reduce_q[sg] = q_sumsq;
reduce_k[sg] = k_sumsq;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float q_total = lane < 4u ? reduce_q[lane] : 0.0f;
float k_total = lane < 4u ? reduce_k[lane] : 0.0f;
q_total = simd_sum(q_total);
k_total = simd_sum(k_total);
q[index] = sq[tid] * rsqrt(q_total + 1.0e-6f) *
0x1.6a09e6p-4f;
k[index] = sk[tid] * rsqrt(k_total + 1.0e-6f);
threadgroup_barrier(mem_flags::mem_threadgroup |
mem_flags::mem_device);
}
}
kernel void kernel_glm53_kda_prefill_recurrence(
constant glm53_kda_args &args,
device const float *q,
device const float *k,
device const float *v,
device const float *decay,
device const float *raw_beta,
device float *state,
device float *out,
uint2 tgpig [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
const uint head = tgpig.x;
const uint value = tgpig.y * 4u + sg;
if (head >= args.n_heads || value >= D) return;
const uint projection = args.n_heads * D;
const uint k0 = lane * 4u;
device float4 *state_ptr = (device float4 *)(
state + ((ulong)head * D + value) * D + k0);
float4 h = *state_ptr;
for (uint token = 0; token < args.n_rows; token++) {
const ulong base = (ulong)token * projection + head * D;
const float4 q4 = *((device const float4 *)(q + base + k0));
const float4 k4 = *((device const float4 *)(k + base + k0));
const float4 decay4 =
*((device const float4 *)(decay + base + k0));
h *= decay4;
const float hk = simd_sum(dot(h, k4));
const float beta = 1.0f /
(1.0f + exp(-raw_beta[(ulong)token * args.n_heads + head]));
const float delta_v = (v[base + value] - hk) * beta;
h = fma(k4, float4(delta_v), h);
const float result = simd_sum(dot(h, q4));
if (lane == 0u) out[base + value] = result;
}
*state_ptr = h;
}
kernel void kernel_glm53_kda_prefill_output(
constant glm53_kda_args &args,
device float *out,
device const float *output_gate,
device const float *output_norm,
threadgroup float *partial [[threadgroup(0)]],
uint2 tgpig [[threadgroup_position_in_grid]],
ushort tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]]) {
constexpr uint D = 128u;
const uint token = tgpig.x;
const uint head = tgpig.y;
if (token >= args.n_rows || head >= args.n_heads) return;
const uint projection = args.n_heads * D;
const ulong base = (ulong)token * projection + head * D;
const float raw = out[base + tid];
float sumsq = simd_sum(raw * raw);
if (lane == 0u) partial[sg] = sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
float total = lane < 4u ? partial[lane] : 0.0f;
total = simd_sum(total);
const float scale = rsqrt(total / (float)D + args.norm_eps);
out[base + tid] = raw * scale * output_norm[tid] /
(1.0f + exp(-output_gate[base + tid]));
}

271
metal/glm53_vision.metal Normal file
View File

@@ -0,0 +1,271 @@
// GLM-5.3 Flash vision operations not covered by the shared BF16 matmuls.
struct glm53_vision_rows_args {
uint width;
uint rows;
float eps;
};
struct glm53_vision_qkv_args {
uint rows;
uint grid_h;
uint grid_w;
float eps;
};
struct glm53_vision_attention_args {
uint rows;
float scale;
};
struct glm53_vision_scatter_args {
uint dst_row;
uint image_row;
uint rows;
uint total_rows;
uint width;
uint hc;
};
static inline float glm53_vision_erf(float x) {
const float sign = x < 0.0f ? -1.0f : 1.0f;
const float a = abs(x);
const float t = 1.0f / (1.0f + 0.3275911f * a);
const float p = (((((1.061405429f * t - 1.453152027f) * t) +
1.421413741f) * t - 0.284496736f) * t +
0.254829592f) * t;
return sign * (1.0f - p * exp(-a * a));
}
kernel void kernel_glm53_vision_add_bias(
constant glm53_vision_rows_args &args,
device float *x,
device const ushort *bias,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
x[(ulong)gid.y * args.width + gid.x] += glm53_bf16_to_f32(bias[gid.x]);
}
kernel void kernel_glm53_vision_rms_bf16(
constant glm53_vision_rows_args &args,
device const float *x,
device const ushort *weight,
device float *out,
threadgroup float *partial,
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
if (row >= args.rows) return;
device const float *xr = x + (ulong)row * args.width;
device float *yr = out + (ulong)row * args.width;
float sum = 0.0f;
for (uint d = tid; d < args.width; d += 256u) sum = fma(xr[d], xr[d], sum);
sum = simd_sum(sum);
if (lane == 0u) partial[sg] = sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = rsqrt(v / (float)args.width + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float inv = partial[0];
for (uint d = tid; d < args.width; d += 256u) {
yr[d] = xr[d] * inv * glm53_bf16_to_f32(weight[d]);
}
}
kernel void kernel_glm53_vision_qkv_rope(
constant glm53_vision_qkv_args &args,
device const float *qkv,
device const ushort *bias,
device const ushort *q_weight,
device const ushort *k_weight,
device float *q,
device float *k,
device float *v,
uint2 group [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint row = group.x;
const uint head = group.y;
if (row >= args.rows || head >= 16u) return;
const ulong qkv_base = (ulong)row * 3072u + (ulong)head * 64u;
const ulong out_base = (ulong)row * 1024u + (ulong)head * 64u;
float q0 = qkv[qkv_base + lane] + glm53_bf16_to_f32(bias[(ulong)head * 64u + lane]);
float q1 = qkv[qkv_base + lane + 32u] +
glm53_bf16_to_f32(bias[(ulong)head * 64u + lane + 32u]);
float k0 = qkv[qkv_base + 1024u + lane] +
glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane]);
float k1 = qkv[qkv_base + 1024u + lane + 32u] +
glm53_bf16_to_f32(bias[1024u + (ulong)head * 64u + lane + 32u]);
const float qsum = simd_sum(fma(q0, q0, q1 * q1));
const float ksum = simd_sum(fma(k0, k0, k1 * k1));
const float qinv = rsqrt(qsum / 64.0f + args.eps);
const float kinv = rsqrt(ksum / 64.0f + args.eps);
q0 *= qinv * glm53_bf16_to_f32(q_weight[lane]);
q1 *= qinv * glm53_bf16_to_f32(q_weight[lane + 32u]);
k0 *= kinv * glm53_bf16_to_f32(k_weight[lane]);
k1 *= kinv * glm53_bf16_to_f32(k_weight[lane + 32u]);
const uint merge_w = args.grid_w / 2u;
const uint group_index = row / 4u;
const uint within = row & 3u;
const uint py = (group_index / merge_w) * 2u + within / 2u;
const uint px = (group_index % merge_w) * 2u + within % 2u;
const uint freq_index = lane & 15u;
const uint pos = lane < 16u ? py : px;
const float inv_freq = powr(10000.0f, -(float)freq_index / 16.0f);
const float angle = (float)pos * inv_freq;
const float cs = cos(angle);
const float sn = sin(angle);
q[out_base + lane] = q0 * cs - q1 * sn;
q[out_base + lane + 32u] = q1 * cs + q0 * sn;
k[out_base + lane] = k0 * cs - k1 * sn;
k[out_base + lane + 32u] = k1 * cs + k0 * sn;
v[out_base + lane] = qkv[qkv_base + 2048u + lane] +
glm53_bf16_to_f32(bias[2048u + (ulong)head * 64u + lane]);
v[out_base + lane + 32u] = qkv[qkv_base + 2048u + lane + 32u] +
glm53_bf16_to_f32(bias[2048u + (ulong)head * 64u + lane + 32u]);
}
/* A simdgroup owns one query/head and keeps its 64 output values in registers.
* This is quadratic in compute, as the model graph requires, but linear in
* memory and never materializes the attention matrix. */
kernel void kernel_glm53_vision_attention(
constant glm53_vision_attention_args &args,
device const float *q,
device const float *k,
device const float *v,
device float *out,
uint2 group [[threadgroup_position_in_grid]],
ushort lane [[thread_index_in_simdgroup]]) {
const uint row = group.x;
const uint head = group.y;
if (row >= args.rows || head >= 16u) return;
const ulong base = (ulong)row * 1024u + (ulong)head * 64u;
const float q0 = q[base + lane];
const float q1 = q[base + lane + 32u];
float acc0 = 0.0f, acc1 = 0.0f;
float max_score = -INFINITY;
float denom = 0.0f;
for (uint key_row = 0; key_row < args.rows; key_row++) {
const ulong kb = (ulong)key_row * 1024u + (ulong)head * 64u;
float score = simd_sum(q0 * k[kb + lane] + q1 * k[kb + lane + 32u]);
score *= args.scale;
const float next_max = max(max_score, score);
const float old_scale = max_score == -INFINITY ? 0.0f : exp(max_score - next_max);
const float new_scale = exp(score - next_max);
denom = denom * old_scale + new_scale;
acc0 = acc0 * old_scale + new_scale * v[kb + lane];
acc1 = acc1 * old_scale + new_scale * v[kb + lane + 32u];
max_score = next_max;
}
out[base + lane] = acc0 / denom;
out[base + lane + 32u] = acc1 / denom;
}
kernel void kernel_glm53_vision_bias_residual(
constant glm53_vision_rows_args &args,
device float *x,
device const ushort *bias,
device const float *residual,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
const ulong off = (ulong)gid.y * args.width + gid.x;
x[off] += glm53_bf16_to_f32(bias[gid.x]) + residual[off];
}
kernel void kernel_glm53_vision_swiglu_bias(
constant glm53_vision_rows_args &args,
device const float *gate,
device const ushort *gate_bias,
device const float *up,
device const ushort *up_bias,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
if (gid.x >= args.width || gid.y >= args.rows) return;
const ulong off = (ulong)gid.y * args.width + gid.x;
const float g = min(gate[off] + glm53_bf16_to_f32(gate_bias[gid.x]), 10.0f);
const float u = clamp(up[off] + glm53_bf16_to_f32(up_bias[gid.x]), -10.0f, 10.0f);
out[off] = (g / (1.0f + exp(-g))) * u;
}
kernel void kernel_glm53_vision_downsample_reorder(
device const float *x,
device float *out,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint row = gid.y;
if (d >= 4096u) return;
const uint channel = d / 4u;
const uint within = d & 3u;
out[(ulong)row * 4096u + d] = x[((ulong)row * 4u + within) * 1024u + channel];
}
kernel void kernel_glm53_vision_layernorm_gelu(
constant glm53_vision_rows_args &args,
device const float *x,
device const ushort *weight,
device const ushort *bias,
device float *out,
threadgroup float *partial,
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
ushort lane [[thread_index_in_simdgroup]],
ushort sg [[simdgroup_index_in_threadgroup]],
ushort nsg [[simdgroups_per_threadgroup]]) {
if (row >= args.rows) return;
device const float *xr = x + (ulong)row * args.width;
device float *yr = out + (ulong)row * args.width;
float sum = 0.0f;
for (uint d = tid; d < args.width; d += 256u) sum += xr[d];
sum = simd_sum(sum);
if (lane == 0u) partial[sg] = sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = v / (float)args.width;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float mean = partial[0];
float var = 0.0f;
for (uint d = tid; d < args.width; d += 256u) {
const float centered = xr[d] - mean;
var = fma(centered, centered, var);
}
var = simd_sum(var);
if (lane == 0u) partial[sg] = var;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sg == 0u) {
float v = lane < nsg ? partial[lane] : 0.0f;
v = simd_sum(v);
if (lane == 0u) partial[0] = rsqrt(v / (float)args.width + args.eps);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const float inv = partial[0];
const float inv_sqrt2 = 0.7071067811865475f;
for (uint d = tid; d < args.width; d += 256u) {
float v = (xr[d] - mean) * inv * glm53_bf16_to_f32(weight[d]) +
glm53_bf16_to_f32(bias[d]);
yr[d] = 0.5f * v * (1.0f + glm53_vision_erf(v * inv_sqrt2));
}
}
kernel void kernel_glm53_vision_scatter_hc(
constant glm53_vision_scatter_args &args,
device float *hc,
device const float *image,
uint2 gid [[thread_position_in_grid]]) {
const uint d = gid.x;
const uint linear_row = gid.y;
if (d >= args.width || linear_row >= args.rows * args.hc) return;
const uint image_delta = linear_row / args.hc;
const uint hc_index = linear_row % args.hc;
const ulong dst = ((ulong)(args.dst_row + image_delta) * args.hc + hc_index) *
args.width + d;
const ulong src = (ulong)(args.image_row + image_delta) * args.width + d;
if (args.dst_row + image_delta < args.total_rows) hc[dst] = image[src];
}

89
metal/moe.metal
View File

@@ -423,6 +423,7 @@ struct ds4_metal_glm_routed_moe_args {
uint32_t n_tokens;
uint32_t mid_token_stride;
uint32_t down_type;
float swiglu_clamp;
/* Expert ownership under tensor parallelism: tp_world 0/1 = full
* compute; otherwise each rank owns a contiguous expert range. */
int32_t tp_rank;
@@ -436,6 +437,14 @@ struct ds4_metal_glm_routed_moe_args {
uint64_t down_row_bytes;
};
static inline float ds4_glm_swiglu(float gate, float up, float limit) {
if (limit > 1.0e-6f) {
gate = min(gate, limit);
up = clamp(up, -limit, limit);
}
return (gate / (1.0f + exp(-gate))) * up;
}
static inline bool ds4_tp_owns_expert(int expert, int n_total,
int tp_rank, int tp_world) {
@@ -712,18 +721,20 @@ kernel void kernel_glm_q4_K_pair_swiglu_f32(
const uint64_t mid_off = (uint64_t)token * args.mid_token_stride +
(uint64_t)slot * args.mid_dim + row;
const int expert = selected[selected_off];
if (expert < 0 || (uint)expert >= args.n_total_expert) {
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tid == 0u) mid[mid_off] = 0.0f;
return;
}
const int local_expert = expert - args.tp_expert_base;
device const block_q4_K *gate_row =
(device const block_q4_K *)(gate +
(uint64_t)(uint)expert * args.gate_expert_bytes +
(uint64_t)(uint)local_expert * args.gate_expert_bytes +
(uint64_t)row * args.gate_row_bytes);
device const block_q4_K *up_row =
(device const block_q4_K *)(up +
(uint64_t)(uint)expert * args.up_expert_bytes +
(uint64_t)(uint)local_expert * args.up_expert_bytes +
(uint64_t)row * args.up_row_bytes);
float acc_gate = 0.0f;
@@ -747,10 +758,8 @@ kernel void kernel_glm_q4_K_pair_swiglu_f32(
}
if (tid == 0u) {
const float g = scratch[0];
const float u = scratch[ntg];
const float sw = g / (1.0f + exp(-g));
mid[mid_off] = sw * u * weights[selected_off];
mid[mid_off] = ds4_glm_swiglu(scratch[0], scratch[ntg],
args.swiglu_clamp) * weights[selected_off];
}
}
@@ -887,8 +896,8 @@ static inline void glm_q2_K_pair_swiglu_simd_f32_impl(
const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]);
if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g));
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off];
mid[mid_base + row0 + (uint)row] =
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
}
}
@@ -944,7 +953,9 @@ kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32(
(uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) {
if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) {
for (short row = 0;
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
@@ -1005,7 +1016,9 @@ kernel void kernel_glm_q2_K_addr_pair_swiglu2_f32_masked(
(uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) {
if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) {
for (short row = 0;
row < N_R0_GLM_Q2_PAIR2_K && row0 + (uint)row < args.mid_dim;
@@ -1191,8 +1204,8 @@ static inline void glm_q4_K_pair_swiglu_simd_f32_impl(
const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]);
if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g));
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off];
mid[mid_base + row0 + (uint)row] =
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
}
}
@@ -1216,9 +1229,12 @@ kernel void kernel_glm_q4_K_pair_swiglu2_f32(
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
const int expert = selected[selected_off];
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR2_K>(
args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, expert, tiisg, sgitg);
tgpig, slot, token, selected_off,
expert - args.tp_expert_base, tiisg, sgitg);
}
kernel void kernel_glm_q4_K_addr_pair_swiglu_f32(
@@ -1245,7 +1261,9 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32(
(uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) {
if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) {
for (short row = 0;
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
@@ -1306,7 +1324,9 @@ kernel void kernel_glm_q4_K_addr_pair_swiglu_f32_masked(
(uint64_t)slot * args.mid_dim;
if (row0 >= args.mid_dim) return;
if (expert < 0 || (uint)expert >= args.n_total_expert) {
if (expert < 0 || (uint)expert >= args.n_total_expert ||
!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) {
if (tiisg == 0u) {
for (short row = 0;
row < N_R0_Q4_K && row0 + (uint)row < args.mid_dim;
@@ -1358,9 +1378,12 @@ kernel void kernel_glm_q4_K_pair_swiglu4_f32(
if (slot >= args.n_expert_used || token >= args.n_tokens) return;
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
const int expert = selected[selected_off];
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_GLM_Q4_PAIR_K>(
args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, expert, tiisg, sgitg);
tgpig, slot, token, selected_off,
expert - args.tp_expert_base, tiisg, sgitg);
}
kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
@@ -1378,6 +1401,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
ushort sgitg [[simdgroup_index_in_threadgroup]]) {
const uint expert = tgpig.z;
if (expert >= args.n_total_expert) return;
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
const uint count = htpe[expert];
const uint map_base = tgpig.y * 32u;
for (uint i = 0; i < 32u; i++) {
@@ -1391,7 +1416,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_f32(
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg);
tgpig, slot, token, selected_off,
(int)expert - args.tp_expert_base, tiisg, sgitg);
}
}
@@ -1411,6 +1437,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32(
const uint expert = tgpig.z;
const uint map_row = tgpig.y;
if (expert >= args.n_total_expert || map_row >= htpe[expert]) return;
if (!ds4_tp_owns_expert((int)expert, args.n_total_expert,
args.tp_rank, args.tp_world)) return;
const int id = hids[(uint64_t)expert * args.n_tokens + map_row];
if (id < 0) return;
const uint token = (uint)id / args.n_expert_used;
@@ -1419,7 +1447,8 @@ kernel void kernel_glm_q4_K_pair_swiglu2_mapped_row_f32(
const uint64_t selected_off = (uint64_t)token * args.n_expert_used + slot;
glm_q4_K_pair_swiglu_simd_f32_impl<N_R0_Q4_K>(
args, gate, up, x, weights, mid, scratch,
tgpig, slot, token, selected_off, (int)expert, tiisg, sgitg);
tgpig, slot, token, selected_off,
(int)expert - args.tp_expert_base, tiisg, sgitg);
}
static inline void glm_q5_K_pair_swiglu_f32_impl(
@@ -1591,8 +1620,8 @@ static inline void glm_q5_K_pair_swiglu_f32_impl(
const float g = simd_sum(sumg[row]);
const float u = simd_sum(sumu[row]);
if (tiisg == 0u) {
const float sw = g / (1.0f + exp(-g));
mid[mid_base + row0 + (uint)row] = sw * u * weights[selected_off];
mid[mid_base + row0 + (uint)row] =
ds4_glm_swiglu(g, u, args.swiglu_clamp) * weights[selected_off];
}
}
@@ -1913,6 +1942,8 @@ kernel void kernel_glm_q2_K_addr_down_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue;
device const block_q2_K *x =
@@ -1996,10 +2027,12 @@ kernel void kernel_glm_q4_K_down_f32(
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
device const block_q4_K *down_row =
(device const block_q4_K *)(down +
(uint64_t)(uint)expert * args.down_expert_bytes +
(uint64_t)(uint)(expert - args.tp_expert_base) *
args.down_expert_bytes +
(uint64_t)row * args.down_row_bytes);
device const float *slot_mid = mid + mid_base + (uint64_t)slot * args.mid_dim;
for (uint k = tid; k < args.mid_dim; k += ntg) {
@@ -2038,6 +2071,8 @@ kernel void kernel_glm_q4_K_addr_down_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue;
device const block_q4_K *down_row =
@@ -2093,11 +2128,13 @@ kernel void kernel_glm_q4_K_down_simd_f32(
const uint64_t mid_base = (uint64_t)token * args.mid_token_stride;
for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
device const block_q4_K *x =
(device const block_q4_K *)(down +
(uint64_t)(uint)expert * args.down_expert_bytes +
(uint64_t)(uint)(expert - args.tp_expert_base) *
args.down_expert_bytes +
(uint64_t)row0 * args.down_row_bytes);
device const float *y = mid + mid_base + (uint64_t)slot * args.mid_dim;
device const float *y4 = y + ix * QK_K + 64 * iq + 8 * ir;
@@ -2197,6 +2234,8 @@ kernel void kernel_glm_q4_K_addr_down_simd_f32(
for (uint slot = 0; slot < args.n_expert_used; slot++) {
const int expert = selected[selected_base + slot];
if (expert < 0 || (uint)expert >= args.n_total_expert) continue;
if (!ds4_tp_owns_expert(expert, args.n_total_expert,
args.tp_rank, args.tp_world)) continue;
const uint64_t down_addr = down_addrs[(uint)expert];
if (down_addr == 0) continue;