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

Save inference parity implementation and evaluation harness

Browse Source
This commit is contained in:
Georg Bauer
2026-09-10 22:45:59 +02:00
parent b99ce2aa10
commit 02db0968ae
198 changed files with 111205 additions and 586 deletions
Download Patch File Download Diff File Expand all files Collapse all files
metal/qwen38.metal
+570 -40
View File
@@ -340,8 +340,8 @@ kernel void kernel_qwen_affine_qmv_fast(
// MLX affine_qmv_wide for the verifier's S=2..4 matrices. M5-class GPUs
// route these shapes here instead of running one independent QMV per row.
template <ushort bits, ushort group_size>
static inline void qwen_affine_qmv_wide_impl(
template <ushort bits, ushort group_size, ushort rows>
static inline void qwen_affine_qmv_wide_rows_impl(
constant qwen_kernel_args &args,
device float *out,
device const float *x,
@@ -361,12 +361,11 @@ static inline void qwen_affine_qmv_wide_impl(
simd_group * outputs_per_simdgroup + simd_row;
const uint in_dim = args.u[0];
const uint out_dim = args.u[1];
const uint rows = args.u[4];
const uint row = min(output_row, out_dim - 1u);
const uint groups_per_row = in_dim / group_size;
const ulong packed_row = (ulong)row * in_dim * bits / 8u;
const ulong parameter_row = (ulong)row * groups_per_row;
float result[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float result[rows] = {0.0f};
for (uint quant_group = k_lane; quant_group < groups_per_row;
quant_group += k_lanes) {
@@ -374,11 +373,13 @@ static inline void qwen_affine_qmv_wide_impl(
scales, args.u[14], parameter_row + quant_group));
const float bias = qwen_bf16(qwen_weight_u16(
biases, args.u[15], parameter_row + quant_group));
#pragma unroll
for (uint chunk = 0u; chunk < group_size / sub; chunk++) {
const uint column = quant_group * group_size + chunk * sub;
const ulong weight_byte = (ulong)args.u[13] + packed_row +
(ulong)column * bits / 8u;
float weights[sub];
#pragma unroll
for (uint index = 0u; index < sub; index++) {
uint quantized;
if constexpr (bits == 4) {
@@ -389,8 +390,10 @@ static inline void qwen_affine_qmv_wide_impl(
}
weights[index] = scale * (float)quantized + bias;
}
#pragma unroll
for (uint vector = 0u; vector < rows; vector++) {
float sum = 0.0f;
#pragma unroll
for (uint index = 0u; index < sub; index++) {
sum += x[(ulong)vector * in_dim + column + index] * weights[index];
}
@@ -411,6 +414,35 @@ static inline void qwen_affine_qmv_wide_impl(
}
}
// MTPLX specializes vecs_per_tg: keeping the accumulator indices constant
// avoids a runtime-indexed register array in the verifier's inner loop.
template <ushort bits, ushort group_size>
static inline void qwen_affine_qmv_wide_impl(
constant qwen_kernel_args &args,
device float *out,
device const float *x,
device const uchar *packed,
device const uchar *scales,
device const uchar *biases,
uint group,
uint simd_group,
uint lane) {
switch (args.u[4]) {
case 2:
qwen_affine_qmv_wide_rows_impl<bits, group_size, 2>(
args, out, x, packed, scales, biases, group, simd_group, lane);
break;
case 3:
qwen_affine_qmv_wide_rows_impl<bits, group_size, 3>(
args, out, x, packed, scales, biases, group, simd_group, lane);
break;
case 4:
qwen_affine_qmv_wide_rows_impl<bits, group_size, 4>(
args, out, x, packed, scales, biases, group, simd_group, lane);
break;
}
}
template <ushort bits, ushort group_size>
kernel void kernel_qwen_affine_qmv_batch_fast(
constant qwen_kernel_args &args [[buffer(0)]],
@@ -583,23 +615,111 @@ kernel qwen_affine_pair_qmv_wide_b8g64 kernel_qwen_affine_pair_qmv_wide<8, 64>;
#ifdef DS4_METAL_HAS_TENSOR
using namespace mpp::tensor_ops;
template <ushort bits, ushort group_size, ushort tile_m>
kernel void kernel_qwen_affine_qmm_mpp(
kernel void kernel_qwen_route_map(
constant qwen_kernel_args &args [[buffer(0)]],
device uint *counts [[buffer(1)]],
device const uint *selected [[buffer(2)]],
device uint *route_map [[buffer(3)]],
device uchar *work [[buffer(4)]],
threadgroup ushort *staged [[threadgroup(0)]],
uint expert [[thread_index_in_threadgroup]],
uint threads [[threads_per_threadgroup]]) {
uint count = 0u;
device uint *expert_map = route_map + (ulong)expert * args.u[4];
for (uint first_row = 0u; first_row < args.u[4]; first_row += threads) {
const uint row = first_row + expert;
if (row < args.u[4]) {
for (uint slot = 0u; slot < args.u[8]; slot++) {
staged[expert * args.u[8] + slot] =
(ushort)selected[(ulong)row * args.u[8] + slot];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
const uint valid = min(threads, args.u[4] - first_row);
for (uint local_row = 0u; local_row < valid; local_row++) {
for (uint slot = 0u; slot < args.u[8]; slot++) {
if (staged[local_row * args.u[8] + slot] == expert) {
expert_map[count++] = (first_row + local_row) * args.u[8] + slot;
break;
}
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
counts[expert] = count;
threadgroup_barrier(mem_flags::mem_device);
uint route_base = 0u;
for (uint other = 0u; other < expert; other++) route_base += counts[other];
device uint *route_offsets = (device uint *)(work + 8);
device uint *inverse = route_offsets + 512;
route_offsets[expert] = route_base;
for (uint index = 0u; index < count; index++) {
inverse[expert_map[index]] = route_base + index;
}
staged[expert] = (ushort)((count + 63u) / 64u);
threadgroup_barrier(mem_flags::mem_threadgroup);
uint work_base = 0u;
for (uint other = 0u; other < expert; other++) {
work_base += staged[other];
}
device uint *work_count = (device uint *)work;
const ulong items_offset =
(8ul + 512ul * 4ul + (ulong)args.u[4] * args.u[8] * 4ul + 7ul) & ~7ul;
device uint2 *work_items = (device uint2 *)(work + items_offset);
for (uint tile = 0u; tile < staged[expert]; tile++) {
work_items[work_base + tile] = uint2(expert, tile * 64u);
}
if (expert + 1u == threads) {
work_count[0] = work_base + staged[expert];
}
}
kernel void kernel_qwen_route_gather_rows(
constant qwen_kernel_args &args [[buffer(0)]],
device bfloat *out [[buffer(1)]],
device const float *x [[buffer(2)]],
device const uchar *work [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]) {
const uint column = gid.x;
const uint route = gid.y;
if (column >= args.u[0] || route >= args.u[4] * args.u[8]) return;
device const uint *inverse = (device const uint *)(work + 8) + 512;
const uint sorted = inverse[route];
out[(ulong)sorted * args.u[0] + column] =
(bfloat)x[(ulong)(route / args.u[8]) * args.u[0] + column];
}
template <ushort bits, ushort group_size>
kernel void kernel_qwen_affine_gather_qmm_mpp(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device const float *x [[buffer(2)]],
device const uint *route_map [[buffer(3)]],
device const uchar *work [[buffer(4)]],
device const uchar *packed [[buffer(5)]],
device const uchar *scales [[buffer(6)]],
device const uchar *biases [[buffer(7)]],
device const uint *counts [[buffer(8)]],
uint2 group [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]) {
constexpr uint tile_m = 64u;
constexpr uint tile_n = 64u;
constexpr uint tile_k = 64u;
constexpr uint threads = 128u;
device const uint *work_count = (device const uint *)work;
if (group.y >= work_count[0]) return;
const ulong items_offset =
(8ul + 512ul * 4ul + (ulong)args.u[4] * args.u[8] * 4ul + 7ul) & ~7ul;
device const uint2 *work_items = (device const uint2 *)(work + items_offset);
const uint2 item = work_items[group.y];
const uint expert = item.x;
const uint first_m = item.y;
if (first_m >= counts[expert]) return;
const uint first_n = group.x * tile_n;
device const uint *expert_map = route_map + (ulong)expert * args.u[4];
threadgroup bfloat xs[tile_m * tile_k];
threadgroup bfloat ws[tile_n * tile_k];
const uint first_m = group.y * tile_m;
const uint first_n = group.x * tile_n;
constexpr auto descriptor = matmul2d_descriptor(
tile_m, tile_n, tile_k, false, true, false,
matmul2d_descriptor::mode::multiply_accumulate);
@@ -609,28 +729,29 @@ kernel void kernel_qwen_affine_qmm_mpp(
tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>,
float>();
#pragma clang loop unroll(full)
for (uint i = 0u; i < accum.get_capacity(); i++) {
accum[i] = 0.0f;
}
for (uint i = 0u; i < accum.get_capacity(); i++) accum[i] = 0.0f;
for (uint first_k = 0u; first_k < args.u[0]; first_k += tile_k) {
for (uint index = tid; index < tile_m * tile_k; index += threads) {
const uint row = index / tile_k;
const uint column = index % tile_k;
const uint input_m = first_m + row;
xs[index] = input_m < args.u[4]
? (bfloat)x[(ulong)input_m * args.u[0] + first_k + column]
const bool valid = first_m + row < counts[expert];
const uint route = valid ? expert_map[first_m + row] : 0u;
const uint input_row = args.u[5] != 0u ? route : route / args.u[8];
xs[index] = valid
? (bfloat)x[(ulong)input_row * args.u[0] + first_k + column]
: bfloat(0.0f);
}
for (uint index = tid; index < tile_n * tile_k; index += threads) {
const uint row = index / tile_k;
const uint column = index % tile_k;
const uint weight_n = first_n + row;
const ulong table_row = (ulong)expert * args.u[1] + weight_n;
ws[index] = weight_n < args.u[1]
? (bfloat)qwen_quant_weight(
packed, scales, biases,
args.u[13], args.u[14], args.u[15],
weight_n, first_k + column,
table_row, first_k + column,
args.u[0], bits, group_size)
: bfloat(0.0f);
}
@@ -648,15 +769,267 @@ kernel void kernel_qwen_affine_qmm_mpp(
const auto index = accum.get_multidimensional_index(i);
const uint output_n = first_n + index[0];
const uint output_m = first_m + index[1];
if (output_n < args.u[1] && output_m < args.u[4]) {
const float value = accum[i];
out[(ulong)output_m * args.u[1] + output_n] = args.u[11] != 0u
? qwen_round_bf16(value)
: value;
if (output_n < args.u[1] && output_m < counts[expert]) {
const uint route = expert_map[output_m];
out[(ulong)route * args.u[1] + output_n] =
qwen_round_bf16(accum[i]);
}
}
}
template <ushort bits, ushort group_size>
kernel void kernel_qwen_affine_sorted_qmm_mpp(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device bfloat *x [[buffer(2)]],
device const uint *route_map [[buffer(3)]],
device const uchar *work [[buffer(4)]],
device const uchar *packed [[buffer(5)]],
device const uchar *scales [[buffer(6)]],
device const uchar *biases [[buffer(7)]],
device const uint *counts [[buffer(8)]],
uint2 group [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]) {
constexpr uint tile_m = 64u;
constexpr uint tile_n = 64u;
constexpr uint tile_k = 32u;
constexpr uint threads = 128u;
device const uint *work_count = (device const uint *)work;
if (group.y >= work_count[0]) return;
const ulong items_offset =
(8ul + 512ul * 4ul + (ulong)args.u[4] * args.u[8] * 4ul + 7ul) & ~7ul;
device const uint2 *work_items = (device const uint2 *)(work + items_offset);
const uint2 item = work_items[group.y];
const uint expert = item.x;
const uint local_m = item.y;
if (local_m >= counts[expert]) return;
device const uint *route_offsets = (device const uint *)(work + 8);
const uint first_m = route_offsets[expert] + local_m;
const uint first_n = group.x * tile_n;
threadgroup bfloat ws[2u * tile_n * tile_k];
auto weights0 = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
ws, dextents<int32_t, 2>(tile_k, tile_n));
auto weights1 = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
ws + tile_n * tile_k, dextents<int32_t, 2>(tile_k, tile_n));
auto activations = tensor<device bfloat, dextents<int32_t, 2>, tensor_inline>(
x, dextents<int32_t, 2>(args.u[0], args.u[6]),
array<int, 2>({1, (int)args.u[0]}));
constexpr auto descriptor = matmul2d_descriptor(
tile_m, tile_n, tile_k, false, true, true,
matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<descriptor, execution_simdgroups<4>> multiply;
auto accum = multiply.template get_destination_cooperative_tensor<
decltype(activations), decltype(weights0), float>();
#pragma clang loop unroll(full)
for (uint i = 0u; i < accum.get_capacity(); i++) accum[i] = 0.0f;
auto stage_weights = [&](uint first_k, threadgroup bfloat *target) {
constexpr uint per_word = 32u / bits;
constexpr uint words_per_row = tile_k / per_word;
for (uint index = tid; index < tile_n * words_per_row; index += threads) {
const uint row = index / words_per_row;
const uint column = (index % words_per_row) * per_word;
const uint weight_n = first_n + row;
const ulong table_row = (ulong)expert * args.u[1] + weight_n;
uint word = 0u;
float scale = 0.0f, bias = 0.0f;
if (weight_n < args.u[1]) {
word = qwen_weight_u32(packed, args.u[13],
table_row * (args.u[0] / per_word) + (first_k + column) / per_word);
const ulong quant_group = table_row * (args.u[0] / group_size)
+ (first_k + column) / group_size;
scale = qwen_bf16(qwen_weight_u16(scales, args.u[14], quant_group));
bias = qwen_bf16(qwen_weight_u16(biases, args.u[15], quant_group));
}
#pragma clang loop unroll(full)
for (uint element = 0u; element < per_word; element++) {
const uint quant = (word >> (element * bits)) & ((1u << bits) - 1u);
target[row * tile_k + column + element] =
(bfloat)fma((float)quant, scale, bias);
}
}
};
stage_weights(0u, ws);
threadgroup_barrier(mem_flags::mem_threadgroup);
uint selected_weights = 0u;
for (uint first_k = 0u; first_k < args.u[0]; first_k += tile_k) {
auto weight_tile = selected_weights ? weights1 : weights0;
auto activation_tile = activations.slice(first_k, first_m);
multiply.run(activation_tile, weight_tile, accum);
const uint next_k = first_k + tile_k;
if (next_k < args.u[0]) {
selected_weights ^= 1u;
stage_weights(next_k, selected_weights ? ws + tile_n * tile_k : ws);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
#pragma clang loop unroll(full)
for (uint i = 0u; i < accum.get_capacity(); i++) {
const auto index = accum.get_multidimensional_index(i);
const uint output_n = first_n + index[0];
const uint local_output_m = local_m + index[1];
if (output_n < args.u[1] && local_output_m < counts[expert]) {
out[(ulong)(first_m + index[1]) * args.u[1] + output_n] =
qwen_round_bf16(accum[i]);
}
}
(void)route_map;
}
typedef decltype(kernel_qwen_affine_sorted_qmm_mpp<4, 32>) qwen_affine_sorted_qmm_mpp_b4g32;
typedef decltype(kernel_qwen_affine_sorted_qmm_mpp<4, 64>) qwen_affine_sorted_qmm_mpp_b4g64;
typedef decltype(kernel_qwen_affine_sorted_qmm_mpp<8, 64>) qwen_affine_sorted_qmm_mpp_b8g64;
template [[host_name("kernel_qwen_affine_sorted_qmm_mpp_b4g32")]]
kernel qwen_affine_sorted_qmm_mpp_b4g32 kernel_qwen_affine_sorted_qmm_mpp<4, 32>;
template [[host_name("kernel_qwen_affine_sorted_qmm_mpp_b4g64")]]
kernel qwen_affine_sorted_qmm_mpp_b4g64 kernel_qwen_affine_sorted_qmm_mpp<4, 64>;
template [[host_name("kernel_qwen_affine_sorted_qmm_mpp_b8g64")]]
kernel qwen_affine_sorted_qmm_mpp_b8g64 kernel_qwen_affine_sorted_qmm_mpp<8, 64>;
typedef decltype(kernel_qwen_affine_gather_qmm_mpp<4, 32>) qwen_affine_gather_qmm_mpp_b4g32;
typedef decltype(kernel_qwen_affine_gather_qmm_mpp<4, 64>) qwen_affine_gather_qmm_mpp_b4g64;
typedef decltype(kernel_qwen_affine_gather_qmm_mpp<8, 64>) qwen_affine_gather_qmm_mpp_b8g64;
template [[host_name("kernel_qwen_affine_gather_qmm_mpp_b4g32")]]
kernel qwen_affine_gather_qmm_mpp_b4g32 kernel_qwen_affine_gather_qmm_mpp<4, 32>;
template [[host_name("kernel_qwen_affine_gather_qmm_mpp_b4g64")]]
kernel qwen_affine_gather_qmm_mpp_b4g64 kernel_qwen_affine_gather_qmm_mpp<4, 64>;
template [[host_name("kernel_qwen_affine_gather_qmm_mpp_b8g64")]]
kernel qwen_affine_gather_qmm_mpp_b8g64 kernel_qwen_affine_gather_qmm_mpp<8, 64>;
template <ushort bits, ushort group_size, ushort tile_m>
kernel void kernel_qwen_affine_qmm_mpp(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device float *x [[buffer(2)]],
device const uchar *packed [[buffer(5)]],
device const uchar *scales [[buffer(6)]],
device const uchar *biases [[buffer(7)]],
uint2 group [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]) {
constexpr uint tile_n = 64u;
constexpr uint tile_k = 32u;
constexpr uint threads = 128u;
threadgroup bfloat ws[2u * tile_n * tile_k];
const uint m_tiles = (args.u[4] + tile_m - 1u) / tile_m;
const uint partition = group.y / m_tiles;
const uint first_m = (group.y % m_tiles) * tile_m;
const uint first_n = group.x * tile_n;
const uint partition_k = args.u[0] / max(args.u[8], 1u);
const uint start_k = partition * partition_k;
const uint end_k = start_k + partition_k;
out += partition * args.u[4] * args.u[1];
auto weights0 = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
ws, dextents<int32_t, 2>(tile_k, tile_n));
auto weights1 = tensor<threadgroup bfloat, dextents<int32_t, 2>, tensor_inline>(
ws + tile_n * tile_k, dextents<int32_t, 2>(tile_k, tile_n));
auto activations = tensor<device float, dextents<int32_t, 2>, tensor_inline>(
x, dextents<int32_t, 2>(args.u[0], args.u[4]),
array<int, 2>({1, (int)args.u[0]}));
constexpr auto descriptor = matmul2d_descriptor(
tile_m, tile_n, tile_k, false, true, true,
matmul2d_descriptor::mode::multiply_accumulate);
matmul2d<descriptor, execution_simdgroups<4>> multiply;
auto accum = multiply.template get_destination_cooperative_tensor<
decltype(activations), decltype(weights0),
float>();
#pragma clang loop unroll(full)
for (uint i = 0u; i < accum.get_capacity(); i++) {
accum[i] = 0.0f;
}
auto stage_weights = [&](uint first_k, threadgroup bfloat *target) {
// Decode a packed word once, reusing its scale/bias across its values,
// as in the sorted QMM loader and MTPLX's quantized matrix loader.
constexpr uint per_word = 32u / bits;
constexpr uint words_per_row = tile_k / per_word;
for (uint index = tid; index < tile_n * words_per_row; index += threads) {
const uint row = index / words_per_row;
const uint column = (index % words_per_row) * per_word;
const uint weight_n = first_n + row;
uint word = 0u;
float scale = 0.0f, bias = 0.0f;
if (weight_n < args.u[1]) {
word = qwen_weight_u32(packed, args.u[13],
(ulong)weight_n * (args.u[0] / per_word) + (first_k + column) / per_word);
const ulong quant_group = (ulong)weight_n * (args.u[0] / group_size)
+ (first_k + column) / group_size;
scale = qwen_bf16(qwen_weight_u16(scales, args.u[14], quant_group));
bias = qwen_bf16(qwen_weight_u16(biases, args.u[15], quant_group));
}
#pragma clang loop unroll(full)
for (uint element = 0u; element < per_word; element++) {
const uint quant = (word >> (element * bits)) & ((1u << bits) - 1u);
target[row * tile_k + column + element] =
(bfloat)fma((float)quant, scale, bias);
}
}
};
stage_weights(start_k, ws);
threadgroup_barrier(mem_flags::mem_threadgroup);
uint selected_weights = 0u;
for (uint first_k = start_k; first_k < end_k; first_k += tile_k) {
auto weight_tile = selected_weights ? weights1 : weights0;
auto activation_tile = activations.slice(first_k, first_m);
multiply.run(activation_tile, weight_tile, accum);
const uint next_k = first_k + tile_k;
if (next_k < end_k) {
selected_weights ^= 1u;
stage_weights(
next_k,
selected_weights ? ws + tile_n * tile_k : ws);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
auto output = tensor<device float, dextents<int32_t, 2>, tensor_inline>(
out, dextents<int32_t, 2>(args.u[1], args.u[4]),
array<int, 2>({1, (int)args.u[1]}));
if (args.u[11] != 0u) {
#pragma clang loop unroll(full)
for (uint i = 0u; i < accum.get_capacity(); i++) {
accum[i] = qwen_round_bf16(accum[i]);
}
}
auto output_tile = output.slice(first_n, first_m);
accum.store(output_tile);
}
// Match MTPLX's BF16 column-reduction order, including intermediate rounding.
// Each simdgroup handles one output; small reductions use eight strided lanes,
// larger reductions use the reference's 32-lane BF16 simd reduction.
kernel void kernel_qwen_affine_splitk_reduce(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device const float *parts [[buffer(2)]],
uint group [[threadgroup_position_in_grid]],
uint simd_group [[simdgroup_index_in_threadgroup]],
uint lane [[thread_index_in_simdgroup]]) {
const uint index = group * 2u + simd_group;
const uint stride = args.u[4] * args.u[1];
if (index >= stride) return;
const uint count = args.u[8];
const uint lanes = count < 32u ? min(count, 8u) : 32u;
bfloat total = bfloat(0.0f);
for (uint p = lane; p < count; p += lanes) {
if (lane < lanes) total = bfloat(float(total) + parts[p * stride + index]);
}
if (count < 32u) {
bfloat sum = total;
for (uint p = 1u; p < lanes; p++) {
const float next = simd_shuffle(float(total), p);
sum = bfloat(float(sum) + float(next));
}
if (lane == 0u) out[index] = float(sum);
} else {
// MTPLX's bfloat16_t overload reduces in float, then rounds once.
const bfloat sum = bfloat(simd_sum(float(total)));
if (lane == 0u) out[index] = float(sum);
}
}
typedef decltype(kernel_qwen_affine_qmm_mpp<4, 32, 32>) qwen_affine_qmm_mpp_b4g32_bm32;
typedef decltype(kernel_qwen_affine_qmm_mpp<4, 64, 32>) qwen_affine_qmm_mpp_b4g64_bm32;
typedef decltype(kernel_qwen_affine_qmm_mpp<8, 64, 32>) qwen_affine_qmm_mpp_b8g64_bm32;
@@ -1209,6 +1582,33 @@ kernel void kernel_qwen_weighted_sum10(
out[index] = value;
}
kernel void kernel_qwen_weighted_sum10_sorted(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device const float *experts [[buffer(2)]],
device const float *weights [[buffer(3)]],
device const uchar *work [[buffer(4)]],
uint index [[thread_position_in_grid]]) {
if (index >= args.u[0]) return;
const uint hidden = args.u[1];
const uint row = index / hidden;
const uint column = index % hidden;
device const uint *inverse = (device const uint *)(work + 8) + 512;
const ulong routes_base = (ulong)row * args.u[8];
float weighted[10];
for (uint slot = 0; slot < 10u; slot++) {
const ulong route = routes_base + slot;
weighted[slot] = qwen_round_bf16(
experts[(ulong)inverse[route] * hidden + column] * weights[route]);
}
float value = qwen_round_bf16(weighted[0] + weighted[8]);
value = qwen_round_bf16(value + qwen_round_bf16(weighted[1] + weighted[9]));
for (uint slot = 2; slot < 8u; slot++) {
value = qwen_round_bf16(value + weighted[slot]);
}
out[index] = value;
}
kernel void kernel_qwen_affine_embedding(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
@@ -1959,7 +2359,9 @@ kernel void kernel_qwen_gdn_conv_norm(
: qwen_round_bf16(normalized);
}
// MTPLX fused_gdn_conv_norm_rows, including its in-window convolution tail.
// MTPLX fused_gdn_conv_norm_rows at S<=6. Larger prefills preserve the
// Conv1d -> SiLU -> L2 BF16 boundaries and parallelize independent rows.
// In the parallel path state_out must not alias the input state.
kernel void kernel_qwen_gdn_conv_norm_rows(
constant qwen_kernel_args &args [[buffer(0)]],
device float *q_out [[buffer(1)]],
@@ -1972,12 +2374,12 @@ kernel void kernel_qwen_gdn_conv_norm_rows(
uint tid [[thread_index_in_threadgroup]],
uint lane [[thread_index_in_simdgroup]],
uint simd_group [[simdgroup_index_in_threadgroup]],
uint group [[threadgroup_position_in_grid]]) {
uint2 group [[threadgroup_position_in_grid]]) {
constexpr uint width = 10240u;
constexpr uint key_width = 2048u;
constexpr uint dim = 128u;
constexpr float inverse_scale = 0.08838834764831845f;
const uint channel = group * 1024u + tid;
const uint channel = group.x * 1024u + tid;
if (channel >= width) return;
threadgroup float values[1024];
@@ -1988,7 +2390,10 @@ kernel void kernel_qwen_gdn_conv_norm_rows(
const float w3 = qwen_bf16(qwen_weight_u16(weight, args.u[14], (ulong)channel * 4u + 3u));
const bool value_channel = channel >= 2u * key_width;
for (uint row = 0u; row < args.u[4]; row++) {
const bool prefill = args.u[4] > 6u;
const uint first_row = prefill ? group.y : 0u;
const uint end_row = prefill ? first_row + 1u : args.u[4];
for (uint row = first_row; row < end_row; row++) {
const float x0 = row < 3u
? qwen_bf16(state[(ulong)row * width + channel])
: qkv[(ulong)(row - 3u) * width + channel];
@@ -2000,7 +2405,10 @@ kernel void kernel_qwen_gdn_conv_norm_rows(
: qkv[(ulong)(row - 1u) * width + channel];
const float x3 = qkv[(ulong)row * width + channel];
const float convolved = w0 * x0 + w1 * x1 + w2 * x2 + w3 * x3;
const float activated = convolved / (1.0f + exp(-convolved));
const float rounded = qwen_round_bf16(convolved);
const float activated = prefill
? qwen_round_bf16(rounded * qwen_silu_sigmoid_bf16(rounded))
: convolved / (1.0f + exp(-convolved));
if (value_channel) {
v_out[(ulong)row * (width - 2u * key_width) + channel - 2u * key_width] =
qwen_round_bf16(activated);
@@ -2013,10 +2421,12 @@ kernel void kernel_qwen_gdn_conv_norm_rows(
const uint first = (simd_group / 4u) * 4u;
sum = partial[first] + partial[first + 1u] +
partial[first + 2u] + partial[first + 3u];
const float normalized = values[tid] * rsqrt(sum + 1.0e-6f);
const float normalized = values[tid] * (prefill
? precise::rsqrt(sum + 1.0e-6f) : rsqrt(sum + 1.0e-6f));
if (channel < key_width) {
q_out[(ulong)row * key_width + channel] =
qwen_round_bf16(normalized * inverse_scale);
prefill ? qwen_round_bf16(qwen_round_bf16(normalized) * 0.08837890625f)
: qwen_round_bf16(normalized * inverse_scale);
} else {
k_out[(ulong)row * key_width + channel - key_width] =
qwen_round_bf16(normalized);
@@ -2024,6 +2434,7 @@ kernel void kernel_qwen_gdn_conv_norm_rows(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (prefill && group.y != 0u) return;
for (uint tail = 0u; tail < 3u; tail++) {
const uint sequence = args.u[4] + tail;
const float value = sequence < 3u
@@ -2609,18 +3020,28 @@ kernel void kernel_qwen_gdn_norm_gate(
}
}
kernel void kernel_qwen_swiglu(
template <typename T>
kernel void kernel_qwen_swiglu_typed(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device T *out [[buffer(1)]],
device const float *gate [[buffer(2)]],
device const float *up [[buffer(3)]],
uint index [[thread_position_in_grid]]) {
if (index >= args.u[0]) return;
const float activated = qwen_round_bf16(
gate[index] * qwen_silu_sigmoid_bf16(gate[index]));
out[index] = qwen_round_bf16(activated * up[index]);
out[index] = (T)qwen_round_bf16(activated * up[index]);
}
template [[host_name("kernel_qwen_swiglu")]]
kernel void kernel_qwen_swiglu_typed<float>(
constant qwen_kernel_args &, device float *, device const float *,
device const float *, uint);
template [[host_name("kernel_qwen_swiglu_bf16")]]
kernel void kernel_qwen_swiglu_typed<bfloat>(
constant qwen_kernel_args &, device bfloat *, device const float *,
device const float *, uint);
kernel void kernel_qwen_unpack_gdn_inputs(
constant qwen_kernel_args &args [[buffer(0)]],
device float *qkv [[buffer(1)]],
@@ -3017,10 +3438,18 @@ kernel void kernel_qwen_qsa_scores(
device float *scores [[buffer(1)]],
device const float *query [[buffer(2)]],
device const ushort *pooled [[buffer(3)]],
uint block [[thread_position_in_grid]]) {
uint2 group [[thread_position_in_grid]]) {
const uint block = group.x;
const uint row = group.y;
const uint dim = args.u[0];
if (block >= args.u[2]) return;
const ulong score_index = (ulong)row * args.u[2] + block;
if (args.u[5] != 0u && block >= (args.u[3] + row) / args.u[5]) {
scores[score_index] = -INFINITY;
return;
}
float score = 0.0f;
query += (ulong)row * args.u[1] * dim;
for (uint head = 0; head < args.u[1]; head++) {
float head_score = 0.0f;
for (uint i = 0; i < dim; i++) {
@@ -3030,14 +3459,14 @@ kernel void kernel_qwen_qsa_scores(
}
score += max(head_score, 0.0f);
}
scores[block] = score * args.f[0];
scores[score_index] = score * args.f[0];
}
kernel void kernel_qwen_qsa_sort_blocks(
constant qwen_kernel_args &args [[buffer(0)]],
device int *selected [[buffer(1)]],
uint gid [[thread_position_in_grid]]) {
if (gid != 0u) return;
uint row [[thread_position_in_grid]]) {
selected += (ulong)row * args.u[0];
for (uint i = 1; i < args.u[0]; i++) {
const int value = selected[i];
uint j = i;
@@ -3049,15 +3478,45 @@ kernel void kernel_qwen_qsa_sort_blocks(
}
}
kernel void kernel_qwen_qsa_mask(
constant qwen_kernel_args &args [[buffer(0)]],
device uchar *mask [[buffer(1)]],
device const int *selected [[buffer(2)]],
uint row [[threadgroup_position_in_grid]],
uint lane [[thread_index_in_threadgroup]]) {
if (row >= args.u[4]) return;
const uint tokens = args.u[0];
const uint visible = args.u[8] + row;
const uint tail_start = (visible / args.u[5]) * args.u[5];
const ulong mask_base = (ulong)row * tokens;
for (uint token = lane; token < tokens; token += args.u[12]) {
mask[mask_base + token] = token >= tail_start && token < visible;
}
threadgroup_barrier(mem_flags::mem_device);
if (lane < args.u[6]) {
const int block = selected[(ulong)row * args.u[6] + lane];
if (block >= 0) {
const uint token0 = (uint)block * args.u[5];
if (token0 < tail_start) {
for (uint within = 0u; within < args.u[5]; within++) {
mask[mask_base + token0 + within] = true;
}
}
}
}
}
kernel void kernel_qwen_sparse_attention(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
device const float *query [[buffer(2)]],
device const ushort *cache [[buffer(3)]],
device const int *selected [[buffer(4)]],
uint head [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint2 group [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]],
uint lane [[thread_index_in_simdgroup]]) {
const uint head = group.x;
const uint row = group.y;
const uint heads = args.u[0];
const uint kv_heads = args.u[1];
const uint dim = args.u[2];
@@ -3071,7 +3530,14 @@ kernel void kernel_qwen_sparse_attention(
threadgroup float sum_exp_scores[simdgroups];
threadgroup float outputs[simdgroups * 256u];
const float attention_scale = precise::rsqrt((float)dim);
q[tid] = query[(ulong)head * dim + tid] * attention_scale;
const ulong query_base = ((ulong)row * heads + head) * dim;
const ulong output_base = query_base;
selected += (ulong)row * args.u[4];
const uint tokens = args.u[8] != 0u ? args.u[8] + row : args.u[3];
const uint tail_start = args.u[8] != 0u
? (tokens / args.u[5]) * args.u[5]
: args.u[6];
q[tid] = query[query_base + tid] * attention_scale;
threadgroup_barrier(mem_flags::mem_threadgroup);
float result[values_per_lane] = {0.0f};
@@ -3101,8 +3567,8 @@ kernel void kernel_qwen_sparse_attention(
max_score = new_max;
}
}
for (uint token = args.u[6] + simd_group;
token < args.u[3];
for (uint token = tail_start + simd_group;
token < tokens;
token += simdgroups) {
const ulong base = (ulong)token * kv_heads * dim * 2u +
(ulong)kv_head * dim + column;
@@ -3142,7 +3608,7 @@ kernel void kernel_qwen_sparse_attention(
merged_sum += sum_exp_scores[group] * weight;
merged_value += outputs[group * dim + tid] * weight;
}
out[(ulong)head * dim + tid] = qwen_round_bf16(merged_value / merged_sum);
out[output_base + tid] = qwen_round_bf16(merged_value / merged_sum);
}
kernel void kernel_qwen_dense_attention_masked(
@@ -3393,6 +3859,52 @@ kernel void kernel_qwen_attention_fallback_softmax(
}
}
kernel void kernel_qwen_attention_fallback_softmax_wide(
constant qwen_kernel_args &args [[buffer(0)]],
device ushort *out [[buffer(1)]],
device const ushort *input [[buffer(2)]],
uint group [[threadgroup_position_in_grid]],
uint lid [[thread_position_in_threadgroup]],
uint lane [[thread_index_in_simdgroup]],
uint simd_group [[simdgroup_index_in_threadgroup]]) {
threadgroup float local_max[32];
threadgroup float local_normalizer[32];
const uint tokens = args.u[3];
const uint threads = args.u[12];
const ulong base = (ulong)group * tokens;
float maximum = -FLT_MAX;
for (uint token = lid; token < tokens; token += threads) {
maximum = max(maximum, qwen_bf16(input[base + token]));
}
maximum = simd_max(maximum);
if (lane == 0u) local_max[simd_group] = maximum;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_group == 0u) {
maximum = simd_max(local_max[lane]);
if (lane == 0u) local_max[0] = maximum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
maximum = local_max[0];
float normalizer = 0.0f;
for (uint token = lid; token < tokens; token += threads) {
normalizer += fast::exp(qwen_bf16(input[base + token]) - maximum);
}
normalizer = simd_sum(normalizer);
if (lane == 0u) local_normalizer[simd_group] = normalizer;
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_group == 0u) {
normalizer = simd_sum(local_normalizer[lane]);
if (lane == 0u) local_normalizer[0] = normalizer;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
normalizer = 1.0f / local_normalizer[0];
for (uint token = lid; token < tokens; token += threads) {
out[base + token] = qwen_to_bf16(
fast::exp(qwen_bf16(input[base + token]) - maximum) * normalizer);
}
}
kernel void kernel_qwen_attention_fallback_output(
constant qwen_kernel_args &args [[buffer(0)]],
device float *out [[buffer(1)]],
@@ -3449,6 +3961,24 @@ kernel void kernel_qwen_prepare_dense_kv(
values[target] = cache[source + width];
}
kernel void kernel_qwen_prepare_dense_kv_f16(
constant qwen_kernel_args &args [[buffer(0)]],
device half *keys [[buffer(1)]],
device const ushort *cache [[buffer(2)]],
device half *values [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]) {
const uint width = args.u[0];
const uint dim = args.u[1];
const uint tokens = args.u[3];
if (gid.x >= width || gid.y >= tokens) return;
const uint head = gid.x / dim;
const uint column = gid.x % dim;
const ulong source = (ulong)gid.y * width * 2u + gid.x;
const ulong target = ((ulong)head * tokens + gid.y) * dim + column;
keys[target] = (half)qwen_bf16(cache[source]);
values[target] = (half)qwen_bf16(cache[source + width]);
}
kernel void kernel_qwen_float_to_bf16(
constant qwen_kernel_args &args [[buffer(0)]],
device ushort *out [[buffer(1)]],